changes for vae mean and std

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -242,8 +242,8 @@ def __init__(
       self.transformer_spatial_patch_size = getattr(self.transformer.config, "patch_size", 1)
       self.transformer_temporal_patch_size = getattr(self.transformer.config, "patch_size_t", 1)
 
-      self.audio_sampling_rate = getattr(self.audio_vae.config, "sample_rate", 16000)
-      self.audio_hop_length = getattr(self.audio_vae.config, "mel_hop_length", 160)
+      self.audio_sampling_rate = getattr(self.audio_vae.config, "sample_rate", 16000) if getattr(self, "audio_vae", None) is not None else 16000
+      self.audio_hop_length = getattr(self.audio_vae.config, "mel_hop_length", 160) if getattr(self, "audio_vae", None) is not None else 160
 
       # Initialize video processor
       self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio)

test_vae_load.py

@@ -6,7 +6,8 @@
 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])
+pipeline = LTX2Pipeline.from_pretrained(pyconfig.config, load_transformer=False)
+print("latents_mean (Video VAE):", pipeline.vae.latents_mean.value[:5])
+print("latents_std (Video VAE):", pipeline.vae.latents_std.value[:5])
+print("latents_mean (Audio VAE):", pipeline.audio_vae.latents_mean.value[:5])
+print("latents_std (Audio VAE):", pipeline.audio_vae.latents_std.value[:5])