generate_audio flag added

Author: prishajain1

src/maxdiffusion/configs/ltx2_3_video.yml

@@ -36,6 +36,7 @@ audio_stg_scale: 1.0
 modality_scale: 3.0
 audio_modality_scale: 3.0
 use_cross_timestep: true
+generate_audio: True
 spatio_temporal_guidance_blocks: [28]
 fps: 24
 pipeline_type: multi-scale

src/maxdiffusion/configs/ltx2_video.yml

@@ -42,6 +42,7 @@ audio_stg_scale: 0.0
 modality_scale: 1.0
 audio_modality_scale: 1.0
 use_cross_timestep: false
+generate_audio: True
 spatio_temporal_guidance_blocks: []
 noise_scale: 1.0
 fps: 24

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -654,8 +654,9 @@ def _create_common_components(cls, config: HyperParameters, vae_only=False):
     components["tokenizer"] = cls.load_tokenizer(config)
     components["text_encoder"] = cls.load_text_encoder(config)
     components["connectors"] = cls.load_connectors(devices_array, mesh, rngs, config)
-    components["audio_vae"] = cls.load_audio_vae(devices_array, mesh, rngs, config)
-    components["vocoder"] = cls.load_vocoder(devices_array, mesh, rngs, config)
+    if getattr(config, "generate_audio", True):
+      components["audio_vae"] = cls.load_audio_vae(devices_array, mesh, rngs, config)
+      components["vocoder"] = cls.load_vocoder(devices_array, mesh, rngs, config)
     components["scheduler"] = cls.load_scheduler(config)
 
     return components
@@ -1325,15 +1326,18 @@ def __call__(
     )
     audio_num_frames = round(duration_s * audio_latents_per_second)
 
-    audio_latents = self.prepare_audio_latents(
-        batch_size=batch_size,
-        num_channels_latents=audio_channels,
-        audio_latent_length=audio_num_frames,
-        noise_scale=noise_scale,
-        dtype=dtype,
-        generator=key_audio,
-        latents=audio_latents,
-    )
+    if getattr(self.config, "generate_audio", True):
+      audio_latents = self.prepare_audio_latents(
+          batch_size=batch_size,
+          num_channels_latents=audio_channels,
+          audio_latent_length=audio_num_frames,
+          noise_scale=noise_scale,
+          dtype=dtype,
+          generator=key_audio,
+          latents=audio_latents,
+      )
+    else:
+      audio_latents = jnp.zeros((batch_size, audio_channels, audio_num_frames), dtype=dtype)
 
     # 5. Prepare Timesteps
     sigmas = jnp.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
@@ -1700,24 +1704,27 @@ 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()
-    # Cache the JITted function on the pipeline so it doesn't recompile on the 2nd run
-    if not hasattr(self, "_jitted_vocoder"):
-      self._jitted_vocoder = nnx.jit(lambda m, x: m(x))
-    
-    audio = self._jitted_vocoder(self.vocoder, generated_mel_spectrograms)
-    jax.block_until_ready(audio)
-    max_logging.log(f"⏱️ BWE Vocoder Execution Time: {time.time() - vocoder_start_time:.4f} seconds")
+    if getattr(self.config, "generate_audio", True) and self.audio_vae is not None:
+      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()
+      # Cache the JITted function on the pipeline so it doesn't recompile on the 2nd run
+      if not hasattr(self, "_jitted_vocoder"):
+        self._jitted_vocoder = nnx.jit(lambda m, x: m(x))
+      
+      audio = self._jitted_vocoder(self.vocoder, generated_mel_spectrograms)
+      jax.block_until_ready(audio)
+      max_logging.log(f"⏱️ BWE Vocoder Execution Time: {time.time() - vocoder_start_time:.4f} seconds")
 
-    # Convert audio to numpy
-    audio = np.array(audio)
+      # Convert audio to numpy
+      audio = np.array(audio)
+    else:
+      audio = None
 
     return LTX2PipelineOutput(frames=video, audio=audio)