Fix

Author: Perseus14

src/maxdiffusion/models/lora_nnx.py

@@ -224,13 +224,34 @@ def merge_lora(model: nnx.Module, state_dict: dict, scale: float, translate_fn=N
 
     assigned_count = 0
     for path, module in nnx.iter_graph(model):
-        if not isinstance(module, (nnx.Linear, nnx.Conv)): 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 = translate_fn(nnx_path_str) if translate_fn else None
 
         if lora_key and lora_key in lora_params:
             weights = lora_params[lora_key]
+
+            # Handle Embeddings
+            if isinstance(module, nnx.Embed):
+                if "diff" in weights and hasattr(module, 'embedding'):
+                    module.embedding.value += np.array(weights["diff"]).reshape(module.embedding.shape).astype(module.embedding.dtype)
+                    assigned_count += 1
+                continue
+            # Handle Norms
+            elif isinstance(module, (nnx.LayerNorm, nnx.RMSNorm)):
+                scale_diff = weights.get("diff", None)
+                bias_diff = weights.get("diff_b", None)
+                updated = False
+                if scale_diff is not None and hasattr(module, 'scale') and module.scale is not None:
+                    module.scale.value += np.array(scale_diff).reshape(module.scale.shape).astype(module.scale.dtype)
+                    updated = True
+                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
+                if updated:
+                    assigned_count += 1
+                continue
 
             # Prepare LoRA terms
             down_w, up_w, current_scale = None, None, None
@@ -298,14 +319,53 @@ def merge_lora_for_scanned(model: nnx.Module, state_dict: dict, scale: float, tr
 
     assigned_count = 0
     for path, module in nnx.iter_graph(model):
-        if not isinstance(module, (nnx.Linear, nnx.Conv)): 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
 
+        # Handle Scanned Embeddings
+        if isinstance(module, nnx.Embed) and hasattr(module, 'embedding') and module.embedding.ndim > 2:
+            num_layers = module.embedding.shape[0]
+            embed_diffs_to_add = np.zeros_like(module.embedding.value)
+            updated = False
+            for i in range(num_layers):
+                lora_key = lora_key_template.format(i)
+                if lora_key in lora_params and "diff" in lora_params[lora_key]:
+                    embed_diffs_to_add[i] = np.array(lora_params[lora_key]["diff"]).reshape(module.embedding.shape[1:])
+                    updated = True
+            if updated:
+                module.embedding.value += embed_diffs_to_add.astype(module.embedding.dtype)
+                assigned_count += 1
+            continue
+
+        # Handle Scanned Norms
+        if isinstance(module, (nnx.LayerNorm, nnx.RMSNorm)) and hasattr(module, 'scale') and module.scale is not None and module.scale.ndim > 1:
+            num_layers = module.scale.shape[0]
+            scale_diffs_to_add = np.zeros_like(module.scale.value)
+            bias_diffs_to_add = np.zeros_like(module.bias.value) if isinstance(module, nnx.LayerNorm) and hasattr(module, 'bias') and module.bias is not None else None
+            updated_scale, updated_bias = False, False
+            for i in range(num_layers):
+                lora_key = lora_key_template.format(i)
+                if lora_key in lora_params:
+                    weights = lora_params[lora_key]
+                    if "diff" in weights:
+                        scale_diffs_to_add[i] = np.array(weights["diff"]).reshape(module.scale.shape[1:])
+                        updated_scale = True
+                    if "diff_b" in weights and bias_diffs_to_add is not None:
+                        bias_diffs_to_add[i] = np.array(weights["diff_b"]).reshape(module.bias.shape[1:])
+                        updated_bias = True
+            if updated_scale:
+                module.scale.value += scale_diffs_to_add.astype(module.scale.dtype)
+            if updated_bias and bias_diffs_to_add is not None:
+                module.bias.value += bias_diffs_to_add.astype(module.bias.dtype)
+            if updated_scale or updated_bias:
+                assigned_count += 1
+            continue
+
         is_linear = isinstance(module, nnx.Linear) and module.kernel.ndim == 3
         is_conv   = isinstance(module, nnx.Conv) and module.kernel.ndim == 5