Automated PR - 2026-06-17

packages/ltx-pipelines/CLAUDE.md

@@ -17,6 +17,7 @@ Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for
 | Pipeline | File | Stages | Model | Sampler | Use case |
 |----------|------|--------|-------|---------|----------|
 | `TI2VidOneStagePipeline` | `ti2vid_one_stage.py` | 1 | Full | Euler | Simple text/image-to-video |
+| `T2AOneStagePipeline` | `t2a_one_stage.py` | 1 | Full | Euler | Text-to-audio (audio-only output, no video branch) |
 | `TI2VidTwoStagesPipeline` | `ti2vid_two_stages.py` | 2 | Full + distilled LoRA | Euler | Production quality |
 | `TI2VidTwoStagesHQPipeline` | `ti2vid_two_stages_hq.py` | 2 | Full + distilled LoRA (both stages) | Res2s | Highest quality, fewer steps |
 | `A2VidPipelineTwoStage` | `a2vid_two_stage.py` | 2 | Full + distilled LoRA | Euler | Audio-conditioned video |

packages/ltx-pipelines/README.md

@@ -58,6 +58,7 @@ Available pipeline modules:
 - `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
 - `ltx_pipelines.ti2vid_two_stages_hq` - Two-stage text/image-to-video (different sampler, better quality).
 - `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
+- `ltx_pipelines.t2a_one_stage` - Single-stage text-to-audio (audio-only output).
 - `ltx_pipelines.distilled` - Fast text/image-to-video pipeline using only the distilled model.
 - `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA.
 - `ltx_pipelines.keyframe_interpolation` - Keyframe interpolation.
@@ -254,6 +255,20 @@ Uses IC-LoRA on a **distilled** checkpoint with a **single** lip-dub IC-LoRA app
 
 ---
 
+### 11. T2AOneStagePipeline
+
+**Best for:** Text-to-audio — generating speech/audio only (no video) from a text prompt, e.g. driving an audio-style LoRA such as an accent LoRA.
+
+**Source**: [`src/ltx_pipelines/t2a_one_stage.py`](src/ltx_pipelines/t2a_one_stage.py)
+
+Single-stage, **audio-only** generation: the video branch is absent (`video=None`), so only the audio modality is denoised and decoded through the audio VAE + vocoder, producing a wave file. Audio duration is derived from `--num-frames` / `--frame-rate` (the same `8k+1` frame convention as video). Audio guidance (CFG/STG) is optional — the `--audio-*` flags default to the model's values; the video→audio cross-modal guidance is disabled since there is no video modality.
+
+**Extra CLI arguments (all optional, with sensible defaults):** `--num-frames`, `--frame-rate`, `--negative-prompt`, `--audio-cfg-guidance-scale`, `--audio-stg-guidance-scale`, `--audio-stg-blocks`, `--audio-rescale-scale`, `--audio-skip-step`. No `--height/--width/--image` (audio has no spatial dimensions).
+
+**Use when:** You need speech/audio from text alone, or to evaluate an audio-only LoRA (accent, voice style) without generating video.
+
+---
+
 ## 🎨 Conditioning Types
 
 Pipelines use different conditioning methods from [`ltx-core`](../ltx-core/) for controlling generation. See the [ltx-core conditioning documentation](../ltx-core/README.md#conditioning--control) for details.
@@ -400,6 +415,69 @@ By default, pipelines clean GPU memory (especially transformer weights) between
 # utils.cleanup_memory()  # Comment out if you have enough VRAM
 ```
 
