test vae load

Author: prishajain1

src/maxdiffusion/models/ltx2/ltx2_utils.py

@@ -315,6 +315,10 @@ def load_vae_weights(
           pt_tuple_key = tuple(pt_list)
 
           flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict)
+          
+          if flax_key in [("latents_mean",), ("latents_std",)]:
+              flax_key = flax_key + ("value",)
+              
           flax_key = _tuple_str_to_int(flax_key)
 
           flax_key_str = [str(x) for x in flax_key]
@@ -545,6 +549,9 @@ def load_audio_vae_weights(
                 flax_key_parts.append(part)
 
         flax_key = tuple(flax_key_parts)
+        
+        if flax_key in [("latents_mean",), ("latents_std",)]:
+            flax_key = flax_key + ("value",)
 
         if "up_stages" in flax_key:
              try:

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1144,9 +1144,6 @@ def __call__(
       )
       audio_num_frames = round(duration_s * audio_latents_per_second)
 
-      # Pad audio sequence length to cleanly divide block sizes for Pallas flash attention on TPUs
-      audio_num_frames = ((audio_num_frames + 127) // 128) * 128
-
       audio_latents = self.prepare_audio_latents(
           batch_size=batch_size,
           num_channels_latents=audio_channels,

test_vae_load.py

@@ -0,0 +1,12 @@
+import jax
+import sys
+import maxdiffusion.pyconfig as pyconfig
+from maxdiffusion.pipelines.ltx2.ltx2_pipeline import LTX2Pipeline
+
+argv = ["", "src/maxdiffusion/configs/ltx2_video.yml"]
+pyconfig.initialize(argv)
+
+pipeline = LTX2Pipeline.from_pretrained(pyconfig.config, vae_only=True)
+print("latents_mean:", pipeline.vae.latents_mean.value[:10])
+print("latents_std:", pipeline.vae.latents_std.value[:10])
+print("audio_latents_mean:", pipeline.audio_vae.latents_mean.value[:10])