fn_registry/python/functions/ml/comfyui_pixelize_image.py

"""comfyui_pixelize_image — post-proceso pixel-perfect de una imagen (Fase 2 pixelart).

Convierte una imagen "pixelart borroso de IA" (o cualquier PNG) en pixelart de
verdad: downscale nearest-neighbor por factor (colapsa cada bloque borroso a un
pixel duro) -> cuantizacion a N colores o a una paleta fija (NES / Game Boy /
PICO-8) -> opcional re-upscale nearest conservando los pixeles duros.

Es la Fase 2 del pipeline pixelart (la Fase 1, generar con SDXL + SDXL_pixel-art
LoRA, vive en otra funcion). Determinista y CPU-only: el nucleo `_pixelize_pil`
es puro (PIL), no toca la GPU ni la red. La funcion publica es impura solo por la
lectura/escritura de disco (mismo patron que comfyui_build_grid).

Por que nearest y no lanczos/cubic: el downscale tiene que colapsar cada bloque
borroso a UN pixel; cualquier interpolacion suave re-difumina el grid. La
cuantizacion (Image.quantize) limita la paleta, que es lo que da la identidad
retro y elimina el ruido de cientos de colores de la difusion.
"""

import os


# Paletas retro fijas (hex sin '#'). Embebidas: cero red, deterministas.
# Fuentes: lospec.com (game-boy, pico-8) + paleta NES clasica reducida.
_BUILTIN_PALETTES = {
    "game-boy": ["0f380f", "306230", "8bac0f", "9bbc0f"],
    "pico-8": [
        "000000", "1d2b53", "7e2553", "008751", "ab5236", "5f574f",
        "c2c3c7", "fff1e8", "ff004d", "ffa300", "ffec27", "00e436",
        "29adff", "83769c", "ff77a8", "ffccaa",
    ],
    # NES: subconjunto representativo de 16 colores de la paleta 2C02.
    "nes": [
        "000000", "fcfcfc", "f8f8f8", "bcbcbc", "7c7c7c", "0000fc",
        "0078f8", "00b800", "b8f818", "f83800", "e40058", "f878f8",
        "fca044", "f8b800", "503000", "00a800",
    ],
}


def _hex_to_rgb(h: str) -> tuple:
    """'1a2b3c' o '#1a2b3c' -> (26, 43, 60)."""
    h = h.strip().lstrip("#")
    if len(h) != 6:
        raise ValueError(f"hex de color invalido: {h!r} (esperado rrggbb)")
    return tuple(int(h[i:i + 2], 16) for i in (0, 2, 4))


def _normalize_palette(palette):
    """palette: None | nombre builtin (str) | lista de hex -> list[(r,g,b)] | None."""
    if palette is None:
        return None
    if isinstance(palette, str):
        key = palette.strip().lower().replace("_", "-")
        if key not in _BUILTIN_PALETTES:
            raise ValueError(
                f"paleta builtin desconocida: {palette!r}. "
                f"Disponibles: {sorted(_BUILTIN_PALETTES)} o pasa una lista de hex."
            )
        hexes = _BUILTIN_PALETTES[key]
    else:
        hexes = list(palette)
        if not hexes:
            raise ValueError("lista de paleta vacia")
    return [_hex_to_rgb(h) for h in hexes]


