Update app.py

app.py

@@ -1,338 +1,247 @@
 """
-Gradio Demo for UniPic-3 DMD Multi-Image Composition
-Hugging Face Space compatible version
-
-Upload up to 6 images and generate a composed result using DMD model with 8-step inference.
+UniPic-3 DMD Multi-Image Composition
+Hugging Face Space 
 """
 
-import gradio as gr
+import os
+import sys
 import torch
+import gradio as gr
 from PIL import Image
-import os
 from spaces import GPU
 
-# Use local pipeline to ensure compatibility
-import sys
+# -----------------------------------------------------------------------------
+# Local imports
+# -----------------------------------------------------------------------------
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
 try:
     from pipeline_qwenimage_edit import QwenImageEditPipeline
 except ImportError:
-    # Fallback to diffusers if local not available
-    try:
-        from diffusers import QwenImageEditPipeline
-    except ImportError:
-        raise ImportError(
-            "QwenImageEditPipeline not found. Please ensure pipeline_qwenimage_edit.py "
-            "is in the same directory or diffusers is installed."
-        )
+    from diffusers import QwenImageEditPipeline
 
-from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageTransformer2DModel, AutoencoderKLQwenImage
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    QwenImageTransformer2DModel,
+    AutoencoderKLQwenImage,
+)
 from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
 
-
-# Global pipeline
+# -----------------------------------------------------------------------------
+# Globals
+# -----------------------------------------------------------------------------
 pipe = None
 
-# Model paths (can be set via environment variables)
-
 MODEL_NAME = os.environ.get("MODEL_NAME", "Skywork/Unipic3-DMD")
-default_transformer = "Skywork/Unipic3-DMD"
-TRANSFORMER_PATH = os.environ.get("TRANSFORMER_PATH", "Skywork/Unipic3-DMD/ema_transformer")
+TRANSFORMER_PATH = os.environ.get(
+    "TRANSFORMER_PATH", "Skywork/Unipic3-DMD/ema_transformer"
+)
 
+# -----------------------------------------------------------------------------
+# Model loader (LAZY)
+# -----------------------------------------------------------------------------
 def load_model():
-    """Load the DMD model and pipeline"""
     global pipe
-    
+
     if pipe is not None:
-        return pipe
-    
-    print(f"Loading model from {TRANSFORMER_PATH}...")
-    
-    # Load scheduler
+        return
+
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "❌ GPU not available. This Space is GPU-only."
+        )
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+
+    print("🚀 Loading UniPic-3 DMD on GPU")
+    print("Device:", device)
+    print("Dtype:", dtype)
+
+    # Scheduler
     scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-        pretrained_model_name_or_path=MODEL_NAME, subfolder='scheduler'
+        MODEL_NAME, subfolder="scheduler"
     )
-    
-    # Load text encoder
+
+    # Text encoder
     text_encoder = AutoModel.from_pretrained(
-        pretrained_model_name_or_path=MODEL_NAME, subfolder='text_encoder',
-        device_map='auto', torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
-    )
-    
-    # Load tokenizer and processor
+        MODEL_NAME,
+        subfolder="text_encoder",
+        torch_dtype=dtype,
+    ).to(device)
+
+    # Tokenizer / Processor
     tokenizer = AutoTokenizer.from_pretrained(
-        pretrained_model_name_or_path=MODEL_NAME, subfolder='tokenizer',
+        MODEL_NAME, subfolder="tokenizer"
     )
     processor = Qwen2VLProcessor.from_pretrained(
-        pretrained_model_name_or_path=MODEL_NAME, subfolder='processor',
+        MODEL_NAME, subfolder="processor"
     )
-    
-    # Load transformer (DMD model)
-    # Handle both local paths and HuggingFace repo paths
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print("device:", device)
-    
+
+    # Transformer (DMD)
     if os.path.exists(TRANSFORMER_PATH):
