fn_registry/python/functions/ml/comfyui_compose_capabilities.py

"""comfyui_compose_capabilities — mezclador de capacidades sobre un workflow base.

Toma un workflow ComfyUI en API format (la base: salida de
comfyui_build_skill_workflow o comfyui_build_txt2img_workflow) y aplica EN ORDEN
las capacidades que se activen, componiendo los inyectores/builders ENCADENABLES
del registry. Cada capacidad es un argumento keyword opcional: None (default) =
desactivada. Asi el mismo dict base se mezcla a la carta y se puede ir mejorando
(activar/desactivar una capacidad cambia el grafo resultante).

Orden de aplicacion (de mas cerca del checkpoint a la salida):

    1. loras        -> comfyui_inject_multi_lora     (cadena MODEL/CLIP)
    2. controlnet   -> comfyui_inject_controlnet     (re-condiciona KSampler.positive)
    3. ipadapter    -> comfyui_inject_ipadapter      (re-condiciona KSampler.model, tras loras)
    4. facedetailer -> comfyui_build_facedetailer_workflow (regenera caras del VAEDecode)
    5. hires        -> comfyui_inject_hires_fix       (UltimateSDUpscale tras el VAEDecode)

Cada capacidad es independiente: se puede activar cualquier subconjunto. Sin
ninguna activada devuelve una copia del base intacta.

Funcion PURA: sin red, sin I/O. No muta el dict de entrada (copia profunda). Solo
compone funciones puras del registry.

Limitacion conocida (piezas actuales): hires y facedetailer NO encadenan entre
si. Ambos toman su imagen del VAEDecode original del render; combinarlos deja a
uno de los dos sin efecto sobre la salida final. Usa uno U otro por workflow, o
encadenalos manualmente fuera del mixer. Ver el .md (## Gotchas).
"""
from __future__ import annotations

import copy
import os
import sys

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

from ml.comfyui_build_facedetailer_workflow import comfyui_build_facedetailer_workflow  # noqa: E402
from ml.comfyui_inject_controlnet import comfyui_inject_controlnet  # noqa: E402
from ml.comfyui_inject_hires_fix import comfyui_inject_hires_fix  # noqa: E402
from ml.comfyui_inject_ipadapter import comfyui_inject_ipadapter  # noqa: E402
from ml.comfyui_inject_multi_lora import comfyui_inject_multi_lora  # noqa: E402


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 _detect_checkpoint(wf: dict) -> str:
    """Nombre del checkpoint del primer CheckpointLoaderSimple, o '' si no hay."""
    for node in wf.values():
        if node.get("class_type") == "CheckpointLoaderSimple":
            return str(node.get("inputs", {}).get("ckpt_name", "")) or ""
    return ""


def _detect_prompts(wf: dict) -> tuple[str, str]:
    """Texto (positivo, negativo) de los dos primeros CLIPTextEncode del workflow.

    En los builders del registry el positivo se inserta antes que el negativo, asi
    que el primer CLIPTextEncode es el positivo y el segundo el negativo.
    """
    texts = [
        str(n.get("inputs", {}).get("text", ""))
        for n in wf.values()
        if n.get("class_type") == "CLIPTextEncode"
    ]
    positive = texts[0] if texts else ""
    negative = texts[1] if len(texts) > 1 else ""
    return positive, negative


def _prune_redundant_saveimages(wf: dict, keep_source_class: str) -> None:
    """Deja un unico SaveImage: el alimentado por un nodo `keep_source_class`.

    Tras encadenar facedetailer queda el SaveImage del render base (que ya no es
    la salida final) ademas del SaveImage del detailer. Se borra el primero para
    que el workflow tenga una sola imagen de salida (la procesada). Muta `wf` in
    situ (el caller ya trabaja sobre una copia). No-op si hay <=1 SaveImage o si
    no se encuentra el SaveImage alimentado por `keep_source_class`.
    """
    saves = [
        (nid, n) for nid, n in wf.items() if n.get("class_type") == "SaveImage"
    ]
    if len(saves) <= 1:
        return
    keep = None
    for nid, node in saves:
        src = node.get("inputs", {}).get("images")
        if _is_link(src) and wf.get(src[0], {}).get("class_type") == keep_source_class:
            keep = nid
            break
    if keep is None:
        return
    for nid, _ in saves:
        if nid != keep:
            del wf[nid]


