Update app.py

app.py

@@ -1,258 +1,629 @@
 """
-UniPic-3 DMD – ZeroGPU Final Architecture
-UI Always-On (CPU) + GPU On-Demand Inference
-with GPU Queue Status Indicator
+UniPic-3 DMD Multi-Image Composition
+Hugging Face Space - UI Persistent + GPU On-Demand Architecture
+
+核心优化：
+1. UI 常驻 - 页面始终可用，无需等待模型加载
+2. GPU on-demand - 仅在推理时调用 GPU，节省资源
+3. 优化的前端界面 - 现代美观的 UI 设计
 """
 
-import os
-import sys
-import time
-import threading
-import torch
 import gradio as gr
+import torch
 from PIL import Image
-from huggingface_hub import snapshot_download
-from spaces import GPU
-
-# =============================================================================
-# Paths & Globals
-# =============================================================================
-
-MODEL_ID = "Skywork/Unipic3-DMD"
-CACHE_ROOT = "./hf_cache"
-LOCAL_MODEL_DIR = os.path.join(CACHE_ROOT, MODEL_ID)
-
-pipe = None
-model_lock = threading.Lock()
-
-# GPU state (for UI display)
-GPU_STATE = {
-    "status": "idle",   # idle | waiting | loading | ready | running | error
-    "message": "UI ready. GPU not requested yet."
-}
-
-# =============================================================================
-# CPU Stage: Pre-cache weights (NO GPU)
-# =============================================================================
-
-def precache_weights():
-    if os.path.exists(LOCAL_MODEL_DIR):
-        print("✅ Weights already cached")
-        return
-
-    print("📦 Pre-caching UniPic-3 DMD weights (CPU stage)...")
-
-    snapshot_download(
-        repo_id=MODEL_ID,
-        local_dir=LOCAL_MODEL_DIR,
-        local_dir_use_symlinks=False,
-        resume_download=True,
-        allow_patterns=[
-            "scheduler/*",
-            "text_encoder/*",
-            "tokenizer/*",
-            "processor/*",
-            "vae/*",
-            "ema_transformer/*",
-        ],
-    )
-
-    print("✅ Pre-cache complete")
-
-# =============================================================================
-# Local imports AFTER cache
-# =============================================================================
-
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-
-from diffusers import (
-    FlowMatchEulerDiscreteScheduler,
-    QwenImageTransformer2DModel,
-    AutoencoderKLQwenImage,
-)
-from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
+import os
+import sys
 
+# Hugging Face Spaces GPU decorator
 try:
-    from pipeline_qwenimage_edit import QwenImageEditPipeline
+    from spaces import GPU
+    HF_SPACES = True
 except ImportError:
-    from diffusers import QwenImageEditPipeline
+    HF_SPACES = False
+    def GPU(duration=60):
+        def decorator(func):
+            return func
+        return decorator
 
