# -*- coding: utf-8 -*-
"""
IMAGEM — Qwen-Image (Qwen/Qwen-Image), via diffusers.

Escolha do dono. Conferido no Hub em 2026-09-15: `Qwen/Qwen-Image` existe,
é `QwenImagePipeline`, Apache-2.0, 316 mil downloads — e ele é GRANDE: 57,7 GB
no total (transformer 41 GB em 9 partes, text-encoder 16,6 GB, VAE 0,25 GB).
Esse número não é curiosidade, é o que decide tudo o que está escrito abaixo.

── POR QUE ELE NÃO PODE ENTRAR NO z_image.py ───────────────────────────────
Trocar só o `IMAGE_MODEL` parecia bastar — o `DiffusionPipeline` resolve a
classe pelo `model_index.json` e carregaria o Qwen-Image sem reclamar. O que
não vem junto é a CALIBRAGEM, e ela é o modelo:

  - o z_image pede 8 passos, que é o ponto doce do Z-Image-**Turbo**. O
    Qwen-Image base é um modelo de 50 passos; com 8 ele devolve borrão.
  - o z_image não passa `true_cfg_scale`. O Qwen-Image usa CFG de verdade
    (4.0 é o do cartão do modelo); sem ele o prompt é seguido pela metade.

É a mesma família do defeito que este arquivo existe para não repetir: dois
modelos no mesmo módulo, com os números de um só, é o modelo novo entregando
o pior de si e ninguém sabendo dizer por quê.

── PASSOS E CFG ANDAM JUNTOS, SEMPRE ───────────────────────────────────────
O par é a coisa mais fácil de errar aqui, então ele mora num lugar só
(`_receita`). Com o LoRA Lightning o modelo é DESTILADO: ele quer poucos
passos e CFG DESLIGADO (1.0) — pôr 4.0 em cima de um destilado queima a
imagem. Sem o LoRA é o contrário. Mexer num sem o outro é o defeito.

── O LIGHTNING É O PADRÃO, E ISSO É CONTA, NÃO GOSTO ───────────────────────
`lightx2v/Qwen-Image-Lightning` (0,85 GB, bf16, 8 passos V2.0) é um LoRA que
roda EM CIMA destes mesmos pesos. Ele é o padrão porque 50 passos de um
modelo de 20B não cabem no prazo de GPU que esta casa tem: o ZeroGPU gratuito
recusou 240 s declarados e 180 é o maior valor provado (ver CLAUDE.md), e
esse prazo ainda precisa cobrir o carregamento do disco. Um modelo que estoura
o prazo não é um modelo mais bonito — é uma geração que falha.

`IMAGE_LIGHTNING=0` desliga e devolve a receita de 50 passos, para o dia em
que houver GPU dedicada. O LoRA é reserva de si mesmo: se ele não baixar, a
receita cai para a completa em vez de gerar com 8 passos sem destilação, que
seria o borrão do primeiro parágrafo por outro caminho.

── O QUE NÃO FOI VALIDADO, E É HONESTO DIZER ───────────────────────────────
A inferência NÃO rodou: não há GPU nem pesos neste ambiente. O que está
provado é o que dá para provar aqui — que o id existe, que a classe do
pipeline é a esperada, e a escolha de módulo por família (t_modelo_imagem.py).
O que conferir no primeiro deploy com GPU está nomeado em `docs/MODELS.md`.
"""
from __future__ import annotations

import io
import logging
import os
from typing import Dict, Optional, Tuple

from . import base

logger = logging.getLogger("worker.qwen_image")

# Com offload. O modelo inteiro em VRAM quer bem mais; o offload troca
# velocidade por caber, e caber ganha — imagem lenta ainda é imagem.
VRAM_MINIMA_GB = 12

# O LoRA de destilação. Repositório e ARQUIVO, porque o repositório tem doze
# arquivos e a escolha entre 4 e 8 passos é dele: 8 passos V2.0 é o meio-termo
# medido pelos autores entre velocidade e fidelidade, e o bf16 (0,85 GB) é
# metade do download do fp32 sem diferença que se veja na tela.
_LORA_REPO = "lightx2v/Qwen-Image-Lightning"
_LORA_ARQUIVO = "Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors"

_pipe = None
_modelo_carregado = ""
_lora_vivo = False          # o LoRA ENTROU? Não é o mesmo que "foi pedido".


def _modelo() -> str:
    """O PADRÃO TEM DE SER O QUE CABE, e isto é conserto de defeito.

    Ele era `Qwen/Qwen-Image` — 57,7 GB, num disco de Space de 50. O backend
    manda o id escolhido dentro do pedido, mas quem lê aqui é a VARIÁVEL de
    ambiente do processo: sem `IMAGE_MODEL` configurado no Space, o padrão
    escrito nesta linha é quem manda, e ele mandava baixar 57,7 GB. O disco
    enche, a geração falha lá dentro, e o app degrada para o gerador de
    reserva sem erro em lugar nenhum — o sintoma que o dono já relatou duas
    vezes ("continua indo pro gerador de sempre").

    Padrão que não cabe é o mesmo que padrão que não existe. Este é o FP8 de
    30,1 GB (ver a tabela no `backend/geracao/registro.py`), o mesmo id que o
    registro usa: dois padrões diferentes para a mesma escolha seriam dois
    donos, e o que fica para trás é sempre o que ninguém abre — e o Space é
    justamente esse lado.
    """
    return os.getenv("IMAGE_MODEL", "Ilus-AI/Qwen-Image-2512-FP8").strip()


def _quer_lightning() -> bool:
    return (os.getenv("IMAGE_LIGHTNING", "1").strip() or "1") != "0"


def _receita(params: Dict) -> Tuple[int, float]:
    """(passos, true_cfg_scale) — O PAR, num lugar só.

    Quem manda é o LoRA que de fato ENTROU (`_lora_vivo`), nunca o que foi
    pedido: o download do LoRA pode falhar, e nesse caso gerar com 8 passos
    sem destilação devolve borrão. Errar para a receita completa custa tempo;
    errar para a curta custa a imagem.
    """
    pedido = params.get("steps")
    if _lora_vivo:
        passos = int(pedido or 8)
        # CFG DESLIGADO: o destilado já traz a orientação nos pesos, e somar
        # CFG por cima satura a imagem. 1.0 é "sem CFG" no Qwen-Image.
        return max(4, min(passos, 12)), 1.0
    # SEM O LORA, OS 50 PASSOS DO CARTÃO NÃO CABEM NO PRAZO — e esta era a
    # metade que faltava do parágrafo do Lightning lá em cima. Ele diz que a
    # falha do LoRA "cai para a receita completa em vez de gerar borrão", e
    # isso está certo sobre a QUALIDADE e errado sobre o DESFECHO: num teto
    # duro de 180 s (ZeroGPU gratuito), 50 passos de um modelo de 20B é a
    # chamada abortada no meio, com a quota do dia gasta e NENHUMA imagem —
    # e o sintoma não aponta para cá, é a geração caindo no gerador de
    # reserva como se a GPU estivesse dormindo.
    #
    # Então o teto do padrão é o que cabe. 24 passos sem destilação é uma
    # imagem pior que 50; é infinitamente melhor que imagem nenhuma — a mesma
    # conta do 480p do `wan22.py`. `GPU_DEDICADA=1` devolve os 50 do cartão,
    # e é a MESMA variável que o vídeo lê, porque o fato é um só.
    # PADRÃO e TETO são dois números, e não um: com GPU de verdade o padrão é
    # o do cartão (50) e o teto dá folga para quem pedir mais (60). Sem ela os
    # dois colapsam no que cabe — não adianta aceitar um pedido de 50 passos
    # que vai ser abortado no meio.
    padrao, teto = (50, 60) if base.gpu_dedicada() else (24, 24)
    return max(20, min(int(pedido or padrao), teto)), 4.0


def baixar(progresso: base.Progresso) -> None:
    base.baixar_repo(_modelo(), progresso)
    if _quer_lightning():
        # A FALHA AQUI NÃO DERRUBA NADA. O LoRA é aceleração, não o modelo;
        # sem ele a `_receita` volta para os 50 passos e a imagem sai igual,
        # só mais devagar. Derrubar o job por causa do acelerador seria
        # trocar "mais lento" por "nenhuma imagem".
        try:
            # SÓ O ARQUIVO, e isto não é detalhe: o repositório do Lightning
            # tem 65,5 GB em doze variantes, e o que se usa aqui pesa 0,85.
            base.baixar_repo(_LORA_REPO, progresso, arquivos=[_LORA_ARQUIVO])
        except Exception as e:
            logger.warning("LoRA Lightning não baixou (%s) — segue sem ele", e)


