fn_registry/python/functions/ml/comfyui_inject_lora.py

"""Inserta un nodo LoraLoader en un workflow ComfyUI ya construido (API format).

Reconecta las salidas model/clip de la fuente actual (el CheckpointLoaderSimple
o un LoraLoader previo) hacia el nuevo LoraLoader, y repunta a los consumidores
(KSampler, CLIPTextEncode) para que pasen por el LoRA. Llamar varias veces sobre
el mismo workflow encadena LoRAs.

Convencion de slots ComfyUI: tanto CheckpointLoaderSimple como LoraLoader
exponen MODEL en el output 0 y CLIP en el output 1.

Funcion pura: no muta el dict de entrada (trabaja sobre una copia profunda).
"""
import copy


def comfyui_inject_lora(
    workflow: dict,
    lora_name: str,
    *,
    strength_model: float = 1.0,
    strength_clip: float = 1.0,
    model_node: str | None = None,
    clip_node: str | None = None,
) -> dict:
    """Devuelve una copia del workflow con un LoraLoader insertado y reconectado.

    Args:
        workflow: dict en API format (ej. salida de
            comfyui_build_txt2img_workflow). No se muta.
        lora_name: nombre del archivo .safetensors del LoRA en models/loras/.
        strength_model: fuerza del LoRA sobre el modelo (UNet). keyword-only.
        strength_clip: fuerza del LoRA sobre el CLIP. keyword-only.
        model_node: node_id cuya salida MODEL (slot 0) alimentara el LoRA. Si
            None, se detecta la fuente que hoy alimenta el KSampler.model (con el
            CheckpointLoaderSimple como fallback). keyword-only.
        clip_node: node_id cuya salida CLIP (slot 1) alimentara el LoRA. Si None,
            se detecta la fuente que hoy alimenta los CLIPTextEncode.clip.
            keyword-only.

    Returns:
        copia del workflow con el LoraLoader insertado. El nuevo node_id es el
        maximo id numerico existente + 1.

    Raises:
        ValueError: si no se puede determinar la fuente model/clip y no se pasan
            model_node/clip_node explicitos.
    """
    wf = copy.deepcopy(workflow)

    def _is_link(v) -> bool:
        return (
            isinstance(v, list)
            and len(v) == 2
            and isinstance(v[0], str)
            and isinstance(v[1], int)
        )

    def _find_class(prefix):
        for nid, node in wf.items():
            if str(node.get("class_type", "")).startswith(prefix):
                return nid
        return None

    ckpt = _find_class("CheckpointLoader")

    # fuente actual de model/clip: la que alimenta KSampler.model y CLIPTextEncode.clip
    model_src = None
    clip_src = None
    for node in wf.values():
        ins = node.get("inputs", {})
        if str(node.get("class_type", "")).endswith("KSampler") and _is_link(ins.get("model")):
            model_src = list(ins["model"])
        if node.get("class_type") == "CLIPTextEncode" and clip_src is None and _is_link(ins.get("clip")):
            clip_src = list(ins["clip"])

    if model_node is not None:
        model_src = [model_node, 0]
    elif model_src is None and ckpt is not None:
        model_src = [ckpt, 0]

    if clip_node is not None:
        clip_src = [clip_node, 1]
    elif clip_src is None and ckpt is not None:
        clip_src = [ckpt, 1]

    if model_src is None or clip_src is None:
        raise ValueError(
            "comfyui_inject_lora: no se pudo determinar la fuente model/clip; "
            "pasa model_node y clip_node explicitos."
        )

    numeric = [int(k) for k in wf.keys() if str(k).isdigit()]
    new_id = str((max(numeric) + 1) if numeric else len(wf) + 1)

    wf[new_id] = {
        "class_type": "LoraLoader",
        "inputs": {
            "lora_name": lora_name,
            "strength_model": strength_model,
            "strength_clip": strength_clip,
            "model": list(model_src),
            "clip": list(clip_src),
        },
    }

    # repuntar consumidores de model_src/clip_src hacia el LoraLoader (no el propio LoRA)
    for nid, node in wf.items():
        if nid == new_id:
            continue
        ins = node.get("inputs", {})
        for k, v in list(ins.items()):
            if _is_link(v) and list(v) == list(model_src):
                ins[k] = [new_id, 0]
            elif _is_link(v) and list(v) == list(clip_src):
                ins[k] = [new_id, 1]

    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 cat")
    wf = comfyui_inject_lora(base, "add_detail.safetensors", strength_model=0.8)
    print(json.dumps(wf, indent=2))