-# =============================================================================
-# GPU Stage: Model loader (NO network)
-# =============================================================================
-
-def load_model_on_gpu():
-    global pipe
-
-    with model_lock:
-        if pipe is not None:
-            return
-
-        GPU_STATE["status"] = "loading"
-        GPU_STATE["message"] = "Loading model onto GPU..."
-
-        device = torch.device("cuda")
-        dtype = torch.bfloat16
-
-        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            LOCAL_MODEL_DIR, subfolder="scheduler"
-        )
-
-        text_encoder = AutoModel.from_pretrained(
-            LOCAL_MODEL_DIR,
-            subfolder="text_encoder",
-            torch_dtype=dtype,
-        ).to(device)
-
-        tokenizer = AutoTokenizer.from_pretrained(
-            LOCAL_MODEL_DIR, subfolder="tokenizer"
-        )
+# Local pipeline import
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
-        processor = Qwen2VLProcessor.from_pretrained(
-            LOCAL_MODEL_DIR, subfolder="processor"
+# Model configuration
+MODEL_NAME = os.environ.get("MODEL_NAME", "Skywork/Unipic3-DMD")
+TRANSFORMER_PATH = os.environ.get("TRANSFORMER_PATH", "Skywork/Unipic3-DMD/ema_transformer")
+
+# ============================================================
+# GPU On-Demand: Model loading happens inside @GPU decorated function
+# ============================================================
+
+def get_device():
+    """Get the appropriate device"""
+    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def get_dtype():
+    """Get the appropriate dtype"""
+    return torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
+@GPU(duration=120)
+def generate_image(
+    images: list[Image.Image],
+    prompt: str,
+    true_cfg_scale: float,
+    seed: int,
+    num_steps: int
+) -> Image.Image:
+    """
+    GPU on-demand inference function.
+    Model is loaded fresh each call to work with ZeroGPU.
+    """
+    # Import dependencies inside GPU function for ZeroGPU compatibility
+    try:
+        from pipeline_qwenimage_edit import QwenImageEditPipeline
+    except ImportError:
+        from diffusers import QwenImageEditPipeline
+    
+    from diffusers import (
+        FlowMatchEulerDiscreteScheduler, 
+        QwenImageTransformer2DModel, 
+        AutoencoderKLQwenImage
+    )
+    from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
+    
+    device = get_device()
+    dtype = get_dtype()
+    
+    print(f"🚀 Loading model on {device}...")
+    
+    # Load scheduler
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        MODEL_NAME, subfolder='scheduler'
+    )
+    
+    # Load text encoder
+    text_encoder = AutoModel.from_pretrained(
+        MODEL_NAME,
+        subfolder='text_encoder',
+        torch_dtype=dtype
+    ).to(device)
+    
+    # Load tokenizer & processor
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
+    processor = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
+    
+    # Load transformer
+    transformer = load_transformer(device, dtype)
+    
+    # Load VAE
+    vae = AutoencoderKLQwenImage.from_pretrained(
+        MODEL_NAME,
+        subfolder='vae',
+        torch_dtype=dtype
+    ).to(device)
+    
+    # Create pipeline
+    pipe = QwenImageEditPipeline(
+        scheduler=scheduler,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        processor=processor,
+        transformer=transformer
+    )
+    
+    print(f"✅ Model loaded! Generating with {len(images)} image(s)...")
+    
+    # Generate
+    with torch.no_grad():
+        generator = torch.Generator(device=device).manual_seed(int(seed))
+        
+        if len(images) == 1:
+            result = pipe(
+                images[0],
+                prompt=prompt,
+                height=1024,
+                width=1024,
+                negative_prompt=' ',
+                num_inference_steps=num_steps,
+                true_cfg_scale=true_cfg_scale,
+                generator=generator
+            ).images[0]
+        else:
+            result = pipe(
+                images=images,
+                prompt=prompt,
+                height=1024,
+                width=1024,
+                negative_prompt=' ',
+                num_inference_steps=num_steps,
+                true_cfg_scale=true_cfg_scale,
+                generator=generator
+            ).images[0]
+    
+    # Cleanup to free VRAM
+    del pipe, transformer, vae, text_encoder
+    torch.cuda.empty_cache()
+    
+    return result
+
+
+def load_transformer(device, dtype):
+    """Load transformer with proper path handling"""
+    from diffusers import QwenImageTransformer2DModel
+    
+    if os.path.exists(TRANSFORMER_PATH):
+        # Local path
+        if os.path.isdir(TRANSFORMER_PATH):
+            config_path = os.path.join(TRANSFORMER_PATH, "config.json")
+            if os.path.exists(config_path):
+                return QwenImageTransformer2DModel.from_pretrained(
+                    TRANSFORMER_PATH,
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).to(device)
+            else:
+                return QwenImageTransformer2DModel.from_pretrained(
+                    TRANSFORMER_PATH,
+                    subfolder='transformer',
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).to(device)
+        raise ValueError(f"Invalid transformer path: {TRANSFORMER_PATH}")
+    else:
+        # HuggingFace repo path
+        path_parts = TRANSFORMER_PATH.split('/')
+        if len(path_parts) >= 3:
+            repo_id = '/'.join(path_parts[:2])
+            subfolder = path_parts[2]
+            return QwenImageTransformer2DModel.from_pretrained(
+                repo_id,
+                subfolder=subfolder,
+                torch_dtype=dtype
+            ).to(device)
+        else:
+            return QwenImageTransformer2DModel.from_pretrained(
+                TRANSFORMER_PATH,
+                subfolder='transformer',
+                torch_dtype=dtype
+            ).to(device)
+
+
+# ============================================================
+# UI Logic (CPU-only, always available)
+# ============================================================
+
+def process_images(
+    img1, img2, img3, img4, img5, img6,
+    prompt: str,
+    cfg_scale: float,
+    seed: int,
+    num_steps: int
+):
+    """Process images - validates input then calls GPU function"""
+    
+    # Filter valid images
+    images = [img for img in [img1, img2, img3, img4, img5, img6] if img is not None]
+    
+    # Validation
+    if len(images) == 0:
+        return None, "❌ Please upload at least one image"
+    
+    if len(images) > 6:
+        return None, f"❌ Maximum 6 images allowed (got {len(images)})"
+    
+    if not prompt or prompt.strip() == "":
+        return None, "❌ Please enter an editing instruction"
+    
+    try:
+        # Convert to RGB
+        images = [img.convert("RGB") for img in images]
+        
+        # Call GPU function
+        result = generate_image(
+            images=images,
+            prompt=prompt,
+            true_cfg_scale=cfg_scale,
+            seed=seed,
+            num_steps=num_steps
         )
+        
+        return result, f"✅ Generated from {len(images)} image(s) in {num_steps} steps"
+        
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return None, f"❌ Error: {str(e)}"
+
+
+def update_image_visibility(num):
+    """Update visibility of image upload slots"""
+    return [gr.update(visible=(i < num)) for i in range(6)]
+
+
+# ============================================================
+# Custom CSS for Beautiful UI
+# ============================================================
+
+CUSTOM_CSS = """
+/* Import distinctive fonts */
+@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');
+
+/* Root variables */
+:root {
+    --primary: #6366f1;
+    --primary-dark: #4f46e5;
+    --accent: #f472b6;
+    --surface: #0f0f23;
+    --surface-light: #1a1a3e;
+    --surface-elevated: #252552;
+    --text: #e2e8f0;
+    --text-muted: #94a3b8;
+    --border: #334155;
+    --success: #10b981;
+    --error: #ef4444;
+    --gradient-1: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+    --gradient-2: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
+    --gradient-hero: linear-gradient(135deg, #0f0f23 0%, #1a1a3e 50%, #252552 100%);
+}
 
-        transformer = QwenImageTransformer2DModel.from_pretrained(
-            LOCAL_MODEL_DIR,
-            subfolder="ema_transformer",
-            torch_dtype=dtype,
-        ).to(device)
-
-        vae = AutoencoderKLQwenImage.from_pretrained(
-            LOCAL_MODEL_DIR,
-            subfolder="vae",
-            torch_dtype=dtype,
-        ).to(device)
-
-        pipe = QwenImageEditPipeline(
-            scheduler=scheduler,
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            processor=processor,
-            transformer=transformer,
-        )
+/* Global styles */
+.gradio-container {
+    font-family: 'Outfit', sans-serif !important;
+    background: var(--gradient-hero) !important;
+    min-height: 100vh;
+}
 
-        pipe.to(device)
+/* Header styling */
+.main-header {
+    text-align: center;
+    padding: 2rem 1rem;
+    background: linear-gradient(180deg, rgba(99, 102, 241, 0.1) 0%, transparent 100%);
+    border-radius: 24px;
+    margin-bottom: 2rem;
+    border: 1px solid rgba(99, 102, 241, 0.2);
+}
 
-        GPU_STATE["status"] = "ready"
-        GPU_STATE["message"] = "GPU ready. Model loaded."
+.main-header h1 {
+    font-size: 2.5rem;
+    font-weight: 700;
+    background: linear-gradient(135deg, #fff 0%, #a5b4fc 50%, #f472b6 100%);
+    -webkit-background-clip: text;
+    -webkit-text-fill-color: transparent;
+    background-clip: text;
+    margin-bottom: 0.5rem;
+}
 
-# =============================================================================
-# GPU On-Demand Inference (THIS is the only @GPU function)
-# =============================================================================
+.main-header p {
+    color: var(--text-muted);
+    font-size: 1.1rem;
+    max-width: 600px;
+    margin: 0 auto;
+}
 
-@GPU
-def run_inference(
-    img1, img2, img3, img4, img5, img6,
-    prompt, cfg, seed, steps
-):
-    global pipe
+/* Feature badges */
+.feature-badges {
+    display: flex;
+    gap: 1rem;
+    justify-content: center;
+    flex-wrap: wrap;
+    margin-top: 1.5rem;
+}
 
-    try:
-        GPU_STATE["status"] = "waiting"
-        GPU_STATE["message"] = "Waiting for GPU..."
+.badge {
+    display: inline-flex;
+    align-items: center;
+    gap: 0.5rem;
+    padding: 0.5rem 1rem;
+    background: rgba(99, 102, 241, 0.15);
+    border: 1px solid rgba(99, 102, 241, 0.3);
+    border-radius: 9999px;
+    color: #a5b4fc;
+    font-size: 0.875rem;
+    font-weight: 500;
+}
 
-        # ZeroGPU will block here until GPU is assigned
-        if not torch.cuda.is_available():
-            return None, "⏳ Waiting for GPU, please retry."
+/* Section headers */
+.section-header {
+    display: flex;
+    align-items: center;
+    gap: 0.75rem;
+    margin-bottom: 1rem;
+    padding-bottom: 0.75rem;
+    border-bottom: 1px solid var(--border);
+}
 
-        load_model_on_gpu()
+.section-header h3 {
+    font-size: 1.125rem;
+    font-weight: 600;
+    color: var(--text);
+    margin: 0;
+}
 
-        GPU_STATE["status"] = "running"
-        GPU_STATE["message"] = "Running inference..."
+/* Card styling */
+.card {
+    background: var(--surface-light) !important;
+    border: 1px solid var(--border) !important;
+    border-radius: 16px !important;
+    padding: 1.5rem !important;
+}
 
-        images = [i for i in [img1, img2, img3, img4, img5, img6] if i is not None]
-        if not images:
-            return None, "❌ Please upload at least one image."
+/* Image upload grid */
+.image-grid {
+    display: grid;
+    grid-template-columns: repeat(3, 1fr);
+    gap: 1rem;
+}
 
-        images = [img.convert("RGB") for img in images]
+/* Button styling */
+.generate-btn {
+    background: var(--gradient-1) !important;
+    border: none !important;
+    border-radius: 12px !important;
+    padding: 1rem 2rem !important;
+    font-size: 1.1rem !important;
+    font-weight: 600 !important;
+    color: white !important;
+    cursor: pointer !important;
+    transition: all 0.3s ease !important;
+    box-shadow: 0 4px 15px rgba(99, 102, 241, 0.4) !important;
+}
 
-        generator = torch.Generator(device="cuda").manual_seed(int(seed))
-
-        with torch.no_grad():
-            if len(images) == 1:
-                out = pipe(
-                    images[0],
-                    prompt=prompt,
-                    height=768,
-                    width=768,
-                    num_inference_steps=steps,
-                    true_cfg_scale=cfg,
-                    generator=generator,
-                ).images[0]
-            else:
-                out = pipe(
-                    images=images,
-                    prompt=prompt,
-                    height=768,
-                    width=768,
-                    num_inference_steps=steps,
-                    true_cfg_scale=cfg,
-                    generator=generator,
-                ).images[0]
+.generate-btn:hover {
+    transform: translateY(-2px) !important;
+    box-shadow: 0 6px 20px rgba(99, 102, 241, 0.5) !important;
+}
 
-        GPU_STATE["status"] = "ready"
-        GPU_STATE["message"] = "Inference complete."
+/* Input styling */
+.gr-textbox textarea,
+.gr-textbox input {
+    background: var(--surface) !important;
+    border: 1px solid var(--border) !important;
+    border-radius: 12px !important;
+    color: var(--text) !important;
+    font-family: 'Outfit', sans-serif !important;
+}
 
-        return out, "✅ Done"
+.gr-textbox textarea:focus,
+.gr-textbox input:focus {
+    border-color: var(--primary) !important;
+    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.2) !important;
+}
 
-    except Exception as e:
-        GPU_STATE["status"] = "error"
-        GPU_STATE["message"] = str(e)
-        return None, f"❌ Error: {e}"
+/* Slider styling */
+.gr-slider input[type="range"] {
+    accent-color: var(--primary) !important;
+}
 
-# =============================================================================
-# UI Helpers (CPU)
-# =============================================================================
+/* Output image */
+.output-image {
+    border-radius: 16px;
+    overflow: hidden;
+    border: 2px solid transparent;
+    background: linear-gradient(var(--surface-light), var(--surface-light)) padding-box,
+                var(--gradient-1) border-box;
+}
 
-def get_gpu_status():
-    return f"**GPU Status:** `{GPU_STATE['status']}`\n\n{GPU_STATE['message']}"
+/* Status text */
+.status-success {
+    color: var(--success) !important;
+    font-weight: 500;
+}
 
-# =============================================================================
-# UI (ALWAYS CPU, ALWAYS ON)
-# =============================================================================
+.status-error {
+    color: var(--error) !important;
+    font-weight: 500;
+}
 
-with gr.Blocks(title="UniPic-3 DMD (ZeroGPU)") as demo:
-    gr.Markdown("# 🔥 UniPic-3 DMD – ZeroGPU Demo")
+/* Accordion */
+.gr-accordion {
+    background: var(--surface-light) !important;
+    border: 1px solid var(--border) !important;
+    border-radius: 12px !important;
+}
 
-    status_box = gr.Markdown(get_gpu_status())
+/* Labels */
+label {
+    color: var(--text) !important;
+    font-weight: 500 !important;
+}
 
-    with gr.Row():
-        with gr.Column():
-            imgs = [gr.Image(type="pil", label=f"Image {i+1}") for i in range(6)]
-            prompt = gr.Textbox(label="Prompt", value="Combine the reference images.")
-            cfg = gr.Slider(1, 8, value=4, label="CFG")
-            seed = gr.Number(42, precision=0, label="Seed")
-            steps = gr.Slider(1, 6, value=6, label="Steps")
+/* Tooltip / info text */
+.gr-info {
+    color: var(--text-muted) !important;
+    font-size: 0.875rem !important;
+}
 
-            btn = gr.Button("🚀 Generate")
+/* Responsive adjustments */
+@media (max-width: 768px) {
+    .image-grid {
+        grid-template-columns: repeat(2, 1fr);
+    }
+    
+    .main-header h1 {
+        font-size: 1.75rem;
+    }
+    
+    .feature-badges {
+        flex-direction: column;
+        align-items: center;
+    }
+}
+"""
 
