Switch to Instruct-Pix2Pix

app.py

@@ -1,299 +1,127 @@
 import torch
 import spaces
 import gradio as gr
-from diffusers import DiffusionPipeline, AutoPipelineForImage2Image
+from diffusers import StableDiffusionInstructPix2PixPipeline, StableDiffusionPipeline
 from diffusers.utils import load_image
-import PIL.Image
 
-# Load the pipeline once at startup
-print("Loading Z-Image-Turbo pipeline...")
-pipe = DiffusionPipeline.from_pretrained(
-    "Tongyi-MAI/Z-Image-Turbo",
-    torch_dtype=torch.bfloat16,
-    low_cpu_mem_usage=True,
-)
-pipe.to("cuda")
-
-try:
-    print("Loading Image-to-Image pipeline...")
-    pipe_i2i = AutoPipelineForImage2Image.from_pipe(pipe)
-except Exception as e:
-    print(f"Warning: Could not load Img2Img pipeline: {e}")
-    pipe_i2i = None
+print("Loading Models...")
+# 1. Text-to-Image Model (Standard Generation)
+pipe_t2i = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16,
+    safety_checker=None
+).to("cuda")
 
-# ======== AoTI compilation + FA3 ========
-# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+# 2. Instruct-Pix2Pix Model (Instruction-Based Editing)
+pipe_edit = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+    "timbrooks/instruct-pix2pix",
+    torch_dtype=torch.float16,
+    safety_checker=None
+).to("cuda")
 
-print("Pipeline loaded!")
+print("Models loaded successfully!")
 
 @spaces.GPU
-def generate_image(prompt, input_image, strength, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
-    """Generate or edit an image from the given prompt and optional input image."""
+def generate_or_edit(prompt, input_image, image_guidance_scale, guidance_scale, 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))
     
+    if not prompt:
+        raise gr.Error("Please enter a prompt or instruction!")
+    
     if input_image is not None:
-        if pipe_i2i is None:
-            raise gr.Error("Image-to-Image is not supported by your current diffusers version for Z-Image.")
-        # Resize input image to match target height/width
-        input_image = input_image.resize((int(width), int(height)))
-        image = pipe_i2i(
-            prompt=prompt,
+        # Edit mode
+        input_image = input_image.convert("RGB")
+        # Resize image for SD1.5 (InstructPix2Pix)
+        input_image = input_image.resize((512, 512))
+        
+        image = pipe_edit(
+            prompt,
             image=input_image,
-            strength=strength,
             num_inference_steps=int(num_inference_steps),
-            guidance_scale=0.0,
+            image_guidance_scale=image_guidance_scale,
+            guidance_scale=guidance_scale,
             generator=generator,
         ).images[0]
     else:
-        image = pipe(
-            prompt=prompt,
-            height=int(height),
-            width=int(width),
+        # Generate mode
+        image = pipe_t2i(
+            prompt,
             num_inference_steps=int(num_inference_steps),
-            guidance_scale=0.0,
+            guidance_scale=guidance_scale,
             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"],
-    ["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)
+# Build the Gradio interface
 custom_theme = gr.themes.Soft(
-    primary_hue="yellow",
-    secondary_hue="amber",
-    neutral_hue="slate",
+    primary_hue="blue",
+    secondary_hue="indigo",
     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
+with gr.Blocks(theme=custom_theme, fill_height=True) as demo:
     gr.Markdown(
         """
-        # 🎨 Z-Image-Turbo
-        **Ultra-fast AI image generation** • Generate stunning images in just 8 steps
-        """,
-        elem_classes="header-text"
+        # 🎨 AI Image Studio (Generation & Instruction Editing)
+        Generate images from scratch, or upload an image and use a prompt like *"Make him wear sunglasses"* to edit it seamlessly.
+        """
     )
     
-    with gr.Row(equal_height=False):
-        # Left column - Input controls
-        with gr.Column(scale=1, min_width=320):
+    with gr.Row():
+        with gr.Column(scale=1):
             prompt = gr.Textbox(
-                label="✨ Your Prompt",
-                placeholder="Describe the image you want to create or edit...",
-                lines=5,
-                max_lines=10,
-                autofocus=True,
+                label="✨ Instruction / Prompt", 
+                lines=3, 
+                placeholder="e.g. 'A futuristic city at night' (for generation) OR 'Turn the daytime into nighttime' (for editing)",
+                autofocus=True
             )
-            
             input_image = gr.Image(
-                label="🖼️ Input Image (Optional, for Image-to-Image Editing)",
-                type="pil",
-                height=256,
+                label="🖼️ Input Image (Optional - Leave blank to generate from scratch)", 
+                type="pil"
             )
             
             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)"
-                )
-                
-                strength = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.5,
-                    step=0.05,
-                    label="Denoising Strength (For Image-to-Image)",
-                    info="Higher values = more changes from input image"
+                num_inference_steps = gr.Slider(minimum=10, maximum=50, value=30, step=1, label="Inference Steps")
+                guidance_scale = gr.Slider(minimum=1.0, maximum=15.0, value=7.5, step=0.5, label="Text Guidance Scale")
+                image_guidance_scale = gr.Slider(
+                    minimum=1.0, 
+                    maximum=3.0, 
+                    value=1.5, 
+                    step=0.1, 
+                    label="Image Guidance Scale (Editing Only)", 
+                    info="Lower values = more changes. Higher values = stays closer to original image."
                 )
                 
                 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 = gr.Checkbox(label="🎲 Random Seed", value=True)
+                    seed = gr.Number(label="Seed", value=42, precision=0, visible=False)
                 
                 randomize_seed.change(
-                    toggle_seed,
-                    inputs=[randomize_seed],
+                    lambda r: gr.Number(visible=not r), 
+                    inputs=[randomize_seed], 
                     outputs=[seed]
                 )
+                
+            generate_btn = gr.Button("🚀 Generate / Edit Image", variant="primary", size="lg")
             
-            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
-    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>
-        """,
-        elem_classes="footer-text"
-    )
-    
-    # Connect the generate button
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Result", type="pil", interactive=False)
+            used_seed = gr.Number(label="Seed Used", interactive=False)
+
+    # Connections
     generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt, input_image, strength, height, width, num_inference_steps, seed, randomize_seed],
-        outputs=[output_image, used_seed],
+        fn=generate_or_edit,
+        inputs=[prompt, input_image, image_guidance_scale, guidance_scale, 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, input_image, strength, height, width, num_inference_steps, seed, randomize_seed],
-        outputs=[output_image, used_seed],
+        fn=generate_or_edit,
+        inputs=[prompt, input_image, image_guidance_scale, guidance_scale, num_inference_steps, seed, randomize_seed],
+        outputs=[output_image, used_seed]
     )
 
 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()