diff --git a/utils/download.py b/utils/download.py index 7dbb39f7..6168a637 100644 --- a/utils/download.py +++ b/utils/download.py @@ -3,6 +3,7 @@ This allows memory-constrained CI/CD runners to build container images with large bundled models. See ../deployments/bundle/ for examples. """ +import os import sys import tempfile @@ -11,6 +12,11 @@ if len(sys.argv) < 3: sys.exit("usage: python download.py REPO_ID LOCAL_DIR [REVISION]") +if os.getenv("TENSOR_FORMAT") == "safetensors": + tensor_format = "*.safetensors" +else: + tensor_format = "*.bin" + with tempfile.TemporaryDirectory() as cache_dir: huggingface_hub.snapshot_download( repo_id=sys.argv[1], @@ -19,5 +25,5 @@ cache_dir=cache_dir, local_dir_use_symlinks=False, resume_download=True, - allow_patterns=["*.bin", "*.json", "*.model", "*.py"], + allow_patterns=[tensor_format, "*.json", "*.model", "*.py"], )