missing keys error

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -276,10 +276,12 @@ def load_base_wan_transformer(
               renamed_pt_key = renamed_pt_key.replace("kernel", "scale")
 
       if "norm_added_q" in renamed_pt_key:
-           # 1. Restore the dot before norm_added_q so JAX sees it as a submodule
-           renamed_pt_key = renamed_pt_key.replace("attn2_norm_added_q", "attn2.norm_added_q")
+           # 1. Restore the hierarchy: Convert flattened "attn2_norm_added_q" back to "attn2.norm_added_q"
+           # We use a broad replace to catch it regardless of what prefix (blocks_X_) is before it.
+           if "attn2_norm_added_q" in renamed_pt_key:
+               renamed_pt_key = renamed_pt_key.replace("attn2_norm_added_q", "attn2.norm_added_q")
 
-           # 2. Force 'weight' -> 'scale' (JAX requirement)
+           # 2. Force 'weight' -> 'scale' (Crucial: JAX explicitly asks for 'scale')
            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.")