-        with gr.Column():
-            out = gr.Image(label="Output")
-            msg = gr.Textbox(label="Result")
+# ============================================================
+# Build Gradio Interface
+# ============================================================
+
+def create_demo():
+    with gr.Blocks(
+        title="UniPic-3 DMD",
+        theme=gr.themes.Base(
+            primary_hue="indigo",
+            secondary_hue="pink",
+            neutral_hue="slate",
+            font=("Outfit", "sans-serif"),
+        ),
+        css=CUSTOM_CSS
+    ) as demo:
+        
+        # Header
+        gr.HTML("""
+        <div class="main-header">
+            <h1>🎨 UniPic-3 DMD</h1>
+            <p>Multi-Image Composition with Distribution-Matching Distillation</p>
+            <div class="feature-badges">
+                <span class="badge">⚡ 8-Step Fast Inference</span>
+                <span class="badge">🖼️ Up to 6 Images</span>
+                <span class="badge">🚀 12.5× Speedup</span>
+            </div>
+        </div>
+        """)
+        
+        with gr.Row(equal_height=True):
+            # Left Column - Inputs
+            with gr.Column(scale=1):
+                
+                # Image Upload Section
+                gr.HTML('<div class="section-header"><span>📸</span><h3>Upload Images</h3></div>')
+                
+                num_images = gr.Slider(
+                    minimum=1,
+                    maximum=6,
+                    value=2,
+                    step=1,
+                    label="Number of Images",
+                    info="Select how many images to compose"
+                )
+                
+                with gr.Row():
+                    img1 = gr.Image(type="pil", label="Image 1", visible=True)
+                    img2 = gr.Image(type="pil", label="Image 2", visible=True)
+                
+                with gr.Row():
+                    img3 = gr.Image(type="pil", label="Image 3", visible=False)
+                    img4 = gr.Image(type="pil", label="Image 4", visible=False)
+                
+                with gr.Row():
+                    img5 = gr.Image(type="pil", label="Image 5", visible=False)
+                    img6 = gr.Image(type="pil", label="Image 6", visible=False)
+                
+                image_inputs = [img1, img2, img3, img4, img5, img6]
+                
+                num_images.change(
+                    fn=update_image_visibility,
+                    inputs=num_images,
+                    outputs=image_inputs
+                )
+                
+                # Prompt Section
+                gr.HTML('<div class="section-header"><span>✍️</span><h3>Editing Instruction</h3></div>')
+                
+                prompt_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="e.g., A man from Image1 standing on a surfboard from Image2, riding ocean waves under a bright blue sky.",
+                    lines=3,
+                    value="Combine the reference images to generate the final result."
+                )
+                
+                # Advanced Settings
+                with gr.Accordion("⚙️ Advanced Settings", open=False):
+                    cfg_scale = gr.Slider(
+                        minimum=1.0,
+                        maximum=10.0,
+                        value=4.0,
+                        step=0.5,
+                        label="CFG Scale",
+                        info="Higher = more prompt alignment"
+                    )
+                    
+                    with gr.Row():
+                        seed = gr.Number(
+                            value=42,
+                            label="Seed",
+                            info="For reproducibility",
+                            precision=0
+                        )
+                        num_steps = gr.Slider(
+                            minimum=1,
+                            maximum=8,
+                            value=8,
+                            step=1,
+                            label="Steps",
+                            info="8 recommended for DMD"
+                        )
+                
+                # Generate Button
+                generate_btn = gr.Button(
+                    "🚀 Generate Image",
+                    variant="primary",
+                    size="lg",
+                    elem_classes=["generate-btn"]
+                )
+            
+            # Right Column - Output
+            with gr.Column(scale=1):
+                gr.HTML('<div class="section-header"><span>🎨</span><h3>Generated Result</h3></div>')
+                
+                output_image = gr.Image(
+                    type="pil",
+                    label="Output",
+                    elem_classes=["output-image"],
+                    show_download_button=True
+                )
+                
+                status_text = gr.Textbox(
+                    label="Status",
+                    value="✨ Ready! Upload images and click Generate.",
+                    interactive=False,
+                    show_copy_button=False
+                )
+                
+                # Tips
+                gr.HTML("""
+                <div style="
+                    margin-top: 1.5rem;
+                    padding: 1rem;
+                    background: rgba(99, 102, 241, 0.1);
+                    border-radius: 12px;
+                    border: 1px solid rgba(99, 102, 241, 0.2);
+                ">
+                    <p style="color: #a5b4fc; font-weight: 600; margin-bottom: 0.5rem;">💡 Tips</p>
+                    <ul style="color: #94a3b8; font-size: 0.9rem; margin: 0; padding-left: 1.25rem;">
+                        <li>Reference images as "Image1", "Image2", etc. in your prompt</li>
+                        <li>Use descriptive prompts for better composition</li>
+                        <li>First run may take longer due to model loading</li>
+                    </ul>
+                </div>
+                """)
+        
+        # Connect generate button
+        generate_btn.click(
+            fn=process_images,
+            inputs=[*image_inputs, prompt_input, cfg_scale, seed, num_steps],
+            outputs=[output_image, status_text]
+        )
+        
+        # Examples
+        gr.HTML('<div class="section-header" style="margin-top: 2rem;"><span>📚</span><h3>Example Prompts</h3></div>')
+        
+        gr.Examples(
+            examples=[
+                ["A person from Image1 wearing the outfit from Image2"],
+                ["Combine Image1 and Image2 into a single cohesive scene"],
+                ["The object from Image1 placed in the environment from Image2"],
+                ["Create a portrait using the face from Image1 and hairstyle from Image2"],
+            ],
+            inputs=[prompt_input],
+            label=""
+        )
+    
+    return demo
 
-    btn.click(
-        run_inference,
-        inputs=[*imgs, prompt, cfg, seed, steps],
-        outputs=[out, msg],
-    )
 
-    # Periodic status refresh (CPU only)
-    demo.load(
-        fn=get_gpu_status,
-        inputs=[],
-        outputs=status_box,
-        every=1.0,
-    )
+# ============================================================
+# Launch
+# ============================================================
 
-# =============================================================================
-# Entry
-# =============================================================================
+demo = create_demo()
 
 if __name__ == "__main__":
-    # CPU phase: cache weights
-    precache_weights()
-
-    # UI always-on
-    demo.launch(ssr_mode=False)
+    demo.launch()