Update app.py

app.py

@@ -28,28 +28,50 @@ pipe.to("cuda")
 
 
 
+# @spaces.GPU
+# def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed):
+#     """Generate an image from the given prompt."""
+#     print(prompt)
+#     if randomize_seed:
+#         seed = torch.randint(0, 2**32 - 1, (1,)).item()
+    
+#     generator = torch.Generator("cuda").manual_seed(int(seed))
+#     image = pipe(
+#         prompt=prompt,
+#         height=int(height),
+#         width=int(width),
+#         num_inference_steps=int(num_inference_steps),
+#         guidance_scale=0.0,  # Guidance should be 0 for Turbo models
+#         generator=generator,
+#         max_sequence_length=1024,
+#         num_images_per_prompt=2
+        
+#     ).images[0]
+    
+#     return image, seed
+
+
+
 @spaces.GPU
 def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed):
-    """Generate an image from the given prompt."""
-    print(prompt)
+    """Generate images from the given prompt."""
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
     
     generator = torch.Generator("cuda").manual_seed(int(seed))
-    image = pipe(
+    result = pipe(
         prompt=prompt,
         height=int(height),
         width=int(width),
         num_inference_steps=int(num_inference_steps),
-        guidance_scale=0.0,  # Guidance should be 0 for Turbo models
+        guidance_scale=0.0,
         generator=generator,
         max_sequence_length=1024,
         num_images_per_prompt=2
-        
-    ).images[0]
+    )
     
-    return image, seed
-
+    # return both images
+    return result.images, seed
 
 # Example prompts
 examples = [
@@ -119,7 +141,7 @@ with gr.Blocks(title="Z-Image-Turbo Demo") as demo:
             generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
         
         with gr.Column(scale=1):
-            output_image = gr.Image(
+            output_image = gr.Gallery(
                 label="Generated Image",
                 type="pil",
             )
@@ -137,18 +159,17 @@ with gr.Blocks(title="Z-Image-Turbo Demo") as demo:
 
      
     # Connect the generate button
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
-        outputs=[output_image, used_seed],
-    )
-    
-    # Also allow generating by pressing Enter in the prompt box
+   generate_btn.click(
+       fn=generate_image,
+       inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
+        outputs=[output_images, used_seed],
+        )
     prompt.submit(
         fn=generate_image,
         inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
-        outputs=[output_image, used_seed],
-    )
+      outputs=[output_images, used_seed],
+       )
+
 
 if __name__ == "__main__":
     demo.launch()