fn_registry/python/functions/ml/comfyui_read_png_metadata.py

"""Lee los parametros de generacion de un PNG generado por ComfyUI.

Extrae el chunk "prompt" (API format) de los chunks de texto del PNG y resume
los parametros de generacion: modelo, seed, steps, cfg, sampler, scheduler,
denoise y los prompts positivo/negativo (siguiendo las conexiones del KSampler).

Impura: lectura de disco. Solo stdlib (struct, zlib, json).
"""
import json
import struct
import zlib


def comfyui_read_png_metadata(png_path: str) -> dict:
    """Devuelve {ok, prompt, parameters, error} de un PNG de ComfyUI.

    Args:
        png_path: ruta del PNG generado por ComfyUI.

    Returns:
        dict con:
          - ok: bool.
          - prompt: el workflow API format embebido (dict), o {}.
          - parameters: resumen {model, seed, steps, cfg, sampler_name,
            scheduler, denoise, positive, negative} extraido del KSampler y los
            nodos conectados, o {}.
          - error: mensaje si algo fallo.
    """
    try:
        with open(png_path, "rb") as f:
            data = f.read()
    except OSError as exc:
        return {"ok": False, "prompt": {}, "parameters": {},
                "error": f"no se pudo leer {png_path!r}: {exc}"}
    try:
        chunks = _png_text_chunks(data)
    except ValueError as exc:
        return {"ok": False, "prompt": {}, "parameters": {}, "error": str(exc)}

    if "prompt" not in chunks:
        return {"ok": False, "prompt": {}, "parameters": {},
                "error": "el PNG no contiene chunk 'prompt' de ComfyUI"}
    try:
        prompt = json.loads(chunks["prompt"])
    except json.JSONDecodeError as exc:
        return {"ok": False, "prompt": {}, "parameters": {},
                "error": f"chunk 'prompt' no es JSON valido: {exc}"}

    return {"ok": True, "prompt": prompt, "parameters": _extract_params(prompt), "error": ""}


def _extract_params(prompt: dict) -> dict:
    params = {}
    ksampler = None
    for node in prompt.values():
        if isinstance(node, dict) and str(node.get("class_type", "")).endswith("KSampler"):
            ksampler = node
            break
    if ksampler:
        ins = ksampler.get("inputs", {})
        for k in ("seed", "steps", "cfg", "sampler_name", "scheduler", "denoise"):
            if k in ins and not isinstance(ins[k], list):
                params[k] = ins[k]
        for slot in ("positive", "negative"):
            link = ins.get(slot)
            if isinstance(link, list) and len(link) == 2:
                tnode = prompt.get(str(link[0]), {})
                txt = tnode.get("inputs", {}).get("text")
                if isinstance(txt, str):
                    params[slot] = txt
    for node in prompt.values():
        if isinstance(node, dict) and str(node.get("class_type", "")).startswith("CheckpointLoader"):
            ck = node.get("inputs", {}).get("ckpt_name")
            if ck:
                params["model"] = ck
            break
    return params


def _png_text_chunks(data: bytes) -> dict:
    """Lee los chunks de texto (tEXt/zTXt/iTXt) de un PNG -> {keyword: texto}."""
    if data[:8] != b"\x89PNG\r\n\x1a\n":
        raise ValueError("no es un PNG valido (firma incorrecta)")
    out = {}
    off = 8
    n = len(data)
    while off + 8 <= n:
        length = struct.unpack(">I", data[off:off + 4])[0]
        ctype = data[off + 4:off + 8]
        body = data[off + 8:off + 8 + length]
        off += 12 + length
        if ctype == b"tEXt":
            kw, _, txt = body.partition(b"\x00")
            out[kw.decode("latin1")] = txt.decode("latin1")
        elif ctype == b"zTXt":
            kw, _, rest = body.partition(b"\x00")
            if rest:
                try:
                    out[kw.decode("latin1")] = zlib.decompress(rest[1:]).decode("latin1")
                except zlib.error:
                    pass
        elif ctype == b"iTXt":
            kw, _, rest = body.partition(b"\x00")
            if len(rest) >= 2:
                comp_flag = rest[0]
                parts = rest[2:].split(b"\x00", 2)
                if len(parts) == 3:
                    text_bytes = parts[2]
                    if comp_flag == 1:
                        try:
                            text_bytes = zlib.decompress(text_bytes)
                        except zlib.error:
                            text_bytes = b""
                    out[kw.decode("latin1")] = text_bytes.decode("utf-8", "replace")
        elif ctype == b"IEND":
            break
    return out


if __name__ == "__main__":
    import sys

    path = sys.argv[1] if len(sys.argv) > 1 else "/tmp/missing.png"
    res = comfyui_read_png_metadata(path)
    print(json.dumps({"ok": res["ok"], "parameters": res["parameters"], "error": res["error"]}, indent=2))