Update app.py

app.py

@@ -1,8 +1,226 @@
+import io
+import tempfile
+import zipfile
+
+import random
 import torch
 import spaces
 import gradio as gr
-from diffusers import DiffusionPipeline
+from diffusers import ZImagePipeline
+MAX_SEED = 2**32 - 1
+
 
+# ===== Custom aesthetic =====
+# Neo-noir dusk palette with cyan + amber accents, glass panels, and subtle grain.
+CUSTOM_CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');
+:root {
+    /* Light Mode (Professional & Clean) */
+    --bg: #fdfdfd;
+    --panel: rgba(255, 255, 255, 0.95);
+    --card: #ffffff;
+    --border: #e5e7eb;
+    --border-hover: #d1d5db;
+    --text: #111827;
+    --text-secondary: #4b5563;
+    --muted: #9ca3af;
+    --accent: #0f172a; /* Dark sleek accent for professionalism */
+    --accent-hover: #1e293b;
+    --accent-text: #ffffff;
+    --primary-gradient: linear-gradient(135deg, #0f172a 0%, #334155 100%);
+    --glow: 0 0 0 transparent;
+    --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
+    --shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.05), 0 2px 4px -1px rgba(0, 0, 0, 0.03);
+    --shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.05), 0 4px 6px -2px rgba(0, 0, 0, 0.03);
+    --radius: 12px;
+    --input-bg: #ffffff;
+    --input-border: #e2e8f0;
+    --checkbox-bg: #f1f5f9;
+    --body-bg: #f8fafc; /* Very subtle cool gray */
+    --font-heading: 'Inter', -apple-system, sans-serif;
+    --font-body: 'Inter', -apple-system, sans-serif;
+}
+.dark {
+    /* Dark Mode (Neo-Noir Polished) */
+    --bg: #05080f;
+    --panel: rgba(12, 18, 32, 0.85);
+    --card: rgba(18, 28, 46, 0.70);
+    --border: rgba(36, 224, 194, 0.15);
+    --border-hover: rgba(36, 224, 194, 0.3);
+    --text: #e9f3ff;
+    --text-secondary: #94a3b8;
+    --muted: #64748b;
+    --accent: #24e0c2;
+    --accent-hover: #18cdb0;
+    --accent-text: #041019;
+    --primary-gradient: linear-gradient(120deg, #24e0c2 0%, #ffb347 100%);
+    --glow: 0 8px 32px rgba(36, 224, 194, 0.12);
+    --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.2);
+    --shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.3);
+    --shadow-lg: 0 20px 40px -5px rgba(0, 0, 0, 0.4);
+    --radius: 16px;
+    --input-bg: rgba(255,255,255,0.03);
+    --input-border: rgba(255,255,255,0.08);
+    --checkbox-bg: #0d1829;
+    --body-bg: radial-gradient(circle at 20% 20%, rgba(36, 224, 194, 0.06), transparent 35%),
+               radial-gradient(circle at 82% 12%, rgba(0, 156, 196, 0.06), transparent 35%),
+               linear-gradient(145deg, #05080f 0%, #080f1e 100%);
+    --font-heading: 'Inter', -apple-system, sans-serif;
+    --font-body: 'Inter', -apple-system, sans-serif;
+}
+body, .gradio-container {
+    font-family: var(--font-body) !important;
+    background: var(--body-bg) !important;
+    color: var(--text);
+    min-height: 100vh;
+}
+/* Titles & Typography */
+.gradio-container .prose h1,
+.gradio-container .prose h2,
+.gradio-container .prose h3 {
+    font-family: var(--font-heading);
+    letter-spacing: -0.025em;
+    font-weight: 700;
+    color: var(--text);
+}
+.gradio-container .prose h1 {
+    font-size: 2.25rem;
+    margin-bottom: 0.5rem;
+    background: var(--primary-gradient);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    display: inline-block;
+}
+.gradio-container * { letter-spacing: -0.01em; }
+/* Panels & Cards */
+.gr-block, .gr-panel, .gr-group {
+    background: var(--panel);
+    border: 1px solid var(--border);
+    border-radius: var(--radius);
+    box-shadow: var(--shadow-sm);
+    backdrop-filter: blur(8px);
+    transition: box-shadow 0.2s ease, border-color 0.2s ease;
+}
+.hero-card {
+    background: var(--card);
+    border: 1px solid var(--border);
+    padding: 24px;
+    border-radius: var(--radius);
+    box-shadow: var(--shadow-md);
+    position: relative;
+    overflow: hidden;
+}
+.tagline {
+    display: inline-flex;
+    align-items: center;
+    gap: 8px;
+    padding: 6px 14px;
+    background: var(--input-bg);
+    border: 1px solid var(--border);
+    border-radius: 999px;
+    font-size: 0.875rem;
+    font-weight: 500;
+    color: var(--text-secondary);
+    margin-bottom: 12px;
+}
+.hero-card p {
+    color: var(--text-secondary);
+    font-size: 1.05rem;
+    line-height: 1.6;
+    max-width: 65ch;
+}
+/* Inputs */
+textarea, input:not([type='checkbox']):not([type='radio']),
+.gr-input, .gr-textbox, .gr-number, .gr-slider input {
+    background: var(--input-bg) !important;
+    border: 1px solid var(--input-border) !important;
+    border-radius: 10px !important;
+    color: var(--text) !important;
+    font-family: var(--font-body);
+    transition: all 0.2s ease;
+}
+textarea:focus, input:focus, .gr-input:focus-within {
+    border-color: var(--text-secondary) !important;
+    box-shadow: 0 0 0 2px rgba(var(--accent), 0.1);
+}
+label, .gr-box label {
+    color: var(--text-secondary) !important;
+    font-weight: 600;
+    font-size: 0.875rem;
+    margin-bottom: 6px;
+    text-transform: none !important;
+}
+/* Sliders */
+.gr-slider input[type='range'] {
+    accent-color: var(--accent);
+}
+/* Buttons */
+.gr-button-primary, button.primary {
+    background: var(--primary-gradient) !important;
+    color: var(--accent-text) !important;
+    font-weight: 600 !important;
+    border: 1px solid rgba(255,255,255,0.1) !important;
+    box-shadow: var(--shadow-md);
+    border-radius: 10px !important;
+    padding: 10px 24px;
+    transition: transform 0.1s, box-shadow 0.2s;
+}
+.gr-button-primary:hover {
+    transform: translateY(-1px);
+    box-shadow: var(--shadow-lg);
+    filter: brightness(1.1);
+}
+.gr-button-secondary, button.secondary, .gr-downloadbutton {
+    background: var(--input-bg) !important;
+    border: 1px solid var(--border) !important;
+    color: var(--text) !important;
+    font-weight: 500;
+    border-radius: 10px !important;
+    box-shadow: var(--shadow-sm);
+}
+.gr-button-secondary:hover {
+    border-color: var(--border-hover) !important;
+    background: var(--card) !important;
+}
+.gr-downloadbutton, .gr-downloadbutton > button { width: 100%; }
+/* Gallery */
+.gr-gallery {
+    background: var(--input-bg);
+    border-radius: var(--radius);
+    border: 1px solid var(--border);
+    padding: 8px;
+}
+.gr-gallery .thumbnail-item {
+    border-radius: 8px;
+    overflow: hidden;
+    box-shadow: var(--shadow-sm);
+    border: 1px solid transparent;
+    transition: all 0.2s;
+}
+.gr-gallery .thumbnail-item:hover {
+    box-shadow: var(--shadow-md);
+    transform: scale(1.02);
+}
+.gr-gallery img { object-fit: cover; }
+/* Footer */
+.footer-note {
+    color: var(--muted);
+    font-size: 0.875rem;
+    text-align: center;
+    margin-top: 2rem;
+    opacity: 0.8;
+}
+.footer-note a {
+    color: var(--text-secondary);
+    text-decoration: none;
+    border-bottom: 1px dotted var(--muted);
+}
+.footer-note a:hover {
+    color: var(--accent);
+    border-bottom-style: solid;
+}
+"""
 # Load the pipeline once at startup
 print("Loading Z-Image-Turbo pipeline...")
 pipe = DiffusionPipeline.from_pretrained(
@@ -19,241 +237,239 @@ pipe.to("cuda")
 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,
+    images_count,
+    num_inference_steps,
+    guidance_scale,
+    seed,
+    randomize_seed,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate N images using a deterministic seed cascade (x1..xN)."""
     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,
