fn_registry/python/functions/pipelines/comfyui_pixelart_real_oneshot.py

"""comfyui_pixelart_real_oneshot — prompt de texto -> sprite pixel-art REAL en disco.

Pipeline one-shot (issue 0087) que materializa el metodo ganador de la
investigacion (report 0215): la difusion NO sabe pintar pixel-perfect (su salida
tiene decenas de miles de colores y bordes con anti-aliasing — pixel-art FALSO),
asi que el pixel-art de verdad es siempre post-proceso en dos ejes: colapsar a un
grid duro y limitar la paleta. El metodo ganador combina:

  1. Generar a alta resolucion con el look pixel-art (SDXL Juggernaut + LoRA
     SDXL_pixel-art), via comfyui_build_pixelart_workflow.
  2. Downscale contrast-aware con PixelOE (pixeloe_downscale): elige el pixel mas
     representativo de cada zona y engrosa contornos -> silueta legible. Es lo que
     distingue un sprite reconocible de una mancha. Solo para sujetos con silueta
     (engine="pixeloe"); para tiles/texturas sin contorno, un downscale nearest
     simple basta (engine="nearest") y es mas barato.
  3. Cuantizacion dura con comfyui_pixelize_image (downscale=1): clava la paleta
     exacta (N colores libres MEDIANCUT, o paleta fija pico-8 / nes / game-boy)
     sobre el grid ya hecho -> 16 colores exactos + 100% grid duro.

Resultado del combo verificado por PIL: grid duro perfecto + paleta limitada +
outline nitido. Sweet-spot: 64px personajes/sprites, 32px iconos/objetos simples.

Compone funciones del registry, no reescribe su logica:
    comfyui_build_pixelart_workflow_py_ml  (workflow SDXL + LoRA pixel-art)
    comfyui_submit_workflow_py_ml          (POST /prompt)
    comfyui_wait_result_py_ml              (poll /history)
    comfyui_fetch_output_image_py_ml       (GET /view -> disco, imagen base)
    pixeloe_downscale_py_ml                (downscale contrast-aware, engine pixeloe)
    comfyui_pixelize_image_py_ml           (cuantizacion dura + nearest fallback)

Pipeline impuro: red (HTTP a ComfyUI) + escritura en disco. No-throw: cualquier
fallo se captura y se devuelve en el dict de estado (campo error).
"""

from __future__ import annotations

import os
import sys

# Importa las funciones del registry (mismo arbol python/functions).
_FUNCTIONS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _FUNCTIONS_ROOT not in sys.path:
    sys.path.insert(0, _FUNCTIONS_ROOT)

from ml.comfyui_build_pixelart_workflow import comfyui_build_pixelart_workflow
from ml.comfyui_fetch_output_image import comfyui_fetch_output_image
from ml.comfyui_pixelize_image import comfyui_pixelize_image
from ml.comfyui_submit_workflow import comfyui_submit_workflow
from ml.comfyui_wait_result import comfyui_wait_result
from ml.crop_to_content import crop_to_content
from ml.pixeloe_downscale import pixeloe_downscale

# Sufijo de encuadre: empuja al sujeto a llenar el frame para que tras el
# downscale conserve detalle por pixel (gotcha del report: un sujeto que ocupa el
# 25% del frame queda diminuto a 64px). Solo se anade si no esta ya presente.
_FRAME_HINT = "full body, centered, fills frame, no margins"


def _frame_subject(subject: str, fill_frame: bool) -> str:
    """Anade el hint de encuadre al subject si fill_frame y no esta ya."""
    if not fill_frame:
        return subject
    low = subject.lower()
    if "fills frame" in low or "full body" in low or "centered" in low:
        return subject
    return f"{subject}, {_FRAME_HINT}"


