Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -324,8 +324,8 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
 
             if lora_key_template:
                 num_layers, in_features, out_features = module.kernel.shape
-                deltas = []
-                has_lora = False
+                kernel_value_updated = module.kernel.value
+                lora_found_in_module = 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]:
@@ -335,15 +335,13 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
                         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
+                        kernel_value_updated = kernel_value_updated.at[i].add(delta_i)
+                        lora_found_in_module = True
+
+                if lora_found_in_module:
+                    module.kernel.value = kernel_value_updated
                     assigned_count += 1
+                    max_logging.log(f"Merged LoRA into scanned layer {nnx_path_str}")
                 else:
                     max_logging.log(f"Scanned layer {nnx_path_str} matched template but no LoRA weights found for any block.")
             else: