scripts/prepare_runtime.py

"""Fetch pinned source/model assets and fx. No inference or API requests are made."""
import hashlib
import io
import json
import os
from pathlib import Path
import subprocess
import shutil
import sys
import importlib
import tarfile
import urllib.request

ROOT = Path(__file__).resolve().parents[1]
RUNTIME = ROOT / ".runtime"
LOCK = json.loads((ROOT / "models.lock.json").read_text())


def fetch(url):
    with urllib.request.urlopen(url, timeout=120) as response:
        return response.read()


def repository(name, revision, destination):
    marker = destination / ".revision"
    if marker.exists() and marker.read_text() == revision:
        return
    data = fetch(f"https://api.github.com/repos/{name}/tarball/{revision}")
    staging = RUNTIME / (destination.name + "-download")
    staging.mkdir(parents=True, exist_ok=True)
    with tarfile.open(fileobj=io.BytesIO(data), mode="r:gz") as archive:
        archive.extractall(staging, filter="data")
    source = next(p for p in staging.iterdir() if p.is_dir())
    if destination.exists():
        shutil.rmtree(destination)
    source.rename(destination)
    shutil.rmtree(staging)
    marker.write_text(revision)


ROLES = ("reference", "geometry", "segment", "texture", "motion")


def worker_role(role=None):
    role = role or os.environ.get("STUDIO_WORKER_ROLE") or "all"
    if role not in (*ROLES, "all"):
        raise ValueError(f"Unknown STUDIO_WORKER_ROLE: {role}")
    return role


def check_imports(*modules):
    for module in modules:
        importlib.import_module(module)
        print(f"Runtime import ready: {module}", flush=True)


def prepare_reference():
    from huggingface_hub import snapshot_download
    check_imports("diffusers")
    model = "black-forest-labs/FLUX.2-klein-4B"
    snapshot_download(model, revision=LOCK[model],
                      allow_patterns=["model_index.json", "scheduler/*", "tokenizer/*",
                                      "text_encoder/*", "transformer/*", "vae/*"])


def prepare_segment():
    from huggingface_hub import hf_hub_download
    repository("nv-tlabs/PartField", LOCK["partfield_code"], RUNTIME / "PartField")
    encoder = RUNTIME / "PartField/partfield/model/PVCNN/encoder_pc.py"
    original = "from torch_scatter import scatter_mean #, scatter_max"
    patched = "from studio.partfield import scatter_mean # native torch, no compiled scatter extension"
    text = encoder.read_text()
    if original not in text and patched not in text:
        raise RuntimeError("Pinned PartField source no longer matches the reviewed scatter adapter.")
    encoder.write_text(text.replace(original, patched))
    sys.path.insert(0, str(RUNTIME / "PartField"))
    check_imports("partfield.model.PVCNN.encoder_pc", "partfield.model.triplane")
    hf_hub_download("mikaelaangel/partfield-ckpt", "model_objaverse.ckpt",
                    revision=LOCK["mikaelaangel/partfield-ckpt"])


def prepare_geometry():
    from huggingface_hub import hf_hub_download, snapshot_download
    snapshot_download("microsoft/TRELLIS.2", repo_type="space", revision=LOCK["trellis_space"],
                      allow_patterns=["trellis2/**"], local_dir=RUNTIME / "trellis")
    extractor = RUNTIME / "trellis/trellis2/modules/image_feature_extractor.py"
    original = "enumerate(self.model.layer)"
    patched = "enumerate(self.model.model.layer)"
    text = extractor.read_text()
    if original not in text and patched not in text:
        raise RuntimeError("Pinned TRELLIS DINOv3 extractor no longer matches the Transformers adapter.")
    extractor.write_text(text.replace(original, patched))
    sys.path.insert(0, str(RUNTIME / "trellis"))
    check_imports("flash_attn", "flex_gemm", "nvdiffrast.torch", "cumesh", "o_voxel",
                  "trellis2.pipelines.trellis2_image_to_3d")
    # Check gated access before downloading the large generation checkpoints.
    for model in ("facebook/dinov3-vitl16-pretrain-lvd1689m", "briaai/RMBG-2.0"):
        hf_hub_download(model, "config.json", revision=LOCK[model])
    # Freeze TRELLIS's three transitive checkpoints as well as the top-level model.
    trellis = Path(snapshot_download("microsoft/TRELLIS.2-4B", revision=LOCK["microsoft/TRELLIS.2-4B"]))
    local = RUNTIME / "trellis-model"
    local.mkdir(exist_ok=True)
    if (local / "ckpts").is_symlink():
        (local / "ckpts").unlink()
    if not (local / "ckpts").exists():
        (local / "ckpts").symlink_to(trellis / "ckpts", target_is_directory=True)
    config = json.loads((trellis / "pipeline.json").read_text())
    for suffix in ("json", "safetensors"):
        source = hf_hub_download("microsoft/TRELLIS-image-large",
                                f"ckpts/ss_dec_conv3d_16l8_fp16.{suffix}",
                                revision=LOCK["microsoft/TRELLIS-image-large"])
        target = local / f"ss_decoder.{suffix}"
        if target.is_symlink():
            target.unlink()
        target.symlink_to(source)
    config["args"]["models"]["sparse_structure_decoder"] = "ss_decoder"
    for key, model in (("image_cond_model", "facebook/dinov3-vitl16-pretrain-lvd1689m"),
                       ("rembg_model", "briaai/RMBG-2.0")):
        config["args"][key]["args"]["model_name"] = snapshot_download(
            model, revision=LOCK[model], allow_patterns=["*.json", "*.py", "*.safetensors"])
    (local / "pipeline.json").write_text(json.dumps(config, indent=2))