-        # Local path - load directly to device (avoid device_map issues with .bin files)
-        if os.path.isdir(TRANSFORMER_PATH):
-            # Check if it's a direct transformer directory or has subfolder
-            if os.path.exists(os.path.join(TRANSFORMER_PATH, "config.json")):
-                transformer = QwenImageTransformer2DModel.from_pretrained(
-                    pretrained_model_name_or_path=TRANSFORMER_PATH,
-                    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-                    use_safetensors=False  # Use .bin file
-                ).to(device)
-            else:
-                transformer = QwenImageTransformer2DModel.from_pretrained(
-                    pretrained_model_name_or_path=TRANSFORMER_PATH,
-                    subfolder='transformer',
-                    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-                    use_safetensors=False
-                ).to(device)
-        else:
-            raise ValueError(f"Transformer path does not exist: {TRANSFORMER_PATH}")
-    else:
-        # HuggingFace repo path
-        # Handle paths like "Skywork/Unipic3-DMD/ema_transformer"
-        path_parts = TRANSFORMER_PATH.split('/')
-        if len(path_parts) >= 3:
-            # Has subfolder: "Skywork/Unipic3-DMD/ema_transformer"
-            repo_id = '/'.join(path_parts[:2])  # "Skywork/Unipic3-DMD"
-            subfolder = path_parts[2]  # "ema_transformer"
-            transformer = QwenImageTransformer2DModel.from_pretrained(
-                pretrained_model_name_or_path=repo_id,
-                subfolder=subfolder,
-                device_map='auto',
-                torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
-            )
-        elif len(path_parts) == 2:
-            # Just repo: "Skywork/Unipic3-DMD"
-            transformer = QwenImageTransformer2DModel.from_pretrained(
-                pretrained_model_name_or_path=TRANSFORMER_PATH,
-                subfolder='transformer',
-                device_map='auto',
-                torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
-            )
-        else:
-            # Single name, assume it's a repo ID
-            transformer = QwenImageTransformer2DModel.from_pretrained(
-                pretrained_model_name_or_path=TRANSFORMER_PATH,
-                subfolder='transformer',
-                device_map='auto',
-                torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
-            )
-    
-    # Load VAE
-    # Get device from transformer (handle both .device and device_map cases)
-    if hasattr(transformer, 'device'):
-        vae_device = transformer.device
-    elif hasattr(transformer, 'hf_device_map'):
-        # If using device_map, get the first device
-        vae_device = device
+        transformer = QwenImageTransformer2DModel.from_pretrained(
+            TRANSFORMER_PATH,
+            torch_dtype=dtype,
+            use_safetensors=False,
+        ).to(device)
     else:
-        vae_device = device
-    
+        # HF repo path: Skywork/Unipic3-DMD/ema_transformer
+        repo_id = "/".join(TRANSFORMER_PATH.split("/")[:2])
+        subfolder = TRANSFORMER_PATH.split("/")[-1]
+        transformer = QwenImageTransformer2DModel.from_pretrained(
+            repo_id,
+            subfolder=subfolder,
+            torch_dtype=dtype,
+        ).to(device)
+
+    # VAE
     vae = AutoencoderKLQwenImage.from_pretrained(
-        pretrained_model_name_or_path=MODEL_NAME, 
-        subfolder='vae', 
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-    ).to(vae_device)
-    
-    # Create pipeline
+        MODEL_NAME,
+        subfolder="vae",
+        torch_dtype=dtype,
+    ).to(device)
+
+    # Pipeline
     pipe = QwenImageEditPipeline(
         scheduler=scheduler,
         vae=vae,
         text_encoder=text_encoder,
         tokenizer=tokenizer,
         processor=processor,
-        transformer=transformer
+        transformer=transformer,
     )
-    
-    print("Model loaded successfully!")
-    # Don't return pipe for demo.load() - it expects no return value
 
+    pipe.to(device)
 
