fn_registry/python/functions/ml/comfyui_inject_controlnet.py

"""Inyecta una rama ControlNet en un workflow ComfyUI ya construido (API format).

Toma un workflow en API format (dict, p.ej. salida de
comfyui_build_txt2img_workflow) que tiene un KSampler cuyo condicionamiento
positivo viene de un CLIPTextEncode, y le encadena la rama de ControlNet:

    LoadImage (imagen de control) ---+
    ControlNetLoader (modelo CN) ----+--> ControlNetApply --> KSampler.positive
    CLIPTextEncode (positivo) -------+

ControlNetApply re-condiciona el positivo con la imagen de control (canny, depth,
pose, scribble, ...) y el KSampler se repunta para tomar ese condicionamiento.

Es la version ENCADENABLE-sobre-dict del builder
comfyui_build_controlnet_workflow, que construye el grafo entero desde cero y NO
encadena. Reusa los mismos class_types/inputs (LoadImage, ControlNetLoader,
ControlNetApply). Pensada para componerse con inject_lora / inject_ipadapter /
inject_hires_fix sobre un mismo dict base (ver comfyui_compose_capabilities).

Funcion pura: sin red, sin I/O. No muta el dict de entrada (copia profunda).
"""
import copy


def _is_link(v) -> bool:
    """True si v es una conexion ComfyUI [node_id(str), output_index(int)]."""
    return (
        isinstance(v, list)
        and len(v) == 2
        and isinstance(v[0], str)
        and isinstance(v[1], int)
    )


def comfyui_inject_controlnet(
    workflow: dict,
    control_image: str,
    cn_name: str,
    *,
    strength: float = 1.0,
    positive_node: str | None = None,
) -> dict:
    """Devuelve una copia del workflow con una rama ControlNet inyectada.

    Localiza el condicionamiento positivo actual del KSampler (lo que hoy
    alimenta su input `positive`), inserta LoadImage + ControlNetLoader +
    ControlNetApply, y repunta el KSampler para que tome el positivo ya
    condicionado por el ControlNet.

    Args:
        workflow: dict en API format (ej. salida de
            comfyui_build_txt2img_workflow). No se muta; se devuelve una copia.
        control_image: nombre del archivo de la imagen de control dentro de la
            carpeta input/ del servidor ComfyUI (lo carga el nodo LoadImage).
            Suele ser un mapa preprocesado (canny/depth/openpose). No puede estar
            vacio.
        cn_name: nombre del modelo ControlNet en models/controlnet/ tal como lo
            lista /object_info para ControlNetLoader (control_net_name).
        strength: fuerza con la que el ControlNet condiciona la generacion
            (0.0 = nula, 1.0 = plena). keyword-only.
        positive_node: node_id cuya salida CONDITIONING (slot 0) se usara como
            positivo de entrada del ControlNetApply. Si None, se detecta la
            fuente que hoy alimenta el KSampler.positive. keyword-only.

    Returns:
        copia del workflow con LoadImage + ControlNetLoader + ControlNetApply
        anadidos (node_ids = max id numerico existente + 1, + 2, + 3) y el
        KSampler.positive repuntado a la salida del ControlNetApply.

    Raises:
        ValueError: si control_image esta vacio, si no se encuentra un KSampler,
            o si no se puede determinar la fuente del condicionamiento positivo
            (y no se pasa positive_node explicito).
    """
    if not control_image:
        raise ValueError(
            "comfyui_inject_controlnet: control_image no puede estar vacio "
            "(ControlNet necesita una imagen de control en input/)."
        )

    wf = copy.deepcopy(workflow)

    ksampler_id = next(
        (nid for nid, n in wf.items() if str(n.get("class_type", "")).endswith("KSampler")),
        None,
    )
    if ksampler_id is None:
        raise ValueError(
            "comfyui_inject_controlnet: no se encontro ningun KSampler en el workflow."
        )

    ks_inputs = wf[ksampler_id].get("inputs", {})
    if positive_node is not None:
        pos_src = [positive_node, 0]
    elif _is_link(ks_inputs.get("positive")):
        pos_src = list(ks_inputs["positive"])
    else:
        raise ValueError(
            "comfyui_inject_controlnet: no se pudo determinar la fuente del "
            "condicionamiento positivo; pasa positive_node explicito."
        )

    numeric = [int(k) for k in wf.keys() if str(k).isdigit()]
    base = (max(numeric) + 1) if numeric else len(wf) + 1
    load_id = str(base)
    loader_id = str(base + 1)
    apply_id = str(base + 2)

    wf[load_id] = {
        "class_type": "LoadImage",
        "inputs": {"image": control_image},
    }
    wf[loader_id] = {
        "class_type": "ControlNetLoader",
        "inputs": {"control_net_name": cn_name},
    }
    wf[apply_id] = {
        "class_type": "ControlNetApply",
        "inputs": {
            "conditioning": list(pos_src),
            "control_net": [loader_id, 0],
            "image": [load_id, 0],
            "strength": strength,
        },
    }

    # Repunta el KSampler para que tome el positivo condicionado por el ControlNet.
    wf[ksampler_id]["inputs"]["positive"] = [apply_id, 0]
    return wf


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

    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow

    base = comfyui_build_txt2img_workflow("IMG_dreamshaper_8.safetensors", "a knight, dramatic")
    wf = comfyui_inject_controlnet(
        base, "pose_canny.png", "control_v11p_sd15_canny.pth", strength=0.8
    )
    print(json.dumps(wf, indent=2))