jit vocoder

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_ltx2.py

@@ -678,54 +678,35 @@ def __init__(
         window_type="hann",
     )
 
+  @nnx.jit
   def __call__(self, mel_spec: Array) -> Array:
-    print(f"=== BWE Vocoder Debug ===")
-    print(f"Input mel_spec - shape: {mel_spec.shape}, min: {mel_spec.min()}, max: {mel_spec.max()}")
-
     x = self.vocoder(mel_spec)
-    print(f"Base vocoder output (x) - shape: {x.shape}, min: {x.min()}, max: {x.max()}")
-
     x = jnp.transpose(x, (0, 2, 1))
     batch_size, num_samples, num_channels = x.shape
-    print(f"Transposed x - shape: {x.shape}")
 
     remainder = num_samples % self.hop_length
     if remainder != 0:
       x = jnp.pad(x, ((0, 0), (0, self.hop_length - remainder), (0, 0)))
-      print(f"Padded x - shape: {x.shape}")
 
     x_flattened = x.transpose(0, 2, 1).reshape(-1, x.shape[1], 1)
-    print(f"x_flattened - shape: {x_flattened.shape}")
-
     log_mel, _, _, _ = self.mel_stft(x_flattened)
-    print(f"MelSTFT output (log_mel) before reshape - shape: {log_mel.shape}, min: {log_mel.min()}, max: {log_mel.max()}")
 
     log_mel = log_mel.reshape(batch_size, num_channels, -1, log_mel.shape[-1])
-    print(f"Reshaped log_mel - shape: {log_mel.shape}")
-
     residual = self.bwe_generator(log_mel, time_last=False)
-    print(f"BWE generator output (residual) - shape: {residual.shape}, min: {residual.min()}, max: {residual.max()}")
-
     skip = self.resampler(x)
-    print(f"Resampler output (skip) - shape: {skip.shape}, min: {skip.min()}, max: {skip.max()}")
 
     residual = jnp.transpose(residual, (0, 2, 1))
 
     if residual.shape[1] < skip.shape[1]:
       residual = jnp.pad(residual, ((0, 0), (0, skip.shape[1] - residual.shape[1]), (0, 0)), mode='edge')
     elif residual.shape[1] > skip.shape[1]:
       residual = residual[:, :skip.shape[1], :]
-    print(f"Matched residual - shape: {residual.shape}")
 
     raw_waveform = residual + skip
-    print(f"Raw waveform (residual + skip) - min: {raw_waveform.min()}, max: {raw_waveform.max()}")
-
     waveform = jnp.clip(raw_waveform, -1, 1)
 
     output_samples = num_samples * self.output_sampling_rate // self.input_sampling_rate
     waveform = waveform[:, :output_samples, :]
     waveform = jnp.transpose(waveform, (0, 2, 1))
-    print(f"Final waveform - shape: {waveform.shape}, min: {waveform.min()}, max: {waveform.max()}")
-    print(f"=========================")
 
     return waveform

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1655,12 +1655,18 @@ def convert_to_vel(lat, x0, sig):
     video = self.video_processor.postprocess_video(torch.from_numpy(video_np), output_type=output_type)
 
     # Decode Audio
+    import time
     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)
+
+    vocoder_start_time = time.time()
     audio = self.vocoder(generated_mel_spectrograms)
+    jax.block_until_ready(audio)
+    vocoder_execution_time = time.time() - vocoder_start_time
+    max_logging.log(f"BWE Vocoder Execution Time: {vocoder_execution_time:.4f} seconds")
 
     # Convert audio to numpy
     audio = np.array(audio)