Update app.py

app.py

@@ -9,11 +9,16 @@ pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
 
 def generate_image(prompt: str, inference_steps: int = 30, prng_seed: int = 0):
     rng = jax.random.PRNGKey(int(prng_seed))
-
-    prompt_ids = pipeline.prepare_inputs(prompt)
+    rng = jax.random.split(rng, jax.device_count())
+    p_params = replicate(params)
+    
+    num_samples = 1
+    prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)
+    prompt_ids = shard(prompt_ids)
+    
     images = pipeline(
         prompt_ids=prompt_ids,
-        params=pipeline_params,
+        params=p_params,
         prng_seed=rng,
         height=128,
         width=128,
@@ -21,8 +26,9 @@ def generate_image(prompt: str, inference_steps: int = 30, prng_seed: int = 0):
         jit=True,
     ).images
 
-    pil_imgs = pipeline.numpy_to_pil(images)
-    return pil_imgs[0]
+    images = images.reshape((num_samples,) + output.shape[-3:])
+    images = pipeline.numpy_to_pil(images)
+    return images
 
 
 prompt_input = gr.inputs.Textbox(