def comfyui_compose_capabilities(
    base_workflow: dict,
    *,
    loras: list[dict] | None = None,
    controlnet: dict | None = None,
    ipadapter: dict | None = None,
    hires: dict | None = None,
    facedetailer: dict | None = None,
) -> dict:
    """Aplica en orden las capacidades activadas sobre un workflow base.

    Args:
        base_workflow: dict en API format (salida de
            comfyui_build_skill_workflow o comfyui_build_txt2img_workflow). No se
            muta; se devuelve una copia.
        loras: lista de dicts {name, strength_model?, strength_clip?} para
            comfyui_inject_multi_lora. None o lista vacia = sin LoRAs. keyword-only.
        controlnet: dict para comfyui_inject_controlnet. Claves: control_image
            (str, obligatoria), cn_name (str, obligatoria), strength (float),
            positive_node (str). None = sin ControlNet. keyword-only.
        ipadapter: dict para comfyui_inject_ipadapter. Claves: ref_image (str,
            obligatoria), mode ('style'|'faceid'), weight (float) y demas
            keyword-only del inyector. None = sin IPAdapter. keyword-only.
        hires: dict de kwargs para comfyui_inject_hires_fix (upscale_by, denoise,
            steps, cfg, seed, upscale_model, ...). {} = hires con defaults. None =
            sin hires. keyword-only.
        facedetailer: dict de overrides para comfyui_build_facedetailer_workflow.
            Claves opcionales: ckpt_name (str; si falta se detecta del workflow),
            positive / negative (str; si faltan se detectan de los CLIPTextEncode),
            y demas params del builder (denoise, steps, cfg, seed, bbox_model, ...).
            {} = facedetailer con detect + defaults. None = sin facedetailer.
            keyword-only.

    Returns:
        copia del base con las capacidades activadas encadenadas en orden. Si no
        se activa ninguna, una copia del base intacta.

    Raises:
        ValueError: si una capacidad activada es incompatible (p.ej. controlnet
            sin control_image, ipadapter sin ref_image): se propaga el ValueError
            del inyector correspondiente con el contexto del fallo.
    """
    wf = copy.deepcopy(base_workflow)

    if loras:
        wf = comfyui_inject_multi_lora(wf, loras)

    if controlnet is not None:
        cn = dict(controlnet)
        control_image = cn.pop("control_image", "")
        cn_name = cn.pop("cn_name", "")
        wf = comfyui_inject_controlnet(wf, control_image, cn_name, **cn)

    if ipadapter is not None:
        ip = dict(ipadapter)
        ref_image = ip.pop("ref_image", "")
        wf = comfyui_inject_ipadapter(wf, ref_image, **ip)

    if facedetailer is not None:
        fd = dict(facedetailer)
        ckpt_name = fd.pop("ckpt_name", None) or _detect_checkpoint(wf)
        det_pos, det_neg = _detect_prompts(wf)
        positive = fd.pop("positive", None)
        if positive is None:
            positive = det_pos
        negative = fd.pop("negative", None)
        if negative is None:
            negative = det_neg
        wf = comfyui_build_facedetailer_workflow(wf, ckpt_name, positive, negative, **fd)
        # facedetailer anade su propio SaveImage; el del render base ya no es la
        # salida final -> dejar solo el del detailer.
        _prune_redundant_saveimages(wf, "FaceDetailer")

    if hires is not None:
        h = dict(hires) if isinstance(hires, dict) else {}
        wf = comfyui_inject_hires_fix(wf, **h)

    return wf


# Alias con el nombre completo del ID para descubrimiento por convencion.
compose_capabilities = comfyui_compose_capabilities


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 hero, 3d render")
    mixed = comfyui_compose_capabilities(
        base,
        loras=[
            {"name": "3d_render_redmond_sd15.safetensors", "strength_model": 0.9},
            {"name": "detail_tweaker_sd15.safetensors", "strength_model": 0.5},
        ],
        facedetailer={"denoise": 0.45},
    )
    print(json.dumps({"base_nodes": list(base), "mixed_nodes": list(mixed)}, indent=2))