fn_registry/python/functions/ml/pixeloe_downscale.py

"""pixeloe_downscale — downscale contrast-aware a un grid de pixel-art (etapa SOTA).

Colapsa una ilustracion a un grid de pixel-art pequeno (p.ej. 64x64) usando la
libreria `pixeloe` de Kohaku (Contrast-Aware Outline Expansion), el metodo SOTA
para preservar contornos/silueta al reducir. Es la etapa de *downscale* del
metodo ganador de pixel-art (ver report 0215): NO cuantiza la paleta — esa dureza
de color la aplica despues otra funcion (`comfyui_pixelize_image`).

Gotcha de entorno (resuelto con un "bridge" de interprete): la lib `pixeloe` solo
esta instalada en el venv de ComfyUI (`~/ComfyUI/.venv`), no en el venv del
registry, y su modulo vive en `pixeloe.legacy.pixelize` (la ruta vieja
`pixeloe.pixelize` ya no existe). Por eso la funcion:

  1. Intenta `import pixeloe` en el interprete actual y ejecuta el nucleo directo.
  2. Si falta (`ModuleNotFoundError`), re-ejecuta este mismo archivo como subprocess
     (`python pixeloe_downscale.py --bridge <json>`) con un interprete que SI la
     tenga, parseando la unica linea JSON que ese hijo imprime a stdout.
  3. Si no hay ningun interprete con pixeloe, devuelve ok=False (sin excepcion);
     el pipeline que la llama hara fallback.

La funcion es no-throw: cualquier error se captura y viaja en el campo `error`.
Determinista; impura solo por la lectura/escritura de disco y el subprocess.
"""

import os


def _resolve_comfy_python(comfy_python):
    """Devuelve el primer interprete candidato que exista como archivo, o None.

    Orden: arg comfy_python -> env COMFY_PYTHON -> ~/ComfyUI/.venv/bin/python3.
    """
    candidates = []
    if comfy_python:
        candidates.append(comfy_python)
    env = os.environ.get("COMFY_PYTHON")
    if env:
        candidates.append(env)
    candidates.append(os.path.expanduser("~/ComfyUI/.venv/bin/python3"))
    for c in candidates:
        if c and os.path.isfile(c):
            return c
    return None


def _run_core(src_path, dst_path, mode, target_size, patch_size, thickness,
              color_matching, no_upscale):
    """Nucleo no-throw: requiere `pixeloe` importable EN ESTE interprete.

    Lee src como RGB uint8 (numpy + PIL, sin cv2), llama
    `pixeloe.legacy.pixelize.pixelize` y guarda el resultado como PNG. Devuelve el
    dict de resultado. NO lanza excepciones: las captura en `error`.
    """
    out = {
        "ok": False,
        "out_path": "",
        "size": [0, 0],
        "mode": mode,
        "target_size": int(target_size),
        "via": "inproc",
        "error": "",
    }

    try:
        import numpy as np
        from PIL import Image
    except Exception as exc:  # noqa: BLE001 - degradacion limpia, no relanzar
        out["error"] = f"numpy/PIL no disponible en este interprete: {exc}"
        return out

    if not os.path.isfile(src_path):
        out["error"] = f"src_path no existe: {src_path!r}"
        return out

    try:
        from pixeloe.legacy.pixelize import pixelize
    except Exception as exc:  # noqa: BLE001
        out["error"] = f"no se pudo importar pixeloe.legacy.pixelize: {exc}"
        return out

    try:
        img = np.array(Image.open(src_path).convert("RGB"))  # HxWx3 uint8
    except Exception as exc:  # noqa: BLE001
        out["error"] = f"no se pudo leer/decodificar {src_path!r}: {exc}"
        return out

    try:
        res = pixelize(
            img,
            mode=mode,
            target_size=int(target_size),
            patch_size=int(patch_size),
            thickness=int(thickness),
            contrast=1.0,
            saturation=1.0,
            color_matching=bool(color_matching),
            no_upscale=bool(no_upscale),
        )
    except TypeError as exc:
        # Firma de pixelize distinta a la esperada: reseñar, no relanzar.
        out["error"] = f"pixelize rechazo los kwargs (firma distinta?): {exc}"
        return out
    except Exception as exc:  # noqa: BLE001
        out["error"] = f"pixelize fallo: {exc}"
        return out

    try:
        arr = np.asarray(res)
        result_img = Image.fromarray(arr)
        dst_dir = os.path.dirname(os.path.abspath(dst_path))
        os.makedirs(dst_dir, exist_ok=True)
        result_img.save(dst_path)
    except Exception as exc:  # noqa: BLE001
        out["error"] = f"no se pudo escribir {dst_path!r}: {exc}"
        return out

    out.update(ok=True, out_path=dst_path, size=list(result_img.size), error="")
    return out


