transposing vae decoded output

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1395,6 +1395,9 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
       # Decode Audio
       audio_latents = audio_latents.astype(self.audio_vae.dtype)
       generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0]
+      
+      # Audio VAE outputs (B, T, F, C), Vocoder expects (B, Channels, Time, MelBins)
+      generated_mel_spectrograms = generated_mel_spectrograms.transpose(0, 3, 1, 2)
       audio = self.vocoder(generated_mel_spectrograms)
 
       # Convert audio to numpy