Upload app.py

app.py

@@ -43,56 +43,88 @@ _PIPE = None
 _MESH = None
 
 def _load():
+    """Load on CPU at import time. Models are moved to GPU inside the @spaces.GPU call,
+    so the timed GPU window is spent on inference, not on multi-GB model loading —
+    which is what caused first-call stalls/timeouts."""
     global _PIPE, _MESH
     if _PIPE is not None:
         return
+    import time
     from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
     from controlnet_aux import MeshGraphormerDetector
-    _MESH = MeshGraphormerDetector.from_pretrained(MESHGRAPHORMER_ID).to("cuda")
+    t0 = time.time()
+    print("[load] starting model load on CPU…", flush=True)
+    _MESH = MeshGraphormerDetector.from_pretrained(MESHGRAPHORMER_ID)
+    print(f"[load] meshgraphormer ok ({time.time()-t0:.0f}s)", flush=True)
     cn = ControlNetModel.from_pretrained(CONTROLNET_ID, torch_dtype=torch.float16)
     pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
         SD_INPAINT_ID, controlnet=cn, torch_dtype=torch.float16, safety_checker=None
     )
     pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-    _PIPE = pipe.to("cuda")
+    try: pipe.enable_attention_slicing()
+    except Exception as e: print("[load] attn-slicing skip:", e, flush=True)
+    try: pipe.enable_vae_tiling()
+    except Exception as e: print("[load] vae-tiling skip:", e, flush=True)
+    _PIPE = pipe
+    print(f"[load] pipeline ready on CPU ({time.time()-t0:.0f}s total)", flush=True)
+
+# preload at import — runs once when the container boots, OUTSIDE any GPU-timed window
+try:
+    _load()
+except Exception as _e:
+    print("[load] preload deferred:", _e, flush=True)
 
 def _fit(img):
     w, h = img.size
     s = min(1.0, MAX_SIDE / max(w, h))
     return img.resize((max(8, int(round(w*s/8))*8), max(8, int(round(h*s/8))*8)), Image.LANCZOS), (w, h)
 
-@spaces.GPU(duration=90)
+@spaces.GPU(duration=120)
 def fix_hands(image, mask_layers, prompt, strength):
-    """ZeroGPU-allocated worker, wired directly into the Gradio Interface below."""
+    """ZeroGPU-allocated worker. Models are already loaded (CPU) at import;
+    here we move them onto the GPU that ZeroGPU just attached, then infer."""
+    import time, traceback
     if image is None:
         raise gr.Error("Upload an image first.")
-    _load()
-    init, (ow, oh) = _fit(image.convert("RGB"))
-    W, H = init.size
+    try:
+        t0 = time.time()
+        _load()  # no-op if already loaded
+        _MESH.to("cuda")
+        _PIPE.to("cuda")
+        print(f"[fix] models on GPU, t={time.time()-t0:.0f}s", flush=True)
+        init, (ow, oh) = _fit(image.convert("RGB"))
+        W, H = init.size
+        print(f"[fix] input fitted to {W}x{H}", flush=True)
 
-    # optional hand-drawn mask from the ImageMask/Sketchpad component
-    sent_mask = None
-    if isinstance(mask_layers, dict):
-        layers = mask_layers.get("layers") or []
-        if layers:
-            m = layers[0].convert("L").resize((W, H), Image.LANCZOS)
-            if m.getbbox() is not None:
-                sent_mask = m
+        # optional hand-drawn mask from the ImageMask component
+        sent_mask = None
+        if isinstance(mask_layers, dict):
+            layers = mask_layers.get("layers") or []
+            if layers:
+                m = layers[0].convert("L").resize((W, H), Image.LANCZOS)
+                if m.getbbox() is not None:
+                    sent_mask = m
 
-    mg = _MESH(init)
-    depth_img, auto_mask = (mg[0], (mg[1] if len(mg) > 1 else None)) if isinstance(mg, tuple) else (mg, None)
-    depth_img = depth_img.convert("RGB").resize((W, H), Image.LANCZOS)
-    mask_img = sent_mask or (auto_mask.convert("L").resize((W, H), Image.LANCZOS) if auto_mask else None)
-    if mask_img is None:
-        raise gr.Error("No hands detected. Paint a mask over the hand and try again.")
+        print("[fix] running MeshGraphormer…", flush=True)
+        mg = _MESH(init)
+        depth_img, auto_mask = (mg[0], (mg[1] if len(mg) > 1 else None)) if isinstance(mg, tuple) else (mg, None)
+        depth_img = depth_img.convert("RGB").resize((W, H), Image.LANCZOS)
+        mask_img = sent_mask or (auto_mask.convert("L").resize((W, H), Image.LANCZOS) if auto_mask else None)
+        if mask_img is None:
+            raise gr.Error("No hands detected. Paint a mask over the hand and try again.")
 
-    mask_img = mask_img.filter(ImageFilter.GaussianBlur(2))
-    out = _PIPE(
-        prompt=prompt or DEFAULT_PROMPT, negative_prompt=NEG, image=init, mask_image=mask_img,
-        control_image=depth_img, num_inference_steps=30, strength=float(strength),
-        guidance_scale=7.5, controlnet_conditioning_scale=0.7,
-    ).images[0]
-    return out.resize((ow, oh), Image.LANCZOS)
+        mask_img = mask_img.filter(ImageFilter.GaussianBlur(2))
+        print("[fix] running diffusion…", flush=True)
+        out = _PIPE(
+            prompt=prompt or DEFAULT_PROMPT, negative_prompt=NEG, image=init, mask_image=mask_img,
+            control_image=depth_img, num_inference_steps=25, strength=float(strength),
+            guidance_scale=7.5, controlnet_conditioning_scale=0.7,
+        ).images[0]
+        print(f"[fix] done, total {time.time()-t0:.0f}s", flush=True)
+        return out.resize((ow, oh), Image.LANCZOS)
+    except Exception as e:
+        print("[fix] ERROR:\n" + traceback.format_exc(), flush=True)
+        raise gr.Error(f"Fix failed: {e}")
 
 with gr.Blocks(title="DARKROOM HandRefiner", theme=gr.themes.Base()) as demo:
     gr.Markdown("## 🖐️ DARKROOM HandRefiner\nUpload AI art with bad hands. It auto-detects hands "