+    print("✅ Model loaded successfully")
+
+
+# -----------------------------------------------------------------------------
+# Inference
+# -----------------------------------------------------------------------------
 def process_images(
     img1, img2, img3, img4, img5, img6,
-    prompt: str,
-    true_cfg_scale: float = 4.0,
-    seed: int = 42,
-    num_steps: int = 8
-) -> tuple:
-    """Process multiple images and generate composed result"""
+    prompt,
+    true_cfg_scale,
+    seed,
+    num_steps,
+):
     global pipe
-    
-    # Ensure model is loaded (should be loaded by demo.load() on startup)
+
     if pipe is None:
-        return None, "⏳ Model is still loading, please wait a moment and try again..."
-    
-    # Filter out None images
-    images = [img for img in [img1, img2, img3, img4, img5, img6] if img is not None]
-    
-    # Validate inputs
+        load_model()
+
+    images = [i for i in [img1, img2, img3, img4, img5, img6] if i is not None]
+
     if len(images) == 0:
-        return None, "❌ Error: Please upload at least one image."
-    
+        return None, "❌ Please upload at least one image."
+
     if len(images) > 6:
-        return None, f"❌ Error: Maximum 6 images allowed. You uploaded {len(images)} images."
-    
-    if not prompt or prompt.strip() == "":
-        return None, "❌ Error: Please enter an editing instruction."
-    
+        return None, "❌ Maximum 6 images allowed."
+
+    if not prompt.strip():
+        return None, "❌ Prompt cannot be empty."
+
+    images = [img.convert("RGB") for img in images]
+
+    generator = torch.Generator(device="cuda").manual_seed(int(seed))
+
     try:
-        # Convert to RGB
-        images = [img.convert("RGB") for img in images]
-        
-        print(f"Processing {len(images)} images with prompt: '{prompt}'")
-        print(f"Steps: {num_steps}, CFG Scale: {true_cfg_scale}, Seed: {seed}")
-        
-        # Generate image
-        # Note: images can be passed as first positional argument or as keyword argument
         with torch.no_grad():
-            # Try positional argument first (as shown in pipeline examples)
             if len(images) == 1:
-                # Single image: pass as first positional argument
                 result = pipe(
                     images[0],
                     prompt=prompt,
                     height=1024,
                     width=1024,
-                    negative_prompt=' ',
+                    negative_prompt=" ",
                     num_inference_steps=num_steps,
                     true_cfg_scale=true_cfg_scale,
-                    generator=torch.manual_seed(int(seed))
+                    generator=generator,
                 ).images[0]
             else:
-                # Multiple images: pass as keyword argument
                 result = pipe(
                     images=images,
                     prompt=prompt,
                     height=1024,
                     width=1024,
-                    negative_prompt=' ',
+                    negative_prompt=" ",
                     num_inference_steps=num_steps,
                     true_cfg_scale=true_cfg_scale,
-                    generator=torch.manual_seed(int(seed))
+                    generator=generator,
                 ).images[0]
-        
-        return result, f"✅ Success! Generated from {len(images)} image(s) in {num_steps} steps."
-    
+
+        return result, f"✅ Generated from {len(images)} image(s)"
+
     except Exception as e:
-        error_msg = f"❌ Error: {str(e)}"
-        print(error_msg)
         import traceback
         traceback.print_exc()
