From 92c486e577de166f4efa7c45c3b01103219796b7 Mon Sep 17 00:00:00 2001 From: christopher5106 Date: Thu, 15 Oct 2026 15:10:12 +0200 Subject: [PATCH] [LTX-2] Add `decode_audio` to skip the audio decode in modular pipelines `LTX2AutoAudioDecoderStep` wraps `LTX2AudioDecoderStep` in a `ConditionalPipelineBlocks` that runs it unless `decode_audio=False`, in every LTX-2 and LTX-2.5 decoder chain. A video-only output (e.g. an IC-LoRA alpha matte) then skips the audio VAE and vocoder and gets `audio=None`; the video is unchanged, since the audio latents are still denoised jointly with it. Co-Authored-By: Claude Opus 5.5 --- .../modular_pipelines/ltx2/decoders.py | 9 ++ .../ltx2/modular_blocks_ltx2.py | 88 +++++++++++++++++-- .../ltx2/modular_blocks_ltx25.py | 30 +++++-- .../ltx2/test_modular_pipeline_ltx2.py | 10 +++ .../ltx2/test_modular_pipeline_ltx25.py | 12 ++- 5 files changed, 131 insertions(+), 18 deletions(-) diff --git a/src/diffusers/modular_pipelines/ltx2/decoders.py b/src/diffusers/modular_pipelines/ltx2/decoders.py index cbeb1a66826d..7f9dfcefde62 100644 --- a/src/diffusers/modular_pipelines/ltx2/decoders.py +++ b/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 diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py index 86428328a5a6..00d4121f7f7e 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py +++ b/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`): diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py index 7c77aba94a74..d183730bda5d 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py +++ b/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`): diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py index 2ee330ece20a..0926046571d1 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py +++ b/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") diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py index 69080095c627..40b719e63cba 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py +++ b/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):