nucleo.py

# -*- coding: utf-8 -*-
"""
O NÚCLEO DO SPACE — tudo do adapter ZeroGPU que NÃO é Gradio.

Por que a separação: o Gradio e o `spaces` só existem DENTRO do Space, e a
lógica que importa (autenticação, duração de GPU por tipo, despacho para os
modelos, tradução dos erros de quota) precisa ser testável nesta casa, onde
nenhum dos dois está instalado. O `app.py` é só a casca que liga isto ao
Gradio; este arquivo é o que se mede no `t_zerogpu.py`.

OS MODELOS NÃO SÃO REESCRITOS. Este arquivo importa os MESMOS módulos de
`services/gpu-worker/modelos/` (e o mesmo `armazenamento.py`): no repositório
ele os acha pelo caminho relativo; no Space publicado eles são enviados JUNTO
(o `publicar.sh` faz isso — uma cópia de DEPLOY, nunca uma segunda
implementação no git). Dois fontes dos mesmos modelos seria o defeito de
sempre: um conserto entra num e o outro fica para trás.

POR QUE ISTO NÃO É O WORKER FASTAPI NUM SPACE: o ZeroGPU não é uma VM CUDA.
A GPU só existe DENTRO de uma chamada decorada com `@spaces.GPU`, com prazo
declarado, e o Space precisa ser Gradio. Então a fila daqui é a do PRÓPRIO
Gradio (cada chamada é um job do backend, síncrona do lado de cá), e quem
continua dono do estado do job é o backend — este processo não guarda nada
que não possa evaporar junto com a GPU.
"""
from __future__ import annotations

import json
import logging
import os
import secrets
import sys
from typing import Callable, Dict, Iterator, Optional

logger = logging.getLogger("zerogpu.nucleo")

# ---------------------------------------------------------------------------
# Onde moram os modelos: no repo, na pasta irmã gpu-worker; no Space, ao lado
# (o publicar.sh os põe lá). A ordem tenta o Space primeiro — é onde roda de
# verdade.
# ---------------------------------------------------------------------------
_AQUI = os.path.dirname(os.path.abspath(__file__))
for _cand in (_AQUI, os.path.join(_AQUI, "..", "gpu-worker")):
    if os.path.isdir(os.path.join(_cand, "modelos")):
        if _cand not in sys.path:
            sys.path.insert(0, _cand)
        break

from modelos import base as mbase        # noqa: E402
from modelos import ace_step, escolha, qwen3_tts  # noqa: E402
import armazenamento                      # noqa: E402

# QUEM ESCOLHE O MÓDULO É O `escolha.py`, o MESMO que o worker CUDA usa. Esta
# tabela e a do `gerenciador.py` eram iguais por sorte enquanto cada tipo
# tinha um módulo só; com dois módulos possíveis para imagem, duas tabelas
# seriam duas respostas para a mesma pergunta — e o Space é justamente o lado
# que ninguém abre para conferir.
_MODULOS = {"image": escolha.modulo_de("image"),
            "speech": qwen3_tts, "voice": qwen3_tts,
            "music": ace_step, "video": escolha.modulo_de("video")}
_carregado: Optional[str] = None

# ---------------------------------------------------------------------------
# DURAÇÃO DE GPU POR TIPO — o prazo que o decorator @spaces.GPU declara.
# Configurável porque é a variável mais sensível do ZeroGPU: curto demais, a
# geração é abortada no meio (e a quota foi gasta mesmo assim); longo demais,
# cada chamada reserva mais quota do que usa. Os padrões saem dos tempos de
# catálogo dos modelos escolhidos, com folga para a carga do modelo frio.
# Este prazo é DA GPU, separado de propósito do timeout HTTP do backend
# (_ESPERA_MAX do adapter): o backend espera mais do que a GPU promete, para
# a falha vir sempre daqui, com o motivo certo, e nunca das duas pontas ao
# mesmo tempo.
#
# A CONTA DO ZEROGPU, medida na primeira validação real: o que ele RESERVA é
# duração × 1,5 (180 declarados viraram "270s requested" na mensagem de
# quota), e o plano gratuito tem um teto POR CHAMADA — 240 declarados (360
# reservados) voltaram "larger than the maximum allowed"; 180 (270) passaram.
# Vídeo era 300 e nunca alocaria numa conta gratuita: 180 é o maior valor
# provado. E o download dos pesos NÃO conta aqui — ele roda fora da GPU
# (nucleo.preparar), então o prazo só cobre carregar do disco + inferir.
# ---------------------------------------------------------------------------
# A IMAGEM SUBIU DE 60 PARA 180, e é por causa do modelo, não por gosto. Os
# 60 s bastavam para o Z-Image-Turbo (8 passos, poucos GB). O Qwen-Image pesa
# 57,7 GB, e o `carregar()` roda DENTRO da chamada de GPU — o download é que
# ficou fora (ver `base.baixar_repo`). A primeira chamada depois de um Space
# frio paga o disco inteiro; as seguintes acham o pipeline residente e não
# pagam nada disso. 180 é o TETO provado nesta casa: o plano gratuito recusou
# 240 declarados ("requested GPU duration (360s) is larger than the maximum
# allowed" — o ZeroGPU reserva duração × 1,5).
#
# Declarar mais do que se usa custa RESERVA, não consumo; declarar menos custa
# a geração inteira, no meio, com a quota já gasta. Entre os dois erros, este
# é o barato.
# O VÍDEO CAIU DE 180 PARA 80 — e é o que faz ele SAIR. Medido em produção em
# 25/09/2026: 180 declarados (270 reservados) voltaram "larger than the maximum
# allowed". A quota do ZeroGPU é POR CHAMADOR (docs do Hugging Face, lidas no
# mesmo dia): chamador sem login tem 2 MINUTOS por dia, conta gratuita 5, PRO
# 40 — e uma reserva maior que o que o chamador tem é recusada inteira. 80
# declarados são 120 reservados: cabem até no chamador anônimo. O Wan cabe
# nisso porque o pedido também encolheu (3 s, 20 passos, 480p — ver
# `wan22.gerar`). Quem tiver mais cota sobe `VIDEO_GPU_SECONDS` no Space.
_DURACAO_PADRAO = {"image": 180, "speech": 60, "voice": 60,
                   "music": 180, "video": 80}
