fn_registry/python/functions/ml/comfyui_inject_ipadapter.py

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

Toma un workflow en API format (dict, p.ej. salida de
comfyui_build_txt2img_workflow, ya con LoRAs encadenadas si las hay) y le injerta
la rama IPAdapter del custom node ComfyUI_IPAdapter_plus (cubiq), repuntando el
KSampler para que su MODEL venga condicionado por una imagen de referencia:

- mode='style': IPAdapterUnifiedLoader + IPAdapter. La imagen de referencia
  transfiere estilo/composicion al resultado (image prompt clasico).
- mode='faceid': IPAdapterUnifiedLoaderFaceID + IPAdapterFaceID. Usa insightface
  para extraer el embedding de la cara de la referencia e imponer un rostro
  consistente en el personaje generado.

La fuente del MODEL es la que HOY alimenta el KSampler.model (tras las LoRAs, no
el checkpoint crudo): asi el IPAdapter se aplica sobre el modelo ya modificado
por los LoRAs, en el orden correcto del mixer.

Es la version ENCADENABLE-sobre-dict del builder
comfyui_build_ipadapter_workflow, que construye el grafo entero desde cero y NO
encadena. Reusa sus constantes de preset/weight_type por defecto. Pensada para
componerse con inject_lora / inject_controlnet / 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). Los
class_type/inputs estan verificados contra /object_info del servidor (IPAdapter
plus), reutilizando exactamente los del builder.
"""
from __future__ import annotations

import copy
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

# Reutiliza los defaults de preset/weight_type del builder para no duplicarlos.
from ml.comfyui_build_ipadapter_workflow import (  # noqa: E402
    _DEFAULT_PRESET,
    _DEFAULT_WEIGHT_TYPE,
)


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_ipadapter(
    workflow: dict,
    ref_image: str,
    *,
    mode: str = "style",
    weight: float = 0.8,
    preset: str | None = None,
    weight_type: str | None = None,
    start_at: float = 0.0,
    end_at: float = 1.0,
    weight_faceidv2: float = 1.0,
    lora_strength: float = 0.6,
    combine_embeds: str = "concat",
    embeds_scaling: str = "V only",
    provider: str = "CPU",
    model_node: str | None = None,
) -> dict:
    """Devuelve una copia del workflow con una rama IPAdapter inyectada.

    Args:
        workflow: dict en API format (ej. salida de
            comfyui_build_txt2img_workflow, posiblemente con LoRAs). No se muta;
            se devuelve una copia.
        ref_image: nombre del archivo de imagen de referencia en el directorio
            input/ del servidor ComfyUI (lo carga un nodo LoadImage). En faceid
            debe contener una cara nitida; en style es la imagen de estilo. No
            puede estar vacio.
        mode: 'style' (transfiere estilo/composicion) o 'faceid' (rostro
            consistente via insightface + FaceID). keyword-only.
        weight: peso de la influencia IPAdapter (0..1+). 0.8 es un buen punto de
            partida; sube para mas parecido, baja para mas libertad del prompt.
        preset: preset del UnifiedLoader. Si None usa el default del modo
            ('STANDARD (medium strength)' para style, 'FACEID PLUS V2' para faceid).
        weight_type: tipo de ponderacion del nodo IPAdapter/FaceID. Si None usa el
            default del modo ('standard' para style, 'linear' para faceid).
        start_at: fraccion del sampling donde empieza a aplicar IPAdapter (0..1).
        end_at: fraccion del sampling donde deja de aplicar (0..1).
        weight_faceidv2: peso del embedding FaceID v2 (solo mode='faceid').
        lora_strength: fuerza de la LoRA FaceID que carga el UnifiedLoaderFaceID
            (solo mode='faceid').
        combine_embeds: como combinar embeddings si hay varias caras
            ('concat'|'add'|'subtract'|'average'|'norm average'). Solo faceid.
        embeds_scaling: escalado de embeddings ('V only'|'K+V'|...). Solo faceid.
        provider: backend de insightface ('CPU'|'CUDA'|...). CPU por defecto para
            no competir por VRAM con el modelo de difusion. Solo faceid.
        model_node: node_id cuya salida MODEL (slot 0) alimentara la rama
            IPAdapter. Si None, se detecta la fuente que hoy alimenta el
            KSampler.model (con el CheckpointLoader como fallback). keyword-only.

    Returns:
        copia del workflow con LoadImage + (UnifiedLoader|UnifiedLoaderFaceID) +
        (IPAdapter|IPAdapterFaceID) anadidos (node_ids = max id numerico + 1/2/3)
        y el KSampler.model repuntado a la salida MODEL de la rama IPAdapter.

    Raises:
        ValueError: si mode no es 'style' ni 'faceid', si ref_image esta vacio, si
            no se encuentra un KSampler, o si no se puede determinar la fuente del
            MODEL (y no se pasa model_node explicito).
    """
    if mode not in ("style", "faceid"):
        raise ValueError(
            f"comfyui_inject_ipadapter: mode debe ser 'style' o 'faceid', no {mode!r}"
        )
    if not ref_image:
        raise ValueError("comfyui_inject_ipadapter: ref_image no puede estar vacio")

    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_ipadapter: no se encontro ningun KSampler en el workflow."
        )

    ks_inputs = wf[ksampler_id].get("inputs", {})
    if model_node is not None:
        model_src = [model_node, 0]
    elif _is_link(ks_inputs.get("model")):
        model_src = list(ks_inputs["model"])
    else:
        ckpt = next(
            (nid for nid, n in wf.items()
             if str(n.get("class_type", "")).startswith("CheckpointLoader")),
            None,
        )
        if ckpt is None:
            raise ValueError(
                "comfyui_inject_ipadapter: no se pudo determinar la fuente del "
                "MODEL; pasa model_node explicito."
            )
        model_src = [ckpt, 0]

    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)

    used_preset = preset if preset is not None else _DEFAULT_PRESET[mode]
    used_wtype = weight_type if weight_type is not None else _DEFAULT_WEIGHT_TYPE[mode]

    wf[load_id] = {
        "class_type": "LoadImage",
        "inputs": {"image": ref_image},
    }

    if mode == "style":
        wf[loader_id] = {
            "class_type": "IPAdapterUnifiedLoader",
            "inputs": {"model": list(model_src), "preset": used_preset},
        }
        wf[apply_id] = {
            "class_type": "IPAdapter",
            "inputs": {
                "model": [loader_id, 0],
                "ipadapter": [loader_id, 1],
                "image": [load_id, 0],
                "weight": weight,
                "start_at": start_at,
                "end_at": end_at,
                "weight_type": used_wtype,
            },
        }
    else:  # faceid
        wf[loader_id] = {
            "class_type": "IPAdapterUnifiedLoaderFaceID",
            "inputs": {
                "model": list(model_src),
                "preset": used_preset,
                "lora_strength": lora_strength,
                "provider": provider,
            },
        }
        wf[apply_id] = {
            "class_type": "IPAdapterFaceID",
            "inputs": {
                "model": [loader_id, 0],
                "ipadapter": [loader_id, 1],
                "image": [load_id, 0],
                "weight": weight,
                "weight_faceidv2": weight_faceidv2,
                "weight_type": used_wtype,
                "combine_embeds": combine_embeds,
                "start_at": start_at,
                "end_at": end_at,
                "embeds_scaling": embeds_scaling,
            },
        }

    # Repunta el KSampler para que tome el MODEL condicionado por IPAdapter.
    wf[ksampler_id]["inputs"]["model"] = [apply_id, 0]
    return wf


if __name__ == "__main__":
    import json

    from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow

    base = comfyui_build_txt2img_workflow("dreamshaper_8.safetensors", "a knight, cinematic")
    wf_style = comfyui_inject_ipadapter(base, "style_ref.png", mode="style", weight=0.8)
    wf_face = comfyui_inject_ipadapter(base, "face_ref.png", mode="faceid", weight=0.9)
    print(json.dumps({"style_nodes": list(wf_style), "faceid_nodes": list(wf_face)}, indent=2))