Update app.py

app.py

@@ -38,12 +38,14 @@ def show_images_save(x):
     grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
     return grid_im
 
-def generate(schedul):
+def generate(schedul, num):
     if schedul == "DDIMScheduler":
         scheduler = DDIMScheduler.from_pretrained(pipeline_name)
     else:
         scheduler = PNDMScheduler.from_pretrained(pipeline_name)
-    scheduler.set_timesteps(num_inference_steps=40)
+    if num <=0 or num >= 1000:
+        num = 40
+    scheduler.set_timesteps(num_inference_steps=num)
     x = torch.randn(1, 4, 64, 64).to(device)
     # Minimal sampling loop
     for i, t in tqdm(enumerate(scheduler.timesteps)):
@@ -54,10 +56,10 @@ def generate(schedul):
     # View the results
     return show_images_save(x)
 
-def ex(scheduler):
+def ex(scheduler, num):
     t = time()
     print(ctime(t))
-    return generate(scheduler), generate(scheduler), generate(scheduler), generate(scheduler)
+    return generate(scheduler, num), generate(scheduler, num), generate(scheduler, num), generate(scheduler, num)
 
 demo = gr.Blocks(css="#img_size {max-height: 128px} .container {max-width: 730px; margin: auto;} .min-h-\[15rem\]{min-height: 5rem !important;}")
 
@@ -77,7 +79,9 @@ with demo:
         """
     )
   with gr.Column():
-    model_name = gr.Dropdown(label="Base Scheduler", choices=[m.name for m in model], value=current_model.name)
+    with gr.Row().style(equal_height=True):
+        model_name = gr.Dropdown(label="Base Scheduler", choices=[m.name for m in model], value=current_model.name)
+        number = gr.Number(value="40", label="number of generation steps (Standard value 40, MAX 1000; The larger the number, the better the quality, but the longer it takes)", show_label=True)
     with gr.Row().style(equal_height=True):
         out = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
         out2 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
@@ -85,7 +89,7 @@ with demo:
         out3 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
         out4 = gr.Image(shape=(64,64), image_mode='RGBA', type='pil', elem_id='img_size')
   greet_btn = gr.Button("Generate")
-  greet_btn.click(fn=ex, inputs=[model_name], outputs=[out, out2, out3, out4])
+  greet_btn.click(fn=ex, inputs=[model_name, number], outputs=[out, out2, out3, out4])
   gr.HTML(
             """
                 <div class="footer">