Update app to 4-image UI with extra controls

app.py

@@ -3,6 +3,8 @@ import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
 
+MAX_SEED = 2**32 - 1
+
 # Load the pipeline once at startup
 print("Loading Z-Image-Turbo pipeline...")
 pipe = DiffusionPipeline.from_pretrained(
@@ -19,22 +21,46 @@ spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant
 print("Pipeline loaded!")
 
 @spaces.GPU
-def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
-    """Generate an image from the given prompt."""
+def generate_image(
+    prompt,
+    negative_prompt,
+    height,
+    width,
+    num_inference_steps,
+    guidance_scale,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate 4 images with seeds: seed, 2x, 3x, 4x (mod MAX_SEED)."""
     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,
-    ).images[0]
-    
-    return image, seed
+        seed = torch.randint(0, MAX_SEED, (1,)).item()
+
+    base_seed = int(seed) % MAX_SEED
+    if base_seed < 0:
+        base_seed += MAX_SEED
+
+    seeds = [(base_seed * i) % MAX_SEED for i in range(1, 5)]
+
+    neg_prompt = None
+    if isinstance(negative_prompt, str) and negative_prompt.strip():
+        neg_prompt = negative_prompt
+
+    images = []
+    for s in seeds:
+        generator = torch.Generator("cuda").manual_seed(int(s))
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=neg_prompt,
+            height=int(height),
+            width=int(width),
+            num_inference_steps=int(num_inference_steps),
+            guidance_scale=float(guidance_scale),  # 0.0 is recommended default for Turbo
+            generator=generator,
+        ).images[0]
+        images.append(image)
+
+    return images, ", ".join(str(s) for s in seeds)
 
 
 # Example prompts
@@ -64,7 +90,13 @@ with gr.Blocks(title="Z-Image-Turbo Demo") as demo:
                 placeholder="Enter your image description...",
                 lines=4,
             )
-            
+
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="Things you don't want in the image...",
+                lines=3,
+            )
+
             with gr.Row():
                 height = gr.Slider(
                     minimum=512,
@@ -90,6 +122,15 @@ with gr.Blocks(title="Z-Image-Turbo Demo") as demo:
                     label="Inference Steps",
                     info="9 steps results in 8 DiT forwards",
                 )
+
+            guidance_scale = gr.Slider(
+                minimum=0.0,
+                maximum=7.0,
+                value=0.0,
+                step=0.1,
+                label="CFG Guidance Scale",
+                info="0 = no CFG (recommended for Turbo models)",
+            )
             
             with gr.Row():
                 seed = gr.Number(
@@ -105,12 +146,14 @@ 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(
-                label="Generated Image",
-                type="pil",
+            output_images = gr.Gallery(
+                label="Generated Images",
+                columns=2,
+                rows=2,
+                preview=True,
             )
-            used_seed = gr.Number(
-                label="Seed Used",
+            used_seeds = gr.Textbox(
+                label="Seeds Used (base, 2x, 3x, 4x)",
                 interactive=False,
             )
     
@@ -126,16 +169,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_image, used_seed],
+        inputs=[prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed],
+        outputs=[output_images, used_seeds],
     )
     
     # Also allow generating by pressing Enter in the prompt box
     prompt.submit(
         fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
-        outputs=[output_image, used_seed],
+        inputs=[prompt, negative_prompt, height, width, num_inference_steps, guidance_scale, seed, randomize_seed],
+        outputs=[output_images, used_seeds],
     )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch(mcp_server=True)