fn_registry/python/functions/ml/comfyui_install_custom_node.py

"""Instala un custom node de ComfyUI: git clone + pip install de sus deps.

Clona el repo en `<comfyui_dir>/custom_nodes/<name>` y, si trae
`requirements.txt`, instala sus dependencias en el venv de ComfyUI
(`<comfyui_dir>/.venv`). El venv de ComfyUI suele crearse con uv y no trae pip;
por eso el instalador se autodetecta en orden: `python -m pip`, luego
`uv pip --python <venv>`, luego el binario `pip` del venv. NO reinicia el
servidor por defecto (restart=False): el nodo no se carga hasta el siguiente
arranque de ComfyUI, asi que reiniciar es una decision explicita del caller (un
restart en caliente corta cualquier generacion en curso).

Impura: ejecuta subprocess (git, pip/uv) y escribe en disco. Solo stdlib.
"""
import os
import shutil
import subprocess

_GIT_TIMEOUT = 300.0
_PIP_TIMEOUT = 600.0


def _pip_install_cmd(base: str, req: str):
    """Comando para instalar requirements en el venv de ComfyUI, o None.

    Prueba en orden: `python -m pip` (si el venv tiene pip), `uv pip` apuntando
    al python del venv (venvs uv sin pip), y por ultimo el binario `pip` del
    venv. Devuelve la lista de args lista para subprocess o None si no hay
    instalador disponible.
    """
    venv_py = os.path.join(base, ".venv", "bin", "python")
    if os.path.isfile(venv_py):
        probe = subprocess.run(
            [venv_py, "-m", "pip", "--version"],
            capture_output=True, text=True,
        )
        if probe.returncode == 0:
            return [venv_py, "-m", "pip", "install", "-r", req]
        if shutil.which("uv"):
            return ["uv", "pip", "install", "-r", req, "--python", venv_py]
    for cand in ("pip", "pip3"):
        pip_bin = os.path.join(base, ".venv", "bin", cand)
        if os.path.isfile(pip_bin):
            return [pip_bin, "install", "-r", req]
    return None


def comfyui_install_custom_node(
    repo_url: str,
    *,
    comfyui_dir: str = "~/ComfyUI",
    pip_install: bool = True,
    restart: bool = False,
) -> dict:
    """Clona un custom node y (opcional) instala sus requirements.

    Args:
        repo_url: URL del repositorio git del custom node
            (ej. "https://github.com/rgthree/rgthree-comfy").
        comfyui_dir: raiz de la instalacion de ComfyUI (se expande ~).
            keyword-only.
        pip_install: si True y el repo trae requirements.txt, instala sus
            dependencias en el venv de ComfyUI. keyword-only.
        restart: NO soportado de forma segura desde aqui (por defecto False).
            El nodo se carga al reiniciar el servidor; hazlo tu cuando no haya
            generaciones en curso. Si se pasa True se anota en el error pero NO
            se reinicia (evita cortar trabajo del servidor). keyword-only.

    Returns:
        dict {ok, path, pip_done, error}. ok=True si el nodo quedo clonado en
        disco (o ya estaba). pip_done=True si se instalaron las dependencias.
        error describe el fallo de git/pip o la advertencia de restart.
    """
    base = os.path.expanduser(comfyui_dir)
    custom_dir = os.path.join(base, "custom_nodes")
    name = os.path.basename(repo_url.rstrip("/"))
    if name.endswith(".git"):
        name = name[:-4]
    if not name:
        return {"ok": False, "path": "", "pip_done": False,
                "error": f"repo_url invalido: {repo_url!r}"}

    dest = os.path.join(custom_dir, name)
    already = os.path.isdir(dest)
    if not already:
        os.makedirs(custom_dir, exist_ok=True)
        try:
            proc = subprocess.run(
                ["git", "clone", "--depth", "1", repo_url, dest],
                capture_output=True, text=True, timeout=_GIT_TIMEOUT,
            )
        except (subprocess.TimeoutExpired, OSError) as exc:
            return {"ok": False, "path": "", "pip_done": False,
                    "error": f"git clone fallo: {exc}"}
        if proc.returncode != 0:
            return {"ok": False, "path": "", "pip_done": False,
                    "error": f"git clone fallo ({proc.returncode}): {proc.stderr.strip()[:300]}"}

    notes = []
    if already:
        notes.append(f"ya existia en {dest} (no se re-clono)")

    pip_done = False
    if pip_install:
        req = os.path.join(dest, "requirements.txt")
        if os.path.isfile(req):
            cmd = _pip_install_cmd(base, req)
            if cmd is None:
                notes.append(f"no se encontro instalador (pip/uv) para {base}/.venv (deps omitidas)")
            else:
                try:
                    pproc = subprocess.run(
                        cmd, capture_output=True, text=True, timeout=_PIP_TIMEOUT,
                    )
                    pip_done = pproc.returncode == 0
                    if not pip_done:
                        notes.append(f"pip install fallo: {pproc.stderr.strip()[:300]}")
                except (subprocess.TimeoutExpired, OSError) as exc:
                    notes.append(f"pip install fallo: {exc}")
        else:
            notes.append("sin requirements.txt (nada que instalar)")

    if restart:
        notes.append(
            "restart=True ignorado: reinicia ComfyUI manualmente cuando no haya "
            "generaciones en curso para cargar el nodo"
        )

    return {"ok": True, "path": dest, "pip_done": pip_done, "error": "; ".join(notes)}


if __name__ == "__main__":
    import json

    out = comfyui_install_custom_node(
        "https://github.com/rgthree/rgthree-comfy", restart=False,
    )
    print(json.dumps(out, ensure_ascii=False, indent=2))