dim mismatch fix

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -256,22 +256,41 @@ 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.proj" in renamed_pt_key:
-          renamed_pt_key = renamed_pt_key.replace("net.0.proj", "net_0")
-        elif "net_0.proj" in renamed_pt_key:
-          renamed_pt_key = renamed_pt_key.replace("net_0.proj", "net_0")
-        if "net.2" in renamed_pt_key:
-          renamed_pt_key = renamed_pt_key.replace("net.2", "net_2")
-        renamed_pt_key = renamed_pt_key.replace("norm1", "norm1.layer_norm")
-        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")
-        elif "net_0" in renamed_pt_key or "net_2" in renamed_pt_key:
-          renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
-      if "norm" in renamed_pt_key and "image_embedder" not in renamed_pt_key:
-          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")
+          # 1. Fix Shape Mismatch: Transpose Dense Layers
+          # CHECK FOR BOTH VERSIONS ("net.0"/"net_0" and "net.2"/"net_2")
+          # This ensures the transpose happens regardless of what rename_key did.
+          if "net.0" in renamed_pt_key or "net_0" in renamed_pt_key or \
+             "net.2" in renamed_pt_key or "net_2" in renamed_pt_key:
+              tensor = tensor.t()  # <--- This will now correctly execute
+
+          # 2. FIX net_0: Strip .proj (Handle both raw and renamed keys)
+          if "net.0.proj" in renamed_pt_key:
+              renamed_pt_key = renamed_pt_key.replace("net.0.proj", "net_0")
+              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+          elif "net_0.proj" in renamed_pt_key:
+              renamed_pt_key = renamed_pt_key.replace("net_0.proj", "net_0")
+              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+          
+          # 3. FIX net_2: Ensure naming and kernel (Handle both raw and renamed keys)
+          if "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")
+          elif "net_2" in renamed_pt_key:
+              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+
+          # 4. FIX Norms: Add .layer_norm and force 'scale'
+          renamed_pt_key = renamed_pt_key.replace("norm1", "norm1.layer_norm")
+          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")
+
+      # --- FIX START: Missing Key 'norm_added_q' Fix ---
+      if "norm_added" in renamed_pt_key:
+           renamed_pt_key = renamed_pt_key.replace("weight", "scale")
+           renamed_pt_key = renamed_pt_key.replace("kernel", "scale")
+           
+           if "attn2_norm_added" in renamed_pt_key:
+               renamed_pt_key = renamed_pt_key.replace("attn2_norm_added", "attn2.norm_added")
       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")