"""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)