def _run_via_bridge(interp, src_path, dst_path, mode, target_size, patch_size,
                    thickness, color_matching, no_upscale):
    """Ejecuta el nucleo en otro interprete (que tiene pixeloe) via subprocess.

    Corre `interp <este_archivo> --bridge <json_args>` y parsea la ultima linea de
    stdout que sea JSON valido (pixeloe puede emitir ruido antes). No-throw.
    """
    import json
    import subprocess

    out = {
        "ok": False,
        "out_path": "",
        "size": [0, 0],
        "mode": mode,
        "target_size": int(target_size),
        "via": "bridge",
        "error": "",
    }

    args = {
        "src_path": src_path,
        "dst_path": dst_path,
        "mode": mode,
        "target_size": int(target_size),
        "patch_size": int(patch_size),
        "thickness": int(thickness),
        "color_matching": bool(color_matching),
        "no_upscale": bool(no_upscale),
    }

    try:
        proc = subprocess.run(
            [interp, os.path.abspath(__file__), "--bridge", json.dumps(args)],
            capture_output=True,
            text=True,
            timeout=600,
        )
    except Exception as exc:  # noqa: BLE001
        out["error"] = f"fallo el subprocess bridge ({interp}): {exc}"
        return out

    if proc.returncode != 0:
        tail = (proc.stderr or "").strip()[-500:]
        out["error"] = f"bridge salio con codigo {proc.returncode}: {tail}"
        return out

    # Parsea de atras hacia delante la primera linea que sea JSON valido.
    parsed = None
    for ln in reversed((proc.stdout or "").splitlines()):
        ln = ln.strip()
        if not ln:
            continue
        try:
            parsed = json.loads(ln)
            break
        except Exception:  # noqa: BLE001 - linea de ruido, sigue probando
            continue

    if parsed is None:
        tail = (proc.stderr or "").strip()[-300:]
        out["error"] = f"bridge no produjo salida JSON. stderr: {tail}"
        return out

    parsed["via"] = "bridge"
    return parsed