def _pixelize_pil(img, downscale, colors, palette_rgb, dither, upscale_back):
    """Nucleo puro PIL: imagen RGB -> imagen RGB pixelizada.

    Args:
        img: PIL.Image de entrada.
        downscale: factor entero de reduccion nearest (>=1).
        colors: numero de colores objetivo si no hay paleta fija.
        palette_rgb: lista [(r,g,b), ...] o None (cuantizacion automatica).
        dither: aplica Floyd-Steinberg al cuantizar si True.
        upscale_back: re-escala nearest al tamano original si True.

    Returns:
        PIL.Image RGB pixelizada.
    """
    from PIL import Image

    img = img.convert("RGB")
    w, h = img.size
    # 1. downscale nearest -> grid real (colapsa bloques borrosos a 1 pixel).
    sw, sh = max(1, w // downscale), max(1, h // downscale)
    small = img.resize((sw, sh), Image.NEAREST)
    d = Image.Dither.FLOYDSTEINBERG if dither else Image.Dither.NONE
    # 2. cuantizar la paleta.
    if palette_rgb:
        pal_img = Image.new("P", (1, 1))
        flat = [c for rgb in palette_rgb for c in rgb][:768]
        # Rellena las 256 entradas repitiendo el ultimo color real (no ceros): asi
        # quantize no puede introducir un color extra (negro) por las entradas vacias.
        if flat:
            last = flat[-3:]
            flat += last * ((768 - len(flat)) // 3)
            flat += [0] * (768 - len(flat))
        pal_img.putpalette(flat)
        small = small.quantize(palette=pal_img, dither=d)
    else:
        n = max(2, min(256, int(colors)))
        small = small.quantize(colors=n, method=Image.Quantize.MEDIANCUT, dither=d)
    out = small.convert("RGB")
    # 3. opcional: re-upscale nearest para preview/entrega (pixeles duros).
    if upscale_back:
        out = out.resize((w, h), Image.NEAREST)
    return out


def comfyui_pixelize_image(
    src_path: str,
    dst_path: str,
    *,
    downscale: int = 8,
    colors: int = 16,
    palette=None,
    dither: bool = False,
    upscale_back: bool = True,
) -> dict:
    """Pixeliza una imagen y la guarda como PNG.

    Args:
        src_path: ruta de la imagen de entrada (PNG/JPG/...).
        dst_path: ruta del PNG de salida.
        downscale: factor entero de reduccion nearest-neighbor; cada bloque de
            downscale x downscale px colapsa a 1 pixel. 1 = solo cuantiza, sin
            colapsar el grid. keyword-only.
        colors: numero de colores objetivo (2..256) cuando palette es None;
            cuantizacion MEDIANCUT determinista. keyword-only.
        palette: None (cuantizacion automatica a `colors`), nombre de paleta
            fija builtin ("game-boy", "pico-8", "nes") o lista de hex
            ("#rrggbb"/"rrggbb"). Una paleta fija ignora `colors`. keyword-only.
        dither: aplica Floyd-Steinberg al cuantizar (por defecto off, pixelart
            limpio). keyword-only.
        upscale_back: re-escala nearest al tamano original (preview con pixeles
            duros). False deja la imagen pequena (sw x sh). keyword-only.

    Returns:
        dict con:
        - ok (bool): True si se pixelizo y guardo.
        - out_path (str): ruta del PNG generado.
        - size (list[int]): [w, h] de la imagen final.
        - n_colors_final (int): numero de colores distintos en el resultado.
        - error (str): mensaje de error; cadena vacia si todo OK.
    """
    out = {"ok": False, "out_path": "", "size": [0, 0], "n_colors_final": 0, "error": ""}

    try:
        from PIL import Image
    except ImportError:
        out["error"] = "PIL (Pillow) no esta instalado en este interprete"
        return out

    if not os.path.isfile(src_path):
        out["error"] = f"src_path no existe: {src_path!r}"
        return out
    if int(downscale) < 1:
        out["error"] = f"downscale debe ser >= 1, recibido {downscale!r}"
        return out

    try:
        palette_rgb = _normalize_palette(palette)
    except ValueError as exc:
        out["error"] = f"paleta invalida: {exc}"
        return out

    try:
        with Image.open(src_path) as src:
            result = _pixelize_pil(
                src, int(downscale), colors, palette_rgb, bool(dither), bool(upscale_back)
            )
    except OSError as exc:
        out["error"] = f"no se pudo leer/decodificar {src_path!r}: {exc}"
        return out

    try:
        dst_dir = os.path.dirname(os.path.abspath(dst_path))
        os.makedirs(dst_dir, exist_ok=True)
        result.save(dst_path)
    except OSError as exc:
        out["error"] = f"no se pudo escribir {dst_path!r}: {exc}"
        return out

    colors_found = result.getcolors(maxcolors=1 << 20)
    n_final = len(colors_found) if colors_found is not None else -1
    out.update(
        ok=True, out_path=dst_path, size=list(result.size), n_colors_final=n_final
    )
    return out


if __name__ == "__main__":
    import json
    import sys

    if len(sys.argv) < 3:
        print("uso: comfyui_pixelize_image.py <src> <dst> [downscale] [colors] [palette]",
              file=sys.stderr)
        sys.exit(2)
    src, dst = sys.argv[1], sys.argv[2]
    ds = int(sys.argv[3]) if len(sys.argv) > 3 else 8
    col = int(sys.argv[4]) if len(sys.argv) > 4 else 16
    pal = sys.argv[5] if len(sys.argv) > 5 else None
    print(json.dumps(comfyui_pixelize_image(src, dst, downscale=ds, colors=col, palette=pal),
                     indent=2))