Update app.py

Add seed and inference steps inputs.

app.py

@@ -6,30 +6,38 @@ pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
     "bguisard/stable-diffusion-nano",
 )
 
-prng_seed = jax.random.PRNGKey(0)
-inference_steps = 50
 
+def generate_image(prompt: str, inference_steps: int, prng_seed: int):
+    rng = jax.random.PRNGKey(int(prng_seed))
 
-def generate_image(prompt: str):
     prompt_ids = pipeline.prepare_inputs(prompt)
     images = pipeline(
         prompt_ids=prompt_ids,
         params=pipeline_params,
-        prng_seed=prng_seed,
+        prng_seed=rng,
         height=128,
         width=128,
-        num_inference_steps=inference_steps,
+        num_inference_steps=int(inference_steps),
         jit=False,
     ).images
+
     pil_imgs = pipeline.numpy_to_pil(images)
     return pil_imgs[0]
 
 
+prompt_input = gr.inputs.Textbox(
+    label="Prompt", placeholder="A watercolor painting of a bird"
+)
+inf_steps_input = gr.inputs.Slider(
+    minimum=1, maximum=100, default=30, step=1, label="Inference Steps"
+)
+seed_input = gr.inputs.Number(default=0, label="Seed")
+
 app = gr.Interface(
     fn=generate_image,
-    inputs="text",
+    inputs=[prompt_input, inf_steps_input, seed_input],
     outputs=gr.Image(shape=(128, 128)),
     examples=[["A watercolor painting of a bird"]],
 )
 
-app.launch()
+app.launch()