fn_registry/python/functions/ml/comfyui_extract_template.py

"""Extrae el grafo de nodos de un workflow template oficial de ComfyUI por su nombre.

Funcion impura: lee disco (el .json del template instalado) ejecutando la API oficial
del paquete comfyui-workflow-templates dentro del interprete de ComfyUI.

Dado el nombre de un template (su template_id, p.ej. "image_sdxl" o
"api_bfl_flux2_max_sofa_swap"), devuelve:
  - graph: el dict completo del .json (formato UI: nodes/links con posiciones).
  - class_types: la lista de tipos de nodo (class_type) que usa, aplanando los
    subgrafos de `definitions` si los hay.
  - format: "ui_graph" (lo normal en los templates) o "api".
  - assets: rutas en disco de los ficheros del template (json + previews .webp).

Opcionalmente (to_api=True) intenta convertir el grafo UI a API format reutilizando
comfyui_import_workflow_json del registry. Esa conversion necesita un servidor ComfyUI
vivo para mapear los widgets a sus claves de input; si no lo hay, se devuelve el grafo
UI + class_types igualmente y se reporta el motivo en api_error (KISS: no se fuerza la
conversion de grafos complejos).

El paquete vive en el venv de ComfyUI (no en el del registry), por eso esta funcion no
lo importa: localiza el interprete de ComfyUI y le pasa un script que usa la API oficial.
"""
import json
import os
import subprocess
import sys
import tempfile

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
if _THIS_DIR not in sys.path:
    sys.path.insert(0, _THIS_DIR)


# Script que corre DENTRO del python de ComfyUI. Resuelve un template por id, vuelca su
# grafo + metadata como JSON. Si no existe, devuelve sugerencias cercanas.
_EXTRACT_SCRIPT = r"""
import json, sys, difflib, re
try:
    import comfyui_workflow_templates_core as core
except Exception as exc:
    print(json.dumps({"__err__": "import", "msg": str(exc)}))
    sys.exit(0)

_UUID_RE = re.compile(r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")

TID = json.loads({tid_json!r})
m = core.load_manifest()
if TID not in m.templates:
    near = [k for k in m.templates if TID.lower() in k.lower()][:8]
    if not near:
        near = difflib.get_close_matches(TID, list(m.templates.keys()), n=8, cutoff=0.6)
    print(json.dumps({"__err__": "not_found", "suggestions": near}))
    sys.exit(0)

entry = m.templates[TID]
json_asset = next((a.filename for a in entry.assets if a.filename.endswith(".json")), None)
if not json_asset:
    print(json.dumps({"__err__": "no_json"}))
    sys.exit(0)

path = core.get_asset_path(TID, json_asset)
with open(path, encoding="utf-8") as fh:
    graph = json.load(fh)

# Detecta formato y extrae class_types.
fmt = "unknown"
class_types = set()
has_subgraphs = False
if isinstance(graph, dict) and isinstance(graph.get("nodes"), list):
    fmt = "ui_graph"
    for n in graph["nodes"]:
        t = n.get("type") if isinstance(n, dict) else None
        if t and not _UUID_RE.match(str(t)):
            class_types.add(t)
    defs = graph.get("definitions")
    if isinstance(defs, dict) and isinstance(defs.get("subgraphs"), list):
        for sg in defs["subgraphs"]:
            for n in (sg.get("nodes") or []) if isinstance(sg, dict) else []:
                if isinstance(n, dict) and n.get("type"):
                    has_subgraphs = True
                    if not _UUID_RE.match(str(n["type"])):
                        class_types.add(n["type"])
elif isinstance(graph, dict):
    fmt = "api"
    for v in graph.values():
        if isinstance(v, dict) and v.get("class_type"):
            class_types.add(v["class_type"])

print(json.dumps({
    "graph": graph,
    "class_types": sorted(class_types),
    "format": fmt,
    "has_subgraphs": has_subgraphs,
    "bundle": entry.bundle,
    "version": entry.version,
    "assets": core.resolve_all_assets(TID),
    "json_path": path,
}))
"""


def _find_comfyui_python(explicit: str | None) -> str | None:
    """Localiza un interprete de ComfyUI con el paquete instalado (ver list_templates)."""
    candidates = []
    if explicit:
        candidates.append(os.path.expanduser(explicit))
    env = os.environ.get("COMFYUI_PYTHON")
    if env:
        candidates.append(os.path.expanduser(env))
    candidates += [
        os.path.expanduser("~/ComfyUI/.venv/bin/python"),
        os.path.expanduser("~/ComfyUI/venv/bin/python"),
        os.path.expanduser("~/comfyui/.venv/bin/python"),
        sys.executable,
    ]
    for c in candidates:
        if c and os.path.isfile(c):
            return c
    return None


