test

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -169,6 +169,10 @@ def _merge_lora_layer(module, weights, scale):
     if bias_diff is not None and isinstance(module, nnx.LayerNorm) and hasattr(module, "bias") and module.bias is not None:
       module.bias.value += np.array(bias_diff).reshape(module.bias.shape).astype(module.bias.dtype)
       updated = True
+  elif isinstance(module, nnx.Param):
+    if "diff" in weights:
+      module.value += np.array(weights["diff"]).reshape(module.shape).astype(module.dtype)
+      updated = True
   elif isinstance(module, (nnx.Linear, nnx.Conv)):
     # Prepare LoRA terms
     down_w, up_w, current_scale = None, None, None
@@ -252,7 +256,7 @@ def merge_lora(model: nnx.Module, state_dict: dict, rank: int, scale: float, tra
 
   assigned_count = 0
   for path, module in nnx.iter_graph(model):
-    if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed)):
+    if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed, nnx.Param)):
       continue
 
     nnx_path_str = ".".join(map(str, path))
@@ -285,23 +289,19 @@ def merge_lora_for_scanned(
 
   assigned_count = 0
   for path, module in nnx.iter_graph(model):
+    if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed, nnx.Param)):
+      continue
 
     nnx_path_str = ".".join(map(str, path))
     lora_key_template = translate_fn(nnx_path_str) if translate_fn else None
-    max_logging.log(f"nnx_path_str: {nnx_path_str}")
-    if 'scale_shift_table' in nnx_path_str:
+
+    if 'adaln_scale_shift_table' in nnx_path_str:
       max_logging.log(f"adaln_scale_shift_table: {nnx_path_str}")
       max_logging.log(f"lora_key_template: {lora_key_template}")
-      max_logging.log(f"Module: {module}")
+      max_logging.log(f"Module ndim: {module.ndim}")
       max_logging.log(f"Module Type: {type(module)}")
       continue
 
-    if not isinstance(module, (nnx.Linear, nnx.Conv, nnx.LayerNorm, nnx.RMSNorm, nnx.Embed)):
-      continue
-
-    nnx_path_str = ".".join(map(str, path))
-    lora_key_template = translate_fn(nnx_path_str) if translate_fn else None
-
     if not lora_key_template:
       continue
 
@@ -315,6 +315,8 @@ def merge_lora_for_scanned(
       is_scanned = module.kernel.ndim == 3
     elif isinstance(module, nnx.Conv):
       is_scanned = module.kernel.ndim == 5
+    elif isinstance(module, nnx.Param):
+      is_scanned = module.ndim > 1  # Heuristic for scanned Param
 
     if not is_scanned:
       lora_key = lora_key_template
@@ -366,6 +368,20 @@ def merge_lora_for_scanned(
         module.bias.value += bias_diffs_to_add.astype(module.bias.dtype)
       if updated_scale or updated_bias:
         assigned_count += 1
+    elif isinstance(module, nnx.Param):
+      num_layers = module.shape[0]
+      param_diffs_to_add = np.zeros_like(module.value)
+      updated = False
+      for i in range(num_layers):
+        lora_key = lora_key_template.format(i)
+        if lora_key in lora_params:
+          matched_keys.add(lora_key)
+          if "diff" in lora_params[lora_key]:
+            param_diffs_to_add[i] = np.array(lora_params[lora_key]["diff"]).reshape(module.shape[1:])
+            updated = True
+      if updated:
+        module.value += param_diffs_to_add.astype(module.dtype)
+        assigned_count += 1
     elif isinstance(module, (nnx.Linear, nnx.Conv)):
       is_linear = isinstance(module, nnx.Linear)
       is_conv = isinstance(module, nnx.Conv)