+### Compilation (`torch.compile`)
+
+Compiling the transformer blocks with `torch.compile` speeds up inference. It is **opt-in and off by default**. The blocks are compiled shape-polymorphically (the sequence dimension is marked dynamic), so one compiled artifact serves any token count without recompiling.
+
+**CLI** — the `--compile` flag maps directly to `CompilationConfig`:
+
+| Form | Result |
+| ---- | ------ |
+| *(flag absent)* | eager, no compilation |
+| `--compile` | compile with defaults |
+| `--compile KEY=VALUE ...` | compile, overriding individual fields |
+
+```bash
+# Defaults
+python -m ltx_pipelines.ti2vid_two_stages --compile --checkpoint-path=...
+
+# reduce-overhead captures CUDA graphs -- the main latency lever for the denoising loop.
+# Off by default because graph capture reserves static memory pools (extra VRAM), so it
+# trades memory for speed; enable it when you have headroom.
+python -m ltx_pipelines.ti2vid_two_stages --compile mode=reduce-overhead --checkpoint-path=...
+
+# Several overrides at once
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile mode=max-autotune fullgraph=true dynamic=true --checkpoint-path=...
+```
+
+| Field | Values | Default | Notes |
+| ----- | ------ | ------- | ----- |
+| `mode` | `none`, `reduce-overhead`, `max-autotune`, … | `none` | `reduce-overhead`/`max-autotune` enable CUDA graphs |
+| `backend` | `inductor`, `eager`, … | `inductor` | |
+| `fullgraph` | `true`/`false` | `false` | |
+| `dynamic` | `auto`/`true`/`false` | `auto` | the seq dim is marked dynamic regardless |
+| `inductor_config` | JSON object or path to a `.json` | `{}` | `torch._inductor.config` overrides |
+| `dynamo_config` | JSON object or path to a `.json` | `{"inline_inbuilt_nn_modules": true, "cache_size_limit": 256}` | `torch._dynamo.config` overrides |
+
+**Controlling inductor / dynamo configs.** `inductor_config` and `dynamo_config` take either an inline JSON object or a path to a `.json` file, applied via `torch._inductor.config.patch(...)` / `torch._dynamo.config.patch(...)` around the compiled forward. They **replace the defaults wholesale — they do not merge**, so when overriding `dynamo_config` re-include any defaults you want to keep:
+
+```bash
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile 'inductor_config={"max_autotune": true}' \
+              'dynamo_config={"inline_inbuilt_nn_modules": true, "cache_size_limit": 256, "recompile_limit": 32}' \
+    --checkpoint-path=...
+```
+
+**Programmatically**, pass a `CompilationConfig` to the pipeline:
+
+```python
+from ltx_core.model.transformer.compiling import CompilationConfig
+
+pipeline = TI2VidTwoStagesPipeline(
+    ...,
+    compilation_config=CompilationConfig(mode="reduce-overhead"),
+)
+```
+
+**Faster cache loads: `unsafe_skip_cache_dynamic_shape_guards` (unsafe, opt-in).** Inductor's FX-graph cache re-checks the dynamic-shape guards stored with each entry on every lookup. Setting this flag skips that re-check (every entry is treated as a guard hit), which speeds up warm and cross-process cache loads. It is **not enabled by default** because it is a correctness hazard: a kernel first compiled at a small sequence length keeps int32 address arithmetic, and reusing it at a larger sequence length (roughly **>58k tokens/rank**) overflows int32 and reads out of bounds — surfacing as a CUDA illegal memory access or silently corrupted output. Only enable it when your token counts stay within the range the cached kernels were compiled for:
+
+```bash
+python -m ltx_pipelines.ti2vid_two_stages \
+    --compile 'inductor_config={"unsafe_skip_cache_dynamic_shape_guards": true}' \
+    --checkpoint-path=...
+```
+
 ### Denoising Loop Optimization
 
 **Gradient Estimation Denoising Loop:**

packages/ltx-pipelines/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ltx-pipelines"
-version = "1.1.5"
+version = "1.1.6"
 description = "Pipelines implementation for Lightricks' LTX-2 model"
 readme = "README.md"
 requires-python = ">=3.10"

packages/ltx-pipelines/src/ltx_pipelines/__init__.py

@@ -2,6 +2,7 @@
 LTX-2 Pipelines: High-level video generation pipelines and utilities.
 This package provides ready-to-use pipelines for video generation:
 - TI2VidOneStagePipeline: Text/image-to-video in a single stage
+- T2AOneStagePipeline: Text-to-audio in a single stage (audio-only output)
 - TI2VidTwoStagesPipeline: Two-stage generation with upsampling
 - DistilledPipeline: Fast distilled two-stage generation
 - ICLoraPipeline: Image/video conditioning with distilled LoRA
@@ -18,6 +19,7 @@ from ltx_pipelines.ic_lora import ICLoraPipeline
 from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
 from ltx_pipelines.lipdub import LipDubPipeline
 from ltx_pipelines.retake import RetakePipeline
+from ltx_pipelines.t2a_one_stage import T2AOneStagePipeline
 from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
 from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
 
@@ -28,6 +30,7 @@ __all__ = [
     "KeyframeInterpolationPipeline",
     "LipDubPipeline",
     "RetakePipeline",
+    "T2AOneStagePipeline",
     "TI2VidOneStagePipeline",
     "TI2VidTwoStagesPipeline",
 ]

packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py

@@ -14,6 +14,7 @@ from ltx_core.types import Audio, VideoPixelShape
 from ltx_pipelines.iclora_utils import (
     append_ic_lora_reference_video_conditionings,
     read_lora_reference_downscale_factor,
+    read_lora_reference_temporal_scale_factor,
 )
 from ltx_pipelines.utils.args import (
     ImageConditioningInput,
@@ -102,10 +103,11 @@ class ICLoraPipeline:
         self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
         self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
 
-        # Read reference downscale factor from LoRA metadata.
-        # IC-LoRAs trained with low-resolution reference videos store this factor
-        # so inference can resize reference videos to match training conditions.
+        # Read reference scale factors from LoRA metadata.
+        # IC-LoRAs trained with scaled reference videos store these factors
+        # so inference can resize/subsample reference videos to match training conditions.
         self.reference_downscale_factor = 1
+        self.reference_temporal_scale_factor = 1
         for lora in loras:
             scale = read_lora_reference_downscale_factor(lora.path)
             if scale != 1:
@@ -116,6 +118,15 @@ class ICLoraPipeline:
                         f"specifies {scale}. Cannot combine LoRAs with different reference scales."
                     )
                 self.reference_downscale_factor = scale
