Update app.py

app.py

@@ -1,31 +1,32 @@
 import torch
 import spaces
 import gradio as gr
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, FluxImg2ImgPipeline
 
-# Load the pipeline once at startup
-print("Loading Z-Image-Turbo pipeline...")
-pipe = DiffusionPipeline.from_pretrained(
+print("Loading pipelines...")
+# Text to image pipeline
+pipe_t2i = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
 )
-pipe.to("cuda")
+pipe_t2i.to("cuda")
 
-# ======== AoTI compilation + FA3 ========
-# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+# Image to image pipeline
+pipe_i2i = FluxImg2ImgPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
+)
+pipe_i2i.to("cuda")
 
-print("Pipeline loaded!")
+print("Pipelines 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_t2i(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
-    
     generator = torch.Generator("cuda").manual_seed(int(seed))
-    image = pipe(
+    image = pipe_t2i(
         prompt=prompt,
         height=int(height),
         width=int(width),
@@ -33,19 +34,30 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         guidance_scale=0.0,
         generator=generator,
     ).images[0]
-    
     return image, seed
 
-# Example prompts
-examples = [
-    ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
-    ["A majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors, detailed fantasy art style"],
-    ["Cozy coffee shop interior, warm lighting, rain on windows, plants on shelves, vintage aesthetic, photorealistic"],
+@spaces.GPU
+def generate_i2i(input_image, prompt, strength, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+    if randomize_seed:
+        seed = torch.randint(0, 2**32 - 1, (1,)).item()
+    generator = torch.Generator("cuda").manual_seed(int(seed))
+    image = pipe_i2i(
+        prompt=prompt,
+        image=input_image,
+        strength=float(strength),
+        num_inference_steps=int(num_inference_steps),
+        generator=generator,
+    ).images[0]
+    return image, seed
+
+examples_t2i = [
+    ["Young Chinese woman in red Hanfu, intricate embroidery, elaborate high bun, golden phoenix headdress"],
+    ["A majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors"],
+    ["Cozy coffee shop interior, warm lighting, rain on windows, plants on shelves, photorealistic"],
     ["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"],
+    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff"],
 ]
 
-# Custom theme with modern aesthetics (Gradio 6)
 custom_theme = gr.themes.Soft(
     primary_hue="yellow",
     secondary_hue="amber",
@@ -60,135 +72,86 @@ custom_theme = gr.themes.Soft(
     block_title_text_weight="600",
 )
 
-# Build the Gradio interface
 with gr.Blocks(fill_height=True) as demo:
-    # Header
     gr.Markdown(
         """
         # 🎨 Z-Image-Turbo
-        **Ultra-fast AI image generation** • Generate stunning images in just 8 steps
+        **Ultra-fast AI image generation & editing** • Text to Image + Image to Image
         """,
         elem_classes="header-text"
     )
-    
-    with gr.Row(equal_height=False):
-        # Left column - Input controls
-        with gr.Column(scale=1, min_width=320):
-            prompt = gr.Textbox(
-                label="✨ Your Prompt",
-                placeholder="Describe the image you want to create...",
-                lines=5,
-                max_lines=10,
-                autofocus=True,
-            )
-            
-            with gr.Accordion("⚙️ Advanced Settings", open=False):
-                with gr.Row():
-                    height = gr.Slider(
-                        minimum=512,
-                        maximum=2048,
-                        value=1024,
-                        step=64,
-                        label="Height",
-                        info="Image height in pixels"
-                    )
-                    width = gr.Slider(
-                        minimum=512,
-                        maximum=2048,
-                        value=1024,
-                        step=64,
-                        label="Width",
-                        info="Image width in pixels"
-                    )
-                
-                num_inference_steps = gr.Slider(
-                    minimum=1,
-                    maximum=20,
-                    value=9,
-                    step=1,
-                    label="Inference Steps",
-                    info="9 steps = 8 DiT forwards (recommended)"
-                )
-                
-                with gr.Row():
-                    randomize_seed = gr.Checkbox(
-                        label="🎲 Random Seed",
-                        value=True,
+
+    with gr.Tabs():
+
+        # ── Tab 1: Text to Image ──────────────────────────────────────
+        with gr.Tab("✨ Text to Image"):
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=1, min_width=320):
+                    t2i_prompt = gr.Textbox(
+                        label="✨ Your Prompt",
+                        placeholder="Describe the image you want to create...",
+                        lines=5,
+                        max_lines=10,
+                        autofocus=True,
                     )
-                    seed = gr.Number(
-                        label="Seed",
-                        value=42,
-                        precision=0,
-                        visible=False,
+                    with gr.Accordion("⚙️ Advanced Settings", open=False):
+                        with gr.Row():
+                            t2i_height = gr.Slider(512, 2048, value=1024, step=64, label="Height")
+                            t2i_width = gr.Slider(512, 2048, value=1024, step=64, label="Width")
+                        t2i_steps = gr.Slider(1, 20, value=9, step=1, label="Inference Steps")
+                        with gr.Row():
+                            t2i_randomize = gr.Checkbox(label="🎲 Random Seed", value=True)
+                            t2i_seed = gr.Number(label="Seed", value=42, precision=0, visible=False)
+                        t2i_randomize.change(
+                            lambda r: gr.Number(visible=not r),
+                            inputs=[t2i_randomize], outputs=[t2i_seed]
+                        )
+                    t2i_btn = gr.Button("🚀 Generate Image", variant="primary", size="lg")
+                    gr.Examples(examples=examples_t2i, inputs=[t2i_prompt], label="💡 Try these prompts")
+
+                with gr.Column(scale=1, min_width=320):
+                    t2i_output = gr.Image(label="Generated Image", type="pil", format="png", show_label=False, height=600)
+                    t2i_used_seed = gr.Number(label="🎲 Seed Used", interactive=False)
+
+            t2i_btn.click(generate_t2i, [t2i_prompt, t2i_height, t2i_width, t2i_steps, t2i_seed, t2i_randomize], [t2i_output, t2i_used_seed])
+            t2i_prompt.submit(generate_t2i, [t2i_prompt, t2i_height, t2i_width, t2i_steps, t2i_seed, t2i_randomize], [t2i_output, t2i_used_seed])
+
+        # ── Tab 2: Image to Image ─────────────────────────────────────
+        with gr.Tab("🖼️ Image to Image"):
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=1, min_width=320):
+                    i2i_input = gr.Image(label="Upload Image", type="pil")
+                    i2i_prompt = gr.Textbox(
+                        label="✨ Edit Instruction",
+                        placeholder="Describe how you want to edit the image...",
+                        lines=4,
                     )
-                
-                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
-            )
-            
-            # Example prompts
-            gr.Examples(
-                examples=examples,
-                inputs=[prompt],
-                label="💡 Try these prompts",
-                examples_per_page=5,
-            )
-        
-        # 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"],
-            )
-            
-            used_seed = gr.Number(
-                label="🎲 Seed Used",
-                interactive=False,
-                container=True,
-            )
-    
-    # Footer credits
+                    with gr.Accordion("⚙️ Advanced Settings", open=False):
+                        i2i_strength = gr.Slider(0.1, 1.0, value=0.75, step=0.05, label="Strength", info="Higher = more change")
+                        i2i_steps = gr.Slider(1, 8, value=4, step=1, label="Inference Steps")
+                        with gr.Row():
+                            i2i_randomize = gr.Checkbox(label="🎲 Random Seed", value=True)
+                            i2i_seed = gr.Number(label="Seed", value=42, precision=0, visible=False)
+                        i2i_randomize.change(
+                            lambda r: gr.Number(visible=not r),
+                            inputs=[i2i_randomize], outputs=[i2i_seed]
+                        )
+                    i2i_btn = gr.Button("🚀 Edit Image", variant="primary", size="lg")
+
+                with gr.Column(scale=1, min_width=320):
+                    i2i_output = gr.Image(label="Result", type="pil", format="png", show_label=False, height=600)
+                    i2i_used_seed = gr.Number(label="🎲 Seed Used", interactive=False)
+
+            i2i_btn.click(generate_i2i, [i2i_input, i2i_prompt, i2i_strength, i2i_steps, i2i_seed, i2i_randomize], [i2i_output, i2i_used_seed])
+
     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) • 
-        <strong>Demo by:</strong> <a href="https://x.com/realmrfakename" target="_blank">@mrfakename</a> • 
-        <strong>Redesign by:</strong> AnyCoder • 
-        <strong>Optimizations:</strong> <a href="https://huggingface.co/multimodalart" target="_blank">@multimodalart</a> (FA3 + AoTI)
+        <div style="text-align: center; opacity: 0.7; font-size: 0.9em;">
+        <strong>T2I Model:</strong> Tongyi-MAI/Z-Image-Turbo • 
+        <strong>I2I Model:</strong> FLUX.1-schnell
         </div>
         """,
-        elem_classes="footer-text"
-    )
-    
-    # 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
-    prompt.submit(
-        fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed],
-        outputs=[output_image, used_seed],
     )
 
 if __name__ == "__main__":
@@ -198,63 +161,15 @@ if __name__ == "__main__":
         .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;
-        }
+        .header-text p { font-size: 1.1rem !important; color: #64748b !important; }
+        .gradio-container { max-width: 1400px !important; margin: 0 auto !important; }
+        button, .gr-button { transition: all 0.2s ease !important; }
+        button:hover { transform: translateY(-1px); box-shadow: 0 4px 12px rgba(0,0,0,0.15) !important; }
         """,
-        footer_links=[
-            "api",
-            "gradio"
-        ],
         mcp_server=True
     )