Update app.py

app.py

@@ -1,28 +1,40 @@
 """
-UniPic-3 DMD – ZeroGPU friendly demo
-- Pre-cache all weights on CPU
-- GPU phase does ZERO network IO
-- SSR disabled
+UniPic-3 DMD – ZeroGPU Final Architecture
+UI Always-On (CPU) + GPU On-Demand Inference
+with GPU Queue Status Indicator
 """
 
 import os
 import sys
+import time
+import threading
 import torch
 import gradio as gr
 from PIL import Image
-from spaces import GPU
 from huggingface_hub import snapshot_download
+from spaces import GPU
+
+# =============================================================================
+# Paths & Globals
+# =============================================================================
 
-# -----------------------------------------------------------------------------
-# Paths
-# -----------------------------------------------------------------------------
 MODEL_ID = "Skywork/Unipic3-DMD"
 CACHE_ROOT = "./hf_cache"
 LOCAL_MODEL_DIR = os.path.join(CACHE_ROOT, MODEL_ID)
 
-# -----------------------------------------------------------------------------
-# Pre-cache weights (CPU ONLY)
-# -----------------------------------------------------------------------------
+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")
@@ -47,10 +59,10 @@ def precache_weights():
 
     print("✅ Pre-cache complete")
 
-
-# -----------------------------------------------------------------------------
+# =============================================================================
 # Local imports AFTER cache
-# -----------------------------------------------------------------------------
+# =============================================================================
+
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from diffusers import (
@@ -65,152 +77,182 @@ try:
 except ImportError:
     from diffusers import QwenImageEditPipeline
 
-# -----------------------------------------------------------------------------
-# Globals
-# -----------------------------------------------------------------------------
-pipe = None
+# =============================================================================
+# GPU Stage: Model loader (NO network)
+# =============================================================================
 
-# -----------------------------------------------------------------------------
-# Load model (GPU stage, NO DOWNLOAD)
-# -----------------------------------------------------------------------------
-def load_model():
+def load_model_on_gpu():
     global pipe
 
-    if pipe is not None:
-        return
+    with model_lock:
+        if pipe is not None:
+            return
 
-    if not torch.cuda.is_available():
-        raise RuntimeError("❌ GPU not available")
+        GPU_STATE["status"] = "loading"
+        GPU_STATE["message"] = "Loading model onto GPU..."
 
-    device = torch.device("cuda")
-    dtype = torch.bfloat16
+        device = torch.device("cuda")
+        dtype = torch.bfloat16
 
-    print("🚀 Loading UniPic-3 DMD from local cache")
-    print("Device:", device)
+        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            LOCAL_MODEL_DIR, subfolder="scheduler"
+        )
 
-    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)
 
-    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"
+        )
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        LOCAL_MODEL_DIR, subfolder="tokenizer"
-    )
+        processor = Qwen2VLProcessor.from_pretrained(
+            LOCAL_MODEL_DIR, subfolder="processor"
+        )
 
-    processor = Qwen2VLProcessor.from_pretrained(
-        LOCAL_MODEL_DIR, subfolder="processor"
-    )
+        transformer = QwenImageTransformer2DModel.from_pretrained(
+            LOCAL_MODEL_DIR,
+            subfolder="ema_transformer",
+            torch_dtype=dtype,
+        ).to(device)
 
-    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,
-    )
+        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,
+        )
 
-    pipe.to(device)
+        pipe.to(device)
 
-    print("✅ Model loaded successfully")
+        GPU_STATE["status"] = "ready"
+        GPU_STATE["message"] = "GPU ready. Model loaded."
 
+# =============================================================================
+# GPU On-Demand Inference (THIS is the only @GPU function)
+# =============================================================================
 
-# -----------------------------------------------------------------------------
-# Inference
-# -----------------------------------------------------------------------------
-def run(
+@GPU
+def run_inference(
     img1, img2, img3, img4, img5, img6,
     prompt, cfg, seed, steps
 ):
     global pipe
 
-    if pipe is None:
-        load_model()
-
-    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"
-
-    images = [img.convert("RGB") for img in images]
-
-    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
-# -----------------------------------------------------------------------------
+    try:
+        GPU_STATE["status"] = "waiting"
+        GPU_STATE["message"] = "Waiting for GPU..."
+
+        # ZeroGPU will block here until GPU is assigned
+        if not torch.cuda.is_available():
+            return None, "⏳ Waiting for GPU, please retry."
+
+        load_model_on_gpu()
+
+        GPU_STATE["status"] = "running"
+        GPU_STATE["message"] = "Running inference..."
+
+        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."
+
+        images = [img.convert("RGB") for img in images]
+
+        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]
+
+        GPU_STATE["status"] = "ready"
+        GPU_STATE["message"] = "Inference complete."
+
+        return out, "✅ Done"
+
+    except Exception as e:
+        GPU_STATE["status"] = "error"
+        GPU_STATE["message"] = str(e)
+        return None, f"❌ Error: {e}"
+
+# =============================================================================
+# UI Helpers (CPU)
+# =============================================================================
+
+def get_gpu_status():
+    return f"**GPU Status:** `{GPU_STATE['status']}`\n\n{GPU_STATE['message']}"
+
+# =============================================================================
+# UI (ALWAYS CPU, ALWAYS ON)
+# =============================================================================
+
 with gr.Blocks(title="UniPic-3 DMD (ZeroGPU)") as demo:
-    gr.Markdown("# 🔥 UniPic-3 DMD (ZeroGPU + Precached)")
+    gr.Markdown("# 🔥 UniPic-3 DMD – ZeroGPU Demo")
 
-    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)
+    status_box = gr.Markdown(get_gpu_status())
 
-    btn = gr.Button("Generate")
-    out = gr.Image()
-    status = gr.Textbox()
+    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")
+
+            btn = gr.Button("🚀 Generate")
+
+        with gr.Column():
+            out = gr.Image(label="Output")
+            msg = gr.Textbox(label="Result")
 
     btn.click(
-        run,
+        run_inference,
         inputs=[*imgs, prompt, cfg, seed, steps],
-        outputs=[out, status],
+        outputs=[out, msg],
     )
 
+    # Periodic status refresh (CPU only)
+    demo.load(
+        fn=get_gpu_status,
+        inputs=[],
+        outputs=status_box,
+        every=1.0,
+    )
 
-# -----------------------------------------------------------------------------
+# =============================================================================
 # Entry
-# -----------------------------------------------------------------------------
-@GPU
-def main():
-    # CPU stage (no GPU time)
+# =============================================================================
+
+if __name__ == "__main__":
+    # CPU phase: cache weights
     precache_weights()
 
-    # Start Gradio (NO SSR)
+    # UI always-on
     demo.launch(ssr_mode=False)
-
-
-if __name__ == "__main__":
-    main()