-        generator=generator,
-    ).images[0]
-    
-    return image, seed
+        seed = random.randint(0, MAX_SEED)
+
+    base_seed = int(seed) % MAX_SEED
+    if base_seed < 0:
+        base_seed += MAX_SEED
+
+    # Cap to prevent excessive VRAM usage / latency spikes on the demo space
+    images_count = max(1, min(int(images_count), 12))
+
+    seeds = [(base_seed * i) % MAX_SEED for i in range(1, images_count + 1)]
+
+    neg_prompt = None
+    if isinstance(negative_prompt, str) and negative_prompt.strip():
+        neg_prompt = negative_prompt
+
+    images = []
+    image_paths = []
+    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)
+        tmp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+        image.save(tmp_img.name, format="PNG")
+        image_paths.append(tmp_img.name)
+
+    return images, ", ".join(str(s) for s in seeds), image_paths, base_seed
+
+
+
+def append_history(new_images, history):
+    """Append new images to the history state."""
+    if history is None:
+        history = []
+    updated_history = history + new_images
+    return updated_history, updated_history
+
+
+def package_zip(image_paths):
+    """Pack the current image list into a ZIP file for download."""
+    if not image_paths:
+        raise gr.Error("No images in history to download.")
+
+    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
+    with zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as zf:
+        for idx, path in enumerate(image_paths, start=1):
+            # Store as image_001.png, image_002.png, ...
+            zf.write(path, arcname=f"image_{idx:03d}.png")
+
+    tmp.flush()
+    return tmp.name
+
 
 # Example prompts
 examples = [
-    
-    
-    
-    
-    
+    ["Astronaut riding a horse on Mars, cinematic lighting, sci-fi concept art, highly detailed"],
+    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff, magical forest background"],
 ]
 
