Automated PR - 2026-06-17

packages/ltx-pipelines/src/ltx_pipelines/utils/args.py

@@ -650,6 +650,62 @@ def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
     return parser
 
 
+def default_1_stage_t2a_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
+    """Argument parser for single-stage text-to-audio pipelines (audio-only)."""
+    audio_guider = params.audio_guider_params
+    parser = basic_arg_parser(params=params)
+    parser.add_argument(
+        "--num-frames",
+        type=int,
+        default=params.num_frames,
+        help="Number of frames used to derive audio duration (num-frames / frame-rate).",
+    )
+    parser.add_argument(
+        "--frame-rate",
+        type=float,
+        default=params.frame_rate,
+        help="Frame rate used with --num-frames to derive the audio duration.",
+    )
+    parser.add_argument(
+        "--negative-prompt",
+        type=str,
+        default=DEFAULT_NEGATIVE_PROMPT,
+        help="Negative prompt to steer audio generation away from artifacts.",
+    )
+    parser.add_argument(
+        "--audio-cfg-guidance-scale",
+        type=float,
+        default=audio_guider.cfg_scale,
+        help=f"Audio CFG scale (default: {audio_guider.cfg_scale}).",
+    )
+    parser.add_argument(
+        "--audio-stg-guidance-scale",
+        type=float,
+        default=audio_guider.stg_scale,
+        help=f"Audio STG scale (default: {audio_guider.stg_scale}).",
+    )
+    parser.add_argument(
+        "--audio-rescale-scale",
+        type=float,
+        default=audio_guider.rescale_scale,
+        help=f"Audio rescale scale (default: {audio_guider.rescale_scale}).",
+    )
+    parser.add_argument(
+        "--audio-stg-blocks",
+        type=int,
+        nargs="*",
+        default=audio_guider.stg_blocks,
+        help=f"Blocks to perturb for Audio STG (default: {audio_guider.stg_blocks}).",
+    )
+    parser.add_argument(
+        "--audio-skip-step",
+        type=int,
+        default=audio_guider.skip_step,
+        help=f"Audio skip step (default: {audio_guider.skip_step}).",
+    )
+    return parser
+
+
 def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
     parser = default_1_stage_arg_parser(params=params)
     parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)

packages/ltx-pipelines/src/ltx_pipelines/utils/blocks.py

@@ -7,7 +7,6 @@ removes the need for :class:`ModelLedger`.
 from __future__ import annotations
 
 import copy
-import dataclasses
 import logging
 from collections.abc import Iterator
 from contextlib import AbstractContextManager, contextmanager
@@ -40,9 +39,9 @@ from ltx_core.model.audio_vae import (
 from ltx_core.model.audio_vae import (
     decode_audio as vae_decode_audio,
 )
+from ltx_core.model.model_protocol import LTXModelProtocol, ModelConfigurator
 from ltx_core.model.transformer import (
     LTXV_MODEL_COMFY_RENAMING_MAP,
-    LTXModel,
     LTXModelConfigurator,
     X0Model,
 )
@@ -144,7 +143,7 @@ def _streaming_model(
     """Build a streaming wrapper, yield it, then tear down and free memory."""
     cpu_slots_count = DISK_CPU_SLOTS if offload_mode == OffloadMode.DISK else None
     wrapped = builder.build(
-        target_device=target_device,
+        device=target_device,
         dtype=dtype,
         cpu_slots_count=cpu_slots_count,
     )
@@ -195,7 +194,7 @@ class DiffusionStage:
     pattern in every pipeline.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         checkpoint_path: str,
         dtype: torch.dtype,
@@ -205,7 +204,9 @@ class DiffusionStage:
         registry: Registry | None = None,
         compilation_config: CompilationConfig | None = None,
         offload_mode: OffloadMode = OffloadMode.NONE,
-        transformer_builder: ModelBuilderProtocol[LTXModel] | None = None,
+        transformer_builder: ModelBuilderProtocol[LTXModelProtocol] | None = None,
+        model_configurator: type[ModelConfigurator] = LTXModelConfigurator,
+        model_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP,
     ) -> None:
         self._checkpoint_path = checkpoint_path
         self._dtype = dtype
@@ -213,10 +214,12 @@ class DiffusionStage:
         self._quantization = quantization
         self._compilation_config = compilation_config
         self._offload_mode = offload_mode
+        # A quantization policy may pin its own configurator; otherwise use the one
+        # provided by the caller (defaults to the audio-video LTXModelConfigurator).
         configurator = (
             quantization.model_configurator
             if quantization is not None and quantization.model_configurator is not None
-            else LTXModelConfigurator
+            else model_configurator
         )
         if transformer_builder is not None:
             self._transformer_builder = transformer_builder
@@ -224,14 +227,12 @@ class DiffusionStage:
             self._transformer_builder = Builder(
                 model_path=checkpoint_path,
                 model_class_configurator=configurator,
-                model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+                model_sd_ops=model_sd_ops,
                 loras=tuple(loras),
                 registry=registry or DummyRegistry(),
             )
 
         if offload_mode != OffloadMode.NONE:
-            if compilation_config is not None:
-                raise ValueError("torch.compile is not supported with layer streaming")
             # WeightsProvider currently only supports plain bf16 + fp8_cast LoRA fusion
             # (no companion-key emission). Quantization policies that emit
             # companion keys (e.g. ``.weight_scale``) cannot be streamed yet.
@@ -240,8 +241,15 @@ class DiffusionStage:
                     "Block streaming is not supported with this quantization policy "
                     "(only bf16 and fp8_cast are currently supported)."
                 )