def comfyui_extract_template(
    name: str,
    comfyui_python: str | None = None,
    to_api: bool = False,
    server: str = "127.0.0.1:8188",
) -> dict:
    """Extrae el grafo y los class_types de un template oficial de ComfyUI por nombre.

    Args:
        name: template_id exacto del template (p.ej. "image_sdxl"). Usa
            comfyui_list_templates para ver los nombres disponibles.
        comfyui_python: ruta al interprete python de ComfyUI con el paquete
            comfyui-workflow-templates. Si None, se autodetecta.
        to_api: si True, intenta convertir el grafo UI a API format reutilizando
            comfyui_import_workflow_json (requiere un servidor ComfyUI vivo en
            `server`). Si la conversion falla, se devuelve el grafo UI igualmente y
            el motivo va en api_error.
        server: host:port del servidor ComfyUI para la conversion to_api.

    Returns:
        dict {ok, name, format, class_types, has_subgraphs, n_nodes, graph,
        api_workflow, api_error, bundle, version, assets, error}:
          - graph: el dict del template en formato UI (o API si ya lo estaba).
          - class_types: lista ordenada de tipos de nodo del grafo (incluye los de
            subgrafos de `definitions`).
          - api_workflow: dict en API format si to_api tuvo exito, si no {}.
        Nunca lanza. Nombre inexistente -> ok=False con error legible + sugerencias.
    """
    py = _find_comfyui_python(comfyui_python)
    base = {
        "ok": False,
        "name": name,
        "format": "",
        "class_types": [],
        "has_subgraphs": False,
        "n_nodes": 0,
        "graph": {},
        "api_workflow": {},
        "api_error": "",
        "bundle": "",
        "version": "",
        "assets": [],
        "error": "",
    }
    if not py:
        base["error"] = (
            "no se encontro un interprete de ComfyUI. Pasa comfyui_python=... o "
            "define COMFYUI_PYTHON. Instala el paquete con: "
            "pip install comfyui-workflow-templates"
        )
        return base

    script = _EXTRACT_SCRIPT.replace("{tid_json!r}", repr(json.dumps(name)))
    try:
        proc = subprocess.run(
            [py, "-c", script],
            capture_output=True,
            text=True,
            timeout=60,
        )
    except Exception as exc:  # noqa: BLE001
        base["error"] = f"fallo al ejecutar el interprete de ComfyUI ({py}): {exc}"
        return base

    if proc.returncode != 0:
        base["error"] = f"el interprete de ComfyUI fallo: {proc.stderr.strip()[:500]}"
        return base

    try:
        data = json.loads(proc.stdout.strip().splitlines()[-1])
    except Exception as exc:  # noqa: BLE001
        base["error"] = f"salida no parseable del interprete de ComfyUI: {exc}"
        return base

    err = data.get("__err__")
    if err == "import":
        base["error"] = (
            f"el paquete comfyui-workflow-templates no esta instalado en {py} "
            f"({data.get('msg', '')}). Instalalo con: "
            "pip install comfyui-workflow-templates"
        )
        return base
    if err == "not_found":
        sug = data.get("suggestions", [])
        hint = f" ¿Quizas: {', '.join(sug)}?" if sug else ""
        base["error"] = f"template '{name}' no existe en el paquete.{hint}"
        return base
    if err == "no_json":
        base["error"] = f"el template '{name}' no tiene asset .json."
        return base

    graph = data.get("graph", {})
    fmt = data.get("format", "")
    nodes = graph.get("nodes") if isinstance(graph, dict) else None
    n_nodes = len(nodes) if isinstance(nodes, list) else (
        len(graph) if fmt == "api" and isinstance(graph, dict) else 0
    )

    out = {
        "ok": True,
        "name": name,
        "format": fmt,
        "class_types": data.get("class_types", []),
        "has_subgraphs": data.get("has_subgraphs", False),
        "n_nodes": n_nodes,
        "graph": graph,
        "api_workflow": {},
        "api_error": "",
        "bundle": data.get("bundle", ""),
        "version": data.get("version", ""),
        "assets": data.get("assets", []),
        "error": "",
    }

    if to_api:
        if fmt == "api":
            out["api_workflow"] = graph
        else:
            out["api_workflow"], out["api_error"] = _convert_to_api(graph, server)

    return out


def _convert_to_api(graph: dict, server: str) -> tuple[dict, str]:
    """Convierte un grafo UI a API format via comfyui_import_workflow_json del registry.

    Requiere un servidor ComfyUI vivo para mapear widgets. Devuelve (workflow, "")
    si tuvo exito o ({}, motivo) si fallo. No lanza.
    """
    try:
        from comfyui_import_workflow_json import comfyui_import_workflow_json
    except Exception as exc:  # noqa: BLE001
        return {}, f"no se pudo importar comfyui_import_workflow_json: {exc}"

    tmp = None
    try:
        with tempfile.NamedTemporaryFile(
            "w", suffix=".json", delete=False, encoding="utf-8"
        ) as fh:
            json.dump(graph, fh)
            tmp = fh.name
        res = comfyui_import_workflow_json(tmp, server=server)
        if res.get("ok"):
            return res.get("workflow", {}), ""
        return {}, (
            res.get("error", "conversion fallida")
            + f" (requiere un servidor ComfyUI vivo en {server})"
        )
    except Exception as exc:  # noqa: BLE001
        return {}, f"conversion to_api fallida: {exc}"
    finally:
        if tmp and os.path.exists(tmp):
            try:
                os.unlink(tmp)
            except OSError:
                pass


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(description="Extrae el grafo de un template ComfyUI")
    ap.add_argument("name", help="template_id (ver comfyui_list_templates)")
    ap.add_argument("--comfyui-python", default=None)
    ap.add_argument("--to-api", action="store_true")
    ap.add_argument("--server", default="127.0.0.1:8188")
    ap.add_argument("--full", action="store_true", help="incluye el grafo entero")
    args = ap.parse_args()

    res = comfyui_extract_template(
        args.name,
        args.comfyui_python,
        to_api=args.to_api,
        server=args.server,
    )
    if args.full or not res["ok"]:
        print(json.dumps(res, indent=2, ensure_ascii=False))
    else:
        slim = {k: v for k, v in res.items() if k != "graph"}
        slim["graph_keys"] = list(res["graph"].keys()) if isinstance(res["graph"], dict) else []
        print(json.dumps(slim, indent=2, ensure_ascii=False))