[LTX-2] Add decode_audio to skip the audio decode in modular pipelines

src/diffusers/modular_pipelines/ltx2/decoders.py

@@ -354,6 +354,15 @@ def inputs(self) -> list[tuple[str, Any]]:
                 description="Number of audio latent frames, used to unpack the audio latent sequence.",
             ),
             InputParam.template("output_type", default="pil"),
+            InputParam(
+                "decode_audio",
+                type_hint=bool,
+                default=True,
+                description=(
+                    "Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves "
+                    "`audio` as `None`, for video-only outputs; the video is unchanged."
+                ),
+            ),
         ]
 
     @property

src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py

@@ -53,6 +53,54 @@
 logger = logging.get_logger(__name__)
 
 
+# auto_docstring
+class LTX2AutoAudioDecoderStep(ConditionalPipelineBlocks):
+    """
+    Conditional audio decoder step: vocodes the audio latents unless `decode_audio=False`.
+       - `LTX2AudioDecoderStep` by default.
+       - Skipped when `decode_audio=False`, leaving `audio` as `None`: for video-only outputs (e.g. an IC-LoRA alpha
+         matte) the audio VAE and vocoder are not run. The video is unchanged, since the audio latents are still
+         denoised jointly with it.
+
+      Components:
+          audio_vae (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`)
+
+      Inputs:
+          audio_latents (`Tensor`, *optional*):
+              Denoised audio latents.
+          audio_num_frames (`int`, *optional*):
+              Number of audio latent frames, used to unpack the audio latent sequence.
+          output_type (`str`, *optional*, defaults to pil):
+              Output format: 'pil', 'np', 'pt'.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
+
+      Outputs:
+          audio (`Tensor`):
+              The generated audio waveform.
+    """
+
+    model_name = "ltx2"
+    block_classes = [LTX2AudioDecoderStep]
+    block_names = ["audio_decode"]
+    block_trigger_inputs = ["decode_audio"]
+
+    def select_block(self, decode_audio=None) -> str | None:
+        # Decoding is the default: only an explicit `decode_audio=False` skips the step.
+        return None if decode_audio is False else "audio_decode"
+
+    @property
+    def description(self):
+        return (
+            "Conditional audio decoder step: vocodes the audio latents unless `decode_audio=False`.\n"
+            " - `LTX2AudioDecoderStep` by default.\n"
+            " - Skipped when `decode_audio=False`, leaving `audio` as `None`: for video-only outputs (e.g. an "
+            "IC-LoRA alpha matte) the audio VAE and vocoder are not run. The video is unchanged, since the audio "
+            "latents are still denoised jointly with it."
+        )
+
+
 # auto_docstring
 class LTX2AutoPromptEnhancerStep(ConditionalPipelineBlocks):
     """
@@ -1064,10 +1112,13 @@ class LTX2DecoderStep(SequentialPipelineBlocks):
               be generated in input step.
           dtype (`dtype`):
               The dtype of the model inputs, can be generated in input step.
-          audio_latents (`Tensor`):
+          audio_latents (`Tensor`, *optional*):
               Denoised audio latents.
-          audio_num_frames (`int`):
+          audio_num_frames (`int`, *optional*):
               Number of audio latent frames, used to unpack the audio latent sequence.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -1077,7 +1128,7 @@ class LTX2DecoderStep(SequentialPipelineBlocks):
     """
 
     model_name = "ltx2"
-    block_classes = [LTX2VaeDecoderStep, LTX2AudioDecoderStep]
+    block_classes = [LTX2VaeDecoderStep, LTX2AutoAudioDecoderStep]
     block_names = ["video_decode", "audio_decode"]
 
     @property
@@ -1127,10 +1178,13 @@ class LTX2ConditionDecoderStep(SequentialPipelineBlocks):
               be generated in input step.
           dtype (`dtype`):
               The dtype of the model inputs, can be generated in input step.
-          audio_latents (`Tensor`):
+          audio_latents (`Tensor`, *optional*):
               Denoised audio latents.