def pixeloe_downscale(
    src_path: str,
    dst_path: str,
    *,
    mode: str = "contrast",
    target_size: int = 64,
    patch_size: int = 16,
    thickness: int = 2,
    color_matching: bool = True,
    no_upscale: bool = True,
    comfy_python: str | None = None,
) -> dict:
    """Downscale contrast-aware de una imagen a un grid de pixel-art (no cuantiza).

    Args:
        src_path: ruta de la imagen de entrada (PNG/JPG/...).
        dst_path: ruta del PNG de salida (se crea el directorio si falta).
        mode: algoritmo de downscale de pixeloe: "contrast" (SOTA, conserva
            silueta), "bicubic", "nearest", "center" o "k-centroid". keyword-only.
        target_size: lado del grid resultante en pixeles (64 personajes, 32
            iconos). keyword-only.
        patch_size: tamano del patch que pixeloe colapsa por celda. keyword-only.
        thickness: grosor de la expansion de contorno (outline). keyword-only.
        color_matching: corrige el color de cada celda contra el original si True.
            keyword-only.
        no_upscale: True devuelve el grid real target_size x target_size (lo
            habitual para luego cuantizar); False re-escala al tamano original con
            pixeles duros (preview). keyword-only.
        comfy_python: ruta a un interprete con `pixeloe` para el bridge cuando el
            actual no la tiene. Si None, se prueba COMFY_PYTHON y luego
            ~/ComfyUI/.venv/bin/python3. keyword-only.

    Returns:
        dict con:
        - ok (bool): True si se hizo el downscale y se guardo el PNG.
        - out_path (str): ruta del PNG generado.
        - size (list[int]): [w, h] de la imagen escrita.
        - mode (str): modo de downscale usado.
        - target_size (int): lado del grid pedido.
        - via (str): "inproc" si pixeloe estaba en este interprete, "bridge" si se
          delego a otro interprete por subprocess.
        - error (str): mensaje de error; cadena vacia si todo OK.
    """
    out = {
        "ok": False,
        "out_path": "",
        "size": [0, 0],
        "mode": mode,
        "target_size": int(target_size),
        "via": "",
        "error": "",
    }

    try:
        if not os.path.isfile(src_path):
            out["error"] = f"src_path no existe: {src_path!r}"
            return out

        # 1. pixeloe disponible en el interprete actual -> nucleo directo.
        has_local = True
        try:
            import pixeloe  # noqa: F401
        except ModuleNotFoundError:
            has_local = False
        except Exception:  # noqa: BLE001 - pixeloe presente pero roto -> bridge
            has_local = False

        if has_local:
            res = _run_core(
                src_path, dst_path, mode, int(target_size), int(patch_size),
                int(thickness), bool(color_matching), bool(no_upscale),
            )
            res["via"] = "inproc"
            return res

        # 2. Bridge a un interprete que tenga pixeloe.
        interp = _resolve_comfy_python(comfy_python)
        if interp is None:
            out["error"] = (
                "pixeloe no disponible: no se encontro ningun interprete con "
                "pixeloe (pasa comfy_python, define COMFY_PYTHON, o instala "
                "~/ComfyUI/.venv)"
            )
            return out

        return _run_via_bridge(
            interp, src_path, dst_path, mode, int(target_size), int(patch_size),
            int(thickness), bool(color_matching), bool(no_upscale),
        )
    except Exception as exc:  # noqa: BLE001 - contrato no-throw
        out["error"] = f"error inesperado: {exc}"
        return out


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

    if "--bridge" in sys.argv:
        # Modo bridge: ejecuta el nucleo y emite UNA linea JSON a stdout.
        _idx = sys.argv.index("--bridge")
        _payload = sys.argv[_idx + 1] if len(sys.argv) > _idx + 1 else "{}"
        try:
            _a = json.loads(_payload)
        except Exception as _exc:  # noqa: BLE001
            print(json.dumps({
                "ok": False, "out_path": "", "size": [0, 0], "mode": "",
                "target_size": 0, "via": "inproc",
                "error": f"payload --bridge invalido: {_exc}",
            }))
            sys.exit(0)
        _res = _run_core(
            _a.get("src_path", ""),
            _a.get("dst_path", ""),
            _a.get("mode", "contrast"),
            _a.get("target_size", 64),
            _a.get("patch_size", 16),
            _a.get("thickness", 2),
            _a.get("color_matching", True),
            _a.get("no_upscale", True),
        )
        print(json.dumps(_res))
        sys.exit(0)

    # Modo CLI normal.
    if len(sys.argv) < 3:
        print("uso: pixeloe_downscale.py <src> <dst> [target_size] [mode]",
              file=sys.stderr)
        sys.exit(2)
    _src, _dst = sys.argv[1], sys.argv[2]
    _ts = int(sys.argv[3]) if len(sys.argv) > 3 else 64
    _md = sys.argv[4] if len(sys.argv) > 4 else "contrast"
    print(json.dumps(pixeloe_downscale(_src, _dst, target_size=_ts, mode=_md),
                     indent=2))