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 _img_has_alpha(img) -> bool:
    """True si la imagen lleva transparencia (RGBA, LA o P con transparency)."""
    return img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info)


def _fill_transparent_with_mode(small_rgb, small_alpha, threshold):
    """Rellena los pixeles transparentes con el color opaco mas frecuente (moda).

    Asi el fondo transparente NO aporta colores nuevos a la cuantizacion: las zonas
    con alpha <= threshold toman un color que ya esta en el sujeto (y por tanto en la
    paleta resultante), sin gastar entradas de la paleta en el color de fondo. El
    color real de esas zonas es irrelevante para la salida porque luego reciben
    alpha 0. Si no hay numpy, cae a no rellenar (degradacion limpia).

    Args:
        small_rgb: PIL.Image RGB ya reducida.
        small_alpha: PIL.Image 'L' del alpha ya reducido (mismo tamano).
        threshold: umbral de alpha (0..255); <= threshold = transparente.

    Returns:
        PIL.Image RGB con el fondo transparente relleno con la moda del sujeto.
    """
    from PIL import Image

    rgb = small_rgb.convert("RGB")
    mask = small_alpha.point(lambda p: 255 if p > threshold else 0).convert("L")
    try:
        import numpy as np
    except ImportError:
        return rgb

    arr = np.asarray(rgb).reshape(-1, 3)
    opaque = np.asarray(mask).reshape(-1) > 0
    if not opaque.any():
        return rgb  # nada opaco: caso degenerado, deja igual
    op_pixels = arr[opaque]
    colors, counts = np.unique(op_pixels, axis=0, return_counts=True)
    fill = tuple(int(x) for x in colors[counts.argmax()])
    bg = Image.new("RGB", rgb.size, fill)
    bg.paste(rgb, (0, 0), mask)  # rgb donde mask=255, fill (moda) donde mask=0
    return bg


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

    Si la imagen de entrada tiene canal alpha y keep_alpha es True, la cuantizacion
    de color se hace SOLO sobre el RGB (con el fondo transparente relleno con la moda
    del sujeto para que no entre en la paleta) y el alpha se downscalea nearest por
    separado y se binariza por `alpha_threshold`, recombinando a RGBA. Asi se
    preserva la transparencia sin que las zonas transparentes contaminen la paleta.
    Para imagenes sin alpha (o keep_alpha False) el comportamiento RGB es identico al
    de antes.

    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.
        keep_alpha: si True y la imagen tiene alpha, preserva la transparencia.
        alpha_threshold: umbral (0..255) para binarizar el alpha (opaco/transparente).

    Returns:
        PIL.Image pixelizada: RGB, o RGBA si se preservo la transparencia.
    """
    from PIL import Image

    has_alpha = bool(keep_alpha) and _img_has_alpha(img)
    if has_alpha:
        rgba = img.convert("RGBA")
        alpha_full = rgba.getchannel("A")
        rgb = rgba.convert("RGB")
    else:
        rgb = img.convert("RGB")
        alpha_full = None

    w, h = rgb.size
    # 1. downscale nearest -> grid real (colapsa bloques borrosos a 1 pixel).
    sw, sh = max(1, w // downscale), max(1, h // downscale)
    small = rgb.resize((sw, sh), Image.NEAREST)
    small_alpha = (
        alpha_full.resize((sw, sh), Image.NEAREST) if alpha_full is not None else None
    )
    # 1b. con alpha: el fondo transparente no debe entrar en la paleta.
    if small_alpha is not None:
        small = _fill_transparent_with_mode(small, small_alpha, int(alpha_threshold))

    d = Image.Dither.FLOYDSTEINBERG if dither else Image.Dither.NONE
    # 2. cuantizar la paleta (siempre sobre RGB).
    if palette_rgb:
        pal_img = Image.new("P", (1, 1))
        flat = [c for rgb_c in palette_rgb for c in rgb_c][: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")

    # 2b. recombinar el alpha (binarizado) -> RGBA con transparencia dura.
    if small_alpha is not None:
        out = out.convert("RGBA")
        hard_alpha = small_alpha.point(lambda p: 255 if p > int(alpha_threshold) else 0)
        out.putalpha(hard_alpha)

    # 3. opcional: re-upscale nearest para preview/entrega (pixeles duros).
    if upscale_back:
        out = out.resize((w, h), Image.NEAREST)
    return out


def _count_colors(result) -> int:
    """Numero de colores RGB distintos en el resultado.

    Para salida RGBA cuenta solo los colores de la zona opaca (alpha > 0), que es lo
    que define el sprite; el transparente no es un "color" del pixel-art. Para RGB
    cuenta todos los colores. Devuelve -1 si no se pudo contar.
    """
    if result.mode == "RGBA":
        try:
            import numpy as np
        except ImportError:
            colors_found = result.convert("RGB").getcolors(maxcolors=1 << 20)
            return len(colors_found) if colors_found is not None else -1
        arr = np.asarray(result)
        opaque = arr[..., 3] > 0
        rgb_op = arr[..., :3][opaque]
        if rgb_op.size == 0:
            return 0
        return int(len(np.unique(rgb_op.reshape(-1, 3), axis=0)))
    colors_found = result.getcolors(maxcolors=1 << 20)
    return len(colors_found) if colors_found is not None else -1


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,
    keep_alpha: bool = True,
    alpha_threshold: int = 128,
) -> 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.
        keep_alpha: si True (default) y la imagen de entrada tiene canal alpha,
            preserva la transparencia: cuantiza solo el RGB y downscalea/binariza el
            alpha por separado, devolviendo PNG RGBA. Las zonas transparentes no
            entran en la paleta de color. Si la imagen no tiene alpha, no tiene
            efecto (sale RGB igual que antes). keyword-only.
        alpha_threshold: umbral (0..255) para binarizar el alpha en opaco (255) o
            transparente (0). Solo aplica cuando se preserva el alpha. 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 RGB distintos en el resultado
          (en la zona opaca si la salida es RGBA).
        - has_alpha (bool): True si la salida es RGBA con transparencia preservada.
        - error (str): mensaje de error; cadena vacia si todo OK.
    """
    out = {
        "ok": False, "out_path": "", "size": [0, 0], "n_colors_final": 0,
        "has_alpha": False, "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), bool(keep_alpha), int(alpha_threshold),
            )
    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

    n_final = _count_colors(result)
    out.update(
        ok=True, out_path=dst_path, size=list(result.size), n_colors_final=n_final,
        has_alpha=(result.mode == "RGBA"),
    )
    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))