diff --git a/scripts/download_weights.py b/scripts/download_weights.py index 4286811..cb01ed2 100755 --- a/scripts/download_weights.py +++ b/scripts/download_weights.py @@ -21,7 +21,7 @@ from huggingface_hub.errors import GatedRepoError, HfHubHTTPError ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) -from trellis2.model_revisions import MODEL_REVISIONS, TRELLIS_REPO +from trellis2.model_revisions import MODEL_FILES, MODEL_REVISIONS, TRELLIS_REPO def main() -> int: @@ -52,6 +52,7 @@ def main() -> int: cache_dir=str(cache_dir), local_files_only=args.offline, max_workers=args.max_workers, + allow_patterns=list(MODEL_FILES[repo_id]), ) break except GatedRepoError as exc: @@ -78,6 +79,7 @@ def main() -> int: "cache_dir": str(cache_dir), "primary_repo": TRELLIS_REPO, "revisions": MODEL_REVISIONS, + "runtime_files": MODEL_FILES, "snapshots": snapshots, "offline": args.offline, "max_workers": args.max_workers, diff --git a/tests/test_model_revisions.py b/tests/test_model_revisions.py index ace6711..2a9e500 100644 --- a/tests/test_model_revisions.py +++ b/tests/test_model_revisions.py @@ -3,6 +3,7 @@ from pathlib import Path from trellis2.model_revisions import ( DINOV3_REPO, DINOV3_REVISION, + MODEL_FILES, MODEL_REVISIONS, RMBG_REPO, RMBG_REVISION, @@ -21,6 +22,9 @@ def test_runtime_revisions_are_full_commits(): assert MODEL_REVISIONS[RMBG_REPO] == RMBG_REVISION assert all(len(revision) == 40 for revision in MODEL_REVISIONS.values()) assert all(len(revision) == 40 for revision in SOURCE_REVISIONS.values()) + assert MODEL_FILES.keys() == MODEL_REVISIONS.keys() + assert all(files for files in MODEL_FILES.values()) + assert all(len(files) == len(set(files)) for files in MODEL_FILES.values()) def test_unknown_repo_keeps_explicit_revision(): diff --git a/trellis2/model_revisions.py b/trellis2/model_revisions.py index 333023b..aa77a28 100644 --- a/trellis2/model_revisions.py +++ b/trellis2/model_revisions.py @@ -19,6 +19,43 @@ MODEL_REVISIONS = { RMBG_REPO: RMBG_REVISION, } +# Exact runtime files for the supported image-to-3D CLI. Keeping this manifest +# avoids caching unrelated encoders, legacy checkpoints, and every ONNX/RMBG +# weight variant while still supporting 512, 1024, and 1024_cascade offline. +MODEL_FILES = { + TRELLIS_REPO: ( + "pipeline.json", + "ckpts/ss_flow_img_dit_1_3B_64_bf16.json", + "ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors", + "ckpts/shape_dec_next_dc_f16c32_fp16.json", + "ckpts/shape_dec_next_dc_f16c32_fp16.safetensors", + "ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.json", + "ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.safetensors", + "ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.json", + "ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.safetensors", + "ckpts/tex_dec_next_dc_f16c32_fp16.json", + "ckpts/tex_dec_next_dc_f16c32_fp16.safetensors", + "ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16.json", + "ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16.safetensors", + "ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16.json", + "ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16.safetensors", + ), + TRELLIS_IMAGE_LARGE_REPO: ( + "ckpts/ss_dec_conv3d_16l8_fp16.json", + "ckpts/ss_dec_conv3d_16l8_fp16.safetensors", + ), + DINOV3_REPO: ( + "config.json", + "model.safetensors", + ), + RMBG_REPO: ( + "config.json", + "BiRefNet_config.py", + "birefnet.py", + "model.safetensors", + ), +} + SOURCE_REVISIONS = { "pedronaugusto/mtlgemm": "867aec8234299a7fe1ede7f802c8debe5a939a82", "pedronaugusto/mtldiffrast": "4668cd91cb6d27f5e264731f94a06841fbf7aab8",