_VAR_DURACAO = {"image": "IMAGE_GPU_SECONDS", "speech": "TTS_GPU_SECONDS",
                "voice": "TTS_GPU_SECONDS", "music": "MUSIC_GPU_SECONDS",
                "video": "VIDEO_GPU_SECONDS"}


def duracao_gpu(tipo: str) -> int:
    try:
        valor = int(os.getenv(_VAR_DURACAO.get(tipo, ""), "") or 0)
    except ValueError:
        valor = 0
    return valor if valor > 0 else _DURACAO_PADRAO.get(tipo, 120)


# ---------------------------------------------------------------------------
# A PORTA. O endpoint de um Space é público por natureza — qualquer pessoa
# que achar a URL pode chamar a API do Gradio. Quem protege a quota é ISTO:
# o segredo é o primeiro argumento de toda chamada, conferido ANTES de a
# função de GPU ser invocada — chamada sem segredo morre aqui, sem alocar um
# segundo de ZeroGPU. Fail closed, como no worker CUDA: sem WORKER_SECRET
# configurado no Space, ninguém gera nada.
# ---------------------------------------------------------------------------
def autenticar(segredo: str) -> Optional[str]:
    """None = pode passar; string = o motivo da recusa (sem eco do que veio)."""
    esperado = os.getenv("WORKER_SECRET", "").strip()
    if not esperado:
        return ("WORKER_SECRET não configurado no Space — configure o secret "
                "com o MESMO valor usado no backend.")
    if not segredo or not secrets.compare_digest(str(segredo), esperado):
        return "segredo inválido"
    return None


# ---------------------------------------------------------------------------
# QUOTA É ESTADO. O `spaces` sinaliza quota esgotada/GPU indisponível com
# exceções próprias cuja MENSAGEM é estável ("GPU quota exceeded", "ZeroGPU");
# a classe muda entre versões do pacote, então o reconhecimento é pelo texto
# — e o resultado leva `erro_gpu`, que o backend traduz para
# GeracaoIndisponivel (degrada ou espera), nunca para falha do pedido.
# ---------------------------------------------------------------------------
_SINAIS_DE_QUOTA = ("quota", "zerogpu", "gpu task aborted", "no gpu",
                    "gpu unavailable", "queue is full",
                    # O PRAZO ACIMA DO TETO também é recusa de ALOCAÇÃO, e não
                    # estava aqui: `The requested GPU duration (270s) is larger
                    # than the maximum allowed` saía como `erro`, e a tarefa de
                    # vídeo FALHAVA em vez de o backend cair no gerador local
                    # (medido em produção, 25/09/2026).
                    "maximum allowed", "requested gpu duration")


def e_erro_de_gpu(e: BaseException) -> bool:
    texto = f"{type(e).__name__}: {e}".lower()
    return any(s in texto for s in _SINAIS_DE_QUOTA)


# ---------------------------------------------------------------------------
# A GERAÇÃO — o corpo que o app.py decora com @spaces.GPU.
#
# É um GERADOR: cada yield intermediário vira um evento `generating` no SSE
# do Gradio (o progresso que a tela do app mostra), e o último yield é o
# resultado. O contrato do dict é o que o adapter do backend lê:
#   {"etapa","progress"}  no meio
#   {"resultado": {...}}  no fim
#   {"erro": ...}         falha do pedido (não adianta tentar em outro lugar)
#   {"erro_gpu": ...}     GPU/quota indisponível (degradar ou esperar)
# ---------------------------------------------------------------------------
def preparar(tipo: str) -> Optional[str]:
    """Baixa os pesos para o disco ANTES da chamada de GPU — download é
    trabalho de rede, e dentro do @spaces.GPU ele queimava o prazo inteiro:
    pedir prazo para caber o download estourou o teto do ZeroGPU ("requested
    GPU duration (360s) is larger than the maximum allowed"). Devolve None
    quando está pronto; string = o motivo da falha (repo inexistente, rede)."""
    modulo = _MODULOS.get(tipo)
    if not modulo:
        return f"tipo desconhecido: {tipo}"
    baixar = getattr(modulo, "baixar", None)
    if not baixar:
        return None
    try:
        baixar(lambda pct, etapa: None)
        return None
    except Exception as e:
        _log(tipo, "download_falhou", e)
        return f"não consegui baixar o modelo: {str(e)[:200]}"


