Update app.py

app.py

@@ -26,37 +26,14 @@ pipe = DiffusionPipeline.from_pretrained(
 
 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 images from the given prompt."""
+def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, num_images):
+    """Generate multiple images from the given prompt."""
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
+
+    # Clamp num_images to max 3
+    num_images = min(max(1, int(num_images)), 3)
     
     generator = torch.Generator("cuda").manual_seed(int(seed))
     result = pipe(
@@ -67,10 +44,9 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         guidance_scale=0.0,
         generator=generator,
         max_sequence_length=1024,
-        num_images_per_prompt=2
+        num_images_per_prompt=num_images
     )
     
-    # return both images
     return result.images, seed
 
 # Example prompts
@@ -116,7 +92,15 @@ with gr.Blocks(title="Z-Image-Turbo Demo") as demo:
                     step=64,
                     label="Width",
                 )
-            
+            with gr.Row():
+                 num_images = gr.Slider(
+                 minimum=1,
+                maximum=3,
+                value=2,
+                  step=1,
+                 label="Number of Images",
+                          )
+
             with gr.Row():
                 num_inference_steps = gr.Slider(
                     minimum=1,
@@ -161,15 +145,16 @@ 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_images, used_seed],
-        )
+        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, num_images],
+         outputs=[output_images, used_seed],
+         )
     prompt.submit(
-        fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
-      outputs=[output_images, used_seed],
+         fn=generate_image,
+        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, num_images],
+          outputs=[output_images, used_seed],
        )
 
 
+
 if __name__ == "__main__":
     demo.launch()