| """Prepare Space resources without importing torch or exposing credentials.""" |
| import argparse |
| import importlib.metadata |
| import json |
| import os |
| from pathlib import Path, PurePosixPath |
| import re |
| import shutil |
| import subprocess |
| import sys |
| import tarfile |
| import tempfile |
| import urllib.request |
|
|
| _HERE = Path(__file__).resolve().parent |
| _MANIFEST = json.loads((_HERE / "source_manifest.json").read_text()) |
| _YOLOX_URL = "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l.pth" |
| _REGRESSOR_URL = "https://raw.githubusercontent.com/zju3dv/GVHMR/main/hmr4d/utils/body_model/smpl_neutral_J_regressor.pt" |
| _BODY = "assets/body_models/smplh/neutral/model.npz" |
| _REGRESSOR = "assets/body_models/smpl_neutral_J_regressor.pt" |
|
|
|
|
| def _exists(path): |
| return path.is_file() and path.stat().st_size > 0 |
|
|
|
|
| def _source_dir(root, name): |
| configured = os.environ.get("DINOV3_REPO_DIR") if name == "DINOv3" else None |
| path = Path(configured) if configured else root / _MANIFEST["sources"][name]["directory"] |
| if name == "DINOv3" and not configured and not path.exists() and (root.parent / "dinov3/hubconf.py").is_file(): |
| path = root.parent / "dinov3" |
| return path |
|
|
|
|
| def _models(enable_camera): |
| models = [("FLOWHMR_MODEL", "fafsaf1/flowhmr-0.46B", "flowhmr_latest/" + f, "checkpoints/flowhmr_latest/" + f) |
| for f in ("flowhmr_latest.ckpt", "config.yml")] |
| models += [("FLOWHMR_SAM", "facebook/sam-3d-body-dinov3", f, "ckpts/sam-3d-body-dinov3/" + f) |
| for f in ("model.ckpt", "model_config.yaml", "assets/mhr_model.pt")] |
| if enable_camera: |
| models.append(("FLOWHMR_VGGT", "facebook/VGGT-Omega", "vggt_omega_1b_512.pt", "ckpts/vggt-omega/vggt_omega_1b_512.pt")) |
| return models |
|
|
|
|
| def _asset_paths(enable_camera): |
| assets = _MANIFEST["assets"] |
| return assets["required"] + (assets["camera"] if enable_camera else []) |
|
|
|
|
| def check(root: Path, enable_camera: bool = True) -> list[str]: |
| """Return missing files; no network, imports of model code, or installations.""" |
| root = Path(root).resolve() |
| paths = _asset_paths(enable_camera) + ["assets/motion_stats.json"] |
| missing = [str(root / p) for p in paths if not _exists(root / p)] |
| for name, spec in _MANIFEST["sources"].items(): |
| if name == "VGGT" and not enable_camera: |
| continue |
| directory = _source_dir(root, name) |
| missing += [str(directory / f) for f in spec["required"] if not _exists(directory / f)] |
| return missing |
|
|
|
|
| def _download(url, destination): |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| with tempfile.TemporaryDirectory(dir=destination.parent, prefix=".download-") as tmp: |
| partial = Path(tmp) / "file" |
| request = urllib.request.Request(url, headers={"User-Agent": "FlowHMR-Space"}) |
| with urllib.request.urlopen(request, timeout=60) as response, partial.open("wb") as output: |
| shutil.copyfileobj(response, output, length=1024 * 1024) |
| if not _exists(partial): |
| raise RuntimeError("Downloaded resource is empty") |
| os.replace(partial, destination) |
|
|
|
|
| def _extract(archive, destination): |
| """Extract contained regular files/directories; omit unneeded archive links.""" |
| with tarfile.open(archive, "r:gz") as source: |
| prefix = None |
| for member in source: |
| path = PurePosixPath(member.name) |
| if path.is_absolute() or ".." in path.parts or not path.parts: |
| raise ValueError("Unsafe source archive entry") |
| if prefix is None: |
| prefix = path.parts[0] |
| if path.parts[0] != prefix: |
| raise ValueError("Source archive must contain one root directory") |
| if len(path.parts) == 1: |
| if not member.isdir(): |
| raise ValueError("Invalid archive root") |
| continue |
| target = destination.joinpath(*path.parts[1:]) |
| if not target.resolve().is_relative_to(destination.resolve()): |
| raise ValueError("Source archive escapes destination") |
| if member.issym() or member.islnk(): |
| continue |
| if not (member.isfile() or member.isdir()): |
| raise ValueError("Unsafe source archive entry") |
| target.parent.mkdir(parents=True, exist_ok=True) |
| if member.isdir(): |
| target.mkdir(exist_ok=True) |
| else: |
| with source.extractfile(member) as content, target.open("xb") as output: |
| shutil.copyfileobj(content, output) |
|
|
|
|
| def _patch_sam(directory): |
| path = directory / "sam_3d_body/models/backbones/dinov3.py" |
| source = path.read_text() |
| pattern = r'(torch\.hub\.load\(\s*)[\'"]facebookresearch/dinov3(?::[^\'"]*)?[\'"]\s*,' |
| patched, count = re.subn(pattern, r'\1os.environ["DINOV3_REPO_DIR"],', source) |
| if count != 1: |
| raise RuntimeError("Pinned SAM DINOv3 loader differs from the expected source") |
| patched, count = re.subn(r'\bsource\s*=\s*([\'"])github\1', 'source="local"', patched) |
| if count != 1: |
| raise RuntimeError("Pinned SAM DINOv3 loader has no unique source=github argument") |
| if not re.search(r"^import os\b", patched, re.MULTILINE): |
| patched = patched.replace("import torch", "import os\n\nimport torch", 1) |
| compile(patched, str(path), "exec") |
| path.write_text(patched) |
|
|
| |
| |
| path = directory / "sam_3d_body/models/heads/mhr_head.py" |
| source = path.read_text() |
| patched, count = re.subn( |
| r'map_location=\("cuda" if torch\.cuda\.is_available\(\) else "cpu"\)', |
| 'map_location="cpu"', source, |
| ) |
| if count != 1: |
| raise RuntimeError("Pinned SAM MHR loader differs from the expected source") |
| compile(patched, str(path), "exec") |
| path.write_text(patched) |
|
|
|
|
| def _sources(root, enable_camera): |
| for name, spec in _MANIFEST["sources"].items(): |
| if name == "VGGT" and not enable_camera: |
| continue |
| directory = _source_dir(root, name) |
| if all(_exists(directory / f) for f in spec["required"]): |
| continue |
| if directory.is_symlink() or (directory.exists() and (not directory.is_dir() or any(directory.iterdir()))): |
| raise RuntimeError(f"Incomplete source directory; provide a complete checkout: {directory}") |
| directory.parent.mkdir(parents=True, exist_ok=True) |
| with tempfile.TemporaryDirectory(dir=directory.parent, prefix=".source-") as tmp: |
| archive, unpacked = Path(tmp) / "source.tar.gz", Path(tmp) / "unpacked" |
| unpacked.mkdir() |
| _download(f'https://codeload.github.com/{spec["repo"]}/tar.gz/{spec["revision"]}', archive) |
| _extract(archive, unpacked) |
| if not all(_exists(unpacked / f) for f in spec["required"]): |
| raise RuntimeError(f"Incomplete downloaded source: {name}") |
| if name == "SAM": |
| _patch_sam(unpacked) |
| if name == "YOLOX": |
| shutil.copytree(unpacked / "exps/default", unpacked / "yolox/exp/default", dirs_exist_ok=True) |
| if directory.exists(): |
| directory.rmdir() |
| os.replace(unpacked, directory) |
|
|
|
|
| def _hub_file(repo, filename, destination, revision="main", repo_type="model"): |
| if _exists(destination): |
| return |
| from huggingface_hub import hf_hub_download |
|
|
| cached = Path(hf_hub_download(repo_id=repo, filename=filename, revision=revision, repo_type=repo_type)).resolve() |
| destination.parent.mkdir(parents=True, exist_ok=True) |
| with tempfile.TemporaryDirectory(dir=destination.parent, prefix=".link-") as tmp: |
| link = Path(tmp) / "cached" |
| link.symlink_to(cached) |
| os.replace(link, destination) |
|
|
|
|
| def _correction_ready(): |
| try: |
| if importlib.metadata.version("motion-correction") != "1.0.0" or importlib.metadata.version("rgbmocap-correction") != "0.1.0": |
| return False |
| except importlib.metadata.PackageNotFoundError: |
| return False |
| probe = "import motion_correction._motion_correction as n; import rgbmocap_correction; assert n.rgb_mocap_anatomical_version == 1" |
| return subprocess.run([sys.executable, "-c", probe], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0 |
|
|
|
|
| def prepare(root: Path, download: bool = True, enable_camera: bool = True, enable_postprocess: bool = True) -> None: |
| """Reuse local resources or fetch missing resources when explicitly enabled.""" |
| root = Path(root).resolve() |
| if download: |
| assets_repo = os.environ.get("FLOWHMR_ASSETS_REPO") |
| if not _exists(root / _BODY) and not assets_repo and not os.environ.get("FLOWHMR_BODY_MODEL_REPO"): |
| raise RuntimeError("SMPL-H neutral is missing. Provide assets/body_models/smplh/neutral/model.npz, or set " |
| "FLOWHMR_ASSETS_REPO to your private assets repository, or FLOWHMR_BODY_MODEL_REPO to your authorized private dataset.") |
| _sources(root, enable_camera) |
| if assets_repo: |
| for local in _asset_paths(enable_camera): |
| _hub_file(assets_repo, local, root / local, os.environ.get("FLOWHMR_ASSETS_REVISION", "main"), |
| os.environ.get("FLOWHMR_ASSETS_REPO_TYPE", "model")) |
| else: |
| for prefix, default_repo, filename, local in _models(enable_camera): |
| _hub_file(os.environ.get(prefix + "_REPO", default_repo), filename, root / local, |
| os.environ.get(prefix + "_REVISION", "main")) |
| if not _exists(root / _BODY): |
| _hub_file(os.environ["FLOWHMR_BODY_MODEL_REPO"], os.environ.get("FLOWHMR_BODY_MODEL_FILE", "smplh/neutral/model.npz"), |
| root / _BODY, os.environ.get("FLOWHMR_BODY_MODEL_REVISION", "main"), os.environ.get("FLOWHMR_BODY_MODEL_REPO_TYPE", "dataset")) |
| for url, local in ((_YOLOX_URL, "ckpts/yolox/yolox_l.pth"), (_REGRESSOR_URL, _REGRESSOR)): |
| if not _exists(root / local): |
| _download(url, root / local) |
| missing = check(root, enable_camera) |
| if missing: |
| raise RuntimeError("Missing Space resources:\n" + "\n".join(missing) + "\nRun bootstrap.py --download after configuring authorized asset access.") |
| os.environ["DINOV3_REPO_DIR"] = str(_source_dir(root, "DINOv3").resolve()) |
| if not os.environ.get("FFMPEG"): |
| import imageio_ffmpeg |
| os.environ["FFMPEG"] = imageio_ffmpeg.get_ffmpeg_exe() |
| if enable_postprocess and not _correction_ready(): |
| if not download: |
| raise RuntimeError("RGBMocap correction is not installed. Install the two bundled correction packages, or disable post-processing.") |
| for package in ("KimodoMotionCorrection", "RGBMocapCorrection"): |
| subprocess.run([sys.executable, "-m", "pip", "install", "--no-build-isolation", "--no-deps", |
| str(root / "third_party" / package)], check=True) |
| if not _correction_ready(): |
| raise RuntimeError("Installed correction package lacks RGBMocap anatomical version 1") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--root", type=Path, default=_HERE / "core" if (_HERE / "core").is_dir() else _HERE.parents[1]) |
| mode = parser.add_mutually_exclusive_group() |
| mode.add_argument("--check", action="store_true", help="Check local resource files only (default)") |
| mode.add_argument("--download", action="store_true", help="Download missing resources and install enabled correction packages") |
| parser.add_argument("--no-camera", action="store_true") |
| parser.add_argument("--no-postprocess", action="store_true") |
| args = parser.parse_args() |
| if args.download: |
| prepare(args.root, True, not args.no_camera, not args.no_postprocess) |
| print("Space resources are ready.") |
| else: |
| missing = check(args.root, not args.no_camera) |
| print("Missing resources:\n" + "\n".join(missing) if missing else "All required local resource files are present.") |
| sys.exit(bool(missing)) |
|
|
