Fix WAN2.1 lora

Author: Perseus14

src/maxdiffusion/loaders/lora_conversion_utils.py

@@ -618,35 +618,58 @@ def translate_wan_nnx_path_to_diffusers_lora(nnx_path_str, scan_layers=False):
     template 'diffusion_model.blocks.{}.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",
+    
+    # Handle embeddings - exact paths
+    if nnx_path_str == "patch_embedding":
+        return "diffusion_model.patch_embedding"
+    if nnx_path_str == 'condition_embedder.text_embedder.linear_1':
+        return 'diffusion_model.text_embedding.0'
+    if nnx_path_str == 'condition_embedder.text_embedder.linear_2':
+        return 'diffusion_model.text_embedding.2'
+    if nnx_path_str == 'condition_embedder.time_embedder.linear_1':
+        return 'diffusion_model.time_embedding.0'
+    if nnx_path_str == 'condition_embedder.time_embedder.linear_2':
+        return 'diffusion_model.time_embedding.2'
+
+    # Translation for Attention and FFN layers
+    attn_ffn_map = {
+        "attn1.query": "self_attn.q",
+        "attn1.key": "self_attn.k",
+        "attn1.value": "self_attn.v",
+        "attn1.proj_attn": "self_attn.o",
+        "attn2.query": "cross_attn.q",
+        "attn2.key": "cross_attn.k",
+        "attn2.value": "cross_attn.v",
+        "attn2.proj_attn": "cross_attn.o",
         "ffn.act_fn.proj": "ffn.0",
         "ffn.proj_out": "ffn.2",
     }
-    suffix_pattern = r"(attn[12]\.(?:query|key|value|proj_attn)|ffn\.(?:act_fn\.proj|proj_out))"
+    # Translation for Norm layers
+    norm_map = {
+        "norm3": "norm3",
+        "attn1.norm_q": "self_attn.norm_q",
+        "attn1.norm_k": "self_attn.norm_k",
+        "attn2.norm_q": "cross_attn.norm_q",
+        "attn2.norm_k": "cross_attn.norm_k",
+    }
+
     if scan_layers:
-        m = re.match(r"^blocks\." + suffix_pattern + "$", nnx_path_str)
-        if not m:
-            return None
-        block_idx_str = "{}"
-        suffix = m.group(1)
+        # Handle scanned attn/ffn: blocks.attn1.query -> diffusion_model.blocks.{}.self_attn.q
+        for k, v in attn_ffn_map.items():
+            if nnx_path_str == f"blocks.{k}":
+                return f"diffusion_model.blocks.{{}}.{v}"
+        # Handle scanned norm: blocks.norm3 -> diffusion_model.blocks.{}.norm3
+        for k, v in norm_map.items():
+            if nnx_path_str == f"blocks.{k}":
+                return f"diffusion_model.blocks.{{}}.{v}"
     else:
-        m = re.match(r"^blocks\.(\d+)\." + suffix_pattern + "$", nnx_path_str)
-        if not m:
-            return None
-        block_idx_str = m.group(1)
-        suffix = 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_str}.{lora_part1}.{lora_part2}"
-    elif suffix in translation_map:
-         return f"diffusion_model.blocks.{block_idx_str}.{translation_map[suffix]}"
+        # Handle non-scanned attn/ffn/norm: blocks.0.attn1.query -> diffusion_model.blocks.0.self_attn.q
+        m = re.match(r"^blocks\.(\d+)\.(.+)$", nnx_path_str)
+        if m:
+            idx, suffix = m.group(1), m.group(2)
+            if suffix in attn_ffn_map:
+                return f"diffusion_model.blocks.{idx}.{attn_ffn_map[suffix]}"
+            if suffix in norm_map:
+                return f"diffusion_model.blocks.{idx}.{norm_map[suffix]}"
+    
     return None