"""Pre-download MMS-TTS VITS models at Docker build time.

This runs during `docker build` so that models are baked into the image.
At runtime, from_pretrained() loads from the local cache (fast) instead
of downloading from HuggingFace (slow cold-start).

Individual failures are non-fatal — the build continues with whatever
models succeed.
"""

from transformers import VitsModel, AutoTokenizer

MODELS = [
    "facebook/mms-tts-yor",   # Yorùbá
    "facebook/mms-tts-swh",   # Swahili
    "facebook/mms-tts-hau",   # Hausa
    "facebook/mms-tts-pcm",   # Pidgin
    "facebook/mms-tts-aka",   # Twi (Akan)
    "facebook/mms-tts-lug",   # Luganda
    "facebook/mms-tts-amh",   # Amharic
    "facebook/mms-tts-som",   # Somali
    "facebook/mms-tts-sna",   # Shona
    "khof312/mms-tts-lin",    # Lingala
    "facebook/mms-tts-ara",   # Arabic
]

ok = 0
for m in MODELS:
    try:
        VitsModel.from_pretrained(m)
        AutoTokenizer.from_pretrained(m)
        print(f"  OK {m}")
        ok += 1
    except Exception as e:
        print(f"  SKIP {m}: {e}")

print(f"{ok}/{len(MODELS)} models cached")