-          audio_num_frames (`int`):
+          audio_num_frames (`int`, *optional*):
               Number of audio latent frames, used to unpack the audio latent sequence.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -1140,7 +1194,7 @@ class LTX2ConditionDecoderStep(SequentialPipelineBlocks):
     """
 
     model_name = "ltx2"
-    block_classes = [LTX2TrimConditionTokensStep, LTX2VaeDecoderStep, LTX2AudioDecoderStep]
+    block_classes = [LTX2TrimConditionTokensStep, LTX2VaeDecoderStep, LTX2AutoAudioDecoderStep]
     block_names = ["trim_condition_tokens", "video_decode", "audio_decode"]
 
     @property
@@ -1195,10 +1249,13 @@ class LTX2AutoDecoderStep(AutoPipelineBlocks):
               be generated in input step.
           dtype (`dtype`):
               The dtype of the model inputs, can be generated in input step.
-          audio_latents (`Tensor`):
+          audio_latents (`Tensor`, *optional*):
               Denoised audio latents.
-          audio_num_frames (`int`):
+          audio_num_frames (`int`, *optional*):
               Number of audio latent frames, used to unpack the audio latent sequence.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -1309,6 +1366,9 @@ class LTX2Blocks(SequentialPipelineBlocks):
               The timestep at which the VAE decodes the final latents.
           decode_noise_scale (`None`, *optional*):
               Noise interpolation factor applied to the latents at the decode timestep.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -1424,6 +1484,9 @@ class LTX2ImageToVideoBlocks(SequentialPipelineBlocks):
               The timestep at which the VAE decodes the final latents.
           decode_noise_scale (`None`, *optional*):
               Noise interpolation factor applied to the latents at the decode timestep.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -1534,6 +1597,9 @@ class LTX2ConditionBlocks(SequentialPipelineBlocks):
               The timestep at which the VAE decodes the final latents.
           decode_noise_scale (`None`, *optional*):
               Noise interpolation factor applied to the latents at the decode timestep.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -1669,6 +1735,9 @@ class LTX2InContextBlocks(SequentialPipelineBlocks):
               The timestep at which the VAE decodes the final latents.
           decode_noise_scale (`None`, *optional*):
               Noise interpolation factor applied to the latents at the decode timestep.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -1840,6 +1909,9 @@ class LTX2AutoBlocks(SequentialPipelineBlocks):
               The timestep at which the VAE decodes the final latents.
           decode_noise_scale (`None`, *optional*):
               Noise interpolation factor applied to the latents at the decode timestep.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):

src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py

@@ -17,11 +17,11 @@
 from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
 from ..modular_pipeline_utils import OutputParam
 from .decoders import (
-    LTX2AudioDecoderStep,
     LTX2DiffusionVaeDecoderStep,
     LTX2TrimConditionTokensStep,
 )
 from .modular_blocks_ltx2 import (
+    LTX2AutoAudioDecoderStep,
     LTX2AutoConditionEncoderStep,
     LTX2AutoCoreDenoiseStep,
     LTX2AutoDurationStep,
@@ -57,10 +57,13 @@ class LTX25DecoderStep(SequentialPipelineBlocks):
               Torch generator for deterministic generation.
           dtype (`dtype`):
               The dtype of the model inputs, can be generated in input step.
-          audio_latents (`Tensor`):
+          audio_latents (`Tensor`, *optional*):
               Denoised audio latents.
-          audio_num_frames (`int`):
+          audio_num_frames (`int`, *optional*):
               Number of audio latent frames, used to unpack the audio latent sequence.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -70,7 +73,7 @@ class LTX25DecoderStep(SequentialPipelineBlocks):
     """
 
     model_name = "ltx2.5"
-    block_classes = [LTX2DiffusionVaeDecoderStep, LTX2AudioDecoderStep]
+    block_classes = [LTX2DiffusionVaeDecoderStep, LTX2AutoAudioDecoderStep]
     block_names = ["video_decode", "audio_decode"]
 
     @property
@@ -115,10 +118,13 @@ class LTX25ConditionDecoderStep(SequentialPipelineBlocks):
               Torch generator for deterministic generation.
           dtype (`dtype`):
               The dtype of the model inputs, can be generated in input step.
-          audio_latents (`Tensor`):
+          audio_latents (`Tensor`, *optional*):
               Denoised audio latents.