-            streaming_sd_ops: SDOps = LTXV_MODEL_COMFY_RENAMING_MAP
+            streaming_sd_ops: SDOps = model_sd_ops
             streaming_module_ops: tuple[ModuleOps, ...] = ()
+            streaming_loras = tuple(loras)
+
+            if compilation_config:
+                number_of_layers = self._transformer_builder.model_config()["transformer"]["num_layers"]
+                streaming_sd_ops, streaming_module_ops, streaming_loras = _apply_compile_ops(
+                    streaming_sd_ops, streaming_module_ops, streaming_loras, number_of_layers
+                )
             if quantization is not None:
                 streaming_sd_ops, streaming_module_ops = _chain_quantization(
                     streaming_sd_ops, streaming_module_ops, quantization
@@ -251,7 +259,7 @@ class DiffusionStage:
                 model_path=checkpoint_path,
                 model_sd_ops=streaming_sd_ops,
                 module_ops=streaming_module_ops,
-                loras=tuple(loras),
+                loras=streaming_loras,
                 registry=registry or DummyRegistry(),
                 fuse_rule=quantization.fuse_rule if quantization is not None else bf16_fuse_rule,
                 blocks_attr="transformer_blocks",
@@ -273,9 +281,8 @@ class DiffusionStage:
             (*self._transformer_builder.module_ops, op),
         )
         if self._offload_mode != OffloadMode.NONE:
