Update app.py

app.py

@@ -1,6 +1,8 @@
 """
-UniPic-3 DMD Multi-Image Composition
-Hugging Face Space 
+UniPic-3 DMD – ZeroGPU friendly demo
+- Pre-cache all weights on CPU
+- GPU phase does ZERO network IO
+- SSR disabled
 """
 
 import os
@@ -9,16 +11,47 @@ import torch
 import gradio as gr
 from PIL import Image
 from spaces import GPU
+from huggingface_hub import snapshot_download
 
 # -----------------------------------------------------------------------------
-# Local imports
+# Paths
 # -----------------------------------------------------------------------------
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+MODEL_ID = "Skywork/Unipic3-DMD"
+CACHE_ROOT = "./hf_cache"
+LOCAL_MODEL_DIR = os.path.join(CACHE_ROOT, MODEL_ID)
 
-try:
-    from pipeline_qwenimage_edit import QwenImageEditPipeline
-except ImportError:
-    from diffusers import QwenImageEditPipeline
+# -----------------------------------------------------------------------------
+# Pre-cache weights (CPU ONLY)
+# -----------------------------------------------------------------------------
+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,
@@ -27,18 +60,18 @@ from diffusers import (
 )
 from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
 
+try:
+    from pipeline_qwenimage_edit import QwenImageEditPipeline
+except ImportError:
+    from diffusers import QwenImageEditPipeline
+
 # -----------------------------------------------------------------------------
 # Globals
 # -----------------------------------------------------------------------------
 pipe = None
 
-MODEL_NAME = os.environ.get("MODEL_NAME", "Skywork/Unipic3-DMD")
-TRANSFORMER_PATH = os.environ.get(
-    "TRANSFORMER_PATH", "Skywork/Unipic3-DMD/ema_transformer"
-)
-
 # -----------------------------------------------------------------------------
-# Model loader (LAZY)
+# Load model (GPU stage, NO DOWNLOAD)
 # -----------------------------------------------------------------------------
 def load_model():
     global pipe
@@ -47,62 +80,44 @@ def load_model():
         return
 
     if not torch.cuda.is_available():
-        raise RuntimeError(
-            "❌ GPU not available. This Space is GPU-only."
-        )
+        raise RuntimeError("❌ GPU not available")
 
     device = torch.device("cuda")
     dtype = torch.bfloat16
 
-    print("🚀 Loading UniPic-3 DMD on GPU")
+    print("🚀 Loading UniPic-3 DMD from local cache")
     print("Device:", device)
-    print("Dtype:", dtype)
 
-    # Scheduler
     scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-        MODEL_NAME, subfolder="scheduler"
+        LOCAL_MODEL_DIR, subfolder="scheduler"
     )
 
-    # Text encoder
     text_encoder = AutoModel.from_pretrained(
-        MODEL_NAME,
+        LOCAL_MODEL_DIR,
         subfolder="text_encoder",
         torch_dtype=dtype,
-    ).to('cpu')
+    ).to(device)
 
-    # Tokenizer / Processor
     tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_NAME, subfolder="tokenizer"
+        LOCAL_MODEL_DIR, subfolder="tokenizer"
     )
+
     processor = Qwen2VLProcessor.from_pretrained(
-        MODEL_NAME, subfolder="processor"
+        LOCAL_MODEL_DIR, subfolder="processor"
     )
 
-    # Transformer (DMD)
-    if os.path.exists(TRANSFORMER_PATH):
-        transformer = QwenImageTransformer2DModel.from_pretrained(
-            TRANSFORMER_PATH,
-            torch_dtype=dtype,
-            use_safetensors=False,
-        ).to(device)
-    else:
-        # 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
+    transformer = QwenImageTransformer2DModel.from_pretrained(
+        LOCAL_MODEL_DIR,
+        subfolder="ema_transformer",
+        torch_dtype=dtype,
+    ).to(device)
+
     vae = AutoencoderKLQwenImage.from_pretrained(
-        MODEL_NAME,
+        LOCAL_MODEL_DIR,
         subfolder="vae",
         torch_dtype=dtype,
     ).to(device)
 
-    # Pipeline
     pipe = QwenImageEditPipeline(
         scheduler=scheduler,
         vae=vae,
@@ -120,12 +135,9 @@ def load_model():
 # -----------------------------------------------------------------------------
 # Inference
 # -----------------------------------------------------------------------------
-def process_images(
+def run(
     img1, img2, img3, img4, img5, img6,
-    prompt,
-    true_cfg_scale,
-    seed,
-    num_steps,
+    prompt, cfg, seed, steps
 ):
     global pipe
 
@@ -133,112 +145,70 @@ def process_images(
         load_model()
 
     images = [i for i in [img1, img2, img3, img4, img5, img6] if i is not None]
-
-    if len(images) == 0:
-        return None, "❌ Please upload at least one image."
-
-    if len(images) > 6:
-        return None, "❌ Maximum 6 images allowed."
-
-    if not prompt.strip():
-        return None, "❌ Prompt cannot be empty."
+    if not images:
+        return None, "❌ Please upload at least one image"
 
     images = [img.convert("RGB") for img in images]
 
-    generator = torch.Generator(device="cuda").manual_seed(int(seed))
-
-    try:
-        with torch.no_grad():
-            if len(images) == 1:
-                result = pipe(
-                    images[0],
-                    prompt=prompt,
-                    height=512,
-                    width=512,
-                    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=512,
-                    width=512,
-                    negative_prompt=" ",
-                    num_inference_steps=num_steps,
-                    true_cfg_scale=true_cfg_scale,
-                    generator=generator,
-                ).images[0]
-
-        return result, f"✅ Generated from {len(images)} image(s)"
-
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        return None, f"❌ Error: {e}"
+    gen = 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=gen,
+            ).images[0]
+        else:
+            out = pipe(
+                images=images,
+                prompt=prompt,
+                height=768,
+                width=768,
+                num_inference_steps=steps,
+                true_cfg_scale=cfg,
+                generator=gen,
+            ).images[0]
+
+    return out, "✅ Done"
 
 
 # -----------------------------------------------------------------------------
 # UI
 # -----------------------------------------------------------------------------
-with gr.Blocks(
-    title="UniPic-3 DMD Multi-Image Composition",
-    theme=gr.themes.Soft(),
-) as demo:
+with gr.Blocks(title="UniPic-3 DMD (ZeroGPU)") as demo:
+    gr.Markdown("# 🔥 UniPic-3 DMD (ZeroGPU + Precached)")
 
-    gr.Markdown(
-        """
-# 🔥 UniPic-3 DMD Multi-Image Composition
+    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)
+    seed = gr.Number(42)
+    steps = gr.Slider(1, 8, value=6)
 
-- **Model**: UniPic-3 DMD  
-- **Inference**: 8-step fast generation  
-"""
-    )
-
-    with gr.Row():
-        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(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",
-                lines=3,
-                value="Combine the reference images to generate the final result.",
-            )
-
-            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 = gr.Button("Generate")
+    out = gr.Image()
+    status = gr.Textbox()
 
     btn.click(
-        process_images,
-        inputs=[*image_inputs, prompt, cfg, seed, steps],
-        outputs=[output, status],
+        run,
+        inputs=[*imgs, prompt, cfg, seed, steps],
+        outputs=[out, status],
     )
 
 
 # -----------------------------------------------------------------------------
-# Entry (IMPORTANT)
+# Entry
 # -----------------------------------------------------------------------------
 @GPU
 def main():
+    # CPU stage (no GPU time)
+    precache_weights()
+
+    # Start Gradio (NO SSR)
     demo.launch(ssr_mode=False)