Fix

Author: Perseus14

src/maxdiffusion/loaders/wan_lora_nnx_loader.py

@@ -16,11 +16,78 @@
 
 from flax import nnx
 import jax
+import re
 from .lora_base import LoRABaseMixin
 from .lora_pipeline import StableDiffusionLoraLoaderMixin
 from ..models import lora_nnx
 from .. import max_logging
 
+def _translate_nnx_path_to_lora_key(nnx_path_str):
+    """
+    Translates NNX path like 'blocks.10.attn1.key' to
+    LoRA path like 'diffusion_model.blocks.10.self_attn.k'.
+    Returns None if no match.
+    """
+    translation_map = {
+        "attn1": "self_attn",
+        "attn2": "cross_attn",
+        "query": "q",
+        "key": "k",
+        "value": "v",
+        "proj_attn": "o",
+        "ffn.act_fn.proj": "ffn.0",
+        "ffn.proj_out": "ffn.2",
+    }
+    # Match paths like blocks.10.attn1.key or blocks.5.ffn.proj_out
+    m = re.match(r"^blocks\.(\d+)\.(attn[12]\.(?:query|key|value|proj_attn)|ffn\.(?:act_fn\.proj|proj_out))$", nnx_path_str)
+    if not m:
+        return None
+
+    block_idx, suffix = m.group(1), m.group(2)
+
+    parts = suffix.split('.')
+    if parts[0] == 'attn1' or parts[0] == 'attn2':
+        lora_part1 = translation_map[parts[0]]
+        lora_part2 = translation_map[parts[1]]
+        return f"diffusion_model.blocks.{block_idx}.{lora_part1}.{lora_part2}"
+    elif suffix in translation_map:
+         return f"diffusion_model.blocks.{block_idx}.{translation_map[suffix]}"
+    return None
+
+
+def _translate_scanned_nnx_path_to_lora_key_template(nnx_path_str):
+    """
+    Translates NNX path like 'blocks.attn1.key' to
+    LoRA path template like 'diffusion_model.blocks.{}.self_attn.k'.
+    Returns None if no match.
+    This version assumes block index is missing from path due to scan.
+    """
+    translation_map = {
+        "attn1": "self_attn",
+        "attn2": "cross_attn",
+        "query": "q",
+        "key": "k",
+        "value": "v",
+        "proj_attn": "o",
+        "ffn.act_fn.proj": "ffn.0",
+        "ffn.proj_out": "ffn.2",
+    }
+    # Match paths like blocks.attn1.key or blocks.ffn.proj_out, missing block index
+    m = re.match(r"^blocks\.(attn[12]\.(?:query|key|value|proj_attn)|ffn\.(?:act_fn\.proj|proj_out))$", nnx_path_str)
+    if not m:
+        return None
+
+    suffix = m.group(1)
+
+    parts = suffix.split('.')
+    if parts[0] == 'attn1' or parts[0] == 'attn2':
+        lora_part1 = translation_map[parts[0]]
+        lora_part2 = translation_map[parts[1]]
+        return f"diffusion_model.blocks.{{}}.{lora_part1}.{lora_part2}"
+    elif suffix in translation_map:
+         return f"diffusion_model.blocks.{{}}.{translation_map[suffix]}"
+    return None
+
 class WanNnxLoraLoader(LoRABaseMixin):
   """
   Handles loading LoRA weights into NNX-based WAN models.
@@ -36,20 +103,24 @@ def load_lora_weights(
       low_noise_weight_name: str,
       rank: int,
       scale: float = 1.0,
+      scan_layers: bool = False,
       **kwargs,
   ):
     """
     Merges LoRA weights into the pipeline from a checkpoint.
     """
     lora_loader = StableDiffusionLoraLoaderMixin()
 
+    merge_fn = lora_nnx.merge_lora_for_scanned if scan_layers else lora_nnx.merge_lora
+    translate_fn = _translate_scanned_nnx_path_to_lora_key_template if scan_layers else _translate_nnx_path_to_lora_key
+
     # Handle high noise model
     if hasattr(pipeline, "high_noise_transformer") and high_noise_weight_name:
         max_logging.log(f"Merging LoRA into high_noise_transformer with rank={rank}")
         h_state_dict, _ = lora_loader.lora_state_dict(
             lora_model_path, weight_name=high_noise_weight_name, **kwargs
         )
-        lora_nnx.merge_lora(pipeline.high_noise_transformer, h_state_dict, scale)
+        merge_fn(pipeline.high_noise_transformer, h_state_dict, scale, translate_fn)
     else:
         max_logging.log("high_noise_transformer not found or no weight name provided for LoRA.")
 
@@ -59,7 +130,7 @@ def load_lora_weights(
         l_state_dict, _ = lora_loader.lora_state_dict(
             lora_model_path, weight_name=low_noise_weight_name, **kwargs
         )
-        lora_nnx.merge_lora(pipeline.low_noise_transformer, l_state_dict, scale)
+        merge_fn(pipeline.low_noise_transformer, l_state_dict, scale, translate_fn)
     else:
         max_logging.log("low_noise_transformer not found or no weight name provided for LoRA.")
 

src/maxdiffusion/models/lora_nnx.py

@@ -200,40 +200,7 @@ def inject_lora(
 
     return model
 
-def _translate_nnx_path_to_lora_key(nnx_path_str):
-    """
-    Translates NNX path like 'blocks.10.attn1.key' to
-    LoRA path like 'diffusion_model.blocks.10.self_attn.k'.
-    Returns None if no match.
-    """
-    translation_map = {
-        "attn1": "self_attn",
-        "attn2": "cross_attn",
-        "query": "q",
-        "key": "k",
-        "value": "v",
-        "proj_attn": "o",
-        "ffn.act_fn.proj": "ffn.0",
-        "ffn.proj_out": "ffn.2",
-    }
-    # Match paths like blocks.10.attn1.key or blocks.5.ffn.proj_out
-    m = re.match(r"^blocks\.(\d+)\.(attn[12]\.(?:query|key|value|proj_attn)|ffn\.(?:act_fn\.proj|proj_out))$", nnx_path_str)
-    if not m:
-        return None
-
-    block_idx, suffix = m.group(1), m.group(2)
-
-    parts = suffix.split('.')
-    if parts[0] == 'attn1' or parts[0] == 'attn2':
-        lora_part1 = translation_map[parts[0]]
-        lora_part2 = translation_map[parts[1]]
-        return f"diffusion_model.blocks.{block_idx}.{lora_part1}.{lora_part2}"
-    elif suffix in translation_map:
-         return f"diffusion_model.blocks.{block_idx}.{translation_map[suffix]}"
-    return None
-
-
-def merge_lora(model: nnx.Module, state_dict: dict, scale: float):
+def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=None):
     """
     Merges weights from a Diffusers-formatted state dict directly
     into the kernel of nnx.Linear and nnx.Conv layers.
@@ -271,15 +238,12 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float):
     assigned_count = 0
     for path, module in nnx.iter_graph(model):
         if not isinstance(module, (nnx.Linear, nnx.Conv)):
-            max_logging.log(f"Skipping non-Linear/Conv layer: {module}")
             continue
 
         nnx_path_str = ".".join(map(str, path))
-        max_logging.log(f"Checking NNX layer: {nnx_path_str}")
-        lora_key = _translate_nnx_path_to_lora_key(nnx_path_str)
+        lora_key = translate_fn(nnx_path_str) if translate_fn else None
 
         if lora_key and lora_key in lora_params:
-            max_logging.log(f"NNX layer '{nnx_path_str}' matched LoRA key '{lora_key}'")
             weights = lora_params[lora_key]
             if "down" in weights and "up" in weights:
                 if isinstance(module, nnx.Linear):
@@ -308,4 +272,85 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float):
         else:
             max_logging.log(f"NNX layer '{nnx_path_str}' could not be translated to a LoRA key.")
 
-    max_logging.log(f"Merged weights into {assigned_count} layers in {type(model).__name__}.")
+    max_logging.log(f"Merged weights into {assigned_count} layers in {type(model).__name__}.")
+
+
+def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, translate_fn=None):
+    """
+    Merges weights from a Diffusers-formatted state dict directly
+    into the kernel of nnx.Linear and nnx.Conv layers.
+    Assumes scan_layers=True, so weights are stacked if layers are scanned
+    (e.g. kernel.ndim=3 for Linear).
+    """
+    lora_params = {}
+    # Parse weights and alphas
+    for k, v in state_dict.items():
+        if k.endswith(".alpha"):
+            key_base = k[:-len(".alpha")]
+            if key_base not in lora_params:
+                lora_params[key_base] = {}
+            lora_params[key_base]["alpha"] = jnp.array(v)
+            continue
+
+        m = re.match(r"^(.*?)_lora\.(down|up)\.weight$", k)
+        if m:
+            key_base, weight_type = m.group(1), m.group(2)
+        else:
+            m = re.match(r"^(.*?)\.lora\.(down|up)\.weight$", k)
+            if m:
+                key_base, weight_type = m.group(1), m.group(2)
+            else:
+                m = re.match(r"^(.*?)\.(lora_down|lora_up)\.weight$", k)
+                if m:
+                    key_base, weight_type = m.group(1), m.group(2).replace("lora_", "")
+                else:
+                    max_logging.log(f"Could not parse LoRA key: {k}")
+                    continue
+        if key_base not in lora_params:
+            lora_params[key_base] = {}
+        lora_params[key_base][weight_type] = jnp.array(v)
+    max_logging.log(f"Parsed {len(lora_params)} unique LoRA module keys for scanned merge.")
+
+    assigned_count = 0
+    for path, module in nnx.iter_graph(model):
+        if not isinstance(module, (nnx.Linear, nnx.Conv)):
+            continue
+
+        nnx_path_str = ".".join(map(str, path))
+        
+        # Handle scanned Linear layers
+        if isinstance(module, nnx.Linear) and module.kernel.ndim == 3:
+            lora_key_template = translate_fn(nnx_path_str) if translate_fn else None
+
+            if lora_key_template:
+                num_layers, in_features, out_features = module.kernel.shape
+                deltas = []
+                has_lora = False
+                for i in range(num_layers):
+                    lora_key = lora_key_template.format(i)
+                    if lora_key in lora_params and "down" in lora_params[lora_key] and "up" in lora_params[lora_key]:
+                        weights = lora_params[lora_key]
+                        down_w, up_w = weights["down"], weights["up"]
+                        rank = down_w.shape[0]
+                        alpha = weights.get("alpha", rank)
+                        current_scale = scale * alpha / rank
+                        delta_i = (down_w.T @ up_w.T).reshape(in_features, out_features) * current_scale
+                        deltas.append(delta_i)
+                        has_lora = True
+                    else:
+                        deltas.append(jnp.zeros((in_features, out_features), dtype=module.kernel.dtype))
+                
+                if has_lora:
+                    stacked_delta = jnp.stack(deltas, axis=0)
+                    module.kernel.value += stacked_delta
+                    assigned_count += 1
+                else:
+                    max_logging.log(f"Scanned layer {nnx_path_str} matched template but no LoRA weights found for any block.")
+            else:
+                 max_logging.log(f"Scanned NNX layer '{nnx_path_str}' could not be translated to a LoRA key template.")
+
+        # Handle scanned Conv layers (ndim=5)
+        elif isinstance(module, nnx.Conv) and module.kernel.ndim == 5:
+             max_logging.warn(f"Merging LoRA into scanned Conv layers not implemented: {nnx_path_str}")
+
+    max_logging.log(f"Merged weights into {assigned_count} scanned layers in {type(model).__name__}.")