modelos/z_image.py

# -*- coding: utf-8 -*-
"""
IMAGEM — Z-Image / Z-Image-Turbo (Tongyi-MAI), via diffusers.

Por que ele: é open-weight, entende prompt em duas línguas e a variante Turbo
resolve em ~8 passos — numa GPU de 24 GB isso é imagem em poucos segundos, que
é o que o Estúdio precisa. O id exato vem do ambiente (IMAGE_MODEL); este
arquivo não prende ninguém a uma versão.

O pipeline fica em bfloat16 e com `enable_model_cpu_offload` quando a VRAM é
curta: offload troca velocidade por caber, e caber ganha — imagem lenta ainda
é imagem; OOM é nada.
"""
from __future__ import annotations

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

from . import base

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

VRAM_MINIMA_GB = 8          # com offload; folgado ele usa mais e voa

_pipe = None
_modelo_carregado = ""


def _modelo() -> str:
    return os.getenv("IMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo").strip()


def baixar(progresso: base.Progresso) -> None:
    base.baixar_repo(_modelo(), progresso)


def carregar(progresso: base.Progresso) -> None:
    global _pipe, _modelo_carregado
    if base.modo_mock():
        _pipe = "mock"
        return
    if _pipe is not None and _modelo_carregado == _modelo():
        return
    progresso(None, "Carregando o modelo de imagem…")
    import torch
    from diffusers import DiffusionPipeline
    _pipe = DiffusionPipeline.from_pretrained(
        _modelo(), torch_dtype=torch.bfloat16,
        token=os.getenv("HF_TOKEN") or None)
    livre = base.vram_livre_gb()
    if livre is not None and livre < 20:
        # Pouca VRAM: o offload deixa os blocos na RAM e sobe um por vez.
        _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 1024)
    altura = int(params.get("height") or 1024)
    # 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 = int(params.get("steps") or 8)   # Turbo: 8 é o ponto doce

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

    saida = _pipe(prompt,
                  negative_prompt=params.get("negative_prompt") or None,
                  width=largura, height=altura,
                  num_inference_steps=passos,
                  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
    _pipe = None
    base.limpar_vram()


def config_menor(params: Dict) -> Optional[Dict]:
    """Depois de um OOM: menos pixels primeiro, depois menos passos. Abaixo
    de 512px e 4 passos não há mais gordura — devolve None e o job falha
    limpo dizendo o porquê."""
    p = dict(params)
    w = int(p.get("width") or 1024)
    if w > 512:
        p["width"] = max(512, w // 2)
        p["height"] = max(512, int(p.get("height") or 1024) // 2)
        return p
    passos = int(p.get("steps") or 8)
    if passos > 4:
        p["steps"] = 4
        return p
    return None


def _do_aspecto(ratio: str):
    tabela = {"1:1": (1024, 1024), "16:9": (1280, 720), "9:16": (720, 1280),
              "4:3": (1152, 864), "3:4": (864, 1152)}
    return tabela.get(str(ratio).strip(), (1024, 1024))