def prepare_texture():
    check_imports("spandrel")
    weights = RUNTIME / "RealESRGAN_x4plus.pth"
    if not weights.exists():
        weights.write_bytes(fetch("https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"))
    with weights.open("rb") as stream:
        if hashlib.file_digest(stream, "sha256").hexdigest() != LOCK["realesrgan_sha256"]:
            raise RuntimeError("Real-ESRGAN checkpoint checksum mismatch")


def prepare_models(role="all"):
    selected = worker_role(role)
    for name in ROLES:
        if selected in ("all", name):
            globals()[f"prepare_{name}"]()


def prepare_motion():
    from huggingface_hub import snapshot_download

    source = RUNTIME / "kimodo"
    repository("nv-tlabs/kimodo", LOCK["kimodo_code"], source)
    marker = source / ".installed"
    if not marker.exists() or marker.read_text() != LOCK["kimodo_code"]:
        subprocess.run(["uv", "pip", "install", "--python", sys.executable,
                        "--no-deps", "--no-build-isolation", str(source)], check=True, timeout=900)
        marker.write_text(LOCK["kimodo_code"])
    for module in ("kimodo.model", "motion_correction", "bvhio"):
        importlib.import_module(module)
        print(f"Runtime import ready: {module}", flush=True)
    model = "nvidia/Kimodo-SOMA-RP-v1.1"
    snapshot_download(model, revision=LOCK[model], local_dir=RUNTIME / "kimodo-model",
                      allow_patterns=["config.yaml", "model.safetensors", "stats/**"])
    base = "meta-llama/Meta-Llama-3-8B-Instruct"
    base_path = snapshot_download(base, revision=LOCK[base],
                                 allow_patterns=["*.json", "*.safetensors", "tokenizer.model"])
    for suffix, folder in (("", "mntp"), ("-supervised", "supervised")):
        model = "McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp" + suffix
        target = RUNTIME / "kimodo-text" / folder
        snapshot_download(model, revision=LOCK[model], local_dir=target,
                          allow_patterns=["*.json", "*.safetensors", "tokenizer.model"])
        adapter = target / "adapter_config.json"
        config = json.loads(adapter.read_text())
        config["base_model_name_or_path"] = str(Path(base_path).resolve())
        adapter.write_text(json.dumps(config, indent=2))


def prepare_tools():
    repository("TheOrcDev/skills", LOCK["skills"], RUNTIME / "orc-skills")
    from scripts.prepare_fx_runtime import prepare as prepare_fx
    prepare_fx()


def prepare(tools_only=False, role=None):
    RUNTIME.mkdir(exist_ok=True)
    selected = worker_role(role)
    if tools_only or selected == "all":
        prepare_tools()
        print("fx and cleanup skills are ready.", flush=True)
    if tools_only:
        return
    if selected == "all":
        from studio.render import check_renderer
        check_renderer()
        print("Blender GLTF imports are ready.", flush=True)
    prepare_models(selected)
    marker = "prepared.json" if selected == "all" else f"prepared-{selected}.json"
    (RUNTIME / marker).write_text(json.dumps(LOCK, indent=2))
    print(f"Pinned {selected} model assets and source adapters are ready.", flush=True)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tools-only", action="store_true", help="Install fx and skills without CUDA model downloads")
    parser.add_argument("--role", choices=[*ROLES, "all"], default=None)
    args = parser.parse_args()
    prepare(tools_only=args.tools_only, role=args.role)