Update app.py

app.py

@@ -1,3 +1,35 @@
 import gradio as gr
+import jax
+from diffusers import FlaxStableDiffusionPipeline
 
-gr.Interface.load("models/bguisard/stable-diffusion-nano").launch()
+pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
+    "bguisard/stable-diffusion-nano",
+)
+
+prng_seed = jax.random.PRNGKey(0)
+inference_steps = 50
+
+
+def generate_image(prompt: str):
+    prompt_ids = pipeline.prepare_inputs(prompt)
+    images = pipeline(
+        prompt_ids=prompt_ids,
+        params=pipeline_params,
+        prng_seed=prng_seed,
+        height=128,
+        width=128,
+        num_inference_steps=inference_steps,
+        jit=False,
+    ).images
+    pil_imgs = pipeline.numpy_to_pil(images)
+    return pil_imgs[0]
+
+
+app = gr.Interface(
+    fn=generate_image,
+    inputs="text",
+    outputs=gr.Image(shape=(128, 128)),
+    examples=[["A watercolor painting of a bird"]],
+)
+
+app.launch()