modelos/escolha.py

# -*- coding: utf-8 -*-
"""QUAL MÓDULO ATENDE CADA TIPO — num lugar só, para o worker e o Space.

POR QUE ISTO EXISTE. Havia duas tabelas com a mesma decisão: o `_MODULOS` do
`gerenciador.py` (worker CUDA) e o `_MODULOS` do `nucleo.py` (Space ZeroGPU).
Enquanto cada tipo tinha um módulo só, elas eram iguais por sorte. No dia em
que a imagem passou a ter DOIS módulos possíveis — Z-Image e Qwen-Image,
escolhidos pelo `IMAGE_MODEL` —, duas tabelas viraram duas respostas para a
mesma pergunta, e a que fica para trás é sempre a que ninguém olha. É o
defeito mais caro desta casa, e ele já cobrou a peneira do navegador, os dois
renderizadores de resumo do estudo e as duas regras de imagem no chat.

A REGRA DE SEMPRE VALE AQUI: a ORDEM é nossa, quem diz o que EXISTE é a
conta — ou, neste caso, a variável. O que está escrito abaixo são MARCAS DE
FAMÍLIA (`qwen-image`, `z-image`), nunca um id completo: um `Qwen-Image-2512`
publicado amanhã cai sozinho no módulo certo, sem ninguém tocar no arquivo.
Lista de ids escrita à mão já mordeu este projeto cinco vezes.

E O DESCONHECIDO TEM DONO. Um id que não bate com marca nenhuma vai para o
módulo PADRÃO do tipo, em vez de virar `None`: devolver nada aqui apagaria a
geração inteira por causa de um nome que ninguém previu — e o pior desfecho de
um palpite é uma imagem com a calibragem errada, que ainda é uma imagem.
"""
from __future__ import annotations

import os
from typing import Dict, List, Tuple

# Marca da família -> nome do módulo dentro de `modelos/`. A ordem importa:
# a primeira marca que aparecer no id ganha.
_FAMILIAS_IMAGEM: List[Tuple[str, str]] = [
    ("qwen-image", "qwen_image"),
    ("z-image",    "z_image"),
]
_FAMILIAS_VIDEO: List[Tuple[str, str]] = [
    ("ltx",  "ltx"),
    ("wan",  "wan22"),
]

# O padrão de cada tipo quando o id não bate com marca nenhuma.
_PADRAO = {"image": "qwen_image", "speech": "qwen3_tts", "voice": "qwen3_tts",
           "music": "ace_step", "video": "wan22"}

_VARIAVEL = {"image": "IMAGE_MODEL", "speech": "TTS_MODEL",
             "voice": "TTS_MODEL", "music": "MUSIC_MODEL",
             "video": "VIDEO_MODEL"}

_POR_TIPO: Dict[str, List[Tuple[str, str]]] = {
    "image": _FAMILIAS_IMAGEM,
    "video": _FAMILIAS_VIDEO,
}


def nome_do_modulo(tipo: str, modelo: str = "") -> str:
    """O nome do módulo que atende este tipo, pela FAMÍLIA do id configurado.

    `modelo` vazio faz a função ler a variável de ambiente do tipo — é como os
    dois consumidores a chamam. Passar o id à mão existe para o teste, que
    precisa medir a escolha sem mexer no ambiente do processo.
    """
    tipo = (tipo or "").strip().lower()
    id_modelo = (modelo or os.getenv(_VARIAVEL.get(tipo, ""), "")).strip().lower()
    for marca, modulo in _POR_TIPO.get(tipo, []):
        if marca in id_modelo:
            return modulo
    return _PADRAO.get(tipo, "")


def modulo_de(tipo: str):
    """O módulo em si, importado na hora.

    O import é TARDIO de propósito, e é a regra do `modelos/__init__.py`: o
    worker sobe e responde `/health` sem torch instalado, e o peso de um
    modelo só entra quando um job pede. Importar os dois módulos de imagem no
    topo obrigaria o processo a resolver as dependências dos DOIS para atender
    um pedido de voz.
    """
    nome = nome_do_modulo(tipo)
    if not nome:
        return None
    from importlib import import_module
    return import_module(f".{nome}", package=__package__)