monster/trellis-2-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix: cache only pinned runtime weights

Browse Source
This commit is contained in:
Jourloy 2026-07-17 09:24:41 +03:00
parent f59e6e5981
commit 2b7f47fc87
3 changed files with 44 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
scripts/download_weights.py
View File

@ -21,7 +21,7 @@ from huggingface_hub.errors import GatedRepoError, HfHubHTTPError
ROOT = Path(__file__).resolve().parents[1] ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT)) 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: def main() -> int:
@ -52,6 +52,7 @@ def main() -> int:
cache_dir=str(cache_dir), cache_dir=str(cache_dir),
local_files_only=args.offline, local_files_only=args.offline,
max_workers=args.max_workers, max_workers=args.max_workers,
allow_patterns=list(MODEL_FILES[repo_id]),
) )
break break
except GatedRepoError as exc: except GatedRepoError as exc:
@ -78,6 +79,7 @@ def main() -> int:
"cache_dir": str(cache_dir), "cache_dir": str(cache_dir),
"primary_repo": TRELLIS_REPO, "primary_repo": TRELLIS_REPO,
"revisions": MODEL_REVISIONS, "revisions": MODEL_REVISIONS,
"runtime_files": MODEL_FILES,
"snapshots": snapshots, "snapshots": snapshots,
"offline": args.offline, "offline": args.offline,
"max_workers": args.max_workers, "max_workers": args.max_workers,

4
tests/test_model_revisions.py
View File

@ -3,6 +3,7 @@ from pathlib import Path
from trellis2.model_revisions import ( from trellis2.model_revisions import (
DINOV3_REPO, DINOV3_REPO,
DINOV3_REVISION, DINOV3_REVISION,
MODEL_FILES,
MODEL_REVISIONS, MODEL_REVISIONS,
RMBG_REPO, RMBG_REPO,
RMBG_REVISION, RMBG_REVISION,
@ -21,6 +22,9 @@ def test_runtime_revisions_are_full_commits():
assert MODEL_REVISIONS[RMBG_REPO] == RMBG_REVISION 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 MODEL_REVISIONS.values())
assert all(len(revision) == 40 for revision in SOURCE_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(): def test_unknown_repo_keeps_explicit_revision():

37
trellis2/model_revisions.py
View File

@ -19,6 +19,43 @@ MODEL_REVISIONS = {
RMBG_REPO: RMBG_REVISION, 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 = { SOURCE_REVISIONS = {
"pedronaugusto/mtlgemm": "867aec8234299a7fe1ede7f802c8debe5a939a82", "pedronaugusto/mtlgemm": "867aec8234299a7fe1ede7f802c8debe5a939a82",
"pedronaugusto/mtldiffrast": "4668cd91cb6d27f5e264731f94a06841fbf7aab8", "pedronaugusto/mtldiffrast": "4668cd91cb6d27f5e264731f94a06841fbf7aab8",