fn_registry/python/functions/ml/comfyui_fetch_output_mesh.py

"""Localiza y descarga la malla 3D producida por un workflow ComfyUI a disco local.

Hermana de comfyui_fetch_output_image, pero para mallas 3D: el nodo SaveGLB de un
workflow Hunyuan3D expone su salida en GET /history/{prompt_id} bajo la clave "3d"
(no "images"), con {filename, subfolder, type}. Esta funcion lee ese history,
localiza el primer archivo de malla (.glb/.obj/.ply/.gltf/.fbx/.stl/.usdz), lo baja
via GET /view a disco local y, opcionalmente, lo escribe en `dest`.

Impura: red (HTTP GET a /history y /view) + escritura en disco. Solo stdlib.
"""
import json
import os
import urllib.error
import urllib.parse
import urllib.request

_MESH_EXTS = (".glb", ".gltf", ".obj", ".ply", ".fbx", ".stl", ".usdz", ".splat")


def _find_mesh_output(outputs: dict) -> dict | None:
    """Busca en los outputs de /history el primer archivo de malla 3D.

    Recorre cada nodo y cada lista de su output; el SaveGLB usa la clave "3d",
    pero se acepta cualquier lista de dicts con "filename" de extension de malla.
    Devuelve {filename, subfolder, type} o None si no hay ninguno.
    """
    # Prioriza la clave canonica "3d"; si no, cualquier lista con filename de malla.
    for prefer in (True, False):
        for node_out in outputs.values():
            if not isinstance(node_out, dict):
                continue
            for key, items in node_out.items():
                if prefer and key != "3d":
                    continue
                if not isinstance(items, list):
                    continue
                for item in items:
                    if not isinstance(item, dict):
                        continue
                    fn = item.get("filename", "")
                    if fn.lower().endswith(_MESH_EXTS):
                        return {
                            "filename": fn,
                            "subfolder": item.get("subfolder", ""),
                            "type": item.get("type", "output"),
                        }
    return None


def _resolve_dest(dest: str | None, filename: str) -> str:
    """Resuelve la ruta local destino a partir de `dest` y el basename remoto."""
    base = os.path.basename(filename)
    if dest is None:
        return os.path.join(os.getcwd(), base)
    expanded = os.path.expanduser(dest)
    if os.path.isdir(expanded) or expanded.endswith(os.sep):
        return os.path.join(expanded, base)
    return expanded


def comfyui_fetch_output_mesh(
    prompt_id: str,
    *,
    server: str = "127.0.0.1:8188",
    dest: str | None = None,
    timeout: float = 120.0,
) -> dict:
    """Descarga la malla 3D de un prompt ComfyUI ya ejecutado a disco local.

    Args:
        prompt_id: id devuelto por comfyui_submit_workflow, de un workflow cuyo
            nodo SaveGLB ya termino (usa comfyui_wait_result antes si dudas).
        server: host:port del servidor ComfyUI (sin esquema). keyword-only.
        dest: ruta destino. Si None, escribe el basename de la malla en el cwd.
            Si es un directorio (o termina en separador), escribe el basename
            dentro. Si es una ruta de archivo, escribe ahi. keyword-only.
        timeout: timeout de cada peticion HTTP en segundos. keyword-only.

    Returns:
        dict {ok, path, format, bytes, error}. path = ruta local del archivo de
        malla guardado; format = extension sin punto (ej. "glb"); bytes = tamano
        descargado. Si falla, ok=False y error explica (sin malla en history,
        HTTP, conexion o escritura).
    """
    hist_url = f"http://{server}/history/{prompt_id}"
    try:
        with urllib.request.urlopen(hist_url, timeout=timeout) as resp:
            hist = json.loads(resp.read())
    except urllib.error.HTTPError as exc:
        body = exc.read().decode(errors="replace")[:200]
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"HTTP {exc.code} en {hist_url}: {body}"}
    except urllib.error.URLError as exc:
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"no se pudo conectar a {hist_url}: {exc.reason}"}
    except json.JSONDecodeError as exc:
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"respuesta no es JSON valido desde {hist_url}: {exc}"}

    entry = hist.get(prompt_id)
    if not entry:
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"prompt_id {prompt_id} no esta en /history (¿no termino o se purgo?)"}
    outputs = entry.get("outputs", {})
    mesh = _find_mesh_output(outputs)
    if mesh is None:
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"sin archivo de malla 3D en los outputs de {prompt_id}"}

    qs = urllib.parse.urlencode({
        "filename": mesh["filename"],
        "subfolder": mesh["subfolder"],
        "type": mesh["type"],
    })
    view_url = f"http://{server}/view?{qs}"
    try:
        with urllib.request.urlopen(view_url, timeout=timeout) as resp:
            blob = resp.read()
    except urllib.error.HTTPError as exc:
        body = exc.read().decode(errors="replace")[:200]
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"HTTP {exc.code} en {view_url}: {body}"}
    except urllib.error.URLError as exc:
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"no se pudo conectar a {view_url}: {exc.reason}"}

    out_path = _resolve_dest(dest, mesh["filename"])
    try:
        parent = os.path.dirname(out_path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        with open(out_path, "wb") as f:
            f.write(blob)
    except OSError as exc:
        return {"ok": False, "path": "", "format": "", "bytes": 0,
                "error": f"no se pudo escribir en {out_path!r}: {exc}"}

    fmt = os.path.splitext(mesh["filename"])[1].lstrip(".").lower()
    return {"ok": True, "path": out_path, "format": fmt, "bytes": len(blob), "error": ""}


if __name__ == "__main__":
    import sys

    pid = sys.argv[1] if len(sys.argv) > 1 else "00000000-0000-0000-0000-000000000000"
    res = comfyui_fetch_output_mesh(pid, dest="/tmp/comfy_mesh")
    print(json.dumps(res, indent=2))