Fix

Author: Perseus14

src/maxdiffusion/loaders/wan_lora_nnx_loader.py

@@ -52,79 +52,22 @@ def load_lora_weights(
 
     # Handle high noise model
     if hasattr(pipeline, "high_noise_transformer") and high_noise_weight_name:
-        max_logging.log(f"Injecting LoRA into high_noise_transformer with rank={rank}")
-        lora_nnx.inject_lora(
-            pipeline.high_noise_transformer, rank=rank, scale=scale, rngs=nnx.Rngs(rng), target_linear=True, target_conv=True
-        )
-        h_state_dict, h_alphas = lora_loader.lora_state_dict(
+        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
         )
-        self._assign_weights_to_nnx_model(pipeline.high_noise_transformer, h_state_dict, h_alphas if h_alphas else {})
+        lora_nnx.merge_lora(pipeline.high_noise_transformer, h_state_dict, scale)
     else:
         max_logging.warning("high_noise_transformer not found or no weight name provided for LoRA.")
 
     # Handle low noise model
     if hasattr(pipeline, "low_noise_transformer") and low_noise_weight_name:
-        max_logging.log(f"Injecting LoRA into low_noise_transformer with rank={rank}")
-        lora_nnx.inject_lora(
-            pipeline.low_noise_transformer, rank=rank, scale=scale, rngs=nnx.Rngs(rng), target_linear=True, target_conv=True
-        )
-        l_state_dict, l_alphas = lora_loader.lora_state_dict(
+        max_logging.log(f"Merging LoRA into low_noise_transformer with rank={rank}")
+        l_state_dict, _ = lora_loader.lora_state_dict(
             lora_model_path, weight_name=low_noise_weight_name, **kwargs
         )
-        self._assign_weights_to_nnx_model(pipeline.low_noise_transformer, l_state_dict, l_alphas if l_alphas else {})
+        lora_nnx.merge_lora(pipeline.low_noise_transformer, l_state_dict, scale)
     else:
         max_logging.warning("low_noise_transformer not found or no weight name provided for LoRA.")
 
     return pipeline
-
-  def _assign_weights_to_nnx_model(self, model: nnx.Module, state_dict: dict, network_alphas: dict):
-    """
-    Assigns weights from a Diffusers-formatted state dict to
-    injected LoRALinear/LoRAConv layers in an NNX model.
-    """
-    lora_params = {}
-    for k, v in state_dict.items():
-      m = re.match(r"^(.*?)_lora\.(down|up)\.weight$", k)
-      if not m:
-        m = re.match(r"^(.*?)\.lora\.(down|up)\.weight$", k)
-
-      if m:
-        module_path_str, weight_type = m.group(1), m.group(2)
-        if module_path_str not in lora_params:
-          lora_params[module_path_str] = {}
-        lora_params[module_path_str][weight_type] = jnp.array(v)
-      else:
-        max_logging.warning(f"Could not parse LoRA key: {k}")
-
-    assigned_count = 0
-    for path, submodule in nnx.iter_graph(model):
-        if isinstance(submodule, (lora_nnx.LoRALinear, lora_nnx.LoRAConv)):
-            nnx_path_str = ".".join(map(str, path))
-            
-            matched_key = None
-            if nnx_path_str in lora_params:
-              matched_key = nnx_path_str
-            else:
-              # Fallback: check if any param key matches end of nnx path
-              for k in lora_params:
-                if nnx_path_str.endswith(k) or k.endswith(nnx_path_str):
-                   matched_key = k
-                   break
-            
-            if matched_key and matched_key in lora_params:
-              weights = lora_params[matched_key]
-              if "down" in weights and "up" in weights:
-                  if isinstance(submodule, lora_nnx.LoRALinear):
-                      submodule.A.value = weights["down"].T
-                      submodule.B.value = weights["up"].T
-                      assigned_count +=1
-                  elif isinstance(submodule, lora_nnx.LoRAConv):
-                      submodule.down.kernel.value = weights["down"]
-                      submodule.up.kernel.value = weights["up"]
-                      assigned_count += 1
-                  
-                  pass
-              else:
-                  max_logging.warning(f"LoRA weights for {matched_key} incomplete.")
-    max_logging.log(f"Assigned weights to {assigned_count} LoRA layers in {type(model)}.")

src/maxdiffusion/models/lora_nnx.py

@@ -15,8 +15,10 @@
 """
 
 from typing import Union, Tuple, Optional
+import re
 import jax.numpy as jnp
 from flax import nnx
+from .. import max_logging
 
 class BaseLoRALayer(nnx.Module):
   """
@@ -196,4 +198,64 @@ def inject_lora(
             wrapper = LoRAConv(base_layer=module, rank=rank, scale=scale, rngs=rngs)
             setattr(parent, attr_name, wrapper)
 
-    return model
+    return model
+
+def merge_lora(model: nnx.Module, state_dict: dict, scale: float):
+    """
+    Merges weights from a Diffusers-formatted state dict directly
+    into the kernel of nnx.Linear and nnx.Conv layers.
+    """
+    lora_params = {}
+    for k, v in state_dict.items():
+      m = re.match(r"^(.*?)_lora\.(down|up)\.weight$", k)
+      if not m:
+        m = re.match(r"^(.*?)\.lora\.(down|up)\.weight$", k)
+
+      if m:
+        module_path_str, weight_type = m.group(1), m.group(2)
+        if module_path_str not in lora_params:
+          lora_params[module_path_str] = {}
+        lora_params[module_path_str][weight_type] = jnp.array(v)
+      else:
+        max_logging.warning(f"Could not parse LoRA key: {k}")
+
+    assigned_count = 0
+    for path, module in nnx.iter_graph(model):
+        nnx_path_str = ".".join(map(str, path))
+
+        matched_key = None
+        if nnx_path_str in lora_params:
+          matched_key = nnx_path_str
+        else:
+          # Fallback: check if any param key matches end of nnx path
+          for k in lora_params:
+            if nnx_path_str.endswith(k):
+                matched_key = k
+                break
+
+        if matched_key and matched_key in lora_params:
+          weights = lora_params[matched_key]
+          if "down" in weights and "up" in weights:
+              if isinstance(module, nnx.Linear):
+                  down_w = weights["down"] # (rank, in_features)
+                  up_w = weights["up"]     # (out_features_flat, rank)
+                  # delta = A@B = down.T @ up.T
+                  delta = (down_w.T @ up_w.T).reshape(module.kernel.shape)
+                  module.kernel.value += delta * scale
+                  assigned_count +=1
+              elif isinstance(module, nnx.Conv):
+                  if module.kernel_size == (1, 1):
+                      down_w = weights["down"] # (1,1,in_c,rank)
+                      up_w = weights["up"]     # (1,1,rank,out_c)
+                      # delta = down @ up for channel dimension
+                      delta = (jnp.squeeze(down_w) @ jnp.squeeze(up_w)).reshape(module.kernel.shape)
+                      module.kernel.value += delta * scale
+                      assigned_count += 1
+                  else:
+                      raise NotImplementedError(
+                          f"Merging LoRA weights for Conv layer {matched_key} "
+                          f"with kernel_size {module.kernel_size} > 1 is not supported."
+                      )
+          else:
+              max_logging.warning(f"LoRA weights for {matched_key} incomplete.")
+    max_logging.log(f"Merged weights into {assigned_count} layers in {type(model).__name__}.")