File size: 1,430 Bytes
35bb6f4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """Pre-download models during Docker build or startup."""
from __future__ import annotations
import os
import sys
def main() -> None:
models_str = os.environ.get("NEUTTS_DEFAULT_MODELS", "neutts-nano-q4-gguf")
codec = os.environ.get("NEUTTS_DEFAULT_CODEC", "neuphonic/neucodec-onnx-decoder")
models = [m.strip() for m in models_str.split(",") if m.strip()]
# Import model config to resolve repos
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
from api.src.core.model_config import get_backbone_info
print(f"Pre-downloading {len(models)} model(s) and codec '{codec}'...")
for model_id in models:
info = get_backbone_info(model_id)
if info is None:
print(f" WARNING: Unknown model '{model_id}', skipping")
continue
print(f" Downloading backbone: {info.repo}")
try:
from huggingface_hub import snapshot_download
snapshot_download(info.repo)
print(f" OK: {info.repo}")
except Exception as e:
print(f" FAILED: {info.repo} - {e}")
print(f" Downloading codec: {codec}")
try:
from huggingface_hub import snapshot_download
snapshot_download(codec)
print(f" OK: {codec}")
except Exception as e:
print(f" FAILED: {codec} - {e}")
print("Download complete.")
if __name__ == "__main__":
main()
|