app.py

"""Entrypoint do Hugging Face Space (SDK Gradio + ZeroGPU).

Regra do ZeroGPU: modelos só podem ir para 'cuda' no startup do módulo (a
emulação CUDA não intercepta cargas em runtime). Por isso todos os modelos
do seletor são carregados aqui, no boot. O conjunto é configurável pela
variável ESPELHO_SPACE_MODELS (lista separada por vírgula) nas Settings do
Space — orçamento: ~50 GB de disco e a RAM do contêiner.
"""

import os
import sys
from pathlib import Path

import spaces  # noqa: F401  (importar antes do torch; requisito ZeroGPU)

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

from espelho.model import load_model  # noqa: E402
from espelho.ui import CSS, TEMA, build_app  # noqa: E402

PADRAO = "google/gemma-2-2b-it,google/gemma-3-4b-it,google/gemma-2-9b-it"
MODELOS = [
    m.strip()
    for m in os.environ.get("ESPELHO_SPACE_MODELS", PADRAO).split(",")
    if m.strip()
]

bundles = {}
for model_id in MODELOS:
    print(f"[boot] carregando {model_id} ...", flush=True)
    tokenizer, model = load_model(model_id)
    bundles[model_id] = [tokenizer, model, None]

demo = build_app(bundles, ativo_id=MODELOS[0])
demo.launch(css=CSS, theme=TEMA)