audio frames padded to nearest multiple of 128

Author: prishajain1

src/maxdiffusion/maxdiffusion_utils.py

@@ -297,6 +297,7 @@ def get_dummy_ltx2_inputs(config, pipeline, batch_size):
       pipeline.audio_sampling_rate / pipeline.audio_hop_length / float(pipeline.audio_vae_temporal_compression_ratio)
   )
   audio_num_frames = round(duration_s * audio_latents_per_second)
+  audio_num_frames = ((audio_num_frames + 127) // 128) * 128
 
   hidden_states = pipeline.prepare_latents(
       batch_size,

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1145,6 +1145,9 @@ 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,