fix

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -340,6 +340,9 @@ def load_vae_weights(
           else:
               flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu)
 
+          # print(f"Loaded VAE Key: {flax_key}")
+
+      print(f"Total VAE keys loaded: {len(flax_state_dict)}")
       validate_flax_state_dict(eval_shapes, flax_state_dict)
       flax_state_dict = unflatten_dict(flax_state_dict)
       del tensors

src/maxdiffusion/tests/test_ltx2_utils.py

@@ -108,7 +108,8 @@ def test_load_vae_weights(self):
         print("Validating VAE Weights...")
         # Filter out dropout/rngs keys from eval_shapes as they are not expected in weights
         filtered_eval_shapes = {}
-        for k, v in eval_shapes.items():
+        flat_eval_shapes = flatten_dict(eval_shapes)
+        for k, v in flat_eval_shapes.items():
             k_str = [str(x) for x in k]
             if "dropout" in k_str or "rngs" in k_str:
                 continue