def carregar(progresso: base.Progresso) -> None:
    global _pipe, _modelo_carregado, _lora_vivo
    if base.modo_mock():
        _pipe = "mock"
        _lora_vivo = _quer_lightning()
        return
    if _pipe is not None and _modelo_carregado == _modelo():
        return
    progresso(None, "Carregando o modelo de imagem…")
    import torch
    from diffusers import DiffusionPipeline
    # `DiffusionPipeline` e não `QwenImagePipeline` importado por nome: a
    # classe sai do `model_index.json` do próprio repositório, então um
    # `Qwen-Image-2512` com pipeline própria carrega sem tocar neste arquivo.
    # Importar a classe por nome quebraria com um ImportError num diffusers
    # mais novo que a renomeasse — e esta casa já pagou por id e classe
    # escritos à mão.
    _pipe = DiffusionPipeline.from_pretrained(
        _modelo(), torch_dtype=torch.bfloat16,
        token=os.getenv("HF_TOKEN") or None)

    _lora_vivo = False
    if _quer_lightning():
        try:
            progresso(None, "Carregando o acelerador…")
            _pipe.load_lora_weights(_LORA_REPO, weight_name=_LORA_ARQUIVO)
            _lora_vivo = True
        except Exception as e:
            logger.warning("LoRA Lightning não entrou (%s) — receita completa", e)

    livre = base.vram_livre_gb()
    if livre is not None and livre < 36:
        # O CORTE ACOMPANHA O MODELO, senão ele vira número solto. Eram 60,
        # que era a leitura certa para os 57,7 GB do Qwen-Image cheio; o
        # padrão de hoje é o fp8 de 30,1 GB, e manter 60 ligaria o offload em
        # QUALQUER placa — inclusive naquelas em que os pesos cabem inteiros.
        # O offload sobe blocos da RAM a cada passo, e num prazo de 180 s isso
        # é a diferença entre terminar e ser abortado (a mesma conta do
        # `wan22.py`). 36 = os 30,1 dos pesos mais folga para ativações.
        _pipe.enable_model_cpu_offload()
    else:
        _pipe.to(base.dispositivo())
    _modelo_carregado = _modelo()


def gerar(params: Dict, progresso: base.Progresso) -> base.Saida:
    if base.modo_mock():
        return base.PNG_1X1, "image/png"

    prompt = (params.get("prompt") or "").strip()
    estilo = (params.get("style") or "").strip()
    if estilo:
        prompt = f"{prompt}, {estilo} style"
    largura = int(params.get("width") or 1328)
    altura = int(params.get("height") or 1328)
    # A tela manda aspect_ratio; quem manda width/height explícito ganha.
    if not params.get("width") and params.get("aspect_ratio"):
        largura, altura = _do_aspecto(params["aspect_ratio"])
    passos, cfg = _receita(params)

    def a_cada_passo(pipe, passo, t, kw):
        progresso(int(100 * (passo + 1) / max(1, passos)), "Gerando…")
        return kw

    saida = _pipe(prompt=prompt,
                  # O Qwen-Image quer uma STRING no negativo, e não None: o
                  # pipeline dele monta o par de embeddings sempre, e " " é o
                  # vazio que o cartão do modelo usa.
                  negative_prompt=(params.get("negative_prompt") or " "),
                  width=largura, height=altura,
                  num_inference_steps=passos,
                  true_cfg_scale=cfg,
                  generator=base.semente(params),
                  callback_on_step_end=a_cada_passo)
    img = saida.images[0]
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue(), "image/png"


def descarregar() -> None:
    global _pipe, _lora_vivo
    _pipe = None
    _lora_vivo = False
    base.limpar_vram()


def config_menor(params: Dict) -> Optional[Dict]:
    """Depois de um OOM: menos pixels primeiro, depois menos passos.

    O piso dos passos depende da receita e por isso NÃO é um número escrito
    aqui: com o Lightning, 4 é o mínimo que ainda desenha; sem ele, cortar
    para 4 devolveria borrão — que é uma falha pior que a falha limpa, porque
    ela parece ter dado certo.
    """
    p = dict(params)
    w = int(p.get("width") or 1328)
    if w > 640:
        p["width"] = max(640, w // 2)
        p["height"] = max(640, int(p.get("height") or 1328) // 2)
        return p
    passos, _ = _receita(p)
    piso = 4 if _lora_vivo else 20
    if passos > piso:
        p["steps"] = piso
        return p
    return None


def _do_aspecto(ratio: str):
    """As medidas OFICIAIS do cartão do Qwen-Image, e não as do Z-Image.

    Elas não são 1024-e-múltiplos: o modelo foi treinado nestas, e sair delas
    custa composição — é o motivo de esta tabela existir em vez de uma conta
    de proporção.
    """
    tabela = {"1:1": (1328, 1328), "16:9": (1664, 928), "9:16": (928, 1664),
              "4:3": (1472, 1140), "3:4": (1140, 1472)}
    return tabela.get(str(ratio).strip(), (1328, 1328))