Validate converted Flax PyTree state dict structure against expected initialization layout

Author: amepas

src/maxdiffusion/models/flux/util.py

@@ -118,21 +118,28 @@ def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict):
   expected_pytree: dict - a pytree that comes from initializing the model.
   new_pytree: dict - a pytree that has been created from pytorch weights.
   """
-  expected_pytree = flatten_dict(expected_pytree)
-  if len(expected_pytree.keys()) != len(new_pytree.keys()):
-    set1 = set(expected_pytree.keys())
-    set2 = set(new_pytree.keys())
+  expected_flat = (
+      flatten_dict(expected_pytree) if not isinstance(next(iter(expected_pytree.keys()), None), tuple) else expected_pytree
+  )
+  new_flat = flatten_dict(new_pytree) if not isinstance(next(iter(new_pytree.keys()), None), tuple) else new_pytree
+
+  if len(expected_flat.keys()) != len(new_flat.keys()):
+    set1 = set(expected_flat.keys())
+    set2 = set(new_flat.keys())
     missing_keys = set1 ^ set2
-    max_logging.log(f"missing keys : {missing_keys}")
-  for key in expected_pytree.keys():
-    if key in new_pytree.keys():
+    max_logging.log(
+        f"Missing or extra parameter keys count mismatch ({len(expected_flat)} expected vs {len(new_flat)} converted): {missing_keys}"
+    )
+
+  for key in expected_flat.keys():
+    if key in new_flat.keys():
       try:
-        expected_pytree_shape = expected_pytree[key].shape
+        expected_pytree_shape = expected_flat[key].shape
       except Exception:
-        expected_pytree_shape = expected_pytree[key].value.shape
-      if expected_pytree_shape != new_pytree[key].shape:
+        expected_pytree_shape = getattr(expected_flat[key], "value", expected_flat[key]).shape
+      if expected_pytree_shape != new_flat[key].shape:
         max_logging.log(
-            f"shape mismatch, expected shape of {expected_pytree[key].shape}, but got shape of {new_pytree[key].shape}"
+            f"shape mismatch for key '{key}': expected shape of {expected_pytree_shape}, but got shape of {new_flat[key].shape}"
         )
     else:
       max_logging.log(f"key: {key} not found...")
@@ -266,6 +273,7 @@ def load_flow_model(name: str, eval_shapes: dict, device: str, hf_download: bool
       jax.clear_caches()
   return flax_state_dict
 
+
 # -----------------------------------------------------------------------------
 # Latent Packing & Unpacking Helpers
 # -----------------------------------------------------------------------------
@@ -431,6 +439,8 @@ def load_and_convert_flux_klein_weights(safetensors_path, params, num_double_lay
 
   max_logging.log("Mapping weights to JAX parameters...")
 
+  expected_pytree = jax.tree_util.tree_map(lambda leaf: leaf, params)
+
   first_leaf = jax.tree_util.tree_leaves(params)[0]
   target_dtype = first_leaf.dtype
 
@@ -440,7 +450,9 @@ def convert_and_transpose_tensor(tensor, transpose=False):
     return jnp.array(tensor, dtype=target_dtype)
 
   # Global layers
-  params["context_embedder"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("context_embedder.weight"), transpose=True)
+  params["context_embedder"]["kernel"] = convert_and_transpose_tensor(
+      pt_state_dict.pop("context_embedder.weight"), transpose=True
+  )
   params["x_embedder"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("x_embedder.weight"), transpose=True)
   params["double_stream_modulation_img"]["kernel"] = convert_and_transpose_tensor(
       pt_state_dict.pop("double_stream_modulation_img.linear.weight"), transpose=True
@@ -454,7 +466,9 @@ def convert_and_transpose_tensor(tensor, transpose=False):
   params["proj_out"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("proj_out.weight"), transpose=True)
 
   # norm_out
-  params["norm_out"]["linear"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop("norm_out.linear.weight"), transpose=True)
+  params["norm_out"]["linear"]["kernel"] = convert_and_transpose_tensor(
+      pt_state_dict.pop("norm_out.linear.weight"), transpose=True
+  )
 
   # time_text_embed (Timestep Embedding)
   if "time_guidance_embed.timestep_embedder.linear_1.weight" in pt_state_dict:
@@ -492,18 +506,30 @@ def convert_and_transpose_tensor(tensor, transpose=False):
     jax_db["attn"]["e_qkv"]["kernel"] = jnp.array(np.concatenate([add_q, add_k, add_v], axis=1), dtype=target_dtype)
 
     # Projections out
-    jax_db["attn"]["i_proj"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.to_out.0.weight"), transpose=True)
-    jax_db["attn"]["e_proj"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.to_add_out.weight"), transpose=True)
+    jax_db["attn"]["i_proj"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.to_out.0.weight"), transpose=True
+    )
+    jax_db["attn"]["e_proj"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.to_add_out.weight"), transpose=True
+    )
 
     # Norm scales
     jax_db["attn"]["query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_q.weight"))
     jax_db["attn"]["key_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_k.weight"))
-    jax_db["attn"]["encoder_query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_added_q.weight"))
-    jax_db["attn"]["encoder_key_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "attn.norm_added_k.weight"))
+    jax_db["attn"]["encoder_query_norm"]["scale"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.norm_added_q.weight")
+    )
+    jax_db["attn"]["encoder_key_norm"]["scale"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "attn.norm_added_k.weight")
+    )
 
     # SwiGLU MLPs
-    jax_db["ff"]["linear_in"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True)
-    jax_db["ff"]["linear_out"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True)
+    jax_db["ff"]["linear_in"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "ff.linear_in.weight"), transpose=True
+    )
+    jax_db["ff"]["linear_out"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(prefix + "ff.linear_out.weight"), transpose=True
+    )
     jax_db["ff_context"]["linear_in"]["kernel"] = convert_and_transpose_tensor(
         pt_state_dict.pop(prefix + "ff_context.linear_in.weight"), transpose=True
     )
@@ -518,8 +544,12 @@ def convert_and_transpose_tensor(tensor, transpose=False):
     s_prefix = f"single_transformer_blocks.{block_idx}."
 
     # Joint projections
-    jax_sb["linear1"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.to_qkv_mlp_proj.weight"), transpose=True)
-    jax_sb["linear2"]["kernel"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.to_out.weight"), transpose=True)
+    jax_sb["linear1"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(s_prefix + "attn.to_qkv_mlp_proj.weight"), transpose=True
+    )
+    jax_sb["linear2"]["kernel"] = convert_and_transpose_tensor(
+        pt_state_dict.pop(s_prefix + "attn.to_out.weight"), transpose=True
+    )
 
     # Norm scales
     jax_sb["attn"]["query_norm"]["scale"] = convert_and_transpose_tensor(pt_state_dict.pop(s_prefix + "attn.norm_q.weight"))
@@ -530,7 +560,9 @@ def convert_and_transpose_tensor(tensor, transpose=False):
   )
   del pt_state_dict
   gc.collect()
-  max_logging.log("Weight conversion complete!")
+  max_logging.log("Validating converted Flax PyTree state dict structure...")
+  validate_flax_state_dict(expected_pytree, params)
+  max_logging.log("Weight conversion complete & verified!")
   return params
 
 
@@ -553,11 +585,15 @@ def get_pytorch_weight_tensor(key):
   max_logging.log("Mapping VAE decoder weights to JAX parameters...")
 
   # post_quant_conv
-  jax_params["post_quant_conv"]["kernel"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0))
+  jax_params["post_quant_conv"]["kernel"] = jnp.array(
+      get_pytorch_weight_tensor("post_quant_conv.weight").transpose(2, 3, 1, 0)
+  )
   jax_params["post_quant_conv"]["bias"] = jnp.array(get_pytorch_weight_tensor("post_quant_conv.bias"))
 
   # decoder.conv_in
-  jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0))
+  jax_params["decoder"]["conv_in"]["kernel"] = jnp.array(
+      get_pytorch_weight_tensor("decoder.conv_in.weight").transpose(2, 3, 1, 0)
+  )
   jax_params["decoder"]["conv_in"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_in.bias"))
 
   # decoder.mid_block
@@ -621,13 +657,17 @@ def get_pytorch_weight_tensor(key):
       upsampler_jax = up_block_jax["upsamplers_0"]
       upsampler_pt = f"{up_block_pt}.upsamplers.0"
 
-      upsampler_jax["conv"]["kernel"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0))
+      upsampler_jax["conv"]["kernel"] = jnp.array(
+          get_pytorch_weight_tensor(f"{upsampler_pt}.conv.weight").transpose(2, 3, 1, 0)
+      )
       upsampler_jax["conv"]["bias"] = jnp.array(get_pytorch_weight_tensor(f"{upsampler_pt}.conv.bias"))
 
   # decoder.conv_norm_out & conv_out
   jax_params["decoder"]["conv_norm_out"]["scale"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.weight"))
   jax_params["decoder"]["conv_norm_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_norm_out.bias"))
-  jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0))
+  jax_params["decoder"]["conv_out"]["kernel"] = jnp.array(
+      get_pytorch_weight_tensor("decoder.conv_out.weight").transpose(2, 3, 1, 0)
+  )
   jax_params["decoder"]["conv_out"]["bias"] = jnp.array(get_pytorch_weight_tensor("decoder.conv_out.bias"))
 
   jax_params = jax.tree_util.tree_map(