wan_utils.py fixed

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -256,16 +256,29 @@ 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:
-          if "net.0" in renamed_pt_key:
-              renamed_pt_key = renamed_pt_key.replace("net.0", "net_0.proj")
-              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")  
+          # 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")
+          
+          # 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
           elif "net.2" in renamed_pt_key:
-              renamed_pt_key = renamed_pt_key.replace("net.2", "net_2.proj")
-              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+              renamed_pt_key = renamed_pt_key.replace("net.2", "net_2")
+
+          # 3. Fix Norm1 (Add .layer_norm wrapper)
           renamed_pt_key = renamed_pt_key.replace("norm1", "norm1.layer_norm")
+
+          # 4. Fix Norm Parameter Names (weight/kernel -> 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")
       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")