fn_registry/python/functions/ml/comfyui_inject_hires_fix.py

"""Inyecta una segunda pasada "hires fix" en un workflow ComfyUI ya construido.

Toma un workflow en API format (dict, p.ej. salida de
comfyui_build_txt2img_workflow) que termina en VAEDecode -> SaveImage y le
encadena una re-difusion por tiles con UltimateSDUpscale + un modelo de upscale
(ESRGAN/Remacri), repuntando el SaveImage para que guarde la imagen ampliada en
vez de la base:

    ... -> VAEDecode ----------------+--> SaveImage   (antes)
                                     |
    ... -> VAEDecode -> UltimateSDUpscale -> SaveImage   (despues)
              UpscaleModelLoader ----^

Es la version ENCADENABLE-sobre-dict del builder comfyui_build_hires_fix_workflow,
que construye el grafo entero desde cero y NO encadena. Reusa exactamente los
mismos class_types e inputs (mode_type 'Linear', mask_blur 8, tile_padding 32,
seam_fix_mode 'None', force_uniform_tiles True, tiled_decode False, etc.).

UltimateSDUpscale ES la segunda pasada de muestreo: re-samplea cada tile con el
checkpoint (de ahi que reciba `model`, `positive`, `negative`, `vae`).

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


def comfyui_inject_hires_fix(
    workflow: dict,
    *,
    upscale_by: float = 1.5,
    denoise: float = 0.4,
    steps: int = 20,
    cfg: float = 7.0,
    seed: int = 0,
    upscale_model: str = "4x_foolhardy_Remacri.pth",
    sampler_name: str = "euler",
    scheduler: str = "normal",
    tile_width: int = 512,
    tile_height: int = 512,
) -> dict:
    """Devuelve una copia del workflow con la segunda pasada hires-fix inyectada.

    Args:
        workflow: dict en API format (ej. salida de
            comfyui_build_txt2img_workflow) que termina en VAEDecode -> SaveImage.
            No se muta; se devuelve una copia.
        upscale_by: factor de ampliacion de UltimateSDUpscale sobre la imagen base
            (1.5 -> 512 pasa a 768). keyword-only.
        denoise: fuerza de re-difusion de la segunda pasada (0.4 por defecto). <1
            conserva la composicion base y solo anade detalle; 1.0 la re-generaria
            entera. keyword-only.
        steps: pasos de sampling de la re-difusion tiled. keyword-only.
        cfg: classifier-free guidance de la re-difusion. keyword-only.
        seed: semilla de UltimateSDUpscale. keyword-only.
        upscale_model: modelo de upscale en models/upscale_models/ que usa
            UltimateSDUpscale para escalar antes de re-difundir (ej.
            "4x_foolhardy_Remacri.pth"). keyword-only.
        sampler_name: sampler de la re-difusion. keyword-only.
        scheduler: scheduler de la re-difusion. keyword-only.
        tile_width: ancho de tile de UltimateSDUpscale (px). Tiles mas pequenos =
            menos VRAM, mas costuras. keyword-only.
        tile_height: alto de tile de UltimateSDUpscale (px). keyword-only.

    Returns:
        copia del workflow con UpscaleModelLoader + UltimateSDUpscale anadidos
        (node_ids = max id numerico existente + 1 y + 2) y el SaveImage repuntado
        a la salida de UltimateSDUpscale. Si no habia SaveImage, se anade uno.

    Raises:
        ValueError: si el workflow no contiene un VAEDecode (fuente de imagen) o
            un CheckpointLoaderSimple (model/vae para la re-difusion).
    """
    wf = copy.deepcopy(workflow)

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

    vaedecode = _find_class("VAEDecode")
    if vaedecode is None:
        raise ValueError(
            "comfyui_inject_hires_fix: no se encontro ningun nodo VAEDecode "
            "(fuente de imagen) en el workflow."
        )

    ckpt = _find_class("CheckpointLoaderSimple")
    if ckpt is None:
        raise ValueError(
            "comfyui_inject_hires_fix: no se encontro ningun nodo "
            "CheckpointLoaderSimple (model/vae para la re-difusion) en el workflow."
        )

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

    # positive/negative: los mismos CLIPTextEncode que alimentan el KSampler.
    pos_src = [ckpt, 0]
    neg_src = [ckpt, 0]
    for node in wf.values():
        if str(node.get("class_type", "")).endswith("KSampler"):
            ins = node.get("inputs", {})
            if _is_link(ins.get("positive")):
                pos_src = list(ins["positive"])
            if _is_link(ins.get("negative")):
                neg_src = list(ins["negative"])
            break

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

    wf[loader_id] = {
        "class_type": "UpscaleModelLoader",
        "inputs": {"model_name": upscale_model},
    }
    wf[upscale_id] = {
        "class_type": "UltimateSDUpscale",
        "inputs": {
            "image": [vaedecode, 0],
            "model": [ckpt, 0],
            "positive": list(pos_src),
            "negative": list(neg_src),
            "vae": [ckpt, 2],
            "upscale_model": [loader_id, 0],
            "upscale_by": upscale_by,
            "seed": seed,
            "steps": steps,
            "cfg": cfg,
            "sampler_name": sampler_name,
            "scheduler": scheduler,
            "denoise": denoise,
            "mode_type": "Linear",
            "tile_width": tile_width,
            "tile_height": tile_height,
            "mask_blur": 8,
            "tile_padding": 32,
            "seam_fix_mode": "None",
            "seam_fix_denoise": 1.0,
            "seam_fix_width": 64,
            "seam_fix_mask_blur": 8,
            "seam_fix_padding": 16,
            "force_uniform_tiles": True,
            "tiled_decode": False,
            "batch_size": 1,
        },
    }

    # repuntar el SaveImage existente al UltimateSDUpscale; si no hay, anadir uno.
    save_id = _find_class("SaveImage")
    if save_id is not None:
        wf[save_id]["inputs"]["images"] = [upscale_id, 0]
    else:
        new_save_id = str(base + 2)
        wf[new_save_id] = {
            "class_type": "SaveImage",
            "inputs": {"filename_prefix": "hires", "images": [upscale_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 cat, detailed")
    wf = comfyui_inject_hires_fix(base, upscale_by=2.0, denoise=0.35, seed=42)
    print(json.dumps(wf, indent=2))