Upload demo.py with huggingface_hub

demo.py

@@ -517,6 +517,343 @@ def run_cli(args):
 TOT = 0
 
 
+@spaces.GPU` worker process with a hard time budget (≈60 s on free tier).
+# Importing the Blender shared object inside that budget burns 30–60 s, so
+# the worker is killed *during* bpy import — manifesting as
+# "GPU task aborted" before any model code runs.
+#
+# We start `bpy_server.py` here, in the always-running main process, so the
+# slow bpy import happens exactly once at Space boot.
+#
+# Communication between the main process / ZeroGPU workers and the bpy
+# worker happens via **file-based IPC** under ``/tmp/bpy_ipc/``, avoiding
+# TCP listening ports (which trigger HF Space proxy issues) and HTTP body-
+# size limits.
+# ---------------------------------------------------------------------------
+
+MODEL_CKPTS = [
+    "experiments/articulation_xl_quantization_256_token_4/grpo_1400.ckpt",
+]
+
+HF_PATHS = [
+    "None",
+]
+
+
+def get_dataloader_workers() -> int:
+    # Always use 0 workers — the bpy IPC server is single-threaded and
+    # must process requests sequentially, so parallel loading provides
+    # no benefit.  More importantly, DataLoader worker processes wrap
+    # exceptions from asset loading in a way that can leave the
+    # DataLoader iterator in an inconsistent state, blocking Gradio's
+    # event loop on subsequent requests.
+    return 0
+
+
+# ---------------------------------------------------------------------------
+# bpy_server lifecycle — lazy start so the heavy import doesn't fight ZeroGPU
+# during module load.  Uses file-based IPC (no TCP listening).
+# ---------------------------------------------------------------------------
+_BPY_CLIENT: "Optional[BpyClient]" = None
+
+
+def ensure_bpy_server_started(force: bool = False):
+    """Start the bpy worker (if not already running) and wait until ready.
+
+    Uses a fast PID-file check to detect an already-running bpy_server
+    (works across processes, including ZeroGPU workers).  Falls back to
+    launching a fresh subprocess when no server is found.
+
+    When *force* is True the existing worker (if any) is killed and a
+    fresh one is started unconditionally — use after a timeout to
+    recover from a hung server.
+    """
+    global _BPY_CLIENT
+
+    if not force:
+        # Fast path: we hold a live client handle (same-process, e.g. main).
+        if _BPY_CLIENT is not None and _BPY_CLIENT.is_alive():
+            return
+
+        # Cross-process fast path: another process already started the worker.
+        if is_bpy_server_alive():
+            return
+
+    # If we had a dead client handle, clean it up.
+    if _BPY_CLIENT is not None:
+        try:
+            _BPY_CLIENT.shutdown()
+        except Exception:
+            pass
+
+    # Launch a fresh bpy worker (this blocks until bpy is imported, ~30–60 s
+    # on a cold container, but the @spaces.GPU duration budget accommodates).
+    _BPY_CLIENT = BpyClient.launch()
+
+
+# ---------------------------------------------------------------------------
+# Lazy model loading.
+# ---------------------------------------------------------------------------
+model = None
+tokenizer = None
+transform = None
+CURRENT_MODEL_CKPT: Optional[str] = None
+CURRENT_HF_PATH: Optional[str] = None
+
+
+def load_model(model_ckpt: str, hf_path: Optional[str]) -> Tuple[str, str]:
+    global model, tokenizer, transform, CURRENT_MODEL_CKPT, CURRENT_HF_PATH
+    if hf_path == "None":
+        hf_path = None
+    if model is not None and model_ckpt == CURRENT_MODEL_CKPT and hf_path == CURRENT_HF_PATH:
+        return ("Model already loaded.", model_ckpt)
+
+    if not model_ckpt:
+        raise RuntimeError("model_ckpt is empty. Please select a checkpoint.")
+
+    print(f"Loading model: {model_ckpt}, hf_path={hf_path}")
+    model = get_model(model_ckpt, hf_path=hf_path)
+    assert model.tokenizer_config is not None
+    tokenizer = get_tokenizer(**model.tokenizer_config)
+    transform = Transform.parse(**model.transform_config["predict_transform"])
+    CURRENT_MODEL_CKPT = model_ckpt
+    CURRENT_HF_PATH = hf_path
+    return ("Model loaded.", model_ckpt)
+
+
+# ---------------------------------------------------------------------------
+# File utilities (CLI-side).
+# ---------------------------------------------------------------------------
+SUPPORTED_EXT = {".obj", ".fbx", ".glb"}
+
+
+def collect_files(input_path: Path) -> List[Path]:
+    if input_path.is_file():
+        return [input_path]
+
+    files = []
+    for p in input_path.rglob("*"):
+        if p.suffix.lower() in SUPPORTED_EXT:
+            files.append(p)
+    return files
+
+
+def map_output_path(in_path: Path, input_root: Path, output_root: Path) -> Path:
+    rel = in_path.relative_to(input_root)
+    return (output_root / rel).with_suffix(".glb")
+
+
+# ---------------------------------------------------------------------------
+# Core inference (shared by CLI and Gradio).
+# ---------------------------------------------------------------------------
+def run_rig(
+    filepaths: List[Path],
+    top_k: int,
+    top_p: float,
+    temperature: float,
+    repetition_penalty: float,
+    num_beams: int,
+    use_skeleton: bool,
+    use_transfer: bool,
+    use_postprocess: bool,
+    output_paths: List[Path],
+    model_ckpt: str,
+    hf_path: Optional[str],
+):
+    assert len(filepaths) == len(output_paths)
+    ensure_bpy_server_started()
+    load_model(model_ckpt, hf_path)
+
+    datapath = {
+        "data_name": None,
+        "loader": "bpy_server",
+        "filepaths": {"articulation": [str(p) for p in filepaths]},
+    }
+
+    dataset_config = DatasetConfig.parse(
+        shuffle=False,
+        batch_size=1,
+        num_workers=get_dataloader_workers(),
+        pin_memory=get_dataloader_workers() > 0,
+        persistent_workers=False,
+        datapath=datapath,
+    ).split_by_cls()
+
+    module = RigDatasetModule(
+        predict_dataset_config=dataset_config,
+        predict_transform=transform,
+        tokenizer=tokenizer,
+        process_fn=model._process_fn,
+    )
+
+    dataloader = module.predict_dataloader()["articulation"]
+
+    results_out = []
+    infer_device = model.device if model is not None else "cuda"
+
+    # Use explicit next() so we can catch errors from asset loading
+    # (which happens inside __getitem__ → next(dataloader)) per-file.
+    dataloader_iter = iter(dataloader)
+    for i in tqdm(range(len(filepaths)), desc="Processing"):
+        filepath = filepaths[i]
+        out_path = output_paths[i]
+        try:
+            batch = next(dataloader_iter)
+        except StopIteration:
+            break
+        except Exception as exc:
+            err_msg = str(exc)
+            print(f"[Error] Failed to load '{filepath}': {err_msg}", flush=True)
+            import traceback
+            traceback.print_exc()
+            # Force-restart if the server timed out (hung) so the next file
+            # gets a clean interpreter instead of hitting the same hang.
+            _restart_bpy_if_dead(force=("timeout" in err_msg.lower()))
+            continue
+
+        try:
+            batch = {
+                k: v.to(infer_device) if isinstance(v, Tensor) else v
+                for k, v in batch.items()
+            }
+
+            if not use_skeleton:
+                batch.pop("skeleton_tokens", None)
+                batch.pop("skeleton_mask", None)
+
+            batch["generate_kwargs"] = dict(
+                max_length=2048,
+                top_k=int(top_k),
+                top_p=float(top_p),
+                temperature=float(temperature),
+                repetition_penalty=float(repetition_penalty),
+                num_return_sequences=1,
+                num_beams=int(num_beams),
+                do_sample=True,
+            )
+
+            if "skeleton_tokens" in batch and "skeleton_mask" in batch:
+                mask = batch["skeleton_mask"][0] == 1
+                skeleton_tokens = batch["skeleton_tokens"][0][mask].cpu().numpy()
+            else:
+                skeleton_tokens = None
+
+            preds: List[TokenRigResult] = model.predict_step(
+                batch,
+                skeleton_tokens=[skeleton_tokens] if skeleton_tokens is not None else None,
+                make_asset=True,
+            )["results"]
+
+            asset = preds[0].asset
+            assert asset is not None
+
+            if use_postprocess:
+                voxel = asset.voxel(resolution=196)
+                asset.skin *= voxel_skin(
+                    grid=0,
+                    grid_coords=voxel.coords,
+                    joints=asset.joints,
+                    vertices=asset.vertices,
+                    faces=asset.faces,
+                    mode="square",
+                    voxel_size=voxel.voxel_size,
+                )
+                asset.normalize_skin()
+
+            out_path.parent.mkdir(parents=True, exist_ok=True)
+
+            if use_transfer:
+                payload = dict(
+                    source_asset=asset,
+                    target_path=asset.path,
+                    export_path=str(out_path),
+                    group_per_vertex=4,
+                )
+                res = bpy_server_request("transfer", payload)
+            else:
+                payload = dict(
+                    asset=asset,
+                    filepath=str(out_path),
+                    group_per_vertex=4,
+                )
+                res = bpy_server_request("export", payload)
+
+            if res != "ok":
+                print(f"[Error] {filepath}: {res}")
+            else:
+                print(f"[OK] Exported: {out_path}")
+                results_out.append(out_path)
+        except Exception as exc:
+            err_msg = str(exc)
+            print(f"[Error] Failed to process '{filepath}': {err_msg}", flush=True)
+            import traceback
+            traceback.print_exc()
+            # Force-restart if the server timed out (hung).
+            _restart_bpy_if_dead(force=("timeout" in err_msg.lower()))
+            continue
+
+    return results_out
+
+
+def _restart_bpy_if_dead(force: bool = False):
+    """If the bpy worker process died or hung, start a new one.
+
+    When *force* is True, the existing worker is killed and replaced
+    unconditionally — use after a request timeout to guarantee a clean
+    bpy interpreter for the next file.
+    """
+    if force:
+        print("[run_rig] bpy worker timed out — force-restarting...", flush=True)
+        try:
+            ensure_bpy_server_started(force=True)
+        except Exception as exc:
+            print(f"[run_rig] Failed to force-restart bpy worker: {exc}", flush=True)
+    elif not is_bpy_server_alive():
+        print("[run_rig] bpy worker is dead — restarting...", flush=True)
+        try:
+            ensure_bpy_server_started()
+        except Exception as exc:
+            print(f"[run_rig] Failed to restart bpy worker: {exc}", flush=True)
+
+
+# ---------------------------------------------------------------------------
+# CLI entry point.
+# ---------------------------------------------------------------------------
+def run_cli(args):
+    input_path = Path(args.input).resolve()
+    output_path = Path(args.output).resolve()
+
+    files = collect_files(input_path)
+    if not files:
+        raise RuntimeError("No valid 3D files found.")
+
+    if len(files) == 1 and output_path.suffix:
+        outputs = [output_path]
+    else:
+        outputs = [map_output_path(f, input_path, output_path) for f in files]
+
+    run_rig(
+        files,
+        args.top_k,
+        args.top_p,
+        args.temperature,
+        args.repetition_penalty,
+        args.num_beams,
+        args.use_skeleton,
+        args.use_transfer,
+        args.use_postprocess,
+        outputs,
+        args.model_ckpt,
+        args.hf_path,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Gradio wrapper (with ZeroGPU duration estimator).
+# ---------------------------------------------------------------------------
+TOT = 0
+
+
 def _gpu_duration(
     files,
     top_k,
@@ -533,7 +870,7 @@ def _gpu_duration(
     # Cold workers spend ~30–60 s importing bpy + loading the model before
     # any GPU work. Give every request a generous 240 s floor.
     file_count = len(files) if files is not None else 1
-    return min(900, max(240, 240 + 60 * file_count))
+    return 190  # patched: stay under free-tier per-call cap
 
 
 @spaces.GPU(duration=_gpu_duration)