modelos/base.py

# -*- coding: utf-8 -*-
"""
O CONTRATO DE UM MODELO DENTRO DO WORKER.

Cada modelo é um módulo com quatro deveres, e o gerenciador só conversa por
eles:

  carregar()        — traz os pesos para a GPU. Import de torch/diffusers é
                      feito AQUI DENTRO, nunca no topo do arquivo: o worker
                      precisa subir e responder /health mesmo numa máquina
                      sem as bibliotecas pesadas instaladas.
  gerar(params, progresso) — produz bytes + mime. `progresso(pct, etapa)` é
                      chamado quando o modelo sabe onde está (a difusão sabe
                      o passo; TTS muitas vezes não — aí não chama).
  descarregar()     — devolve a VRAM. É o gerenciador que decide QUANDO.
  config_menor(params) — a versão do pedido que cabe em menos VRAM (menos
                      passos, resolução menor, duração menor). Devolve None
                      quando não há o que reduzir — aí o OOM vira falha limpa
                      em vez de tentativa infinita.

E um modo de MENTIRA declarada: com WORKER_MOCK=1 os modelos devolvem
arquivos minúsculos válidos (PNG 1x1, WAV de silêncio, MP4 de um quadro) sem
tocar em GPU nenhuma. Ele existe para os testes automatizados e para subir a
pilha inteira num notebook sem CUDA — e o resultado DIZ que é mock no
metadado, porque fingir inferência real é o tipo de mentira que este projeto
não conta.
"""
from __future__ import annotations

import os
import struct
from typing import Callable, Dict, Optional, Tuple

Progresso = Callable[[Optional[int], str], None]

# bytes, mime
Saida = Tuple[bytes, str]


def modo_mock() -> bool:
    return os.getenv("WORKER_MOCK", "").strip() == "1"


# ---------------------------------------------------------------------------
# Os arquivos mínimos do modo mock — válidos de verdade (abrem em qualquer
# visualizador), só que vazios de conteúdo.
# ---------------------------------------------------------------------------
PNG_1X1 = bytes.fromhex(
    "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489"
    "0000000d4944415478da63fcffff3f0300050201f34d31b20000000049454e44ae426082")


def wav_silencio(segundos: float = 0.2, taxa: int = 16000) -> bytes:
    n = int(segundos * taxa)
    dados = b"\x00\x00" * n
    return (b"RIFF" + struct.pack("<I", 36 + len(dados)) + b"WAVEfmt "
            + struct.pack("<IHHIIHH", 16, 1, 1, taxa, taxa * 2, 2, 16)
            + b"data" + struct.pack("<I", len(dados)) + dados)


def baixar_repo(repo_id: str, progresso: Progresso,
                revision: Optional[str] = None,
                arquivos: Optional[list] = None) -> None:
    """Traz os pesos para o CACHE local (disco), sem tocar em GPU nenhuma.

    Existe por causa do ZeroGPU: lá a GPU só vive dentro da chamada decorada,
    com prazo — e o primeiro `from_pretrained` gastava o prazo inteiro
    BAIXANDO, que é trabalho de rede. Medido: pedir 240s de duração para
    caber o download estourou o teto do ZeroGPU ("requested GPU duration
    (360s) is larger than the maximum allowed"). Com o download feito antes,
    o `carregar()` de dentro da GPU só lê do disco. No worker CUDA chamar
    isto é inócuo: o from_pretrained acha o cache e não baixa de novo.

    `arquivos` RECORTA o que desce, e ele não é conveniência. Um repositório
    do Hub costuma guardar várias variantes do mesmo peso — o do LoRA
    Lightning tem doze arquivos e 65,5 GB para entregar os 0,85 GB de um
    deles. Sem recorte, pedir o acelerador baixaria setenta vezes o que ele
    pesa, e num Space de disco efêmero isso não é lentidão: é o disco
    acabando antes de o modelo principal caber."""
    if modo_mock():
        return
    progresso(None, "Baixando o modelo…")
    from huggingface_hub import snapshot_download
    snapshot_download(repo_id, revision=revision,
                      allow_patterns=arquivos or None,
                      token=os.getenv("HF_TOKEN") or None)


def vram_livre_gb() -> Optional[float]:
    """Quanta VRAM sobra agora. None = sem CUDA (aí ninguém decide por VRAM)."""
    try:
        import torch
        if not torch.cuda.is_available():
            return None
        livre, _total = torch.cuda.mem_get_info()
        return livre / (1024 ** 3)
    except Exception:
        return None


def limpar_vram() -> None:
    """Solta o que der: cache de CUDA e lixo do Python. Chamado pelo
    gerenciador depois de descarregar e depois de um OOM — nas duas horas em
    que 'quase liberado' e 'liberado' são a diferença entre caber e não."""
    import gc
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()
    except Exception:
        pass


def dispositivo() -> str:
    try:
        import torch
        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"


def gpu_dedicada() -> bool:
    """Há prazo de GPU de verdade? `GPU_DEDICADA=1` diz que sim.

    UM FATO, UMA VARIÁVEL. O ZeroGPU gratuito corta a chamada em 180 segundos
    (é o maior valor que ele aceita — ver CLAUDE.md), e esse teto é DURO:
    estourá-lo não é OOM, então nenhuma escada de recuperação roda — a chamada
    morre no meio com a quota do dia gasta e nada saindo. Isso obriga CADA
    módulo a ter um padrão que cabe: o vídeo na resolução, a imagem no número
    de passos.

    Os dois padrões nascem do MESMO fato, então eles leem a MESMA variável.
    Uma por módulo (`VIDEO_720P`, `IMAGE_PASSOS_CHEIOS`…) seriam vários donos
    de uma decisão só, e no dia em que houver GPU dedicada metade ficaria para
    trás — com o sintoma de sempre: nada quebra, só entrega pior.

    Nasce DESLIGADA de propósito: produção é o ZeroGPU gratuito, e um padrão
    que assume folga é uma geração que aborta para todo mundo.
    """
    return (os.getenv("GPU_DEDICADA", "0").strip() or "0") != "0"


def semente(params: Dict):
    """torch.Generator com a seed pedida — ou None para o aleatório de
    sempre. Seed existe para 'gera de novo IGUAL', então ela é respeitada em
    todo modelo que aceita generator."""
    seed = params.get("seed")
    if seed in (None, ""):
        return None
    try:
        import torch
        return torch.Generator(device=dispositivo()).manual_seed(int(seed))
    except Exception:
        return None