scripts/fetch_models.py

"""Open-weights fetcher. Downloads the models once so the runtime can run locally:
- the LLM GGUF (Qwen2.5-1.5B-Instruct, Q4_K_M) into models/llm/;
- the Supertonic ONNX TTS model into its cache.

Run once locally (or it is invoked automatically at app startup if the GGUF is missing).

    python scripts/fetch_models.py
    python scripts/fetch_models.py --llm-only
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "src"))

LLM_DIR = ROOT / "models" / "llm"
# The single model the game runs on: Qwen2.5-1.5B (fast, <=4B Tiny Titan).
LLM_CANDIDATES = (
    ("Qwen/Qwen2.5-1.5B-Instruct-GGUF", "qwen2.5-1.5b-instruct-q4_k_m.gguf",
     LLM_DIR / "qwen2.5-1.5b-instruct-q4_k_m.gguf"),
)

def _prefetch_supertonic() -> bool:
    """Cache the Supertonic ONNX model so the runtime never downloads (off-grid)."""
    try:
        from supertonic import TTS

        TTS(auto_download=True)
        print("[fetch] Supertonic model cached")
        return True
    except Exception as exc:
        print(f"[fetch] Supertonic prefetch failed: {exc}", file=sys.stderr)
        return False


def _download(repo: str, filename: str, dest: Path) -> bool:
    from huggingface_hub import hf_hub_download

    try:
        path = hf_hub_download(repo_id=repo, filename=filename)
        dest.parent.mkdir(parents=True, exist_ok=True)
        dest.write_bytes(Path(path).read_bytes())
        print(f"[fetch] {repo}/{filename} -> {dest} ({dest.stat().st_size} bytes)")
        return True
    except Exception as exc:
        print(f"[fetch] FAILED {repo}/{filename}: {exc}", file=sys.stderr)
        return False


def main() -> int:
    parser = argparse.ArgumentParser(description="Fetch open weights for Case Zero")
    parser.add_argument("--llm-only", action="store_true")
    parser.add_argument("--tts-only", action="store_true")
    args = parser.parse_args()

    ok = True
    if not args.tts_only:
        llm_ok = False
        for repo, filename, dest in LLM_CANDIDATES:
            if _download(repo, filename, dest):
                llm_ok = True
                break
        ok &= llm_ok
    if not args.llm_only:
        ok &= _prefetch_supertonic()
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())