"""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, }, } # 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("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))