def gerar(tipo: str, params_json: str) -> Iterator[Dict]:
    global _carregado
    try:
        params = json.loads(params_json or "{}")
        if not isinstance(params, dict):
            raise ValueError("params não é um objeto")
    except ValueError as e:
        yield {"erro": f"parâmetros ilegíveis: {str(e)[:80]}"}
        return

    modulo = _MODULOS.get(tipo)
    if not modulo:
        yield {"erro": f"tipo desconhecido: {tipo}"}
        return

    progresso_atual: Dict = {}

    def progresso(pct, etapa):
        # O gerador não pode "yield" de dentro do callback; o valor fica
        # anotado e sai no próximo passo do laço de fases abaixo.
        if pct is not None:
            progresso_atual["progress"] = max(0, min(100, int(pct)))
        if etapa:
            progresso_atual["etapa"] = etapa

    yield {"etapa": "Carregando o modelo…", "progress": None}
    tentativa_oom = 0
    while True:
        try:
            # A troca de residente vale aqui também — na mesma chamada podem
            # sobrar pesos da chamada anterior NA RAM (o processo do Space
            # sobrevive entre alocações de GPU; a GPU é que evapora). Nunca
            # dependa de estado NA GPU entre chamadas: os módulos usam
            # cpu_offload, que sobe os blocos por passada — é exatamente o
            # desenho que o ZeroGPU pede.
            if _carregado and _carregado != tipo:
                antigo = _MODULOS.get(_carregado)
                if antigo:
                    try:
                        antigo.descarregar()
                    except Exception:
                        pass
                mbase.limpar_vram()
                _carregado = None
            modulo.carregar(progresso)
            _carregado = tipo

            yield {"etapa": "Gerando…", "progress": progresso_atual.get("progress")}
            saida, mime = modulo.gerar(params, progresso)

            yield {"etapa": "Finalizando…", "progress": 97}
            resultado = armazenamento.como_url_ou_base64(saida, mime)
            resultado["mock"] = mbase.modo_mock()
            _log(tipo, "completed")
            yield {"resultado": resultado}
            return
        except BaseException as e:            # BaseException: o abort do
            if e_erro_de_gpu(e):               # spaces nem sempre é Exception
                _log(tipo, "gpu_indisponivel", e)
                yield {"erro_gpu": f"ZeroGPU indisponível agora: {str(e)[:120]}"}
                return
            if isinstance(e, ImportError):
                # Biblioteca que não está neste runtime (o ACE-Step saiu do
                # requirements por cravar transformers==4.50.0 contra o
                # >=4.51 do Z-Image). Não é falha DO PEDIDO: é este worker
                # não fazendo este tipo — erro_gpu, para o backend degradar
                # ao gerador local em vez de devolver erro à pessoa.
                _log(tipo, "sem_biblioteca", e)
                yield {"erro_gpu": f"este runtime não tem a biblioteca de "
                                   f"{tipo}: {str(e)[:120]}"}
                return
            if "out of memory" in str(e).lower():
                tentativa_oom += 1
                mbase.limpar_vram()
                menor = modulo.config_menor(params)
                if menor is not None and tentativa_oom <= 2:
                    params = menor
                    yield {"etapa": "Memória curta: tentando uma versão menor…",
                           "progress": None}
                    continue
                _log(tipo, "failed_oom", e)
                yield {"erro": "a GPU não teve memória para este pedido"}
                return
            if isinstance(e, (KeyboardInterrupt, SystemExit)):
                raise
            _log(tipo, "failed", e)
            yield {"erro": str(e)[:300]}
            return


def saude() -> Dict:
    """O retrato para o /api/health/gpu do backend. No ZeroGPU não há GPU
    residente para medir — CUDA só existe dentro de uma chamada decorada — e
    o retrato DIZ isso em vez de fingir um número."""
    return {"worker_status": "ready", "runtime": "zerogpu",
            "mock": mbase.modo_mock(),
            "cuda": None,   # None de propósito: "só se sabe dentro da chamada"
            "duracoes": {t: duracao_gpu(t)
                         for t in ("image", "speech", "music", "video")},
            "modelo_residente": _carregado}


def _log(tipo: str, status: str, erro: BaseException = None) -> None:
    # O mesmo log estruturado do worker CUDA — e igualmente SEM o prompt:
    # prompt é conteúdo da pessoa, e log de Space é ainda mais público que o
    # nosso (aparece no painel do Hugging Face).
    logger.info("geracao runtime=zerogpu tipo=%s status=%s%s", tipo, status,
                f" erro={type(erro).__name__}" if erro else "")