Upload demo.py with huggingface_hub

demo.py

@@ -517,343 +517,6 @@ 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,
@@ -870,7 +533,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 190  # patched: stay under free-tier per-call cap
+    return 190  # patched: stay under free-tier per-call duration cap
 
 
 @spaces.GPU(duration=_gpu_duration)