94 lines
3.2 KiB
Python
Executable File
94 lines
3.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Populate the dedicated Hugging Face cache with every pinned runtime input."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
# The Xet transport has produced non-resumable response-body decode failures
|
|
# on large model shards on macOS. Plain HTTP keeps .incomplete files in the
|
|
# selected HF cache and resumes them on the next attempt.
|
|
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
|
|
|
|
from huggingface_hub import snapshot_download
|
|
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
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--cache-dir",
|
|
type=Path,
|
|
default=Path.home() / ".cache" / "trellis2" / "huggingface",
|
|
)
|
|
parser.add_argument("--offline", action="store_true", help="verify that the pinned cache is complete")
|
|
parser.add_argument("--max-workers", type=int, default=2)
|
|
parser.add_argument("--retries", type=int, default=3)
|
|
args = parser.parse_args()
|
|
if args.max_workers < 1:
|
|
parser.error("--max-workers must be at least 1")
|
|
if args.retries < 1:
|
|
parser.error("--retries must be at least 1")
|
|
cache_dir = args.cache_dir.expanduser().resolve()
|
|
cache_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
snapshots = {}
|
|
for repo_id, revision in MODEL_REVISIONS.items():
|
|
for attempt in range(1, args.retries + 1):
|
|
try:
|
|
snapshots[repo_id] = snapshot_download(
|
|
repo_id=repo_id,
|
|
revision=revision,
|
|
cache_dir=str(cache_dir),
|
|
local_files_only=args.offline,
|
|
max_workers=args.max_workers,
|
|
)
|
|
break
|
|
except GatedRepoError as exc:
|
|
raise RuntimeError(
|
|
f"access to {repo_id}@{revision} is gated; accept its terms and run "
|
|
"`hf auth login`, then retry"
|
|
) from exc
|
|
except (OSError, RuntimeError, HfHubHTTPError) as exc:
|
|
if args.offline or attempt == args.retries:
|
|
raise RuntimeError(
|
|
f"failed to cache {repo_id}@{revision} after {attempt} attempt(s): {exc}"
|
|
) from exc
|
|
delay = min(5, attempt * 2)
|
|
print(
|
|
f"Retrying {repo_id}@{revision} after transient download error "
|
|
f"({attempt}/{args.retries}): {exc}",
|
|
file=sys.stderr,
|
|
)
|
|
time.sleep(delay)
|
|
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"cache_dir": str(cache_dir),
|
|
"primary_repo": TRELLIS_REPO,
|
|
"revisions": MODEL_REVISIONS,
|
|
"snapshots": snapshots,
|
|
"offline": args.offline,
|
|
"max_workers": args.max_workers,
|
|
},
|
|
indent=2,
|
|
sort_keys=True,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|