missing keys error

Author: prishajain1

src/maxdiffusion/models/wan/wan_utils.py

@@ -255,6 +255,8 @@ def load_base_wan_transformer(
     del flattened_dict
     for pt_key, tensor in tensors.items():
       renamed_pt_key = rename_key(pt_key)
+      if "norm_added_q" in pt_key:
+          debug_original = renamed_pt_key
       if "image_embedder" in renamed_pt_key:
           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:
@@ -276,16 +278,27 @@ def load_base_wan_transformer(
               renamed_pt_key = renamed_pt_key.replace("kernel", "scale")
 
       if "norm_added_q" in renamed_pt_key:
-              renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
-              tensor = tensor.T
+          renamed_pt_key = renamed_pt_key.replace("weight", "kernel")
+          tensor = tensor.T
       renamed_pt_key = renamed_pt_key.replace("blocks_", "blocks.")
+
 
       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")
       renamed_pt_key = renamed_pt_key.replace("ffn.net_2", "ffn.proj_out")
       renamed_pt_key = renamed_pt_key.replace("ffn.net_0", "ffn.act_fn")
       renamed_pt_key = renamed_pt_key.replace("norm2", "norm2.layer_norm")
+      if "norm_added_q" in pt_key:
+          print(f"DEBUG REPORT for {pt_key}:")
+          print(f"  1. After rename_key : {debug_original}")
+          print(f"  2. Final Key String : {renamed_pt_key}")
+          
+          # Test parsing
+          pt_tuple_key = tuple(renamed_pt_key.split("."))
+          flax_key, _ = get_key_and_value(pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers)
+          print(f"  3. Parsed Flax Key  : {flax_key}")
+          print("-" * 20)
       pt_tuple_key = tuple(renamed_pt_key.split("."))
       flax_key, flax_tensor = get_key_and_value(pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers)
       flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu)