scripts/prefetch_models.py

#!/usr/bin/env python
"""Baixa pesos para o cache HF — funciona no BUILD Docker e no runtime Spaces.

BUILD (AKIRA-SOFTEDGE Docker):  HF_HOME=/tmp/hf_cache (sem /data no build)
RUNTIME (Akiragpu Gradio):      /data/hf_cache/hub persistente (ZeroGPU: só leitura
de disco dentro da janela GPU, não download)
"""
import os, sys
if os.path.isdir("/data") and os.access("/data", os.W_OK):
    HF_CACHE = "/data/hf_cache/hub"
    os.environ.setdefault("HF_HOME", "/data/hf_cache")
    os.environ.setdefault("HF_HUB_CACHE", HF_CACHE)
else:
    HF_CACHE = os.getenv("HF_HOME", "/tmp/hf_cache") + "/hub"
    os.environ.setdefault("HF_HOME", os.getenv("HF_HOME", "/tmp/hf_cache"))

MODELS = [
    ("MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7", "emotion"),
    ("rufimelo/bert-large-portuguese-cased-sts", "embedding"),
    (os.getenv("AKIRA_LOCAL_MODEL","Orenguteng/Llama-3.1-8B-Lexi-Uncensored-V2"), "local-llm"),
    (os.getenv("AKIRA_LOCAL_VLM","microsoft/Phi-3.5-vision-instruct"), "vision"),
    # FLUX é GATED — só baixa se o token tiver licença aceite; senão skip gracioso
    (os.getenv("AKIRA_LOCAL_IMAGE_GEN","black-forest-labs/FLUX.1-schnell"), "flux"),
]
ok=[]
token = (
    os.getenv("HF_TOKEN") or os.getenv("ANN_HF_TOKEN") or os.getenv("ISAAC_HF_TOKEN")
    or os.getenv("GITAKIRA_HF_TOKEN") or os.getenv("NETFLIX_HF_TOKEN")
    or os.getenv("FUGAKUSAYKU_HF_TOKEN")
)
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
for mid, tag in MODELS:
    try:
        from huggingface_hub import snapshot_download
        print(f"[{tag}] -> {mid} -> {HF_CACHE} ...", flush=True)
        snapshot_download(mid, cache_dir=HF_CACHE, token=token)
        ok.append(mid)
        print(f"[{tag}] OK", flush=True)
    except Exception as e:
        print(f"[{tag}] FAIL {mid}: {str(e)[:200]}", file=sys.stderr, flush=True)
print(f"prefetch done: {len(ok)}/{len(MODELS)} -> {HF_CACHE}", flush=True)