debug_audio_vae

Author: prishajain1

debug_audio_vae.py

@@ -236,7 +236,7 @@ def flatten(d, parent_key=()):
 
     # Also check if validation function itself is behaving as expected
     from flax.traverse_util import unflatten_dict, flatten_dict
-    from maxdiffusion.modeling_flax_pytorch_utils import validate_flax_state_dict
+    from maxdiffusion.models.modeling_flax_pytorch_utils import validate_flax_state_dict
 
     # Construct a dummy flax_state_dict with only the keys we found
     # We need to map our final_keys back to a dict

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -577,6 +577,8 @@ def load_audio_vae_weights(
                          # Assuming 3 stages (0, 1, 2)
                          # 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}")
                          flax_key_parts[up_stages_idx + 1] = new_stage_idx
                          flax_key = tuple(flax_key_parts)
              except ValueError: