steps

app.py

@@ -53,7 +53,7 @@ NUM_INFERENCE_STEPS = 4
 
 
 #@spaces.GPU
-def infer(prompt, seed, randomize_seed):
+def infer(prompt, steps, seed, randomize_seed):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -62,7 +62,7 @@ def infer(prompt, seed, randomize_seed):
     image = pipe(
         prompt=prompt,
         guidance_scale=0,
-        num_inference_steps=NUM_INFERENCE_STEPS,
+        num_inference_steps=steps,
         generator=generator,
     ).images[0]
 
@@ -100,7 +100,7 @@ with gr.Blocks(css=css) as demo:
 #        # ⚡ Flash Diffusion: FlashPixart ⚡
 #        This is an interactive demo of [Flash Diffusion](https://gojasper.github.io/flash-diffusion-project/), a diffusion distillation method proposed in [Flash Diffusion: Accelerating Any Conditional
 #        Diffusion Model for Few Steps Image Generation](http://arxiv.org/abs/2406.02347) *by Clément Chadebec, Onur Tasar, Eyal Benaroche and Benjamin Aubin.*
-#        [This model](https://huggingface.co/jasperai/flash-pixart) is a **66.5M** LoRA distilled version of [Pixart-α](https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS) model that is able to generate 1024x1024 images in **4 steps**. 
+#        [This model](https://huggingface.co/jasperai/flash-pixart) is a **66.5M** LoRA distilled version of [Pixart-α](https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS) model that is able to generate 1024x1024 images in **4 steps**.
 #        Currently running on {power_device}.
 #        """
 #        )
@@ -132,6 +132,12 @@ with gr.Blocks(css=css) as demo:
                 step=1,
                 value=0,
             )
+            steps = gr.Number(
+                label="Steps",
+                minimum=1,
+                step=1,
+                value=4
+            )
 
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
@@ -144,7 +150,7 @@ with gr.Blocks(css=css) as demo:
     gr.on(
         [run_button.click, seed.change, randomize_seed.change, prompt.submit],
         fn=infer,
-        inputs=[prompt, seed, randomize_seed],
+        inputs=[prompt, steps, seed, randomize_seed],
         outputs=[result],
         show_progress="minimal",
         show_api=False,