-          audio_num_frames (`int`):
+          audio_num_frames (`int`, *optional*):
               Number of audio latent frames, used to unpack the audio latent sequence.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -128,7 +134,7 @@ class LTX25ConditionDecoderStep(SequentialPipelineBlocks):
     """
 
     model_name = "ltx2.5"
-    block_classes = [LTX2TrimConditionTokensStep, LTX2DiffusionVaeDecoderStep, LTX2AudioDecoderStep]
+    block_classes = [LTX2TrimConditionTokensStep, LTX2DiffusionVaeDecoderStep, LTX2AutoAudioDecoderStep]
     block_names = ["trim_condition_tokens", "video_decode", "audio_decode"]
 
     @property
@@ -176,10 +182,13 @@ class LTX25AutoDecoderStep(AutoPipelineBlocks):
               Torch generator for deterministic generation.
           dtype (`dtype`):
               The dtype of the model inputs, can be generated in input step.
-          audio_latents (`Tensor`):
+          audio_latents (`Tensor`, *optional*):
               Denoised audio latents.
-          audio_num_frames (`int`):
+          audio_num_frames (`int`, *optional*):
               Number of audio latent frames, used to unpack the audio latent sequence.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):
@@ -332,6 +341,9 @@ class LTX25AutoBlocks(SequentialPipelineBlocks):
               VAE-encoded reference-image latents used for image-to-video conditioning.
           output_type (`str`, *optional*, defaults to pil):
               Output format: 'pil', 'np', 'pt'.
+          decode_audio (`bool`, *optional*, defaults to True):
+              Whether to decode the audio. `False` skips this step (`LTX2AutoAudioDecoderStep`) and leaves `audio` as
+              `None`, for video-only outputs; the video is unchanged.
 
       Outputs:
           videos (`list`):

tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py

@@ -25,6 +25,7 @@
 from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
 from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
 
+from ...testing_utils import assert_tensors_close
 from ..testing_utils import (
     BaseModularPipelineTesterConfig,
     ModularLoadingTesterMixin,
@@ -154,6 +155,15 @@ def test_audio_output(self):
         assert audio.shape[1] == pipe.vocoder.config.out_channels
         assert torch.isnan(audio).sum() == 0
 
+    def test_decode_audio_false_skips_the_audio_only(self):
+        pipe = self.get_pipeline().to("cpu")
+
+        with_audio = pipe(**self.get_dummy_inputs(), output=["videos", "audio"])
+        without_audio = pipe(**self.get_dummy_inputs(), decode_audio=False, output=["videos", "audio"])
+
+        assert without_audio["audio"] is None
+        assert_tensors_close(without_audio["videos"], with_audio["videos"], atol=0, rtol=0)
+
     def test_auto_duration_predicts_a_grid_valid_frame_count(self):
         pipe = self.get_pipeline().to("cpu")
 

tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py

@@ -22,6 +22,7 @@
 from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
 from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
 
+from ...testing_utils import assert_tensors_close
 from ..testing_utils import (
     BaseModularPipelineTesterConfig,
     ModularLoadingTesterMixin,
@@ -263,7 +264,16 @@ def get_dummy_inputs(self, seed=0):
 class TestLTX25InContextModularPipelineFast(
     LTX25InContextModularPipelineTesterConfig, LTX25ModularPipelineFastTesterMixin
 ):
-    pass
+    def test_decode_audio_false_skips_the_audio_only(self):
+        # A video-only IC-LoRA (e.g. an alpha-matte LoRA) has no use for the soundtrack decoded alongside it.
+        pipe = self.get_pipeline().to("cpu")
+
+        with_audio = pipe(**self.get_dummy_inputs(), output=["videos", "audio"])
+        without_audio = pipe(**self.get_dummy_inputs(), decode_audio=False, output=["videos", "audio"])
+
+        assert with_audio["audio"] is not None
+        assert without_audio["audio"] is None
+        assert_tensors_close(without_audio["videos"], with_audio["videos"], atol=0, rtol=0)
 
 
 class TestLTX25InContextModularPipelineLoading(LTX25InContextModularPipelineTesterConfig, ModularLoadingTesterMixin):