-            new._streaming_builder = dataclasses.replace(
-                self._streaming_builder,
-                module_ops=(*self._streaming_builder.module_ops, op),
+            new._streaming_builder = self._streaming_builder.with_module_ops(
+                (*self._streaming_builder.module_ops, op),
             )
         return new
 
@@ -531,7 +538,7 @@ class PromptEncoder:
                 prompts[0] = generate_enhanced_prompt(
                     text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed
                 )
-            raw_outputs = [text_encoder.encode(p) for p in prompts]
+            raw_outputs = text_encoder.encode(prompts)
         logger.info("Text encoder done, building embeddings processor from %s", self._checkpoint_path)
 
         with gpu_model(self._build_embeddings_processor()) as embeddings_processor:

packages/ltx-pipelines/src/ltx_pipelines/utils/denoisers.py

@@ -182,6 +182,8 @@ def _guided_denoise(  # noqa: PLR0913,PLR0915
     r = dict(zip(pass_names, zip(splits_v, splits_a, strict=True), strict=True))
 
     cond_v, cond_a = r["cond"]
+    cond_v = cond_v if isinstance(cond_v, torch.Tensor) else torch.tensor(cond_v)
+    cond_a = cond_a if isinstance(cond_a, torch.Tensor) else torch.tensor(cond_a)
     uncond_v, uncond_a = r.get("uncond", (0.0, 0.0))
     ptb_v, ptb_a = r.get("ptb", (0.0, 0.0))
     mod_v, mod_a = r.get("mod", (0.0, 0.0))

packages/ltx-pipelines/src/ltx_pipelines/utils/media_io.py

@@ -391,6 +391,24 @@ def encode_video(
     logger.info(f"Video saved to {output_path}")
 
 
+def encode_audio(audio: Audio, output_path: str) -> None:
+    """Save an audio waveform as a 16-bit PCM ``.wav`` file at the source sampling rate.
+    Reuses :func:`_write_audio` (the same muxing path used by :func:`encode_video`);
+    the only difference is a PCM (``pcm_s16le``) stream in a WAV container instead of
+    the AAC stream used for muxed video.
+    """
+    container = av.open(output_path, mode="w")
+    audio_stream = container.add_stream("pcm_s16le", rate=audio.sampling_rate)
+    audio_stream.codec_context.sample_rate = audio.sampling_rate
+    audio_stream.codec_context.layout = "stereo"
+    audio_stream.codec_context.time_base = Fraction(1, audio.sampling_rate)
+    try:
+        _write_audio(container, audio_stream, audio)
+    finally:
+        container.close()
+    logger.info(f"Audio saved to {output_path}")
+
+
 def _encode_chunks_threaded(
     container: av.container.Container,
     stream: av.video.stream.VideoStream,

packages/ltx-pipelines/src/ltx_pipelines/utils/samplers.py

@@ -436,7 +436,7 @@ def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915,PLR0912
     return video_state, audio_state
 
 
-def euler_cfg_pp_denoising_loop(
+def euler_cfg_pp_denoising_loop(  # noqa: PLR0912
     sigmas: torch.Tensor,
     video_state: LatentState | None,
     audio_state: LatentState | None,
@@ -514,9 +514,15 @@ def euler_cfg_pp_denoising_loop(
             )
 
         if video_state is not None and denoised_video is not None:
-            denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
+            denoised_video = post_process_latent(
+                denoised_video.float(), video_state.denoise_mask, video_state.clean_latent
+            )
+            noisy_video = video_state.latent.float()
         if audio_state is not None and denoised_audio is not None:
-            denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
+            denoised_audio = post_process_latent(
+                denoised_audio.float(), audio_state.denoise_mask, audio_state.clean_latent
+            )
+            noisy_audio = audio_state.latent.float()
 
         if sigmas[step_idx + 1] == 0:
             if video_state is not None and denoised_video is not None:
@@ -525,30 +531,31 @@ def euler_cfg_pp_denoising_loop(
                 audio_state = replace(audio_state, latent=denoised_audio.to(model_dtype))
             return video_state, audio_state
 
-        # Draw noise consecutively from the same generator: video first, audio second.
-        noise_video = new_noise_fn(video_state.latent, generator) if (video_state is not None and draw_noise) else None
-        noise_audio = new_noise_fn(audio_state.latent, generator) if (audio_state is not None and draw_noise) else None
-
         if video_state is not None and denoised_video is not None:
+            video_noise = new_noise_fn(video_state.latent, generator) if draw_noise else None
             x_next = stepper.step(
-                sample=video_state.latent,
+                sample=noisy_video,
                 denoised_sample=denoised_video,
                 sigmas=sigmas,
                 step_index=step_idx,
                 uncond_denoised=uncond_video,
-                noise=noise_video,
+                noise=video_noise,
             )
+            if draw_noise:
+                x_next = post_process_latent(x_next, video_state.denoise_mask, video_state.clean_latent)
             video_state = replace(video_state, latent=x_next.to(model_dtype))
 
         if audio_state is not None and denoised_audio is not None:
+            audio_noise = new_noise_fn(audio_state.latent, generator) if draw_noise else None
             x_next = stepper.step(
-                sample=audio_state.latent,
+                sample=noisy_audio,
                 denoised_sample=denoised_audio,
                 sigmas=sigmas,
                 step_index=step_idx,
                 uncond_denoised=uncond_audio,
-                noise=noise_audio,
+                noise=audio_noise,
             )
+            if draw_noise:
+                x_next = post_process_latent(x_next, audio_state.denoise_mask, audio_state.clean_latent)
             audio_state = replace(audio_state, latent=x_next.to(model_dtype))
-
     return video_state, audio_state