104 lines
3.3 KiB
Python
104 lines
3.3 KiB
Python
"""sdcpp_python_generate — genera una imagen con stable-diffusion-cpp-python a partir de GenerationConfig."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import os
|
|
import time
|
|
from typing import Any
|
|
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
|
|
from generation_config import GenerationConfig
|
|
from image_gen_result import ImageGenResult
|
|
|
|
# Mapa de sampler del registry al nombre que espera stable-diffusion-cpp-python
|
|
_SAMPLER_MAP: dict[str, str] = {
|
|
"euler": "euler",
|
|
"euler_a": "euler_a",
|
|
"dpm++2m": "dpmpp2m",
|
|
"dpm++2m_v2": "dpmpp2mv2",
|
|
"heun": "heun",
|
|
"dpm2": "dpm2",
|
|
"lcm": "lcm",
|
|
}
|
|
|
|
|
|
def sdcpp_python_generate(sd: Any, cfg: GenerationConfig) -> ImageGenResult:
|
|
"""Genera una imagen con un objeto StableDiffusion usando GenerationConfig como contrato.
|
|
|
|
Mapea los campos del GenerationConfig canonico a los parametros de
|
|
StableDiffusion.generate_image(). Mide la duracion total de la llamada.
|
|
Retorna un ImageGenResult con la primera imagen del batch, metadata del backend
|
|
y duracion en milisegundos. VRAM no se mide (None).
|
|
|
|
Args:
|
|
sd: Instancia StableDiffusion cargada via sdcpp_python_load.
|
|
cfg: Contrato de parametros de generacion. Todos los campos son leidos.
|
|
cfg.sampler debe ser uno de los valores del SamplerName del registry.
|
|
|
|
Returns:
|
|
ImageGenResult con image=PIL.Image, meta con backend/modelo/sampler/seed/wtype,
|
|
duration_ms medido via time.perf_counter(), vram_peak_mb=None.
|
|
|
|
Raises:
|
|
KeyError: Si cfg.sampler no tiene correspondencia en _SAMPLER_MAP.
|
|
ImportError: Si stable_diffusion_cpp no esta instalado.
|
|
RuntimeError: Si generate_image retorna lista vacia o None.
|
|
"""
|
|
try:
|
|
from stable_diffusion_cpp import StableDiffusion # noqa: F401 — verifica disponibilidad
|
|
except ImportError as exc:
|
|
raise ImportError(
|
|
"sdcpp_python_generate requiere stable-diffusion-cpp-python. "
|
|
"Instalar con: pip install stable-diffusion-cpp-python"
|
|
) from exc
|
|
|
|
sample_method = _SAMPLER_MAP.get(cfg.sampler)
|
|
if sample_method is None:
|
|
raise KeyError(
|
|
f"Sampler '{cfg.sampler}' no tiene correspondencia en sdcpp_python_generate. "
|
|
f"Valores soportados: {list(_SAMPLER_MAP.keys())}"
|
|
)
|
|
|
|
# wtype del objeto sd (para metadata)
|
|
wtype = getattr(sd, "wtype", "unknown")
|
|
|
|
t0 = time.perf_counter()
|
|
|
|
images = sd.generate_image(
|
|
prompt=cfg.prompt,
|
|
negative_prompt=cfg.negative_prompt or "",
|
|
cfg_scale=cfg.cfg_scale,
|
|
sample_method=sample_method,
|
|
sample_steps=cfg.steps,
|
|
seed=cfg.seed,
|
|
width=cfg.width,
|
|
height=cfg.height,
|
|
clip_skip=cfg.clip_skip if cfg.clip_skip is not None else -1,
|
|
batch_count=1,
|
|
)
|
|
|
|
duration_ms = int((time.perf_counter() - t0) * 1000)
|
|
|
|
if not images:
|
|
raise RuntimeError(
|
|
"sdcpp_python_generate: generate_image retorno lista vacia o None."
|
|
)
|
|
|
|
meta: dict[str, Any] = {
|
|
"backend": "sdcpp_python",
|
|
"model": cfg.model.name,
|
|
"sampler": cfg.sampler,
|
|
"actual_steps": cfg.steps,
|
|
"seed": cfg.seed,
|
|
"wtype": wtype,
|
|
}
|
|
|
|
return ImageGenResult(
|
|
image=images[0],
|
|
meta=meta,
|
|
duration_ms=duration_ms,
|
|
vram_peak_mb=None,
|
|
)
|