Update app.py

app.py

@@ -1,226 +1,8 @@
-import io
-import tempfile
-import zipfile
-
-import random
 import torch
 import spaces
 import gradio as gr
-from diffusers import ZImagePipeline
-MAX_SEED = 2**32 - 1
-
+from diffusers import DiffusionPipeline
 
-# ===== 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(
@@ -237,239 +19,241 @@ pipe.to("cuda")
 print("Pipeline loaded!")
 
 @spaces.GPU
-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)."""
+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."""
     if randomize_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
-
+        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
 
 # 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
-# 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([])
+with gr.Blocks(fill_height=True) as demo:
+    # Header
     gr.Markdown(
         """
-        <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>
+        # Z-Image-Turbo
         """,
-        sanitize_html=False,
+        elem_classes="header-text"
     )
     
-    with gr.Row():
-        with gr.Column(scale=1):
+    with gr.Row(equal_height=False):
+        # Left column - Input controls
+        with gr.Column(scale=1, min_width=600):
             prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="e.g. bioluminescent reef city at dusk, cinematic, anamorphic glow",
-                lines=4,
-            )
-
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                placeholder="noise, blur, extra limbs, text watermark",
-                lines=3,
+                label="✨ Your Prompt",
+                placeholder="Describe the image you want to create...",
+                lines=5,
+                max_lines=10,
+                autofocus=True,
             )
-
-            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():
+            
+            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________________________"
+                    )
+                
                 num_inference_steps = gr.Slider(
                     minimum=1,
                     maximum=20,
                     value=9,
                     step=1,
                     label="Inference Steps",
-                    info="9 steps → 8 DiT forwards",
-                )
-
-            images_count = gr.Slider(
-                minimum=1,
-                maximum=12,
-                value=4,
-                step=1,
-                label="Images",
-                info="1–12 (higher counts use more VRAM)",
-            )
-
-            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,
+                    info="9 steps = 8 DiT (recommended)"
                 )
-                randomize_seed = gr.Checkbox(
-                    label="Randomize",
-                    value=True,
-                    interactive=True,
+                
+                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", variant="primary", size="lg")
-        
-        with gr.Column(scale=1):
-            output_images = gr.Gallery(
-                label="Generated Grid",
-                columns=4,
-                rows=None,
-                preview=True,
+            generate_btn = gr.Button(
+                "🚀 Generate Image",
+                variant="primary",
+                size="lg",
+                scale=1
             )
-            used_seeds = gr.Textbox(
-                label="Seed Cascade (x1 · x2 · ... · xN)",
-                interactive=False,
+            
+            # Example prompts
+            gr.Examples(
+                examples=examples,
+                inputs=[prompt],
+                label=" ",
+                examples_per_page=5,
             )
-            history_gallery = gr.Gallery(
-                label="History",
-                columns=6,
-                rows=None,
-                preview=True,
-                object_fit="cover"
+        
+        # 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"],
             )
-            download_btn = gr.DownloadButton(
-                label="📦 Download All History (ZIP)",
+            
+            used_seed = gr.Number(
+                label="🎲 Seed Used",
+                interactive=False,
+                container=True,
             )
     
-    gr.Markdown("### 💡 Quick Prompts")
-    gr.Examples(
-        examples=examples,
-        inputs=[prompt],
-        cache_examples=False,
-    )
-
+    # Footer credits
     gr.Markdown(
         """
-        <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 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>
         """,
-        sanitize_html=False,
+        elem_classes="footer-text"
     )
     
     # Connect the generate button
     generate_btn.click(
         fn=generate_image,
-        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],
+        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
     prompt.submit(
         fn=generate_image,
-        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],
+        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
+        outputs=[output_image, used_seed],
     )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True, show_error=True)
+    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
+    )