-# Custom theme with modern aesthetics (Gradio 6)
-custom_theme = gr.themes.Soft(
-    primary_hue="green",
-    secondary_hue="amber",
-    neutral_hue="slate",
-    font=gr.themes.GoogleFont("Inter"),
-    text_size="lg",
-    spacing_size="md",
-    radius_size="lg"
-).set(
-    button_primary_background_fill="*primary_500",
-    button_primary_background_fill_hover="*primary_600",
-    block_title_text_weight="600",
-)
-
 # Build the Gradio interface
-with gr.Blocks(fill_height=True) as demo:
-    # Header
+# Build the Gradio interface
+with gr.Blocks(title="Z-Image-Turbo Demo", css=CUSTOM_CSS, analytics_enabled=False) as demo:
+    image_state = gr.State([])
+    history_state = gr.State([])
     gr.Markdown(
         """
-        # Z-Image-Turbo
+        <div class="hero-card">
+          <div class="tagline">⚡ Turbo diffusion · 8 steps · CUDA ready</div>
+          <h1>Z‑Image Turbo Studio</h1>
+          <p>Draft up to twelve stylized candidates in one pass. Neo‑noir gradients, glass panels, and crisp typography keep the tooling out of your way while you explore ideas.</p>
+        </div>
         """,
-        elem_classes="header-text"
+        sanitize_html=False,
     )
     
-    with gr.Row(equal_height=False):
-        # Left column - Input controls
-        with gr.Column(scale=1, min_width=600):
+    with gr.Row():
+        with gr.Column(scale=1):
             prompt = gr.Textbox(
-                label="✨ Your Prompt",
-                placeholder="Describe the image you want to create...",
-                lines=5,
-                max_lines=10,
-                autofocus=True,
+                label="Prompt",
+                placeholder="e.g. bioluminescent reef city at dusk, cinematic, anamorphic glow",
+                lines=4,
             )
-            
-            with gr.Accordion("⚙️ Advanced Settings", open=True):
-                with gr.Row():
-                    height = gr.Slider(
-                        minimum=512,
-                        maximum=2048,
-                        value=1024,
-                        step=64,
-                        label="Height",
-                        info="Pixels_______________________"
-                    )
-                    width = gr.Slider(
-                        minimum=512,
-                        maximum=2048,
-                        value=1024,
-                        step=64,
-                        label="Width",
-                        info="Pixels________________________"
-                    )
-                
+
+            negative_prompt = gr.Textbox(
+                label="Negative Prompt",
+                placeholder="noise, blur, extra limbs, text watermark",
+                lines=3,
+            )
+
+            with gr.Row():
+                height = gr.Slider(
+                    minimum=512,
+                    maximum=2048,
+                    value=1024,
+                    step=64,
+                    label="Height",
+                )
+                width = gr.Slider(
+                    minimum=512,
+                    maximum=2048,
+                    value=1024,
+                    step=64,
+                    label="Width",
+                )
+
+            with gr.Row():
                 num_inference_steps = gr.Slider(
                     minimum=1,
                     maximum=20,
                     value=9,
                     step=1,
                     label="Inference Steps",
-                    info="9 steps = 8 DiT (recommended)"
+                    info="9 steps → 8 DiT forwards",
                 )
-                
-                with gr.Row():
-                    randomize_seed = gr.Checkbox(
-                        label="🎲 Random Seed",
-                        value=True,
-                    )
-                    seed = gr.Number(
-                        label="Seed",
-                        value=42,
-                        precision=0,
-                        visible=False,
-                    )
-                
-                def toggle_seed(randomize):
-                    return gr.Number(visible=not randomize)
-                
-                randomize_seed.change(
-                    toggle_seed,
-                    inputs=[randomize_seed],
-                    outputs=[seed]
-                )
-            
-            generate_btn = gr.Button(
-                "🚀 Generate Image",
-                variant="primary",
-                size="lg",
-                scale=1
+
+            images_count = gr.Slider(
+                minimum=1,
+                maximum=12,
+                value=4,
+                step=1,
+                label="Images",
+                info="1–12 (higher counts use more VRAM)",
             )
-            
-            # Example prompts
-            gr.Examples(
-                examples=examples,
-                inputs=[prompt],
-                label=" ",
-                examples_per_page=5,
+
+            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(
+                    label="Base Seed",
+                    value=42,
+                    precision=0,
+                )
+                randomize_seed = gr.Checkbox(
+                    label="Randomize",
+                    value=True,
+                    interactive=True,
+                )
+            
+            generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
         
-        # Right column - Output
-        with gr.Column(scale=1, min_width=320):
-            output_image = gr.Image(
-                label="Generated Image",
-                type="pil",
-                format="png",
-                show_label=False,
-                height=600,
-                buttons=["download", "share"],
+        with gr.Column(scale=1):
+            output_images = gr.Gallery(
+                label="Generated Grid",
+                columns=4,
+                rows=None,
+                preview=True,
             )
-            
-            used_seed = gr.Number(
-                label="🎲 Seed Used",
+            used_seeds = gr.Textbox(
+                label="Seed Cascade (x1 · x2 · ... · xN)",
                 interactive=False,
-                container=True,
+            )
+            history_gallery = gr.Gallery(
+                label="History",
+                columns=6,
+                rows=None,
+                preview=True,
+                object_fit="cover"
+            )
+            download_btn = gr.DownloadButton(
+                label="📦 Download All History (ZIP)",
             )
     
-    # Footer credits
+    gr.Markdown("### 💡 Quick Prompts")
+    gr.Examples(
+        examples=examples,
+        inputs=[prompt],
+        cache_examples=False,
+    )
+
     gr.Markdown(
         """
-        ---
-        <div style="text-align: center; opacity: 0.7; font-size: 0.9em; margin-top: 1rem;">
-        <strong>Model:</strong> <a href="https://huggingface.co/Tongyi-MAI/Z-Image-Turbo" target="_blank">Tongyi-MAI/Z-Image-Turbo</a> (Apache 2.0 License) • 
-        
-        
-        
+        <div class="footer-note">
+        Model: Tongyi-MAI/Z-Image-Turbo (Apache 2.0). Demo by <a href="https://z-image-turbo.tech" target="_blank">https://z-image-turbo.tech</a>
         </div>
         """,
-        elem_classes="footer-text"
+        sanitize_html=False,
     )
     
     # 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, images_count, num_inference_steps, guidance_scale, seed, randomize_seed],
+        outputs=[output_images, used_seeds, image_state, seed],
+    ).success(
+        fn=append_history,
+        inputs=[image_state, history_state],
+        outputs=[history_state, history_gallery],
     )
     
     # 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, images_count, num_inference_steps, guidance_scale, seed, randomize_seed],
+        outputs=[output_images, used_seeds, image_state, seed],
+    ).success(
+        fn=append_history,
+        inputs=[image_state, history_state],
+        outputs=[history_state, history_gallery],
+    )
+
+    download_btn.click(
+        fn=package_zip,
+        inputs=[history_state],
+        outputs=[download_btn],
     )
 
 if __name__ == "__main__":
-    demo.launch(
-        theme=custom_theme,
-        css="""
-        .header-text h1 {
-            font-size: 2.5rem !important;
-            font-weight: 700 !important;
-            margin-bottom: 0.5rem !important;
-            background: linear-gradient(135deg, #fbbf24 0%, #f59e0b 100%);
-            -webkit-background-clip: text;
-            -webkit-text-fill-color: transparent;
-            background-clip: text;
-        }
-        
-        .header-text p {
-            font-size: 1.1rem !important;
-            color: #64748b !important;
-            margin-top: 0 !important;
-        }
-        
-        .footer-text {
-            padding: 1rem 0;
-        }
-        
-        .footer-text a {
-            color: #f59e0b !important;
-            text-decoration: none !important;
-            font-weight: 500;
-        }
-        
-        .footer-text a:hover {
-            text-decoration: underline !important;
-        }
-        
-        /* Mobile optimizations */
-        @media (max-width: 768px) {
-            .header-text h1 {
-                font-size: 1.8rem !important;
-            }
-            
-            .header-text p {
-                font-size: 1rem !important;
-            }
-        }
-        
-        /* Smooth transitions */
-        button, .gr-button {
-            transition: all 0.2s ease !important;
-        }
-        
-        button:hover, .gr-button:hover {
-            transform: translateY(-1px);
-            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.15) !important;
-        }
-        
-        /* Better spacing */
-        .gradio-container {
-            max-width: 1400px !important;
-            margin: 0 auto !important;
-        }
-        """,
-        footer_links=[
-            "api",
-            "gradio"
-        ],
-        mcp_server=True
-    )
+    demo.launch(mcp_server=True, show_error=True)