-        return None, error_msg
-
-
-# Create Gradio interface
-with gr.Blocks(title="UniPic-3 DMD Multi-Image Composition", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    # 🔥 UniPic-3 DMD Multi-Image Composition
-    
-    Upload up to **6 images** and provide an editing instruction to generate a composed result.
-    
-    **Model**: DMD (Distribution-Matching Distillation) - **8-step fast inference (12.5× speedup)**
-    
-    **Features**:
-    - Support 1-6 input images
-    - Fast 8-step inference
-    - High-quality multi-image composition
-    """)
-    
+        return None, f"❌ Error: {e}"
+
+
+# -----------------------------------------------------------------------------
+# UI
+# -----------------------------------------------------------------------------
+with gr.Blocks(
+    title="UniPic-3 DMD Multi-Image Composition",
+    theme=gr.themes.Soft(),
+) as demo:
+
+    gr.Markdown(
+        """
+# 🔥 UniPic-3 DMD Multi-Image Composition
+
+- **Model**: UniPic-3 DMD  
+- **Inference**: 8-step fast generation  
+- **GPU-only Hugging Face Space**
+"""
+    )
+
     with gr.Row():
-        with gr.Column(scale=1):
-            gr.Markdown("### 📸 Upload Images (1-6 images)")
+        with gr.Column():
             image_inputs = [
                 gr.Image(type="pil", label=f"Image {i+1}", visible=(i < 2))
                 for i in range(6)
             ]
-            
-            num_images = gr.Slider(
-                minimum=1,
-                maximum=6,
-                value=2,
-                step=1,
-                label="Number of Images",
-                info="Select how many images you want to upload"
-            )
-            
-            def update_image_visibility(num):
-                return [gr.update(visible=(i < num)) for i in range(6)]
-            
-            num_images.change(
-                fn=update_image_visibility,
-                inputs=num_images,
-                outputs=image_inputs
-            )
-            
-            gr.Markdown("### ✍️ Editing Instruction")
-            prompt_input = gr.Textbox(
+
+            num_images = gr.Slider(1, 6, value=2, step=1, label="Number of Images")
+
+            def update_visibility(n):
+                return [gr.update(visible=i < n) for i in range(6)]
+
+            num_images.change(update_visibility, num_images, image_inputs)
+
+            prompt = gr.Textbox(
                 label="Prompt",
-                placeholder="e.g., A man from Image1 is standing on a surfboard from Image2, riding the ocean waves under a bright blue sky.",
                 lines=3,
-                value="Combine the reference images to generate the final result."
-            )
-            
-            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 values make the output more aligned with the prompt"
-                )
-                seed = gr.Number(
-                    value=42,
-                    label="Seed",
-                    info="Random seed for reproducibility",
-                    precision=0
-                )
-                num_steps = gr.Slider(
-                    minimum=1,
-                    maximum=8,
-                    value=8,
-                    step=1,
-                    label="Inference Steps",
-                    info="Number of denoising steps (8 is recommended for DMD)"
-                )
-            
-            generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
-        
-        with gr.Column(scale=1):
-            gr.Markdown("### 🎨 Generated Result")
-            output_image = gr.Image(type="pil", label="Output Image")
-            status_text = gr.Textbox(
-                label="Status",
-                value="Ready. Upload images and enter a prompt, then click Generate.",
-                interactive=False
+                value="Combine the reference images to generate the final result.",
             )
-    
-    # Load model on startup
-    def load_model_wrapper():
-        """Wrapper to load model without returning value"""
-        load_model()
-        return None
-    
-    demo.load(
-        fn=load_model_wrapper,
-        inputs=[],
-        outputs=[],
-        show_progress=True
-    )
-    
-    # Generate button
-    generate_btn.click(
-        fn=process_images,
-        inputs=[*image_inputs, prompt_input, cfg_scale, seed, num_steps],
-        outputs=[output_image, status_text]
+
+            cfg = gr.Slider(1.0, 10.0, value=4.0, step=0.5, label="CFG Scale")
+            seed = gr.Number(value=42, precision=0, label="Seed")
+            steps = gr.Slider(1, 8, value=8, step=1, label="Steps")
+
+            btn = gr.Button("🚀 Generate", variant="primary")
+
+        with gr.Column():
+            output = gr.Image(label="Output")
+            status = gr.Textbox(label="Status", interactive=False)
+
+    btn.click(
+        process_images,
+        inputs=[*image_inputs, prompt, cfg, seed, steps],
+        outputs=[output, status],
     )
 
-# @GPU
-# def main():
-#     demo.launch()
-    
-if __name__ == "__main__":
+
+# -----------------------------------------------------------------------------
+# Entry (IMPORTANT)
+# -----------------------------------------------------------------------------
+@GPU
+def main():
     demo.launch()
 
+
+if __name__ == "__main__":
+    main()