wan_utils.py fixed

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -256,29 +256,33 @@ def load_base_wan_transformer(
     for pt_key, tensor in tensors.items():
       renamed_pt_key = rename_key(pt_key)
       if "image_embedder" in renamed_pt_key:
-          # 1. Handle Layer 0 (HAS .proj in PyTorch -> Remove it for JAX)
-          # Source: condition_embedder.image_embedder.ff.net.0.proj.weight
-          # Target: condition_embedder.image_embedder.ff.net_0.kernel
-          if "net.0.proj" in renamed_pt_key:
-              renamed_pt_key = renamed_pt_key.replace("net.0.proj", "net_0")
+          # 1. Handle Layer 0: "net.0" -> "net_0" (PRESERVE .proj)
+          # Source: ...ff.net.0.proj.weight -> Target: ...ff.net_0.proj.kernel
+          if "net.0" in renamed_pt_key:
+              renamed_pt_key = renamed_pt_key.replace("net.0", "net_0")
+              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
 
-          # 2. Handle Layer 2 (NO .proj in PyTorch -> Standard rename)
-          # Source: condition_embedder.image_embedder.ff.net.2.weight
-          # Target: condition_embedder.image_embedder.ff.net_2.kernel
+          # 2. Handle Layer 2: "net.2" -> "net_2"
+          # Source: ...ff.net.2.weight -> Target: ...ff.net_2.kernel
           elif "net.2" in renamed_pt_key:
               renamed_pt_key = renamed_pt_key.replace("net.2", "net_2")
+              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
 
-          # 3. Fix Norm1 (Add .layer_norm wrapper)
+          # 3. Handle Norm1: "norm1" -> "norm1.layer_norm"
           renamed_pt_key = renamed_pt_key.replace("norm1", "norm1.layer_norm")
 
-          # 4. Fix Norm Parameter Names (weight/kernel -> scale)
+          # 4. Force Norms to use "scale"
           if "norm1" in renamed_pt_key or "norm2" in renamed_pt_key:
               renamed_pt_key = renamed_pt_key.replace("weight", "scale")
               renamed_pt_key = renamed_pt_key.replace("kernel", "scale")
-              
-          # 5. Ensure Dense Weights use 'kernel' (Fixes weight->kernel mapping if missed)
-          if "net_0" in renamed_pt_key or "net_2" in renamed_pt_key:
-              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+      
+      # 5. Global Norm Fix (Fixes 'norm_added_q', 'norm_k', etc.)
+      # Any key containing 'norm' that ends in 'weight'/'kernel' should likely be 'scale'
+      if "norm" in renamed_pt_key and ("weight" in renamed_pt_key or "kernel" in renamed_pt_key):
+           # Exclude 'norm2' if it's already handled, or specific dense layers that might be named norm (unlikely)
+           if "norm_added" in renamed_pt_key or "norm_k" in renamed_pt_key or "norm_q" in renamed_pt_key:
+                renamed_pt_key = renamed_pt_key.replace("weight", "scale")
+                renamed_pt_key = renamed_pt_key.replace("kernel", "scale")
       renamed_pt_key = renamed_pt_key.replace("blocks_", "blocks.")
       renamed_pt_key = renamed_pt_key.replace(".scale_shift_table", ".adaln_scale_shift_table")
       renamed_pt_key = renamed_pt_key.replace("to_out_0", "proj_attn")