def comfyui_pixelart_real_oneshot(
    subject: str,
    *,
    size: int = 64,
    colors: int = 16,
    engine: str = "pixeloe",
    palette=None,
    server: str = "127.0.0.1:8188",
    dest_dir: str = "~/ComfyUI/output",
    seed: int = 0,
    negative: str | None = None,
    mode: str = "contrast",
    patch_size: int = 16,
    thickness: int = 2,
    fill_frame: bool = True,
    transparent: bool = True,
    autocrop: bool = True,
    crop_pad_ratio: float = 0.06,
    rembg_model: str = "u2net",
    upscale_preview: int = 512,
    keep_base: bool = True,
    comfy_python: str | None = None,
    wait_timeout: float = 300.0,
    filename_prefix: str = "pixelart_real",
    **gen_kwargs,
) -> dict:
    """Genera un sprite pixel-art REAL desde un prompt de texto, end-to-end.

    Args:
        subject: prompt positivo (lo que se quiere ver: "pixel art knight, full
            body, side view", etc.). No puede estar vacio.
        size: lado del grid final en pixeles (64 personajes/sprites, 32 iconos).
            keyword-only.
        colors: numero de colores de la paleta libre cuando palette es None
            (cuantizacion MEDIANCUT). keyword-only.
        engine: "pixeloe" (downscale contrast-aware, para sujetos con silueta:
            personajes/criaturas/iconos) o "nearest" (downscale nearest simple,
            mas barato, para tiles/texturas/fondos sin contorno). Si "pixeloe"
            falla o la lib no esta disponible, cae automaticamente a "nearest" y
            lo reporta en engine_used. keyword-only.
        palette: None (paleta libre a `colors`), nombre builtin ("pico-8", "nes",
            "game-boy") o lista de hex. Una paleta fija ignora `colors`.
            keyword-only.
        server: host:port del servidor ComfyUI (sin esquema). keyword-only.
        dest_dir: directorio donde guardar los PNG (se expande ~). keyword-only.
        seed: semilla del KSampler. keyword-only.
        negative: prompt negativo; None usa el default de build_pixelart
            (evita blur/gradientes/anti-alias). keyword-only.
        mode: modo de downscale de PixelOE ("contrast" SOTA, "k-centroid",
            "nearest", "center", "bicubic"); solo aplica con engine="pixeloe".
            keyword-only.
        patch_size: tamano de patch de PixelOE (default 16). keyword-only.
        thickness: grosor del outline expansion de PixelOE (default 2).
            keyword-only.
        fill_frame: si True, anade un hint de encuadre al subject para que el
            sujeto llene el frame (mejor detalle por pixel tras el downscale).
            keyword-only.
        transparent: si True (default) genera con fondo recortado (rembg en el
            workflow) y produce un sprite RGBA con transparencia real. Para
            tiles/texturas que NO quieren alpha, pasar transparent=False (el sprite
            sale RGB sobre fondo opaco). keyword-only.
        autocrop: si True (default) recorta la imagen base al bounding box de su
            contenido y la cuadra antes del downscale, para que el sujeto llene el
            frame (evita el sprite diminuto). Usa el alpha si transparent, o el color
            de fondo si no. keyword-only.
        crop_pad_ratio: margen relativo que deja el autocrop alrededor del sujeto
            (0.06 = 6% del lado). keyword-only.
        rembg_model: modelo Rembg para recortar el fondo ('u2net' general,
            'isnet-anime' para anime). Solo aplica si transparent. keyword-only.
        upscale_preview: si > 0, escribe ademas un PNG re-escalado nearest a
            ese lado (preview con pixeles duros, p.ej. 512). 0 lo desactiva.
            keyword-only.
        keep_base: si True conserva el PNG base de alta resolucion; si False lo
            borra tras pixelizar. keyword-only.
        comfy_python: ruta al interprete de ComfyUI (con la lib pixeloe); None
            autodetecta. keyword-only.
        wait_timeout: segundos maximos esperando al server. keyword-only.
        filename_prefix: prefijo de los archivos de salida. keyword-only.
        **gen_kwargs: params extra para comfyui_build_pixelart_workflow
            (width, height, ckpt_name, lora_strength, use_lcm, steps, cfg, ...).

    Returns:
        dict con:
        - ok (bool): True si se produjo el PNG final pixelizado.
        - out_path (str): ruta del PNG final size x size.
        - out_path_upscaled (str): ruta del preview re-escalado, o "" si off.
        - base_path (str): ruta del PNG base de alta resolucion (o "" si se borro).
        - size (int): lado real del PNG final.
        - colors_final (int): numero de colores distintos en el resultado (en la
          zona opaca si es RGBA).
        - engine_used (str): "pixeloe" o "nearest" (refleja el fallback).
        - has_alpha (bool): True si el PNG final es RGBA con transparencia.
        - autocrop_applied (bool): True si el autocrop recorto la imagen base.
        - prompt_id (str): id del trabajo en ComfyUI.
        - error (str): mensaje de error; vacio si OK.
    """
    out = {
        "ok": False, "out_path": "", "out_path_upscaled": "", "base_path": "",
        "size": int(size), "colors_final": 0, "engine_used": engine,
        "has_alpha": False, "autocrop_applied": False,
        "prompt_id": "", "error": "",
    }

    if not subject or not subject.strip():
        out["error"] = "subject vacio"
        return out
    if int(size) < 1:
        out["error"] = f"size debe ser >= 1, recibido {size!r}"
        return out
    if engine not in ("pixeloe", "nearest"):
        out["error"] = f"engine invalido: {engine!r} (usa 'pixeloe' o 'nearest')"
        return out

    dest = os.path.expanduser(dest_dir)
    try:
        os.makedirs(dest, exist_ok=True)
    except OSError as exc:
        out["error"] = f"no se pudo crear dest_dir {dest!r}: {exc}"
        return out

    # --- Fase 1: generar la imagen base de alta resolucion (look pixel-art) ---
    positive = _frame_subject(subject, fill_frame)
    try:
        if negative is None:
            workflow = comfyui_build_pixelart_workflow(
                positive, seed=seed, transparent=bool(transparent),
                rembg_model=rembg_model,
                filename_prefix=f"{filename_prefix}_base", **gen_kwargs,
            )
        else:
            workflow = comfyui_build_pixelart_workflow(
                positive, negative, seed=seed, transparent=bool(transparent),
                rembg_model=rembg_model,
                filename_prefix=f"{filename_prefix}_base", **gen_kwargs,
            )
    except (ValueError, TypeError) as exc:
        out["error"] = f"build workflow fallo: {exc}"
        return out

    try:
        sub = comfyui_submit_workflow(workflow, server=server)
        prompt_id = sub["prompt_id"]
        out["prompt_id"] = prompt_id
    except (RuntimeError, KeyError, OSError) as exc:
        out["error"] = f"submit fallo (server {server} responde?): {exc}"
        return out

    try:
        outputs = comfyui_wait_result(prompt_id, server=server, timeout=wait_timeout)
    except (TimeoutError, RuntimeError, OSError) as exc:
        out["error"] = f"wait fallo: {exc}"
        return out

    img = None
    for node_out in outputs.values():
        images = node_out.get("images") if isinstance(node_out, dict) else None
        if images:
            img = images[0]
            break
    if img is None:
        out["error"] = f"el workflow no produjo imagenes (outputs={list(outputs)})"
        return out

    fetched = comfyui_fetch_output_image(
        img["filename"], subfolder=img.get("subfolder", ""),
        type_=img.get("type", "output"), server=server, dest_dir=dest,
    )
    if not fetched.get("ok"):
        out["error"] = f"fetch de imagen base fallo: {fetched.get('error')}"
        return out
    base_path = fetched["path"]
    out["base_path"] = base_path

    # --- Fase 1b (opcional): autocrop al contenido + cuadrar (sujeto llena el frame). ---
    # La imagen sobre la que se hace el downscale: la recortada si autocrop, o la base.
    pre_ds_path = base_path
    crop_path = ""
    if autocrop:
        crop_path = os.path.join(dest, f"{filename_prefix}_{size}px_crop.png")
        try:
            from PIL import Image
            with Image.open(base_path) as base_im:
                src_im = base_im.convert("RGBA") if transparent else base_im.convert("RGB")
                before = src_im.size
                cropped = crop_to_content(
                    src_im, pad_ratio=float(crop_pad_ratio), square=True,
                )
            cropped.save(crop_path)
            pre_ds_path = crop_path
            out["autocrop_applied"] = cropped.size != before
        except (ImportError, OSError, ValueError) as exc:
            # Autocrop es best-effort: si falla, se sigue con la base sin recortar.
            crop_path = ""
            pre_ds_path = base_path
            if not out["error"]:
                out["error"] = f"autocrop fallo (no critico): {exc}"

    # --- Fase 2a: downscale a un grid `size` x `size` (mid). ---
    mid_path = os.path.join(dest, f"{filename_prefix}_{size}px_mid.png")
    engine_used = engine

    if engine == "pixeloe":
        ds = pixeloe_downscale(
            pre_ds_path, mid_path, mode=mode, target_size=int(size),
            patch_size=patch_size, thickness=thickness, no_upscale=True,
            comfy_python=comfy_python,
        )
        if not ds.get("ok"):
            # Fallback limpio: PixelOE no disponible / fallo -> nearest.
            engine_used = "nearest"
            out["error"] = (
                f"pixeloe fallo ({ds.get('error')}); fallback a nearest"
            )

    if engine_used == "nearest":
        # Downscale nearest simple a size x size (PIL en el venv del registry).
        # nearest preserva el alpha por canal: si transparent, conserva la silueta.
        try:
            from PIL import Image
            with Image.open(pre_ds_path) as src:
                target_mode = "RGBA" if transparent else "RGB"
                small = src.convert(target_mode).resize(
                    (int(size), int(size)), Image.NEAREST
                )
                small.save(mid_path)
        except (ImportError, OSError) as exc:
            out["error"] = f"downscale nearest fallo: {exc}"
            return out

    if not os.path.isfile(mid_path):
        out["error"] = "no se genero la imagen intermedia (mid)"
        return out

    # --- Fase 2a-bis: recombinar alpha tras pixeloe (PixelOE trabaja en RGB). ---
    # El nucleo de PixelOE convierte a RGB: el grid `mid` sale sin transparencia. Se
    # downscalea el alpha de la imagen pre-downscale por separado (nearest al mismo
    # size) y se reaplica al grid para no perder el recorte ni la transparencia.
    if transparent and engine_used == "pixeloe":
        try:
            from PIL import Image
            with Image.open(pre_ds_path) as src_im:
                alpha = src_im.convert("RGBA").getchannel("A").resize(
                    (int(size), int(size)), Image.NEAREST
                )
            with Image.open(mid_path) as mid_im:
                mid_rgba = mid_im.convert("RGBA")
                mid_rgba.putalpha(alpha)
                mid_rgba.save(mid_path)
        except (ImportError, OSError) as exc:
            if not out["error"]:
                out["error"] = f"recombinacion de alpha fallo (no critico): {exc}"

    # --- Fase 2b: cuantizacion dura (paleta exacta) sobre el grid ya hecho. ---
    final_tag = palette if isinstance(palette, str) else f"q{colors}"
    final_path = os.path.join(
        dest, f"{filename_prefix}_{size}px_{engine_used}_{final_tag}.png"
    )
    quant = comfyui_pixelize_image(
        mid_path, final_path, downscale=1, colors=int(colors),
        palette=palette, upscale_back=False, keep_alpha=bool(transparent),
    )
    if not quant.get("ok"):
        out["error"] = f"cuantizacion fallo: {quant.get('error')}"
        return out

    out["out_path"] = final_path
    out["size"] = quant["size"][0] if quant.get("size") else int(size)
    out["colors_final"] = quant.get("n_colors_final", 0)
    out["has_alpha"] = bool(quant.get("has_alpha", False))
    out["engine_used"] = engine_used

    # --- Fase 3 (opcional): preview re-escalado nearest a pixeles duros. ---
    if int(upscale_preview) > 0:
        up_path = os.path.join(
            dest, f"{filename_prefix}_{size}px_{engine_used}_{final_tag}_up.png"
        )
        try:
            from PIL import Image
            with Image.open(final_path) as fin:
                prev_mode = "RGBA" if transparent else "RGB"
                up = fin.convert(prev_mode).resize(
                    (int(upscale_preview), int(upscale_preview)), Image.NEAREST
                )
                up.save(up_path)
            out["out_path_upscaled"] = up_path
        except (ImportError, OSError) as exc:
            # El preview es opcional: no invalida el resultado.
            out["out_path_upscaled"] = ""
            if not out["error"]:
                out["error"] = f"preview upscale fallo (no critico): {exc}"

    # Limpieza de intermedios (mid + crop temporal).
    for tmp in (mid_path, crop_path):
        if tmp:
            try:
                os.remove(tmp)
            except OSError:
                pass
    if not keep_base:
        try:
            os.remove(base_path)
            out["base_path"] = ""
        except OSError:
            pass

    out["ok"] = True
    return out


if __name__ == "__main__":
    import json

    res = comfyui_pixelart_real_oneshot(
        "pixel art knight, full body, centered, game sprite",
        size=64, colors=16, engine="pixeloe", seed=42,
        transparent=True, autocrop=True,
        dest_dir="/tmp/comfy_pixelart_real",
    )
    print(json.dumps(res, indent=2))