fix

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -578,7 +578,8 @@ def load_audio_vae_weights(
                          # Map 0 -> 2, 1 -> 1, 2 -> 0
                          new_stage_idx = 2 - stage_idx
                          if "upsample" in flax_key:
-                              print(f"DEBUG REVERSAL: {flax_key} -> stage_idx={stage_idx} -> new={new_stage_idx}")
+                              # print(f"DEBUG REVERSAL: {flax_key} -> stage_idx={stage_idx} -> new={new_stage_idx}")
+                              pass
                          flax_key_parts[up_stages_idx + 1] = new_stage_idx
                          flax_key = tuple(flax_key_parts)
              except ValueError:
@@ -599,43 +600,5 @@ def load_audio_vae_weights(
               continue
           filtered_eval_shapes[k] = v
 
-    print(f"DEBUG: Initial eval_shapes count: {len(flattened_eval)}")
-    print(f"DEBUG: Filtered eval_shapes count: {len(filtered_eval_shapes)}")
-    
-    # Check if any rngs remain in filtered
-    rngs_count = 0
-    for k in filtered_eval_shapes:
-        k_str = [str(x) for x in k]
-        for ks in k_str:
-            if "rngs" in ks or "dropout" in ks:
-                 rngs_count += 1
-                 break
-    print(f"DEBUG: Remaining rngs/dropout keys in Expected: {rngs_count}")
-
-    # Check flax_state_dict for rngs (New)
-    rngs_new_count = 0
-    for k in flax_state_dict:
-        k_str = [str(x) for x in k]
-        for ks in k_str:
-            if "rngs" in ks or "dropout" in ks:
-                 rngs_new_count += 1
-                 break
-    print(f"DEBUG: rngs/dropout keys in New (loaded): {rngs_new_count}")
-
-    # Explicit Set Diffs
-    expected_keys = set(filtered_eval_shapes.keys())
-    new_keys = set(flax_state_dict.keys())
-    
-    missing_keys = expected_keys - new_keys
-    extra_keys = new_keys - expected_keys
-    
-    print(f"DEBUG: Truly Missing Keys (in Expected but not New): {len(missing_keys)}")
-    if len(missing_keys) > 0:
-        print(f"DEBUG: Sample Missing: {list(missing_keys)[:5]}")
-        
-    print(f"DEBUG: Truly Extra Keys (in New but not Expected): {len(extra_keys)}")
-    if len(extra_keys) > 0:
-        print(f"DEBUG: Sample Extra: {list(extra_keys)[:5]}")
-
     validate_flax_state_dict(unflatten_dict(filtered_eval_shapes), flax_state_dict)
     return unflatten_dict(flax_state_dict)

src/maxdiffusion/models/modeling_flax_pytorch_utils.py

@@ -36,17 +36,6 @@ def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict):
   new_pytree: dict - a pytree that has been created from pytorch weights.
   """
   expected_pytree = flatten_dict(expected_pytree)
-  
-  # DEBUG PRINTS
-  print(f"DEBUG: validate_flax_state_dict called.")
-  print(f"DEBUG: expected_pytree keys: {len(expected_pytree)}")
-  print(f"DEBUG: new_pytree keys: {len(new_pytree)}")
-  
-  dropout_in_expected = [k for k in expected_pytree.keys() if "dropout" in str(k)]
-  print(f"DEBUG: dropout keys in expected_pytree: {len(dropout_in_expected)}")
-  
-  dropout_in_new = [k for k in new_pytree.keys() if "dropout" in str(k)]
-  print(f"DEBUG: dropout keys in new_pytree: {len(dropout_in_new)}")
 
   if len(expected_pytree.keys()) != len(new_pytree.keys()):
     set1 = set(expected_pytree.keys())

src/maxdiffusion/tests/test_ltx2_utils.py

@@ -223,7 +223,20 @@ def test_load_audio_vae_weights(self):
         )
 
         print("Validating Audio VAE Weights...")
-        validate_flax_state_dict(eval_shapes, flatten_dict(loaded_weights))
+        # Filter eval_shapes for validation as load_audio_vae_weights returns filtered weights
+        filtered_eval_shapes = {}
+        flat_eval = flatten_dict(eval_shapes)
+        for k, v in flat_eval.items():
+            k_str = [str(x) for x in k]
+            is_stat = False
+            for ks in k_str:
+                if "dropout" in ks or "rngs" in ks:
+                    is_stat = True
+                    break
+            if not is_stat:
+                filtered_eval_shapes[k] = v
+                
+        validate_flax_state_dict(unflatten_dict(filtered_eval_shapes), flatten_dict(loaded_weights))
         print("Audio VAE Weights Validated Successfully!")
 
 if __name__ == "__main__":