+            temporal = read_lora_reference_temporal_scale_factor(lora.path)
+            if temporal != 1:
+                if self.reference_temporal_scale_factor not in (1, temporal):
+                    raise ValueError(
+                        f"Conflicting reference_temporal_scale_factor values in LoRAs: "
+                        f"already have {self.reference_temporal_scale_factor}, but {lora.path} "
+                        f"specifies {temporal}. Cannot combine LoRAs with different temporal scales."
+                    )
+                self.reference_temporal_scale_factor = temporal
 
     def __call__(  # noqa: PLR0913
         self,
@@ -318,6 +329,7 @@ class ICLoraPipeline:
             dtype=self.dtype,
             device=self.device,
             reference_downscale_factor=self.reference_downscale_factor,
+            reference_temporal_scale_factor=self.reference_temporal_scale_factor,
             conditioning_attention_strength=conditioning_attention_strength,
             conditioning_attention_mask=conditioning_attention_mask,
             tiling_config=None,

packages/ltx-pipelines/src/ltx_pipelines/iclora_utils.py

@@ -31,6 +31,17 @@ def read_lora_reference_downscale_factor(lora_path: str) -> int:
         return 1
 
 
+def read_lora_reference_temporal_scale_factor(lora_path: str) -> int:
+    """Read ``reference_temporal_scale_factor`` from LoRA safetensors metadata (default 1)."""
+    try:
+        with safe_open(lora_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+            return int(metadata.get("reference_temporal_scale_factor", 1))
+    except Exception as e:
+        logging.warning("Failed to read metadata from LoRA file '%s': %s", lora_path, e)
+        return 1
+
+
 def downsample_mask_video_to_latent(
     mask: torch.Tensor,
     target_latent_shape: VideoLatentShape,
@@ -66,6 +77,12 @@ def downsample_mask_video_to_latent(
     return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
 
 
+def temporal_subsample(video: torch.Tensor, temporal_scale_factor: int) -> torch.Tensor:
+    """VAE-aligned temporal subsampling: keep frame 0, then every Nth frame."""
+    indices = [0, *list(range(1, video.shape[2], temporal_scale_factor))]
+    return video[:, :, indices]
+
+
 def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
     conditionings: list[ConditioningItem],
     video_conditioning: list[tuple[str, float]],
@@ -77,6 +94,7 @@ def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
     dtype: torch.dtype,
     device: torch.device,
     reference_downscale_factor: int,
+    reference_temporal_scale_factor: int = 1,
     conditioning_attention_strength: float,
     conditioning_attention_mask: torch.Tensor | None,
     tiling_config: TilingConfig | None = None,
@@ -93,6 +111,8 @@ def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
     for video_path, strength in video_conditioning:
         frame_gen = decode_video_by_frame(path=video_path, frame_cap=num_frames, device=device)
         video = video_preprocess(frame_gen, ref_height, ref_width, dtype, device)
+        if reference_temporal_scale_factor > 1:
+            video = temporal_subsample(video, reference_temporal_scale_factor)
         if tiling_config is not None:
             encoded_video = video_encoder.tiled_encode(video, tiling_config)
         else:
@@ -113,6 +133,7 @@ def append_ic_lora_reference_video_conditionings(  # noqa: PLR0913
         cond = VideoConditionByReferenceLatent(
             latent=encoded_video,
             downscale_factor=scale,
+            temporal_scale_factor=reference_temporal_scale_factor,
             strength=strength,
         )
         if attn_mask is not None:

packages/ltx-pipelines/src/ltx_pipelines/t2a_one_stage.py

@@ -0,0 +1,192 @@
+import logging
+
+import torch
+
+from ltx_core.components.guiders import (
+    MultiModalGuiderFactory,
+    MultiModalGuiderParams,
+    create_multimodal_guider_factory,
+)
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.model.transformer import LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP, LTXAudioOnlyModelConfigurator
+from ltx_core.model.transformer.compiling import CompilationConfig
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.types import Audio
+from ltx_pipelines.utils import get_device
+from ltx_pipelines.utils.args import (
+    default_1_stage_t2a_arg_parser,
+    detect_checkpoint_path,
+)
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    PromptEncoder,
+)
+from ltx_pipelines.utils.constants import detect_params
+from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser
+from ltx_pipelines.utils.media_io import encode_audio
+from ltx_pipelines.utils.types import ModalitySpec, OffloadMode
+
+# Placeholder pixel dimensions used for ``VideoPixelShape`` construction.
+# Audio-only generation reads ``frames`` and ``fps`` from the pixel shape via
+# ``AudioLatentShape.from_video_pixel_shape`` (height/width are unused).
+_AUDIO_ONLY_PLACEHOLDER_RES = 512
+
+
+class T2AOneStagePipeline:
+    """
+    Single-stage text-to-audio generation pipeline.
+    Generates audio at the target duration in a single diffusion pass with
+    classifier-free guidance (CFG) on the audio modality only. The video
+    modality is fully absent — the transformer runs audio-only by passing
+    ``video=None`` to the ``DiffusionStage``.
+    Assumes full non distilled model is provided in the checkpoint_path.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        compilation_config: CompilationConfig | None = None,
+        offload_mode: OffloadMode = OffloadMode.NONE,
+    ):
+        self.dtype = torch.bfloat16
+        self.device = device or get_device()
+        self._scheduler = LTX2Scheduler()
+        self.prompt_encoder = PromptEncoder(
+            checkpoint_path=checkpoint_path,
+            gemma_root=gemma_root,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+            offload_mode=offload_mode,
+        )
+        # Audio-only: build an audio-only transformer (model_configurator) so the video
+        # weights are never instantiated, plus a use-case-specific SDOps that restricts
+        # checkpoint reads to the audio model's keys, so the video weights are never even
+        # read from disk (the loader skips any key the SDOps maps to None).
+        self.stage = DiffusionStage(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            loras=tuple(loras),
+            quantization=quantization,
+            registry=registry,
+            compilation_config=compilation_config,
+            offload_mode=offload_mode,
+            model_configurator=LTXAudioOnlyModelConfigurator,
+            model_sd_ops=LTXV_AUDIO_ONLY_MODEL_COMFY_RENAMING_MAP,
+        )
+        self.audio_decoder = AudioDecoder(
+            checkpoint_path=checkpoint_path,
+            dtype=self.dtype,
+            device=self.device,
+            registry=registry,
+        )
+
+    def __call__(
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        enhance_prompt: bool = False,
+        max_batch_size: int = 1,
+        sigmas: torch.Tensor | None = None,
+    ) -> Audio:
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        ctx_p, ctx_n = self.prompt_encoder(
+            [prompt, negative_prompt],
+            enhance_first_prompt=enhance_prompt,
+            enhance_prompt_image=None,
+            enhance_prompt_seed=seed,
+        )
+        a_context_p = ctx_p.audio_encoding
+        a_context_n = ctx_n.audio_encoding
+
+        sigmas = (sigmas if sigmas is not None else self._scheduler.execute(steps=num_inference_steps)).to(
+            dtype=torch.float32, device=self.device
+        )
+
+        # Normalize to a guider factory. Plain ``MultiModalGuiderParams`` (the default /
+        # CLI case) becomes a simple sigma-independent guider, but callers may also pass
+        # their own factory for sigma-dependent guidance; ``FactoryGuidedDenoiser`` always
+        # consumes a factory.
+        audio_guider_factory = create_multimodal_guider_factory(
+            params=audio_guider_params,
+            negative_context=a_context_n,
+        )
+
+        _, audio_state = self.stage(
+            denoiser=FactoryGuidedDenoiser(
+                v_context=None,
+                a_context=a_context_p,
+                video_guider_factory=None,
+                audio_guider_factory=audio_guider_factory,
+            ),
+            sigmas=sigmas,
+            noiser=noiser,
+            width=_AUDIO_ONLY_PLACEHOLDER_RES,
+            height=_AUDIO_ONLY_PLACEHOLDER_RES,
+            frames=num_frames,
+            fps=frame_rate,
+            video=None,
+            audio=ModalitySpec(context=a_context_p),
+            max_batch_size=max_batch_size,
+        )
+
+        return self.audio_decoder(audio_state.latent)
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.basicConfig(level=logging.INFO)
+    checkpoint_path = detect_checkpoint_path()
+    params = detect_params(checkpoint_path)
+    parser = default_1_stage_t2a_arg_parser(params=params)
+    args = parser.parse_args()
+    pipeline = T2AOneStagePipeline(
+        checkpoint_path=args.checkpoint_path,
+        gemma_root=args.gemma_root,
+        loras=tuple(args.lora) if args.lora else (),
+        quantization=args.quantization,
+        compilation_config=args.compile,
+        offload_mode=args.offload_mode,
+    )
+    audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        audio_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.audio_cfg_guidance_scale,
+            stg_scale=args.audio_stg_guidance_scale,
+            rescale_scale=args.audio_rescale_scale,
+            # Audio-only generation has no video modality, so the video->audio
+            # (v2a) cross-modal guidance is meaningless here. 1.0 disables it.
+            modality_scale=1.0,
+            skip_step=args.audio_skip_step,
+            stg_blocks=args.audio_stg_blocks,
+        ),
+        max_batch_size=args.max_batch_size,
+    )
+
+    encode_audio(audio=audio, output_path=args.output_path)
+
+
+if __name__ == "__main__":
+    main()

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