Upload wan_lora_nodes.py

wan_lora_nodes.py

@@ -0,0 +1,542 @@
+import hashlib
+import json
+import os
+import re
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Iterable, Optional, Tuple
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+
+
+# ============================================================
+# Helpers
+# ============================================================
+
+
+def _clean_path(path: str) -> str:
+    if path is None:
+        return ""
+    return str(path).strip().strip('"').strip("'")
+
+
+
+def _ensure_metadata_str_dict(metadata: Optional[dict]) -> Dict[str, str]:
+    if not metadata:
+        return {}
+    return {str(k): str(v) for k, v in metadata.items()}
+
+
+
+def _file_signature(path: str) -> str:
+    p = Path(path)
+    stat = p.stat()
+    payload = f"{p.resolve()}|{stat.st_size}|{stat.st_mtime_ns}"
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+
+
+def _load_safetensors(path: str) -> Tuple[Dict[str, torch.Tensor], Dict[str, str]]:
+    tensors: Dict[str, torch.Tensor] = {}
+    metadata: Dict[str, str] = {}
+    with safe_open(path, framework="pt") as f:
+        metadata = _ensure_metadata_str_dict(f.metadata())
+        for key in f.keys():
+            tensors[key] = f.get_tensor(key)
+    return tensors, metadata
+
+
+
+def _detect_format(keys: Iterable[str]) -> str:
+    keys = list(keys)
+    if not keys:
+        return "empty"
+    sample = keys[0]
+    if sample.startswith("lora_unet_"):
+        return "kohya"
+    if sample.startswith("diffusion_model."):
+        return "diffusers"
+    if sample.startswith("transformer."):
+        return "transformer"
+    if sample.startswith("blocks."):
+        return "blocks"
+    if sample.startswith("base_model.model."):
+        return "peft"
+    return "other"
+
+
+
+def _maybe_compensate_rs_lora(sd: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """
+    Handle rank-stabilized PEFT-like LoRAs similarly to WanVideoWrapper.
+    When these files omit alpha, add a compensated alpha tensor.
+    """
+    probe_key = "base_model.model.blocks.0.cross_attn.k.lora_A.weight"
+    if probe_key not in sd:
+        return sd
+
+    rank = int(sd[probe_key].shape[0])
+    # Mirrors the idea used in WanVideoWrapper's compensate_rs_lora_format.
+    alpha_value = rank * rank // max(int(rank ** 0.5), 1)
+    alpha_tensor = torch.tensor(alpha_value, dtype=torch.int32)
+
+    out: Dict[str, torch.Tensor] = {}
+    for k, v in sd.items():
+        out[k] = v
+        if k.endswith(".lora_A.weight"):
+            alpha_key = k.replace(".lora_A.weight", ".alpha")
+            out.setdefault(alpha_key, alpha_tensor)
+    return out
+
+
+
+def _standardize_key(key: str) -> Optional[str]:
+    """
+    Normalize several LoRA naming variants into the WAN2.2/Comfy diffusers-style
+    `diffusion_model.*` key space.
+
+    Supported sources:
+    - Kohya/Fun style: lora_unet_* and lora_unet__*
+    - Diffusers variants: transformer.*, blocks.*, base_model.model.*
+    - Already standardized: diffusion_model.*
+    """
+    k = key
+
+    # Already standardized.
+    if k.startswith("diffusion_model."):
+        return k
+
+    # Light normalization adopted from current WanVideoWrapper loading logic.
+    if k.startswith("transformer."):
+        k = k.replace("transformer.", "diffusion_model.", 1)
+    if k.startswith("pipe.dit."):
+        k = k.replace("pipe.dit.", "diffusion_model.", 1)
+    if k.startswith("blocks."):
+        k = k.replace("blocks.", "diffusion_model.blocks.", 1)
+    if k.startswith("vace_blocks."):
+        k = k.replace("vace_blocks.", "diffusion_model.vace_blocks.", 1)
+    if k.startswith("base_model.model."):
+        k = k.replace("base_model.model.", "diffusion_model.", 1)
+
+    k = k.replace(".default.", ".")
+    k = k.replace(".diff_m", ".modulation.diff")
+
+    if k.startswith("diffusion.model."):
+        k = k.replace("diffusion.model.", "diffusion_model.", 1)
+
+    if ".attn1." in k:
+        k = k.replace(".attn1.", ".cross_attn.")
+        k = k.replace(".to_k.", ".k.")
+        k = k.replace(".to_q.", ".q.")
+        k = k.replace(".to_v.", ".v.")
+        k = k.replace(".to_out.0.", ".o.")
+    elif ".attn2." in k:
+        k = k.replace(".attn2.", ".cross_attn.")
+        k = k.replace(".to_k.", ".k.")
+        k = k.replace(".to_q.", ".q.")
+        k = k.replace(".to_v.", ".v.")
+        k = k.replace(".to_out.0.", ".o.")
+
+    k = k.replace("img_attn.proj", "img_attn_proj")
+    k = k.replace("img_attn.qkv", "img_attn_qkv")
+    k = k.replace("txt_attn.proj", "txt_attn_proj")
+    k = k.replace("txt_attn.qkv", "txt_attn_qkv")
+
+    # AIToolkit/LyCORIS-ish shorthand.
+    if k.startswith("lycoris_blocks_"):
+        k = k.replace("lycoris_blocks_", "diffusion_model.blocks.", 1)
+        k = k.replace("_cross_attn_", ".cross_attn.")
+        k = k.replace("_self_attn_", ".self_attn.")
+        k = k.replace("_ffn_net_0_proj", ".ffn.0")
+        k = k.replace("_ffn_net_2", ".ffn.2")
+        k = k.replace("to_out_0", "o")
+        return k
+
+    # Kohya / Fun LoRA style with double underscore.
+    if k.startswith("lora_unet__"):
+        parts = k.split(".")
+        main_part = parts[0]
+        weight_type = ".".join(parts[1:]) if len(parts) > 1 else ""
+
+        if "blocks_" in main_part:
+            components = main_part[len("lora_unet__"):].split("_")
+            new_key = "diffusion_model"
+            if len(components) >= 2 and components[0] == "blocks":
+                new_key += f".blocks.{components[1]}"
+                idx = 2
+                if idx < len(components):
+                    if idx + 1 < len(components) and components[idx] == "self" and components[idx + 1] == "attn":
+                        new_key += ".self_attn"
+                        idx += 2
+                    elif idx + 1 < len(components) and components[idx] == "cross" and components[idx + 1] == "attn":
+                        new_key += ".cross_attn"
+                        idx += 2
+                    elif components[idx] == "ffn":
+                        new_key += ".ffn"
+                        idx += 1
+
+                    if idx < len(components):
+                        component = components[idx]
+                        idx += 1
+                        if idx < len(components) and components[idx] == "img":
+                            component += "_img"
+                        new_key += f".{component}"
+
+                if weight_type == "alpha":
+                    return new_key + ".alpha"
+                if weight_type in {"lora_down.weight", "lora_down"}:
+                    return new_key + ".lora_A.weight"
+                if weight_type in {"lora_up.weight", "lora_up"}:
+                    return new_key + ".lora_B.weight"
+                if weight_type:
+                    return new_key + f".{weight_type}"
+                return new_key
+
+        # Fallback for remaining lora_unet__ patterns.
+        new_key = main_part.replace("lora_unet__", "diffusion_model.", 1)
+        new_key = new_key.replace("_self_attn", ".self_attn")
+        new_key = new_key.replace("_cross_attn", ".cross_attn")
+        new_key = new_key.replace("_ffn", ".ffn")
+        new_key = new_key.replace("blocks_", "blocks.")
+        new_key = new_key.replace("head_head", "head.head")
+        new_key = new_key.replace("text_embedding", "text.embedding")
+        new_key = new_key.replace("time_embedding", "time.embedding")
+        new_key = new_key.replace("time_projection", "time.projection")
+
+        rebuilt_parts = []
+        for part in new_key.split("."):
+            if part in {"img_emb", "self_attn", "cross_attn"}:
+                rebuilt_parts.append(part)
+            else:
+                rebuilt_parts.append(part.replace("_", "."))
+        new_key = ".".join(rebuilt_parts)
+
+        special_components = {
+            "time.projection": "time_projection",
+            "img.emb": "img_emb",
+            "text.emb": "text_emb",
+            "time.emb": "time_emb",
+        }
+        for old, new in special_components.items():
+            new_key = new_key.replace(old, new)
+
+        if weight_type == "alpha":
+            return new_key + ".alpha"
+        if weight_type in {"lora_down.weight", "lora_down"}:
+            return new_key + ".lora_A.weight"
+        if weight_type in {"lora_up.weight", "lora_up"}:
+            return new_key + ".lora_B.weight"
+        if weight_type:
+            return new_key + f".{weight_type}"
+        return new_key
+
+    # Kohya style from the user's original converter.
+    if k.startswith("lora_unet_"):
+        # alpha support
+        m = re.match(r"lora_unet_blocks_(\d+)_(cross_attn|self_attn)_(\w+)\.alpha$", k)
+        if m:
+            block_num, attn_type, sub_layer = m.groups()
+            return f"diffusion_model.blocks.{block_num}.{attn_type}.{sub_layer}.alpha"
+
+        m = re.match(r"lora_unet_blocks_(\d+)_ffn_(\d+)\.alpha$", k)
+        if m:
+            block_num, ffn_num = m.groups()
+            return f"diffusion_model.blocks.{block_num}.ffn.{ffn_num}.alpha"
+
+        m = re.match(r"lora_unet_blocks_(\d+)_(\w+)_(\w+)\.alpha$", k)
+        if m:
+            block_num, layer1, layer2 = m.groups()
+            return f"diffusion_model.blocks.{block_num}.{layer1}.{layer2}.alpha"
+
+        m = re.match(
+            r"lora_unet_blocks_(\d+)_(cross_attn|self_attn)_(\w+)\.(lora_down|lora_up)\.weight$",
+            k,
+        )
+        if m:
+            block_num, attn_type, sub_layer, matrix = m.groups()
+            matrix_new = "lora_A" if matrix == "lora_down" else "lora_B"
+            return f"diffusion_model.blocks.{block_num}.{attn_type}.{sub_layer}.{matrix_new}.weight"
+
+        m = re.match(r"lora_unet_blocks_(\d+)_ffn_(\d+)\.(lora_down|lora_up)\.weight$", k)
+        if m:
+            block_num, ffn_num, matrix = m.groups()
+            matrix_new = "lora_A" if matrix == "lora_down" else "lora_B"
+            return f"diffusion_model.blocks.{block_num}.ffn.{ffn_num}.{matrix_new}.weight"
+
+        m = re.match(r"lora_unet_blocks_(\d+)_(\w+)_(\w+)\.(lora_down|lora_up)\.weight$", k)
+        if m:
+            block_num, layer1, layer2, matrix = m.groups()
+            matrix_new = "lora_A" if matrix == "lora_down" else "lora_B"
+            return f"diffusion_model.blocks.{block_num}.{layer1}.{layer2}.{matrix_new}.weight"
+
+        return None
+
+    # If earlier normalization got us into the target namespace, keep it.
+    if k.startswith("diffusion_model."):
+        return k
+
+    return None
+
+
+
+def _should_drop_key(raw_key: str, standardized_key: Optional[str], filter_img: bool, extra_patterns: Iterable[str]) -> bool:
+    raw_lower = raw_key.lower()
+    std_lower = (standardized_key or raw_key).lower()
+
+    if filter_img:
+        img_markers = (
+            "_img",
+            ".img_",
+            "img_emb",
+            "img_attn",
+            "clip_vision",
+            "clip.visual",
+            "clip_visual",
+        )
+        if any(marker in raw_lower for marker in img_markers) or any(marker in std_lower for marker in img_markers):
+            return True
+
+    for pattern in extra_patterns:
+        p = pattern.strip().lower()
+        if not p:
+            continue
+        if p in raw_lower or p in std_lower:
+            return True
+    return False
+
+
+
+def _bake_strength_linear(state_dict: Dict[str, torch.Tensor], strength: float) -> Dict[str, torch.Tensor]:
+    """
+    Bake a *linear* LoRA strength by scaling only the up/B side.
+    Scaling both A and B would square the effective strength.
+    """
+    if strength == 1.0:
+        return dict(state_dict)
+
+    baked: Dict[str, torch.Tensor] = {}
+    for key, tensor in state_dict.items():
+        if key.endswith(".lora_B.weight") or key.endswith(".lora_up.weight"):
+            scaled = tensor.to(torch.float32) * float(strength)
+            baked[key] = scaled.to(tensor.dtype)
+        else:
+            baked[key] = tensor
+    return baked
+
+
+
+def _convert_and_save(
+    input_path: str,
+    output_path: str,
+    baked_strength: float,
+    filter_img: bool,
+    extra_exclude: str = "",
+) -> Dict[str, object]:
+    src_sd, src_meta = _load_safetensors(input_path)
+    src_sd = _maybe_compensate_rs_lora(src_sd)
+
+    converted: Dict[str, torch.Tensor] = {}
+    filtered = 0
+    skipped_unmapped = 0
+    preserved = 0
+
+    extra_patterns = [p.strip() for p in extra_exclude.split(",") if p.strip()]
+
+    for raw_key, tensor in src_sd.items():
+        std_key = _standardize_key(raw_key)
+        if _should_drop_key(raw_key, std_key, filter_img=filter_img, extra_patterns=extra_patterns):
+            filtered += 1
+            continue
+
+        if std_key is None:
+            # Preserve already non-LoRA or uncommon keys only if they already live in target namespace.
+            # Otherwise skip because arbitrary passthrough keys are more likely to break WAN2.2 loading.
+            skipped_unmapped += 1
+            continue
+
+        if std_key in converted:
+            # Prefer the first occurrence; duplicates usually indicate multiple source aliases.
+            continue
+
+        converted[std_key] = tensor
+        preserved += 1
+
+    if not converted:
+        raise ValueError(
+            "No convertible WAN LoRA keys were produced. The file may not be a WAN2.1 LoRA in a supported format."
+        )
+
+    baked = _bake_strength_linear(converted, baked_strength)
+
+    meta = dict(src_meta)
+    meta.update(
+        {
+            "wan_toolkit.source_file": Path(input_path).name,
+            "wan_toolkit.converted_for": "WAN2.2",
+            "wan_toolkit.filter_img": str(bool(filter_img)).lower(),
+            "wan_toolkit.extra_exclude": extra_exclude,
+            "wan_toolkit.baked_strength": str(baked_strength),
+            "wan_toolkit.generated_at": datetime.utcnow().replace(microsecond=0).isoformat() + "Z",
+        }
+    )
+
+    output_path = str(Path(output_path))
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    save_file(baked, output_path, metadata=_ensure_metadata_str_dict(meta))
+
+    return {
+        "output_path": output_path,
+        "source_keys": len(src_sd),
+        "saved_keys": len(baked),
+        "filtered_keys": filtered,
+        "skipped_unmapped": skipped_unmapped,
+        "preserved_keys": preserved,
+        "baked_strength": baked_strength,
+        "detected_format": _detect_format(src_sd.keys()),
+    }
+
+
+# ============================================================
+# ComfyUI Node
+# ============================================================
+
+
+class WAN21ToWAN22HighLowConverter:
+    """
+    Convert a WAN2.1 LoRA from path input and emit baked WAN2.2 high/low files.
+    """
+
+    CATEGORY = "WAN/LoRA"
+    FUNCTION = "convert"
+    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
+    RETURN_NAMES = ("high_path", "low_path", "output_dir", "report")
+    OUTPUT_NODE = True
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "input_lora_path": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "multiline": False,
+                        "placeholder": "/full/path/to/wan21_lora.safetensors",
+                    },
+                ),
+                "high_strength": (
+                    "FLOAT",
+                    {"default": 1.75, "min": 0.0, "max": 10.0, "step": 0.05},
+                ),
+                "low_strength": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05},
+                ),
+            },
+            "optional": {
+                "output_dir": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "multiline": False,
+                        "placeholder": "leave empty = same folder as input",
+                    },
+                ),
+                "filter_img_keys": (
+                    "BOOLEAN",
+                    {"default": True},
+                ),
+                "extra_exclude": (
+                    "STRING",
+                    {
+                        "default": "",
+                        "multiline": False,
+                        "placeholder": "comma-separated substrings to drop",
+                    },
+                ),
+            },
+        }
+
+    @classmethod
+    def IS_CHANGED(
+        cls,
+        input_lora_path,
+        high_strength,
+        low_strength,
+        output_dir="",
+        filter_img_keys=True,
+        extra_exclude="",
+    ):
+        path = _clean_path(input_lora_path)
+        if not path or not Path(path).exists():
+            return f"missing:{path}|{high_strength}|{low_strength}|{output_dir}|{filter_img_keys}|{extra_exclude}"
+        return (
+            f"{_file_signature(path)}|{high_strength}|{low_strength}|"
+            f"{_clean_path(output_dir)}|{bool(filter_img_keys)}|{extra_exclude}"
+        )
+
+    def convert(
+        self,
+        input_lora_path,
+        high_strength,
+        low_strength,
+        output_dir="",
+        filter_img_keys=True,
+        extra_exclude="",
+    ):
+        input_lora_path = _clean_path(input_lora_path)
+        output_dir = _clean_path(output_dir)
+
+        if not input_lora_path:
+            raise ValueError("input_lora_path is empty")
+
+        src = Path(input_lora_path)
+        if not src.exists():
+            raise FileNotFoundError(f"Input file not found: {input_lora_path}")
+        if src.suffix.lower() != ".safetensors":
+            raise ValueError("Input file must be a .safetensors file")
+
+        if not output_dir:
+            output_dir = str(src.parent)
+
+        out_dir = Path(output_dir)
+        out_dir.mkdir(parents=True, exist_ok=True)
+
+        base = src.stem
+        hi_path = str(out_dir / f"{base}_HI.safetensors")
+        lo_path = str(out_dir / f"{base}_LO.safetensors")
+
+        hi = _convert_and_save(
+            input_path=str(src),
+            output_path=hi_path,
+            baked_strength=float(high_strength),
+            filter_img=bool(filter_img_keys),
+            extra_exclude=extra_exclude,
+        )
+        lo = _convert_and_save(
+            input_path=str(src),
+            output_path=lo_path,
+            baked_strength=float(low_strength),
+            filter_img=bool(filter_img_keys),
+            extra_exclude=extra_exclude,
+        )
+
+        report = {
+            "source": str(src),
+            "output_dir": str(out_dir),
+            "high": hi,
+            "low": lo,
+        }
+
+        return (hi_path, lo_path, str(out_dir), json.dumps(report, indent=2))
+
+
+NODE_CLASS_MAPPINGS = {
+    "WAN21ToWAN22HighLowConverter": WAN21ToWAN22HighLowConverter,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "WAN21ToWAN22HighLowConverter": "WAN 2.1 → 2.2 LoRA Converter (HI/LO)",
+}