diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d46cd5e --- /dev/null +++ b/__init__.py @@ -0,0 +1,15 @@ +# __init__.py — Registro nodi IAMCCS (solo nuovi nodi LoRA) +from .iamccs_wan_lora_stack import ( + IAMCCS_WanLoRAStack, + IAMCCS_ModelWithLoRA +) + +NODE_CLASS_MAPPINGS = { + "IAMCCS_WanLoRAStack": IAMCCS_WanLoRAStack, + "IAMCCS_ModelWithLoRA": IAMCCS_ModelWithLoRA, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "IAMCCS_WanLoRAStack": "LoRA Stack (WAN-style remap)", + "IAMCCS_ModelWithLoRA": "Apply LoRA to MODEL (Native)", +} diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..8ae060d Binary files /dev/null and b/__pycache__/__init__.cpython-312.pyc differ diff --git a/__pycache__/iamccs_wan_lora_stack.cpython-312.pyc b/__pycache__/iamccs_wan_lora_stack.cpython-312.pyc new file mode 100644 index 0000000..c0ab13f Binary files /dev/null and b/__pycache__/iamccs_wan_lora_stack.cpython-312.pyc differ diff --git a/iamccs_wan_lora_stack.py b/iamccs_wan_lora_stack.py new file mode 100644 index 0000000..79b0508 --- /dev/null +++ b/iamccs_wan_lora_stack.py @@ -0,0 +1,131 @@ +# iamccs_wan_lora_stack.py +# =============================================================== +# IAMCCS_WanLoRAStack / IAMCCS_ModelWithLoRA +# Versione multi-LoRA (4 slots con strength dedicato) +# =============================================================== + +import logging +import comfy.utils +import comfy.sd +import folder_paths + + +# --- Normalizzatore chiavi WAN (stesso del precedente) --- +def standardize_wan_lora_keys(sd: dict) -> dict: + new_sd = {} + for k, v in sd.items(): + nk = k + if nk.startswith("transformer."): + nk = nk.replace("transformer.", "diffusion_model.") + if nk.startswith("pipe.dit."): + nk = nk.replace("pipe.dit.", "diffusion_model.") + if nk.startswith("blocks."): + nk = nk.replace("blocks.", "diffusion_model.blocks.") + if ".attn1." in nk or ".attn2." in nk: + tgt = ".cross_attn." if ".attn2." in nk else ".self_attn." + nk = nk.replace(".attn1.", tgt).replace(".attn2.", tgt) + nk = nk.replace(".to_k.", ".k.").replace(".to_q.", ".q.") + nk = nk.replace(".to_v.", ".v.").replace(".to_out.0.", ".o.") + if nk.startswith("lora_unet__"): + core, *rest = nk.split(".") + core = core.replace("lora_unet__", "diffusion_model.") + core = core.replace("_self_attn", ".self_attn").replace("_cross_attn", ".cross_attn").replace("blocks_", "blocks.") + nk = ".".join([core] + rest) + nk = nk.replace("img_attn.proj", "img_attn_proj") + nk = nk.replace("img_attn.qkv", "img_attn_qkv") + nk = nk.replace("txt_attn.proj", "txt_attn_proj") + nk = nk.replace("txt_attn.qkv", "txt_attn_qkv") + new_sd[nk] = v + return new_sd + + +# =============================================================== +# Nodo 1 — LoRA Stack (fino a 4 LoRA con strength) +# =============================================================== +class IAMCCS_WanLoRAStack: + @classmethod + def INPUT_TYPES(cls): + lora_list = folder_paths.get_filename_list("loras") + return { + "required": { + "lora1": (lora_list,), + "strength1": ("FLOAT", {"default": 1.0, "min": -5.0, "max": 5.0, "step": 0.01}), + "lora2": (lora_list,), + "strength2": ("FLOAT", {"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01}), + "lora3": (lora_list,), + "strength3": ("FLOAT", {"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01}), + "lora4": (lora_list,), + "strength4": ("FLOAT", {"default": 0.0, "min": -5.0, "max": 5.0, "step": 0.01}), + "model_type": (["wan2x", "flow", "standard"], {"default": "flow"}), + } + } + + RETURN_TYPES = ("LORA",) + FUNCTION = "load_and_standardize" + CATEGORY = "IAMCCS/LoRA" + + def load_and_standardize(self, lora1, strength1, + lora2, strength2, + lora3, strength3, + lora4, strength4, + model_type="flow"): + + loras = [] + for name, strength in [ + (lora1, strength1), + (lora2, strength2), + (lora3, strength3), + (lora4, strength4), + ]: + if name and strength != 0.0: + path = folder_paths.get_full_path_or_raise("loras", name) + sd = comfy.utils.load_torch_file(path, safe_load=True) + if model_type != "standard": + sd = standardize_wan_lora_keys(sd) + loras.append({"name": name, "strength": strength, "state_dict": sd}) + + logging.info(f"[IAMCCS_WanLoRAStack] ✅ Caricate {len(loras)} LoRA") + return (loras,) + + +# =============================================================== +# Nodo 2 — Apply LoRA to MODEL (ponte, senza strength) +# =============================================================== +class IAMCCS_ModelWithLoRA: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "lora": ("LORA",), + } + } + + RETURN_TYPES = ("MODEL",) + FUNCTION = "apply" + CATEGORY = "IAMCCS/LoRA" + + def apply(self, model, lora): + if not lora: + return (model,) + model_out = model + for entry in lora: + sd = entry["state_dict"] + strength = entry["strength"] + model_out, _ = comfy.sd.load_lora_for_models(model_out, None, sd, strength, 0) + logging.info(f"[IAMCCS_ModelWithLoRA] ✅ '{entry['name']}' strength={strength}") + return (model_out,) + + +# =============================================================== +# Registrazione +# =============================================================== +NODE_CLASS_MAPPINGS = { + "IAMCCS_WanLoRAStack": IAMCCS_WanLoRAStack, + "IAMCCS_ModelWithLoRA": IAMCCS_ModelWithLoRA, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "IAMCCS_WanLoRAStack": "LoRA Stack (WAN-style remap)", + "IAMCCS_ModelWithLoRA": "Apply LoRA to MODEL (Native)", +} diff --git a/icon.png b/icon.png new file mode 100644 index 0000000..698c262 Binary files /dev/null and b/icon.png differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..2e14f76 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "iamccs-wan-lora-fixer-stack" +version = "1.1.0" +description = "Nodo IAMCCS per ComfyUI che applica fino a 5 LoRA (WAN 2.1 / 2.2 compatibili) con fix automatico e input/output MODEL." +license = { text = "MIT" } +authors = [ + { name = "Carmine Cristallo Scalzi (IAMCCS)", email = "info@carminecristalloscalzi.com" } +] + +[tool.comfy] +publisher_id = "iamccs" +display_name = "IAMCCS WAN LoRA Fixer Stack" +icon = "icon.png" +requires_comfyui = ">=3.0.0" diff --git a/version.json b/version.json new file mode 100644 index 0000000..5ecb857 --- /dev/null +++ b/version.json @@ -0,0 +1,6 @@ +{ + "name": "iamccs-wan-lora-fixer-stack", + "version": "1.1.0", + "author": "Carmine Cristallo Scalzi (IAMCCS)", + "description": "Nodo IAMCCS per ComfyUI che applica fino a 5 LoRA in stack, compatibile con WAN 2.1 e 2.2 (Flow)." +}