Update app.py

app.py

@@ -756,7 +756,7 @@ print(result)
 @spaces.GPU
 @torch.no_grad()
 def generate_image(
-    prompt, width, height, guidance, seed,
+    prompt, width, height, guidance, inference_steps, seed,
     do_img2img, init_image, image2image_strength, resize_img,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -786,7 +786,7 @@ def generate_image(
     generator = torch.Generator(device=device).manual_seed(seed)
     x = torch.randn(1, 16, 2 * math.ceil(height / 16), 2 * math.ceil(width / 16), device=device, dtype=torch.bfloat16, generator=generator) 
     
-    num_steps = 16
+    num_steps = inference_steps
     timesteps = get_schedule(num_steps, (x.shape[-1] * x.shape[-2]) // 4, shift=True)
 
     if do_img2img and init_image is not None:
@@ -876,6 +876,13 @@ def create_demo():
                 width = gr.Slider(minimum=128, maximum=2048, step=64, label="Width", value=1360)
                 height = gr.Slider(minimum=128, maximum=2048, step=64, label="Height", value=768)
                 guidance = gr.Slider(minimum=1.0, maximum=5.0, step=0.1, label="Guidance", value=3.5)
+                inference_steps = gr.Slider(
+                    label="Inference steps",
+                    minimum=1,
+                    maximum=30,
+                    step=1,
+                    value=16,
+                )
                 seed = gr.Number(label="Seed", precision=-1)
                 do_img2img = gr.Checkbox(label="Image to Image", value=False)
                 init_image = gr.Image(label="Input Image", visible=False)
@@ -895,7 +902,7 @@ def create_demo():
 
         generate_button.click(
             fn=generate_image,
-            inputs=[prompt, width, height, guidance, seed, do_img2img, init_image, image2image_strength, resize_img],
+            inputs=[prompt, width, height, guidance, inference_steps, seed, do_img2img, init_image, image2image_strength, resize_img],
             outputs=[output_image, output_seed]
         )