Update app.py

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
 import jax
+from diffusers import FlaxStableDiffusionPipeline
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard
-from diffusers import FlaxStableDiffusionPipeline
 
 pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
     "bguisard/stable-diffusion-nano",
@@ -13,11 +13,11 @@ def generate_image(prompt: str, inference_steps: int = 30, prng_seed: int = 0):
     rng = jax.random.PRNGKey(int(prng_seed))
     rng = jax.random.split(rng, jax.device_count())
     p_params = replicate(pipeline_params)
-    
+
     num_samples = 1
     prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)
     prompt_ids = shard(prompt_ids)
-    
+
     images = pipeline(
         prompt_ids=prompt_ids,
         params=p_params,
@@ -30,7 +30,7 @@ def generate_image(prompt: str, inference_steps: int = 30, prng_seed: int = 0):
 
     images = images.reshape((num_samples,) + images.shape[-3:])
     images = pipeline.numpy_to_pil(images)
-    return images
+    return images[0]
 
 
 prompt_input = gr.inputs.Textbox(
@@ -44,7 +44,7 @@ seed_input = gr.inputs.Number(default=0, label="Seed")
 app = gr.Interface(
     fn=generate_image,
     inputs=[prompt_input, inf_steps_input, seed_input],
-    outputs=gr.Image(shape=(128, 128)),
+    outputs="image",
     title="Stable Diffusion Nano",
     description=(
         "Based on stable diffusion and fine-tuned on 128x128 images, "