wan_utils.py fixed

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -256,27 +256,38 @@ 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: PyTorch has "net.0.proj" -> JAX wants "net_0.proj"
-          # We just replace the separator "net.0" -> "net_0"
-          if "net.0" in renamed_pt_key:
-              renamed_pt_key = renamed_pt_key.replace("net.0", "net_0")
-          
-          # 2. Handle Layer 2: PyTorch has "net.2" (NO proj) -> JAX likely wants "net_2.proj"
-          # We force the addition of ".proj" to match the symmetric JAX structure
-          elif "net.2" in renamed_pt_key:
-              renamed_pt_key = renamed_pt_key.replace("net.2", "net_2.proj")
-
-          # 3. Handle Norm1: "norm1" -> "norm1.layer_norm"
-          renamed_pt_key = renamed_pt_key.replace("norm1", "norm1.layer_norm")
-
-          # 4. Fix Parameter Names:
-          # Norms (norm1, norm2) -> force '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")
-          # Dense Layers (net_0, net_2) -> force 'kernel'
-          elif "net_0" in renamed_pt_key or "net_2" in renamed_pt_key:
-               renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+        # 1. Handle Layer 0: PyTorch "net.0.proj" -> JAX "net_0"
+        # The JAX model likely defines net_0 as a Dense layer directly, not a block with a 'proj' sub-layer.
+        # We must strip ".proj" out.
+        if "net.0.proj" in renamed_pt_key:
+            renamed_pt_key = renamed_pt_key.replace("net.0.proj", "net_0")
+        elif "net.0" in renamed_pt_key:
+            renamed_pt_key = renamed_pt_key.replace("net.0", "net_0")
+
+        # 2. Handle Layer 2: PyTorch "net.2" -> JAX "net_2"
+        # Do NOT add ".proj" here unless your JAX definition explicitly has it. 
+        # The shape mismatch error suggests the model found 'net_2' but the shape was wrong.
+        if "net.2" in renamed_pt_key:
+            renamed_pt_key = renamed_pt_key.replace("net.2", "net_2")
+
+        # 3. Handle Norms: "norm1" -> "norm1.layer_norm" (Keeping your logic if JAX uses this structure)
+        # If standard Flax LayerNorm is used, usually just "norm1" is sufficient.
+        # Ensure "weight" becomes "scale".
+        if "norm" in renamed_pt_key:
+            renamed_pt_key = renamed_pt_key.replace("norm1", "norm1.layer_norm")
+            renamed_pt_key = renamed_pt_key.replace("norm2", "norm2.layer_norm")
+            renamed_pt_key = renamed_pt_key.replace("weight", "scale")
+
+        # 4. Handle Dense Layers (Kernels) and Fix Shape Mismatch
+        if "net_0" in renamed_pt_key or "net_2" in renamed_pt_key:
+            # Rename weight to kernel
+            if "weight" in renamed_pt_key:
+                renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+                
+                # CRITICAL FIX: Transpose the weights
+                # PyTorch Linear is (Out, In), JAX Dense is (In, Out).
+                # Ensure 'pt_tensor' is the variable holding your weight tensor.
+                pt_tensor = pt_tensor.T
 
       # 5. Fix for 'norm_added_q' which showed up in your missing keys list
       # The error said 'kernel' was missing, implying this specific norm might act like a dense layer