Upload 2 files

app.py

@@ -32,8 +32,29 @@ import torch
 from PIL import Image, ImageFilter
 import gradio as gr
 
+# ---------------------------------------------------------------------------
+# transformers compatibility shim (fixes MeshGraphormer import on new transformers)
+# Newer transformers removed prune_linear_layer / Conv1D from transformers.modeling_utils,
+# which is exactly what breaks the vendored MeshGraphormer (ComfyUI issue #578).
+# Re-expose them so the legacy import succeeds.
+# ---------------------------------------------------------------------------
+def _patch_transformers():
+    try:
+        import transformers.modeling_utils as mu
+        need = ("prune_linear_layer", "Conv1D", "prune_layer")
+        if all(hasattr(mu, n) for n in need):
+            return
+        from transformers import pytorch_utils as pu
+        for n in need:
+            if not hasattr(mu, n) and hasattr(pu, n):
+                setattr(mu, n, getattr(pu, n))
+        print("[shim] transformers symbols patched", flush=True)
+    except Exception as e:
+        print("[shim] transformers patch skipped:", e, flush=True)
+_patch_transformers()
+
 SD_INPAINT_ID = "runwayml/stable-diffusion-inpainting"
-CONTROLNET_ID = "lllyasviel/control_v11f1p_sd15_depth"   # -> hr16/ControlNet-HandRefiner-pruned for best
+CONTROLNET_ID = "lllyasviel/control_v11f1p_sd15_depth"
 TILE_CN_ID = "lllyasviel/control_v11f1e_sd15_tile"        # detail-regeneration ControlNet
 SD_BASE_ID = "runwayml/stable-diffusion-v1-5"             # base SD for img2img detail pass
 MESHGRAPHORMER_ID = "hr16/ControlNet-HandRefiner-pruned"
@@ -46,21 +67,34 @@ DETAIL_NEG = "blurry, soft, out of focus, jpeg artifacts, low quality, smudged,
 _PIPE = None
 _MESH = None
 _DETAIL = None
+_MESH_OK = False
+_MESH_ERR = None
+
+def _make_mesh_detector():
+    """controlnet_aux==0.0.6 ships MeshGraphormerDetector at the top level.
+    (Newer versions dropped it — that's why the pin matters.)"""
+    from controlnet_aux import MeshGraphormerDetector as MGD
+    return MGD.from_pretrained(MESHGRAPHORMER_ID)
 
 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
+    """Load SD inpaint + ControlNet (always works, diffusers-only) and attempt
+    MeshGraphormer (optional). If MeshGraphormer fails, the Space still runs;
+    hand auto-detect is then unavailable but manual-mask + detail pass work."""
+    global _PIPE, _MESH, _MESH_OK, _MESH_ERR
     if _PIPE is not None:
         return
     import time
     from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
-    from controlnet_aux import MeshGraphormerDetector
     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)
+    # MeshGraphormer is optional — isolate it so it can't crash the container
+    try:
+        _MESH = _make_mesh_detector()
+        _MESH_OK = True
+        print(f"[load] meshgraphormer ok ({time.time()-t0:.0f}s)", flush=True)
+    except Exception as e:
+        _MESH = None; _MESH_OK = False; _MESH_ERR = str(e)
+        print("[load] meshgraphormer UNAVAILABLE (manual mask still works):", e, 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
@@ -122,7 +156,10 @@ def fix_hands(image, mask_layers, prompt, strength):
         _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)
+        if _MESH_OK and _MESH is not None:
+            try: _MESH.to("cuda")
+            except Exception: pass
+        print(f"[fix] models on GPU, t={time.time()-t0:.0f}s (mesh={_MESH_OK})", flush=True)
         init, (ow, oh) = _fit(image.convert("RGB"))
         W, H = init.size
         print(f"[fix] input fitted to {W}x{H}", flush=True)
@@ -136,14 +173,29 @@ def fix_hands(image, mask_layers, prompt, strength):
                 if m.getbbox() is not None:
                     sent_mask = m
 
-        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)
+        depth_img = None
+        auto_mask = None
+        if _MESH_OK and _MESH is not None:
+            print("[fix] running MeshGraphormer…", flush=True)
+            try:
+                mg = _MESH(init)
+                depth_img, auto_mask = (mg[0], (mg[1] if len(mg) > 1 else None)) if isinstance(mg, tuple) else (mg, None)
+                if depth_img is not None:
+                    depth_img = depth_img.convert("RGB").resize((W, H), Image.LANCZOS)
+            except Exception as e:
+                print("[fix] mesh inference failed, falling back to mask:", e, flush=True)
+
         mask_img = sent_mask or (auto_mask.convert("L").resize((W, H), Image.LANCZOS) if auto_mask else None)
         if mask_img is None:
+            if not _MESH_OK:
+                raise gr.Error("Auto hand-detection isn't available on this Space build. "
+                               "Paint a mask over the bad hand (use the brush on the image) and run again.")
             raise gr.Error("No hands detected. Paint a mask over the hand and try again.")
 
+        # if we have no depth (no mesh), use the masked region of the image as a soft control
+        if depth_img is None:
+            depth_img = init  # tile/identity-style guidance keeps structure from the source
+
         mask_img = mask_img.filter(ImageFilter.GaussianBlur(2))
         print("[fix] running diffusion…", flush=True)
         out = _PIPE(

requirements.txt

@@ -1,10 +1,17 @@
+# ZeroGPU needs modern gradio + spaces
 spaces
 gradio==5.49.1
-torch
-diffusers
-transformers
-accelerate
-controlnet_aux
+# --- pinned stack that ships a WORKING MeshGraphormerDetector ---
+# (from hysts's official controlnet Space; these versions still export it)
+controlnet_aux==0.0.6
+diffusers==0.18.2
+transformers==4.30.2
+accelerate==0.21.0
+mediapipe==0.10.1
+huggingface-hub==0.16.4
+safetensors==0.3.1
+# torch is provided by the ZeroGPU base image; do not pin it here
 pillow
 numpy
 scipy
+einops