Spaces:
Running on Zero
Running on Zero
File size: 3,840 Bytes
c7f05a9 550d35b c7f05a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | # -*- 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))
|