Automated PR - 2026-03-04

README.md

@@ -63,15 +63,18 @@ Download the following models from the [LTX-2 HuggingFace repository](https://hu
 ### Available Pipelines
 
 * **[TI2VidTwoStagesPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)** - Production-quality text/image-to-video with 2x upsampling (recommended)
+* **[TI2VidTwoStagesRes2sPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_res2s.py)** - Same two-stage flow as above but uses the res_2s second-order sampler (fewer steps, different quality/speed trade-off)
 * **[TI2VidOneStagePipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py)** - Single-stage generation for quick prototyping
 * **[DistilledPipeline](packages/ltx-pipelines/src/ltx_pipelines/distilled.py)** - Fastest inference with 8 predefined sigmas
-* **[ICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py)** - Video-to-video and image-to-video transformations
+* **[ICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py)** - Video-to-video and image-to-video transformations (uses distilled model.)
 * **[KeyframeInterpolationPipeline](packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py)** - Interpolate between keyframe images
+* **[A2VidPipelineTwoStage](packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py)** - Audio-to-video generation conditioned on an input audio file
+* **[RetakePipeline](packages/ltx-pipelines/src/ltx_pipelines/retake.py)** - Regenerate a specific time region of an existing video
 
 ### ⚡ Optimization Tips
 
 * **Use DistilledPipeline** - Fastest inference with only 8 predefined sigmas (8 steps stage 1, 4 steps stage 2)
-* **Enable FP8 transformer** - Enables lower memory footprint: `--enable-fp8` (CLI) or `fp8transformer=True` (Python)
+* **Enable FP8 quantization** - Enables lower memory footprint: `--quantization fp8-cast` (CLI) or `quantization=QuantizationPolicy.fp8_cast()` (Python). For Hopper GPUs with TensorRT-LLM, use `--quantization fp8-scaled-mm` for FP8 scaled matrix multiplication.
 * **Install attention optimizations** - Use xFormers (`uv sync --extra xformers`) or [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) for Hopper GPUs
 * **Use gradient estimation** - Reduce inference steps from 40 to 20-30 while maintaining quality (see [pipeline documentation](packages/ltx-pipelines/README.md#denoising-loop-optimization))
 * **Skip memory cleanup** - If you have sufficient VRAM, disable automatic memory cleanup between stages for faster processing

packages/ltx-core/README.md

@@ -56,6 +56,58 @@ pip install -e packages/ltx-core
 - **Loader** ([`loader/`](src/ltx_core/loader/)): Model loading from `.safetensors`, LoRA fusion, weight remapping, and memory management
 - **Quantization** ([`quantization/`](src/ltx_core/quantization/)): FP8 quantization backends for reduced memory footprint and faster inference
 
+### Loader
+
+The `loader/` module provides `SingleGPUModelBuilder`, a frozen dataclass that loads a PyTorch model from `.safetensors` checkpoints and optionally fuses one or more LoRA adapters.
+
+#### Basic usage
+
+```python
+from ltx_core.loader import SingleGPUModelBuilder
+
+builder = SingleGPUModelBuilder(
+    model_class_configurator=MyModelConfigurator,
+    model_path="/path/to/model.safetensors",
+)
+model = builder.build(device=torch.device("cuda"))
+```
+
+#### Loading LoRA adapters
+
+Use the `.lora()` method to attach one or more LoRA adapters before calling `.build()`:
+
+```python
+builder = (
+    SingleGPUModelBuilder(
+        model_class_configurator=MyModelConfigurator,
+        model_path="/path/to/model.safetensors",
+    )
+    .lora("/path/to/lora_a.safetensors", strength=0.8)
+    .lora("/path/to/lora_b.safetensors", strength=0.5)
+)
+model = builder.build(device=torch.device("cuda"))
+```
+
+#### Memory-efficient LoRA loading (`lora_load_device`)
+
+By default, LoRA weights are loaded onto the **CPU** (`lora_load_device=torch.device("cpu")`).  This means each LoRA adapter is kept in CPU memory and transferred to the GPU sequentially during weight fusion, which keeps peak GPU memory low even when fusing large adapters.
+
+If all adapters fit comfortably in GPU memory you can skip the CPU staging by setting `lora_load_device` to the target CUDA device:
+
+```python
+import torch
+from ltx_core.loader import SingleGPUModelBuilder
+
+# Load LoRA weights directly onto the GPU (faster, but uses more GPU memory)
+builder = SingleGPUModelBuilder(
+    model_class_configurator=MyModelConfigurator,
+    model_path="/path/to/model.safetensors",
+    lora_load_device=torch.device("cuda"),
+).lora("/path/to/lora.safetensors", strength=1.0)
+
+model = builder.build(device=torch.device("cuda"))
+```
+
 ### Quantization
 
 The `quantization/` module provides FP8 quantization support for the LTX-2 transformer, significantly reducing memory usage while maintaining quality. Two backends are available:

packages/ltx-core/src/ltx_core/components/diffusion_steps.py

@@ -12,7 +12,7 @@ class EulerDiffusionStep(DiffusionStepProtocol):
     """
 
     def step(
-        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
+        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int, **_kwargs
     ) -> torch.Tensor:
         sigma = sigmas[step_index]
         sigma_next = sigmas[step_index + 1]
@@ -20,3 +20,76 @@ class EulerDiffusionStep(DiffusionStepProtocol):
         velocity = to_velocity(sample, sigma, denoised_sample)
 
         return (sample.to(torch.float32) + velocity.to(torch.float32) * dt).to(sample.dtype)
+
+
+class Res2sDiffusionStep(DiffusionStepProtocol):
+    """
+    Second-order diffusion step for res_2s sampling with SDE noise injection.
+    Used by the res_2s denoising loop. Advances the sample from the current
+    sigma to the next by mixing a deterministic update (from the denoised
+    prediction) with injected noise via ``get_sde_coeff``, producing
+    variance-preserving transitions.
+    """
+
+    @staticmethod
+    def get_sde_coeff(
+        sigma_next: torch.Tensor,
+        sigma_up: torch.Tensor | None = None,
+        sigma_down: torch.Tensor | None = None,
+        sigma_max: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Compute SDE coefficients (alpha_ratio, sigma_down, sigma_up) for the step.
+        Given either ``sigma_down`` or ``sigma_up``, returns the mixing
+        coefficients used for variance-preserving noise injection. If
+        ``sigma_up`` is provided, ``sigma_down`` and ``alpha_ratio`` are
+        derived; if ``sigma_down`` is provided, ``sigma_up`` and
+        ``alpha_ratio`` are derived.
+        """
+        if sigma_down is not None:
+            alpha_ratio = (1 - sigma_next) / (1 - sigma_down)
+            sigma_up = (sigma_next**2 - sigma_down**2 * alpha_ratio**2).clamp(min=0) ** 0.5
+        elif sigma_up is not None:
+            # Fallback to avoid sqrt(neg_num)
+            sigma_up.clamp_(max=sigma_next * 0.9999)
+            sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next)
+            sigma_signal = sigmax - sigma_next
+            sigma_residual = (sigma_next**2 - sigma_up**2).clamp(min=0) ** 0.5
+            alpha_ratio = sigma_signal + sigma_residual
+            sigma_down = sigma_residual / alpha_ratio
+        else:
+            alpha_ratio = torch.ones_like(sigma_next)
+            sigma_down = sigma_next
+            sigma_up = torch.zeros_like(sigma_next)
+
+        sigma_up = torch.nan_to_num(sigma_up if sigma_up is not None else torch.zeros_like(sigma_next), 0.0)
+        # Replace NaNs in sigma_down with corresponding sigma_next elements (float32)
+        nan_mask = torch.isnan(sigma_down)
+        sigma_down[nan_mask] = sigma_next[nan_mask].to(sigma_down.dtype)
+        alpha_ratio = torch.nan_to_num(alpha_ratio, 1.0)
+
+        return alpha_ratio, sigma_down, sigma_up
+
+    def step(
+        self,
+        sample: torch.Tensor,
+        denoised_sample: torch.Tensor,
+        sigmas: torch.Tensor,
+        step_index: int,
+        noise: torch.Tensor,
+    ) -> torch.Tensor:
+        """Advance one step with SDE noise injection via get_sde_coeff."""
+        sigma = sigmas[step_index]
+        sigma_next = sigmas[step_index + 1]
+        alpha_ratio, sigma_down, sigma_up = self.get_sde_coeff(sigma_next, sigma_up=sigma_next * 0.5)
+        output_dtype = denoised_sample.dtype
+        if torch.any(sigma_up == 0) or torch.any(sigma_next == 0):
+            return denoised_sample
+
+        # Extract epsilon prediction
+        eps_next = (sample - denoised_sample) / (sigma - sigma_next)
+        denoised_next = sample - sigma * eps_next
+
+        # Mix deterministic and stochastic components
+        x_noised = alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise
+        return x_noised.to(output_dtype)

packages/ltx-core/src/ltx_core/components/guiders.py

@@ -1,4 +1,5 @@
 import math
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass, field
 
 import torch
@@ -210,10 +211,31 @@ class MultiModalGuiderParams:
     "Skip step controlling how often the model skips the step."
 
 
+def _params_for_sigma_from_sorted_dict(
+    sigma: float, params_by_sigma: Sequence[tuple[float, MultiModalGuiderParams]]
+) -> MultiModalGuiderParams:
+    """
+    Return params for the given sigma from a sorted (sigma_upper_bound -> params) structure.
+    Keys are sorted descending (bin upper bounds). Bin i is (key_{i+1}, key_i].
+    Get all keys >= sigma; use last in list (smallest such key = upper bound of bin containing sigma),
+    or last entry in the sequence if list is empty (sigma above max key).
+    """
+    if not params_by_sigma:
+        raise ValueError("params_by_sigma must be non-empty")
+    sigma = float(sigma)
+    keys_desc = [k for k, _ in params_by_sigma]
+    keys_ge_sigma = [k for k in keys_desc if k >= sigma]
+    # sigma above all keys: use first bin (max key)
+    key = keys_ge_sigma[-1] if keys_ge_sigma else keys_desc[0]
+    return next(p for k, p in params_by_sigma if k == key)
+
+
 @dataclass(frozen=True)
 class MultiModalGuider:
     """
-    Multi-modal guider.
+    Multi-modal guider with constant params per instance.
+    For sigma-dependent params, use MultiModalGuiderFactory.build_from_sigma(sigma) to
+    obtain a guider for each step.
     """
 
     params: MultiModalGuiderParams
@@ -246,33 +268,93 @@ class MultiModalGuider:
         return pred
 
     def do_unconditional_generation(self) -> bool:
-        """
-        Returns True if the guider is doing unconditional generation.
-        """
+        """Returns True if the guider is doing unconditional generation."""
         return not math.isclose(self.params.cfg_scale, 1.0)
 
     def do_perturbed_generation(self) -> bool:
-        """
-        Returns True if the guider is doing perturbed generation.
-        """
+        """Returns True if the guider is doing perturbed generation."""
         return not math.isclose(self.params.stg_scale, 0.0)
 
     def do_isolated_modality_generation(self) -> bool:
-        """
-        Returns True if the guider is doing isolated modality generation.
-        """
+        """Returns True if the guider is doing isolated modality generation."""
         return not math.isclose(self.params.modality_scale, 1.0)
 
     def should_skip_step(self, step: int) -> bool:
-        """
-        Returns True if the guider should skip the step.
-        """
+        """Returns True if the guider should skip the step."""
         if self.params.skip_step == 0:
             return False
-
         return step % (self.params.skip_step + 1) != 0
 
 
+@dataclass(frozen=True)
+class MultiModalGuiderFactory:
+    """
+    Factory that creates a MultiModalGuider for a given sigma.
+    Single source of truth: _params_by_sigma (schedule). Use constant() for
+    one params for all sigma, from_dict() for sigma-binned params.
+    """
+
+    negative_context: torch.Tensor | None = None
+    _params_by_sigma: tuple[tuple[float, MultiModalGuiderParams], ...] = ()
+
+    @classmethod
+    def constant(
+        cls,
+        params: MultiModalGuiderParams,
+        negative_context: torch.Tensor | None = None,
+    ) -> "MultiModalGuiderFactory":
+        """Build a factory with constant params (same guider for all sigma)."""
+        return cls(
+            negative_context=negative_context,
+            _params_by_sigma=((float("inf"), params),),
+        )
+
+    @classmethod
+    def from_dict(
+        cls,
+        sigma_to_params: Mapping[float, MultiModalGuiderParams],
+        negative_context: torch.Tensor | None = None,
+    ) -> "MultiModalGuiderFactory":
+        """
+        Build a factory from a dict of sigma_value -> MultiModalGuiderParams.
+        Keys are sorted descending and used for bin lookup in params(sigma).
+        """
+        if not sigma_to_params:
+            raise ValueError("sigma_to_params must be non-empty")
+        sorted_items = tuple(sorted(sigma_to_params.items(), key=lambda x: x[0], reverse=True))
+        return cls(negative_context=negative_context, _params_by_sigma=sorted_items)
+
+    def params(self, sigma: float | torch.Tensor) -> MultiModalGuiderParams:
+        """Return params effective for the given sigma (getter; single source of truth)."""
+        sigma_val = float(sigma.item() if isinstance(sigma, torch.Tensor) else sigma)
+        return _params_for_sigma_from_sorted_dict(sigma_val, self._params_by_sigma)
+
+    def build_from_sigma(self, sigma: float | torch.Tensor) -> MultiModalGuider:
+        """Return a MultiModalGuider with params effective for the given sigma."""
+        return MultiModalGuider(
+            params=self.params(sigma),
+            negative_context=self.negative_context,
+        )
+
+
+def create_multimodal_guider_factory(
+    params: MultiModalGuiderParams | MultiModalGuiderFactory,
+    negative_context: torch.Tensor | None = None,
+) -> MultiModalGuiderFactory:
+    """
+    Create or return a MultiModalGuiderFactory. Pass constant params for a
+    single-params factory (uses MultiModalGuiderFactory.constant), or an existing
+    MultiModalGuiderFactory. When given a factory, returns it as-is unless
+    negative_context is provided. For sigma-dependent params use
+    MultiModalGuiderFactory.from_dict(...) and pass that as params.
+    """
+    if isinstance(params, MultiModalGuiderFactory):
+        if negative_context is not None and params.negative_context is not negative_context:
+            return MultiModalGuiderFactory.from_dict(dict(params._params_by_sigma), negative_context=negative_context)
+        return params
+    return MultiModalGuiderFactory.constant(params, negative_context=negative_context)
+
+
 def projection_coef(to_project: torch.Tensor, project_onto: torch.Tensor) -> torch.Tensor:
     batch_size = to_project.shape[0]
     positive_flat = to_project.reshape(batch_size, -1)

packages/ltx-core/src/ltx_core/components/protocols.py

@@ -97,5 +97,5 @@ class DiffusionStepProtocol(Protocol):
     """
 
     def step(
-        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int
+        self, sample: torch.Tensor, denoised_sample: torch.Tensor, sigmas: torch.Tensor, step_index: int, **kwargs
     ) -> torch.Tensor: ...

packages/ltx-core/src/ltx_core/components/schedulers.py

@@ -26,9 +26,10 @@ class LTX2Scheduler(SchedulerProtocol):
         base_shift: float = 0.95,
         stretch: bool = True,
         terminal: float = 0.1,
+        default_number_of_tokens: int = MAX_SHIFT_ANCHOR,
         **_kwargs,
     ) -> torch.FloatTensor:
-        tokens = math.prod(latent.shape[2:]) if latent is not None else MAX_SHIFT_ANCHOR
+        tokens = math.prod(latent.shape[2:]) if latent is not None else default_number_of_tokens
         sigmas = torch.linspace(1.0, 0.0, steps + 1)
 
         x1 = BASE_SHIFT_ANCHOR

packages/ltx-core/src/ltx_core/conditioning/__init__.py

@@ -3,6 +3,7 @@
 from ltx_core.conditioning.exceptions import ConditioningError
 from ltx_core.conditioning.item import ConditioningItem
 from ltx_core.conditioning.types import (
+    ConditioningItemAttentionStrengthWrapper,
     VideoConditionByKeyframeIndex,
     VideoConditionByLatentIndex,
     VideoConditionByReferenceLatent,
@@ -11,6 +12,7 @@ from ltx_core.conditioning.types import (
 __all__ = [
     "ConditioningError",
     "ConditioningItem",
+    "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
     "VideoConditionByReferenceLatent",

packages/ltx-core/src/ltx_core/conditioning/mask_utils.py

@@ -0,0 +1,210 @@
+"""Utilities for building 2D self-attention masks for conditioning items."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from ltx_core.types import LatentState
+
+
+def resolve_cross_mask(
+    attention_mask: float | int | torch.Tensor,
+    num_new_tokens: int,
+    batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Convert an attention_mask (scalar or tensor) to a (B, M) cross_mask tensor.
+    Args:
+        attention_mask: Scalar value applied uniformly, 1D tensor of shape (M,)
+            broadcast across batch, or 2D tensor of shape (B, M).
+        num_new_tokens: Number of new conditioning tokens M.
+        batch_size: Batch size B.
+        device: Device for the output tensor.
+        dtype: Data type for the output tensor.
+    Returns:
+        Cross-mask tensor of shape (B, M).
+    """
+    if isinstance(attention_mask, (int, float)):
+        return torch.full(
+            (batch_size, num_new_tokens),
+            fill_value=float(attention_mask),
+            device=device,
+            dtype=dtype,
+        )
+    mask = attention_mask.to(device=device, dtype=dtype)
+
+    # Handle scalar (0-D) tensor like a Python scalar.
+    if mask.dim() == 0:
+        return torch.full(
+            (batch_size, num_new_tokens),
+            fill_value=float(mask.item()),
+            device=device,
+            dtype=dtype,
+        )
+
+    if mask.dim() == 1:
+        if mask.shape[0] != num_new_tokens:
+            raise ValueError(
+                f"1-D attention_mask length must equal num_new_tokens ({num_new_tokens}), got shape {tuple(mask.shape)}"
+            )
+        mask = mask.unsqueeze(0).expand(batch_size, -1)
+    elif mask.dim() == 2:
+        b, m = mask.shape
+        if m != num_new_tokens:
+            raise ValueError(
+                f"2-D attention_mask second dimension must equal num_new_tokens ({num_new_tokens}), "
+                f"got shape {tuple(mask.shape)}"
+            )
+        if b not in (batch_size, 1):
+            raise ValueError(
+                f"2-D attention_mask batch dimension must equal batch_size ({batch_size}) or 1, "
+                f"got shape {tuple(mask.shape)}"
+            )
+        if b == 1 and batch_size > 1:
+            mask = mask.expand(batch_size, -1)
+    else:
+        raise ValueError(
+            f"attention_mask tensor must be 0-D, 1-D, or 2-D, got {mask.dim()}-D with shape {tuple(mask.shape)}"
+        )
+    return mask
+
+
+def update_attention_mask(
+    latent_state: LatentState,
+    attention_mask: float | torch.Tensor | None,
+    num_noisy_tokens: int,
+    num_new_tokens: int,
+    batch_size: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor | None:
+    """Build or update the self-attention mask for newly appended conditioning tokens.
+    If *attention_mask* is ``None`` and no existing mask is present, returns
+    ``None``.  If *attention_mask* is ``None`` but an existing mask is present,
+    the mask is expanded with full attention (1s) for the new tokens so that
+    its dimensions stay consistent with the growing latent sequence.  Otherwise,
+    resolves *attention_mask* to a per-token cross-mask and expands the 2-D
+    attention mask via :func:`build_attention_mask`.
+    Args:
+        latent_state: Current latent state (provides the existing mask and total
+            existing-token count).
+        attention_mask: Per-token attention weight. Scalar, 1-D ``(M,)``, 2-D
+            ``(B, M)`` tensor, or ``None`` (no-op).
+        num_noisy_tokens: Number of original noisy tokens (from
+            ``latent_tools.target_shape.token_count()``).
+        num_new_tokens: Number of new conditioning tokens being appended.
+        batch_size: Batch size.
+        device: Device for the output tensor.
+        dtype: Data type for the output tensor.
+    Returns:
+        Updated attention mask of shape ``(B, N+M, N+M)``, or ``None`` if no
+        masking is needed.
+    """
+    if attention_mask is None:
+        if latent_state.attention_mask is None:
+            return None
+        # Existing mask present but no new mask requested: pad with 1s (full
+        # attention) so the mask dimensions stay consistent with the growing
+        # latent sequence.
+        cross_mask = torch.ones(batch_size, num_new_tokens, device=device, dtype=dtype)
+        return build_attention_mask(
+            existing_mask=latent_state.attention_mask,
+            num_noisy_tokens=num_noisy_tokens,
+            num_new_tokens=num_new_tokens,
+            num_existing_tokens=latent_state.latent.shape[1],
+            cross_mask=cross_mask,
+            device=device,
+            dtype=dtype,
+        )
+
+    cross_mask = resolve_cross_mask(attention_mask, num_new_tokens, batch_size, device, dtype)
+    return build_attention_mask(
+        existing_mask=latent_state.attention_mask,
+        num_noisy_tokens=num_noisy_tokens,
+        num_new_tokens=num_new_tokens,
+        num_existing_tokens=latent_state.latent.shape[1],
+        cross_mask=cross_mask,
+        device=device,
+        dtype=dtype,
+    )
+
+
+def build_attention_mask(
+    existing_mask: torch.Tensor | None,
+    num_noisy_tokens: int,
+    num_new_tokens: int,
+    num_existing_tokens: int,
+    cross_mask: torch.Tensor,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """
+    Expand the attention mask to include newly appended conditioning tokens.
+    Each conditioning item appends M new reference tokens to the sequence. This function
+    builds a (B, N+M, N+M) attention mask with the following block structure:
+                     noisy      prev_ref    new_ref
+                   (N_noisy)   (N-N_noisy)    (M)
+                 ┌───────────┬───────────┬───────────┐
+        noisy    │           │           │           │
+       (N_noisy) │  existing │  existing │   cross   │
+                 │           │           │           │
+                 ├───────────┼───────────┼───────────┤
+       prev_ref  │           │           │           │
+      (N-N_noisy)│  existing │  existing │     0     │
+                 │           │           │           │
+                 ├───────────┼───────────┼───────────┤
+       new_ref   │           │           │           │
+         (M)     │   cross   │     0     │     1     │
+                 │           │           │           │
+                 └───────────┴───────────┴───────────┘
+    Where:
+      - **existing**: preserved from the previous mask (or 1.0 if first conditioning)
+      - **cross**: values from *cross_mask* (shape B, M), in [0, 1]
+      - **0**: no attention between different reference groups
+    Args:
+        existing_mask: Current attention mask of shape (B, N, N), or None if no mask exists yet.
+            When None, the top-left NxN block is filled with 1s (full attention between all
+            existing tokens including any prior reference tokens that had no mask).
+        num_noisy_tokens: Number of original noisy tokens (always at positions [0:num_noisy_tokens]).
+        num_new_tokens: Number of new conditioning tokens M being appended.
+        num_existing_tokens: Total number of current tokens N (noisy + any prior conditioning tokens).
+        cross_mask: Per-token attention weight of shape (B, M) controlling attention between
+            new reference tokens and noisy tokens. Values in [0, 1].
+        device: Device for the output tensor.
+        dtype: Data type for the output tensor.
+    Returns:
+        Attention mask of shape (B, N+M, N+M) with values in [0, 1].
+    """
+    batch_size = cross_mask.shape[0]
+    total = num_existing_tokens + num_new_tokens
+
+    # Start with zeros
+    mask = torch.zeros((batch_size, total, total), device=device, dtype=dtype)
+
+    # Top-left: preserve existing mask or fill with 1s for noisy tokens
+    if existing_mask is not None:
+        mask[:, :num_existing_tokens, :num_existing_tokens] = existing_mask
+    else:
+        mask[:, :num_existing_tokens, :num_existing_tokens] = 1.0
+
+    # Bottom-right: new reference tokens fully attend to themselves
+    mask[:, num_existing_tokens:, num_existing_tokens:] = 1.0
+
+    # Cross-attention between noisy tokens and new reference tokens
+    # cross_mask shape: (B, M) -> broadcast to (B, N_noisy, M) and (B, M, N_noisy)
+
+    # Noisy tokens attending to new reference tokens: [0:N_noisy, N:N+M]
+    # Each column j in this block gets cross_mask[:, j]
+    mask[:, :num_noisy_tokens, num_existing_tokens:] = cross_mask.unsqueeze(1)
+
+    # New reference tokens attending to noisy tokens: [N:N+M, 0:N_noisy]
+    # Each row i in this block gets cross_mask[:, i]
+    mask[:, num_existing_tokens:, :num_noisy_tokens] = cross_mask.unsqueeze(2)
+
+    # [N_noisy:N, N:N+M] and [N:N+M, N_noisy:N] remain 0 (no cross-ref attention)
+
+    return mask

packages/ltx-core/src/ltx_core/conditioning/types/__init__.py

@@ -1,10 +1,12 @@
 """Conditioning type implementations."""
 
+from ltx_core.conditioning.types.attention_strength_wrapper import ConditioningItemAttentionStrengthWrapper
 from ltx_core.conditioning.types.keyframe_cond import VideoConditionByKeyframeIndex
 from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
 from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
 
 __all__ = [
+    "ConditioningItemAttentionStrengthWrapper",
     "VideoConditionByKeyframeIndex",
     "VideoConditionByLatentIndex",
     "VideoConditionByReferenceLatent",

packages/ltx-core/src/ltx_core/conditioning/types/attention_strength_wrapper.py

@@ -0,0 +1,71 @@
+"""Wrapper conditioning item that adds attention masking to any inner conditioning."""
+
+from dataclasses import replace
+
+import torch
+
+from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.mask_utils import update_attention_mask
+from ltx_core.tools import LatentTools
+from ltx_core.types import LatentState
+
+
+class ConditioningItemAttentionStrengthWrapper(ConditioningItem):
+    """Wraps a conditioning item to add an attention mask for its tokens.
+    Separates the *attention-masking* concern from the underlying conditioning
+    logic (token layout, positional encoding, denoise strength).  The inner
+    conditioning item appends tokens to the latent sequence as usual, and this
+    wrapper then builds or updates the self-attention mask so that the newly
+    added tokens interact with the noisy tokens according to *attention_mask*.
+    Args:
+        conditioning: Any conditioning item that appends tokens to the latent.
+        attention_mask: Per-token attention weight controlling how strongly the
+            new conditioning tokens attend to/from noisy tokens.  Can be a
+            scalar (float) applied uniformly, or a tensor of shape ``(B, M)``
+            for spatial control, where ``M = F * H * W`` is the number of
+            patchified conditioning tokens.  Values in ``[0, 1]``.
+    Example::
+        cond = ConditioningItemAttentionStrengthWrapper(
+            VideoConditionByReferenceLatent(latent=ref, strength=1.0),
+            attention_mask=0.5,
+        )
+        state = cond.apply_to(latent_state, latent_tools)
+    """
+
+    def __init__(
+        self,
+        conditioning: ConditioningItem,
+        attention_mask: float | torch.Tensor,
+    ):
+        self.conditioning = conditioning
+        self.attention_mask = attention_mask
+
+    def apply_to(
+        self,
+        latent_state: LatentState,
+        latent_tools: LatentTools,
+    ) -> LatentState:
+        """Apply inner conditioning, then build the attention mask for its tokens."""
+        # Snapshot the original state for mask building
+        original_state = latent_state
+
+        # Inner conditioning appends tokens (positions, denoise mask, etc.)
+        new_state = self.conditioning.apply_to(latent_state, latent_tools)
+
+        num_new_tokens = new_state.latent.shape[1] - original_state.latent.shape[1]
+        if num_new_tokens == 0:
+            return new_state
+
+        # Build the attention mask using the *original* state as the reference
+        # so that the block structure is computed correctly.
+        new_attention_mask = update_attention_mask(
+            latent_state=original_state,
+            attention_mask=self.attention_mask,
+            num_noisy_tokens=latent_tools.target_shape.token_count(),
+            num_new_tokens=num_new_tokens,
+            batch_size=new_state.latent.shape[0],
+            device=new_state.latent.device,
+            dtype=new_state.latent.dtype,
+        )
+
+        return replace(new_state, attention_mask=new_attention_mask)

packages/ltx-core/src/ltx_core/conditioning/types/keyframe_cond.py

@@ -2,6 +2,7 @@ import torch
 
 from ltx_core.components.patchifiers import get_pixel_coords
 from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.mask_utils import update_attention_mask
 from ltx_core.tools import VideoLatentTools
 from ltx_core.types import LatentState, VideoLatentShape
 
@@ -11,6 +12,11 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
     Conditions video generation on keyframe latents at a specific frame index.
     Appends keyframe tokens to the latent state with positions offset by frame_idx,
     and sets denoise strength according to the strength parameter.
+    To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
+    Args:
+        keyframes: Keyframe latents [B, C, F, H, W].
+        frame_idx: Frame index offset for positional encoding.
+        strength: Conditioning strength (1.0 = clean, 0.0 = fully denoised).
     """
 
     def __init__(self, keyframes: torch.Tensor, frame_idx: int, strength: float):
@@ -45,9 +51,20 @@ class VideoConditionByKeyframeIndex(ConditioningItem):
             dtype=self.keyframes.dtype,
         )
 
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=latent_tools.target_shape.token_count(),
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            device=self.keyframes.device,
+            dtype=self.keyframes.dtype,
+        )
+
         return LatentState(
             latent=torch.cat([latent_state.latent, tokens], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
+            attention_mask=new_attention_mask,
         )

packages/ltx-core/src/ltx_core/conditioning/types/reference_video_cond.py

@@ -4,6 +4,7 @@ import torch
 
 from ltx_core.components.patchifiers import get_pixel_coords
 from ltx_core.conditioning.item import ConditioningItem
+from ltx_core.conditioning.mask_utils import update_attention_mask
 from ltx_core.tools import VideoLatentTools
 from ltx_core.types import LatentState, VideoLatentShape
 
@@ -19,6 +20,7 @@ class VideoConditionByReferenceLatent(ConditioningItem):
     `downscale_factor` scales reference positions to match target coordinates, preserving
     the learned positional relationships. This must match the factor used during training
     (stored in LoRA metadata).
+    To add attention masking, wrap with :class:`ConditioningItemAttentionStrengthWrapper`.
     Args:
         latent: Reference video latents [B, C, F, H, W]
         downscale_factor: Target/reference resolution ratio (e.g., 2 = half-resolution
@@ -70,9 +72,20 @@ class VideoConditionByReferenceLatent(ConditioningItem):
             dtype=self.latent.dtype,
         )
 
+        new_attention_mask = update_attention_mask(
+            latent_state=latent_state,
+            attention_mask=None,
+            num_noisy_tokens=latent_tools.target_shape.token_count(),
+            num_new_tokens=tokens.shape[1],
+            batch_size=tokens.shape[0],
+            device=self.latent.device,
+            dtype=self.latent.dtype,
+        )
+
         return LatentState(
             latent=torch.cat([latent_state.latent, tokens], dim=1),
             denoise_mask=torch.cat([latent_state.denoise_mask, denoise_mask], dim=1),
             positions=torch.cat([latent_state.positions, positions], dim=2),
             clean_latent=torch.cat([latent_state.clean_latent, tokens], dim=1),
+            attention_mask=new_attention_mask,
         )

packages/ltx-core/src/ltx_core/loader/fuse_loras.py

@@ -53,8 +53,11 @@ def _prepare_deltas(
     for lsd, coef in lora_sd_and_strengths:
         if key_a not in lsd.sd or key_b not in lsd.sd:
             continue
-        product = torch.matmul(lsd.sd[key_b] * coef, lsd.sd[key_a])
-        deltas.append(product.to(dtype=dtype, device=device))
+        a = lsd.sd[key_a].to(device=device)
+        b = lsd.sd[key_b].to(device=device)
+        product = torch.matmul(b * coef, a)
+        del a, b
+        deltas.append(product.to(dtype=dtype))
     if len(deltas) == 0:
         return None
     elif len(deltas) == 1:

packages/ltx-core/src/ltx_core/loader/single_gpu_model_builder.py

@@ -26,6 +26,19 @@ logger: logging.Logger = logging.getLogger(__name__)
 class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType], LoRAAdaptableProtocol):
     """
     Builder for PyTorch models residing on a single GPU.
+    Attributes:
+        model_class_configurator: Class responsible for constructing the model from a config dict.
+        model_path: Path (or tuple of shard paths) to the model's `.safetensors` checkpoint(s).
+        model_sd_ops: Optional state-dict operations applied when loading the model weights.
+        module_ops: Sequence of module-level mutations applied to the meta model before weight loading.
+        loras: Sequence of LoRA adapters (path, strength, optional sd_ops) to fuse into the model.
+        model_loader: Strategy for loading state dicts from disk. Defaults to
+            :class:`SafetensorsModelStateDictLoader`.
+        registry: Cache for already-loaded state dicts. Defaults to :class:`DummyRegistry` (no caching).
+        lora_load_device: Device used when loading LoRA weight tensors from disk. Defaults to
+            ``torch.device("cpu")``, which keeps LoRA weights in CPU memory and transfers them to
+            the target GPU sequentially during fusion, reducing peak GPU memory usage compared to
+            loading all LoRA weights directly onto the GPU at once.
     """
 
     model_class_configurator: type[ModelConfigurator[ModelType]]
@@ -35,6 +48,7 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
     loras: tuple[LoraPathStrengthAndSDOps, ...] = field(default_factory=tuple)
     model_loader: StateDictLoader = field(default_factory=SafetensorsModelStateDictLoader)
     registry: Registry = field(default_factory=DummyRegistry)
+    lora_load_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
 
     def lora(self, lora_path: str, strength: float = 1.0, sd_ops: SDOps | None = None) -> "SingleGPUModelBuilder":
         return replace(self, loras=(*self.loras, LoraPathStrengthAndSDOps(lora_path, strength, sd_ops)))
@@ -85,7 +99,8 @@ class SingleGPUModelBuilder(Generic[ModelType], ModelBuilderProtocol[ModelType],
             return self._return_model(meta_model, device)
 
         lora_state_dicts = [
-            self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=device) for lora in self.loras
+            self.load_sd([lora.path], sd_ops=lora.sd_ops, registry=self.registry, device=self.lora_load_device)
+            for lora in self.loras
         ]
         lora_sd_and_strengths = [
             LoraStateDictWithStrength(sd, strength)

packages/ltx-core/src/ltx_core/model/audio_vae/__init__.py

@@ -1,6 +1,6 @@
 """Audio VAE model components."""
 
-from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_audio
+from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder, decode_audio, encode_audio
 from ltx_core.model.audio_vae.model_configurator import (
     AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
     AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
@@ -10,7 +10,7 @@ from ltx_core.model.audio_vae.model_configurator import (
     VocoderConfigurator,
 )
 from ltx_core.model.audio_vae.ops import AudioProcessor
-from ltx_core.model.audio_vae.vocoder import Vocoder
+from ltx_core.model.audio_vae.vocoder import Vocoder, VocoderWithBWE
 
 __all__ = [
     "AUDIO_VAE_DECODER_COMFY_KEYS_FILTER",
@@ -23,5 +23,7 @@ __all__ = [
     "AudioProcessor",
     "Vocoder",
     "VocoderConfigurator",
+    "VocoderWithBWE",
     "decode_audio",
+    "encode_audio",
 ]

packages/ltx-core/src/ltx_core/model/audio_vae/audio_vae.py

@@ -8,12 +8,12 @@ from ltx_core.model.audio_vae.attention import AttentionType, make_attn
 from ltx_core.model.audio_vae.causal_conv_2d import make_conv2d
 from ltx_core.model.audio_vae.causality_axis import CausalityAxis
 from ltx_core.model.audio_vae.downsample import build_downsampling_path
-from ltx_core.model.audio_vae.ops import PerChannelStatistics
+from ltx_core.model.audio_vae.ops import AudioProcessor, PerChannelStatistics
 from ltx_core.model.audio_vae.resnet import ResnetBlock
 from ltx_core.model.audio_vae.upsample import build_upsampling_path
 from ltx_core.model.audio_vae.vocoder import Vocoder
 from ltx_core.model.common.normalization import NormType, build_normalization_layer
-from ltx_core.types import AudioLatentShape
+from ltx_core.types import Audio, AudioLatentShape
 
 LATENT_DOWNSAMPLE_FACTOR = 4
 
@@ -245,6 +245,34 @@ class AudioEncoder(torch.nn.Module):
         return self.patchifier.unpatchify(latent_normalized, latent_shape)
 
 
+def encode_audio(
+    audio: Audio,
+    audio_encoder: AudioEncoder,
+    audio_processor: AudioProcessor | None = None,
+) -> torch.Tensor:
+    """Encode audio waveform into latent representation.
+    Args:
+        audio: Audio container with waveform tensor of shape (batch, channels, samples) and sampling rate.
+        audio_encoder: Audio encoder model
+        audio_processor: Audio processor model (optional, if not provided, it will be created from the audio encoder)
+    """
+    dtype = next(audio_encoder.parameters()).dtype
+    device = next(audio_encoder.parameters()).device
+
+    if audio_processor is None:
+        audio_processor = AudioProcessor(
+            target_sample_rate=audio_encoder.sample_rate,
+            mel_bins=audio_encoder.mel_bins,
+            mel_hop_length=audio_encoder.mel_hop_length,
+            n_fft=audio_encoder.n_fft,
+        ).to(device=device)
+
+    mel_spectrogram = audio_processor.waveform_to_mel(audio.to(device=device))
+
+    latent = audio_encoder(mel_spectrogram.to(dtype=dtype))
+    return latent
+
+
 class AudioDecoder(torch.nn.Module):
     """
     Symmetric decoder that reconstructs audio spectrograms from latent features.
@@ -465,7 +493,7 @@ class AudioDecoder(torch.nn.Module):
         return torch.tanh(h) if self.tanh_out else h
 
 
-def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> torch.Tensor:
+def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "Vocoder") -> Audio:
     """
     Decode an audio latent representation using the provided audio decoder and vocoder.
     Args:
@@ -473,8 +501,8 @@ def decode_audio(latent: torch.Tensor, audio_decoder: "AudioDecoder", vocoder: "
         audio_decoder: Model to decode the latent to waveform features.
         vocoder: Model to convert decoded features to audio waveform.
     Returns:
-        Decoded audio as a float tensor.
+        Decoded audio with waveform and sampling rate.
     """
     decoded_audio = audio_decoder(latent)
-    decoded_audio = vocoder(decoded_audio).squeeze(0).float()
-    return decoded_audio
+    waveform = vocoder(decoded_audio).squeeze(0).float()
+    return Audio(waveform=waveform, sampling_rate=vocoder.output_sampling_rate)

packages/ltx-core/src/ltx_core/model/audio_vae/model_configurator.py

@@ -1,30 +1,107 @@
-from ltx_core.loader.sd_ops import SDOps
+import torch
+
+from ltx_core.loader.sd_ops import KeyValueOperationResult, SDOps
 from ltx_core.model.audio_vae.attention import AttentionType
 from ltx_core.model.audio_vae.audio_vae import AudioDecoder, AudioEncoder
 from ltx_core.model.audio_vae.causality_axis import CausalityAxis
-from ltx_core.model.audio_vae.vocoder import Vocoder
+from ltx_core.model.audio_vae.vocoder import MelSTFT, Vocoder, VocoderWithBWE
 from ltx_core.model.common.normalization import NormType
 from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.utils import check_config_value
+
+
+def _vocoder_from_config(
+    cfg: dict,
+    apply_final_activation: bool = True,
+    output_sampling_rate: int | None = None,
+) -> Vocoder:
+    """Instantiate a Vocoder from a flat config dict.
+    Args:
+        cfg: Vocoder config dict (keys match Vocoder constructor args).
+        apply_final_activation: Whether to apply tanh/clamp at the output.
+        output_sampling_rate: Explicit override for the output sample rate.
+            When None, reads from cfg["output_sampling_rate"] (default 24000).
+    """
+    return Vocoder(
+        resblock_kernel_sizes=cfg.get("resblock_kernel_sizes", [3, 7, 11]),
+        upsample_rates=cfg.get("upsample_rates", [6, 5, 2, 2, 2]),
+        upsample_kernel_sizes=cfg.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]),
+        resblock_dilation_sizes=cfg.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]),
+        upsample_initial_channel=cfg.get("upsample_initial_channel", 1024),
+        resblock=cfg.get("resblock", "1"),
+        output_sampling_rate=(
+            output_sampling_rate if output_sampling_rate is not None else cfg.get("output_sampling_rate", 24000)
+        ),
+        activation=cfg.get("activation", "snake"),
+        use_tanh_at_final=cfg.get("use_tanh_at_final", True),
+        apply_final_activation=apply_final_activation,
+        use_bias_at_final=cfg.get("use_bias_at_final", True),
+    )
 
 
 class VocoderConfigurator(ModelConfigurator[Vocoder]):
+    """Configurator that auto-detects the checkpoint format.
+    Returns a plain Vocoder for pre-ltx-2.3 checkpoints (flat config) or a
+    VocoderWithBWE for ltx-2.3+ checkpoints (nested "vocoder" + "bwe" config).
+    """
+
     @classmethod
-    def from_config(cls: type[Vocoder], config: dict) -> Vocoder:
-        config = config.get("vocoder", {})
-        return Vocoder(
-            resblock_kernel_sizes=config.get("resblock_kernel_sizes", [3, 7, 11]),
-            upsample_rates=config.get("upsample_rates", [6, 5, 2, 2, 2]),
-            upsample_kernel_sizes=config.get("upsample_kernel_sizes", [16, 15, 8, 4, 4]),
-            resblock_dilation_sizes=config.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]]),
-            upsample_initial_channel=config.get("upsample_initial_channel", 1024),
-            stereo=config.get("stereo", True),
-            resblock=config.get("resblock", "1"),
-            output_sample_rate=config.get("output_sample_rate", 24000),
+    def from_config(cls: type[Vocoder], config: dict) -> Vocoder | VocoderWithBWE:
+        cfg = config.get("vocoder", {})
+
+        if "bwe" not in cfg:
+            check_config_value(cfg, "resblock", "1")
+            check_config_value(cfg, "stereo", True)
+            return _vocoder_from_config(cfg)
+
+        vocoder_cfg = cfg.get("vocoder", {})
+        bwe_cfg = cfg["bwe"]
+
+        check_config_value(vocoder_cfg, "resblock", "AMP1")
+        check_config_value(vocoder_cfg, "stereo", True)
+        check_config_value(vocoder_cfg, "activation", "snakebeta")
+        check_config_value(bwe_cfg, "resblock", "AMP1")
+        check_config_value(bwe_cfg, "stereo", True)
+        check_config_value(bwe_cfg, "activation", "snakebeta")
+
+        vocoder = _vocoder_from_config(
+            vocoder_cfg,
+            output_sampling_rate=bwe_cfg["input_sampling_rate"],
+        )
+        bwe_generator = _vocoder_from_config(
+            bwe_cfg,
+            apply_final_activation=False,
+            output_sampling_rate=bwe_cfg["output_sampling_rate"],
+        )
+        mel_stft = MelSTFT(
+            filter_length=bwe_cfg["n_fft"],
+            hop_length=bwe_cfg["hop_length"],
+            win_length=bwe_cfg["n_fft"],
+            n_mel_channels=bwe_cfg["num_mels"],
+        )
+        return VocoderWithBWE(
+            vocoder=vocoder,
+            bwe_generator=bwe_generator,
+            mel_stft=mel_stft,
+            input_sampling_rate=bwe_cfg["input_sampling_rate"],
+            output_sampling_rate=bwe_cfg["output_sampling_rate"],
+            hop_length=bwe_cfg["hop_length"],
         )
 
 
+def _strip_vocoder_prefix(key: str, value: torch.Tensor) -> list[KeyValueOperationResult]:
+    """Strip the leading 'vocoder.' prefix exactly once.
+    Uses removeprefix instead of str.replace so that BWE keys like
+    'vocoder.vocoder.conv_pre' become 'vocoder.conv_pre' (not 'conv_pre').
+    Works identically for legacy keys like 'vocoder.conv_pre' → 'conv_pre'.
+    """
+    return [KeyValueOperationResult(key.removeprefix("vocoder."), value)]
+
+
 VOCODER_COMFY_KEYS_FILTER = (
-    SDOps("VOCODER_COMFY_KEYS_FILTER").with_matching(prefix="vocoder.").with_replacement("vocoder.", "")
+    SDOps("VOCODER_COMFY_KEYS_FILTER")
+    .with_matching(prefix="vocoder.")
+    .with_kv_operation(operation=_strip_vocoder_prefix, key_prefix="vocoder.")
 )
 
 

packages/ltx-core/src/ltx_core/model/audio_vae/ops.py

@@ -2,26 +2,28 @@ import torch
 import torchaudio
 from torch import nn
 
+from ltx_core.types import Audio
+
 
 class AudioProcessor(nn.Module):
     """Converts audio waveforms to log-mel spectrograms with optional resampling."""
 
     def __init__(
         self,
-        sample_rate: int,
+        target_sample_rate: int,
         mel_bins: int,
         mel_hop_length: int,
         n_fft: int,
     ) -> None:
         super().__init__()
-        self.sample_rate = sample_rate
+        self.target_sample_rate = target_sample_rate
         self.mel_transform = torchaudio.transforms.MelSpectrogram(
-            sample_rate=sample_rate,
+            sample_rate=target_sample_rate,
             n_fft=n_fft,
             win_length=n_fft,
             hop_length=mel_hop_length,
             f_min=0.0,
-            f_max=sample_rate / 2.0,
+            f_max=target_sample_rate / 2.0,
             n_mels=mel_bins,
             window_fn=torch.hann_window,
             center=True,
@@ -31,25 +33,20 @@ class AudioProcessor(nn.Module):
             norm="slaney",
         )
 
-    def resample_waveform(
-        self,
-        waveform: torch.Tensor,
-        source_rate: int,
-        target_rate: int,
-    ) -> torch.Tensor:
-        """Resample waveform to target sample rate if needed."""
-        if source_rate == target_rate:
-            return waveform
-        resampled = torchaudio.functional.resample(waveform, source_rate, target_rate)
-        return resampled.to(device=waveform.device, dtype=waveform.dtype)
+    def resample_audio(self, audio: Audio) -> Audio:
+        """Resample audio to the processor's target sample rate if needed."""
+        if audio.sampling_rate == self.target_sample_rate:
+            return audio
+        resampled = torchaudio.functional.resample(audio.waveform, audio.sampling_rate, self.target_sample_rate)
+        resampled = resampled.to(device=audio.waveform.device, dtype=audio.waveform.dtype)
+        return Audio(waveform=resampled, sampling_rate=self.target_sample_rate)
 
     def waveform_to_mel(
         self,
-        waveform: torch.Tensor,
-        waveform_sample_rate: int,
+        audio: Audio,
     ) -> torch.Tensor:
         """Convert waveform to log-mel spectrogram [batch, channels, time, n_mels]."""
-        waveform = self.resample_waveform(waveform, waveform_sample_rate, self.sample_rate)
+        waveform = self.resample_audio(audio).waveform
 
         mel = self.mel_transform(waveform)
         mel = torch.log(torch.clamp(mel, min=1e-5))

packages/ltx-core/src/ltx_core/model/audio_vae/vocoder.py

@@ -6,7 +6,266 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 
-from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1, ResBlock2
+from ltx_core.model.audio_vae.resnet import LRELU_SLOPE, ResBlock1
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+# ---------------------------------------------------------------------------
+# Anti-aliased resampling helpers (kaiser-sinc filters) for BigVGAN v2
+# Adopted from https://github.com/NVIDIA/BigVGAN
+# ---------------------------------------------------------------------------
+
+
+def _sinc(x: torch.Tensor) -> torch.Tensor:
+    return torch.where(
+        x == 0,
+        torch.tensor(1.0, device=x.device, dtype=x.dtype),
+        torch.sin(math.pi * x) / math.pi / x,
+    )
+
+
+def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+    delta_f = 4 * half_width
+    amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if amplitude > 50.0:
+        beta = 0.1102 * (amplitude - 8.7)
+    elif amplitude >= 21.0:
+        beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+    time = torch.arange(-half_size, half_size) + 0.5 if even else torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time)
+        filter_ /= filter_.sum()
+    return filter_.view(1, 1, kernel_size)
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff: float = 0.5,
+        half_width: float = 0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ) -> None:
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, n_channels, _ = x.shape
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        return F.conv1d(x, self.filter.expand(n_channels, -1, -1), stride=self.stride, groups=n_channels)
+
+
+class UpSample1d(nn.Module):
+    def __init__(
+        self,
+        ratio: int = 2,
+        kernel_size: int | None = None,
+        persistent: bool = True,
+        window_type: str = "kaiser",
+    ) -> None:
+        super().__init__()
+        self.ratio = ratio
+        self.stride = ratio
+
+        if window_type == "hann":
+            # Hann-windowed sinc filter equivalent to torchaudio.functional.resample
+            rolloff = 0.99
+            lowpass_filter_width = 6
+            width = math.ceil(lowpass_filter_width / rolloff)
+            self.kernel_size = 2 * width * ratio + 1
+            self.pad = width
+            self.pad_left = 2 * width * ratio
+            self.pad_right = self.kernel_size - ratio
+            time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff
+            time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width)
+            window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2
+            sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view(1, 1, -1)
+        else:
+            # Kaiser-windowed sinc filter (BigVGAN default).
+            self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+            self.pad = self.kernel_size // ratio - 1
+            self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+            self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+            sinc_filter = kaiser_sinc_filter1d(
+                cutoff=0.5 / ratio,
+                half_width=0.6 / ratio,
+                kernel_size=self.kernel_size,
+            )
+
+        self.register_buffer("filter", sinc_filter, persistent=persistent)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        _, n_channels, _ = x.shape
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        filt = self.filter.to(dtype=x.dtype, device=x.device).expand(n_channels, -1, -1)
+        x = self.ratio * F.conv_transpose1d(x, filt, stride=self.stride, groups=n_channels)
+        return x[..., self.pad_left : -self.pad_right]
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio: int = 2, kernel_size: int | None = None) -> None:
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lowpass(x)
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation: nn.Module,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ) -> None:
+        super().__init__()
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.upsample(x)
+        x = self.act(x)
+        return self.downsample(x)
+
+
+class Snake(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        alpha: float = 1.0,
+        alpha_trainable: bool = True,
+        alpha_logscale: bool = True,
+    ) -> None:
+        super().__init__()
+        self.alpha_logscale = alpha_logscale
+        self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha)
+        self.alpha.requires_grad = alpha_trainable
+        self.eps = 1e-9
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        return x + (1.0 / (alpha + self.eps)) * torch.sin(x * alpha).pow(2)
+
+
+class SnakeBeta(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        alpha: float = 1.0,
+        alpha_trainable: bool = True,
+        alpha_logscale: bool = True,
+    ) -> None:
+        super().__init__()
+        self.alpha_logscale = alpha_logscale
+        self.alpha = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha)
+        self.alpha.requires_grad = alpha_trainable
+        self.beta = nn.Parameter(torch.zeros(in_features) if alpha_logscale else torch.ones(in_features) * alpha)
+        self.beta.requires_grad = alpha_trainable
+        self.eps = 1e-9
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2)
+
+
+class AMPBlock1(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple[int, int, int] = (1, 3, 5),
+        activation: str = "snake",
+    ) -> None:
+        super().__init__()
+        act_cls = SnakeBeta if activation == "snakebeta" else Snake
+        self.convs1 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[0],
+                    padding=get_padding(kernel_size, dilation[0]),
+                ),
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[1],
+                    padding=get_padding(kernel_size, dilation[1]),
+                ),
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    1,
+                    dilation=dilation[2],
+                    padding=get_padding(kernel_size, dilation[2]),
+                ),
+            ]
+        )
+
+        self.convs2 = nn.ModuleList(
+            [
+                nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)),
+                nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)),
+                nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)),
+            ]
+        )
+
+        self.acts1 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs1))])
+        self.acts2 = nn.ModuleList([Activation1d(act_cls(channels)) for _ in range(len(self.convs2))])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2, strict=True):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = x + xt
+        return x
 
 
 class Vocoder(torch.nn.Module):
@@ -23,28 +282,33 @@ class Vocoder(torch.nn.Module):
                                This value is read from the checkpoint at `config.vocoder.resblock_dilation_sizes`.
         upsample_initial_channel: Initial number of channels for the upsampling layers.
                                This value is read from the checkpoint at `config.vocoder.upsample_initial_channel`.
-        stereo: Whether to use stereo output.
-                               This value is read from the checkpoint at `config.vocoder.stereo`.
-        resblock: Type of residual block to use.
+        resblock: Type of residual block to use ("1", "2", or "AMP1").
                                 This value is read from the checkpoint at `config.vocoder.resblock`.
-        output_sample_rate: Waveform sample rate.
-                               This value is read from the checkpoint at `config.vocoder.output_sample_rate`.
+        output_sampling_rate: Waveform sample rate.
+                               This value is read from the checkpoint at `config.vocoder.output_sampling_rate`.
+        activation: Activation type for BigVGAN v2 ("snake" or "snakebeta"). Only used when resblock="AMP1".
+        use_tanh_at_final: Apply tanh at the output (when apply_final_activation=True).
+        apply_final_activation: Whether to apply the final tanh/clamp activation.
+        use_bias_at_final: Whether to use bias in the final conv layer.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         resblock_kernel_sizes: List[int] | None = None,
         upsample_rates: List[int] | None = None,
         upsample_kernel_sizes: List[int] | None = None,
         resblock_dilation_sizes: List[List[int]] | None = None,
         upsample_initial_channel: int = 1024,
-        stereo: bool = True,
         resblock: str = "1",
-        output_sample_rate: int = 24000,
-    ):
+        output_sampling_rate: int = 24000,
+        activation: str = "snake",
+        use_tanh_at_final: bool = True,
+        apply_final_activation: bool = True,
+        use_bias_at_final: bool = True,
+    ) -> None:
         super().__init__()
 
-        # Initialize default values if not provided. Note that mutable default values are not supported.
+        # Mutable default values are not supported as default arguments.
         if resblock_kernel_sizes is None:
             resblock_kernel_sizes = [3, 7, 11]
         if upsample_rates is None:
@@ -54,36 +318,60 @@ class Vocoder(torch.nn.Module):
         if resblock_dilation_sizes is None:
             resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
 
-        self.output_sample_rate = output_sample_rate
+        self.output_sampling_rate = output_sampling_rate
         self.num_kernels = len(resblock_kernel_sizes)
         self.num_upsamples = len(upsample_rates)
-        in_channels = 128 if stereo else 64
-        self.conv_pre = nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
-        resblock_class = ResBlock1 if resblock == "1" else ResBlock2
+        self.use_tanh_at_final = use_tanh_at_final
+        self.apply_final_activation = apply_final_activation
+        self.is_amp = resblock == "AMP1"
 
-        self.ups = nn.ModuleList()
-        for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)):
-            self.ups.append(
-                nn.ConvTranspose1d(
-                    upsample_initial_channel // (2**i),
-                    upsample_initial_channel // (2 ** (i + 1)),
-                    kernel_size,
-                    stride,
-                    padding=(kernel_size - stride) // 2,
-                )
+        # All production checkpoints are stereo: 128 input channels (2 stereo channels x 64 mel
+        # bins each), 2 output channels.
+        self.conv_pre = nn.Conv1d(
+            in_channels=128,
+            out_channels=upsample_initial_channel,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+        )
+        resblock_cls = ResBlock1 if resblock == "1" else AMPBlock1
+
+        self.ups = nn.ModuleList(
+            nn.ConvTranspose1d(
+                upsample_initial_channel // (2**i),
+                upsample_initial_channel // (2 ** (i + 1)),
+                kernel_size,
+                stride,
+                padding=(kernel_size - stride) // 2,
             )
+            for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True))
+        )
 
+        final_channels = upsample_initial_channel // (2 ** len(upsample_rates))
         self.resblocks = nn.ModuleList()
-        for i, _ in enumerate(self.ups):
+
+        for i in range(len(upsample_rates)):
             ch = upsample_initial_channel // (2 ** (i + 1))
             for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
-                self.resblocks.append(resblock_class(ch, kernel_size, dilations))
+                if self.is_amp:
+                    self.resblocks.append(resblock_cls(ch, kernel_size, dilations, activation=activation))
+                else:
+                    self.resblocks.append(resblock_cls(ch, kernel_size, dilations))
 
-        out_channels = 2 if stereo else 1
-        final_channels = upsample_initial_channel // (2**self.num_upsamples)
-        self.conv_post = nn.Conv1d(final_channels, out_channels, 7, 1, padding=3)
+        if self.is_amp:
+            self.act_post: nn.Module = Activation1d(SnakeBeta(final_channels))
+        else:
+            self.act_post = nn.LeakyReLU()
 
-        self.upsample_factor = math.prod(layer.stride[0] for layer in self.ups)
+        # All production checkpoints are stereo: this final conv maps `final_channels` to 2 output channels (stereo).
+        self.conv_post = nn.Conv1d(
+            in_channels=final_channels,
+            out_channels=2,
+            kernel_size=7,
+            stride=1,
+            padding=3,
+            bias=use_bias_at_final,
+        )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -104,7 +392,8 @@ class Vocoder(torch.nn.Module):
         x = self.conv_pre(x)
 
         for i in range(self.num_upsamples):
-            x = F.leaky_relu(x, LRELU_SLOPE)
+            if not self.is_amp:
+                x = F.leaky_relu(x, LRELU_SLOPE)
             x = self.ups[i](x)
             start = i * self.num_kernels
             end = start + self.num_kernels
@@ -116,8 +405,171 @@ class Vocoder(torch.nn.Module):
                 [self.resblocks[idx](x) for idx in range(start, end)],
                 dim=0,
             )
-
             x = block_outputs.mean(dim=0)
 
-        x = self.conv_post(F.leaky_relu(x))
-        return torch.tanh(x)
+        x = self.act_post(x)
+        x = self.conv_post(x)
+
+        if self.apply_final_activation:
+            x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1)
+
+        return x
+
+
+class _STFTFn(nn.Module):
+    """Implements STFT as a convolution with precomputed DFT x Hann-window bases.
+    The DFT basis rows (real and imaginary parts interleaved) multiplied by the causal
+    Hann window are stored as buffers and loaded from the checkpoint. Using the exact
+    bfloat16 bases from training ensures the mel values fed to the BWE generator are
+    bit-identical to what it was trained on.
+    """
+
+    def __init__(self, filter_length: int, hop_length: int, win_length: int) -> None:
+        super().__init__()
+        self.hop_length = hop_length
+        self.win_length = win_length
+        n_freqs = filter_length // 2 + 1
+        self.register_buffer("forward_basis", torch.zeros(n_freqs * 2, 1, filter_length))
+        self.register_buffer("inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length))
+
+    def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Compute magnitude and phase spectrogram from a batch of waveforms.
+        Applies causal (left-only) padding of win_length - hop_length samples so that
+        each output frame depends only on past and present input — no lookahead.
+        Args:
+            y: Waveform tensor of shape (B, T).
+        Returns:
+            magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames).
+            phase:     Phase spectrogram in radians, shape (B, n_freqs, T_frames).
+        """
+        if y.dim() == 2:
+            y = y.unsqueeze(1)  # (B, 1, T)
+        left_pad = max(0, self.win_length - self.hop_length)  # causal: left-only
+        y = F.pad(y, (left_pad, 0))
+        spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0)
+        n_freqs = spec.shape[1] // 2
+        real, imag = spec[:, :n_freqs], spec[:, n_freqs:]
+        magnitude = torch.sqrt(real**2 + imag**2)
+        phase = torch.atan2(imag.float(), real.float()).to(real.dtype)
+        return magnitude, phase
+
+
+class MelSTFT(nn.Module):
+    """Causal log-mel spectrogram module whose buffers are loaded from the checkpoint.
+    Computes a log-mel spectrogram by running the causal STFT (_STFTFn) on the input
+    waveform and projecting the linear magnitude spectrum onto the mel filterbank.
+    The module's state dict layout matches the 'mel_stft.*' keys stored in the checkpoint
+    (mel_basis, stft_fn.forward_basis, stft_fn.inverse_basis).
+    """
+
+    def __init__(
+        self,
+        filter_length: int,
+        hop_length: int,
+        win_length: int,
+        n_mel_channels: int,
+    ) -> None:
+        super().__init__()
+        self.stft_fn = _STFTFn(filter_length, hop_length, win_length)
+
+        # Initialized to zeros; load_state_dict overwrites with the checkpoint's
+        # exact bfloat16 filterbank (vocoder.mel_stft.mel_basis, shape [n_mels, n_freqs]).
+        n_freqs = filter_length // 2 + 1
+        self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs))
+
+    def mel_spectrogram(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Compute log-mel spectrogram and auxiliary spectral quantities.
+        Args:
+            y: Waveform tensor of shape (B, T).
+        Returns:
+            log_mel:   Log-compressed mel spectrogram, shape (B, n_mel_channels, T_frames).
+            magnitude: Linear amplitude spectrogram, shape (B, n_freqs, T_frames).
+            phase:     Phase spectrogram in radians, shape (B, n_freqs, T_frames).
+            energy:    Per-frame energy (L2 norm over frequency), shape (B, T_frames).
+        """
+        magnitude, phase = self.stft_fn(y)
+        energy = torch.norm(magnitude, dim=1)
+        mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude)
+        log_mel = torch.log(torch.clamp(mel, min=1e-5))
+        return log_mel, magnitude, phase, energy
+
+
+class VocoderWithBWE(nn.Module):
+    """Vocoder with bandwidth extension (BWE) upsampling.
+    Chains a mel-to-wav vocoder with a BWE module that upsamples the output
+    to a higher sample rate. The BWE computes a mel spectrogram from the
+    vocoder output, runs it through a second generator to predict a residual,
+    and adds it to a sinc-resampled skip connection.
+    """
+
+    def __init__(
+        self,
+        vocoder: Vocoder,
+        bwe_generator: Vocoder,
+        mel_stft: MelSTFT,
+        input_sampling_rate: int,
+        output_sampling_rate: int,
+        hop_length: int,
+    ) -> None:
+        super().__init__()
+        self.vocoder = vocoder
+        self.bwe_generator = bwe_generator
+        self.mel_stft = mel_stft
+        self.input_sampling_rate = input_sampling_rate
+        self.output_sampling_rate = output_sampling_rate
+        self.hop_length = hop_length
+        # Compute the resampler on CPU so the sinc filter is materialized even when
+        # the model is constructed on meta device (SingleGPUModelBuilder pattern).
+        # The filter is not stored in the checkpoint (persistent=False).
+        with torch.device("cpu"):
+            self.resampler = UpSample1d(
+                ratio=output_sampling_rate // input_sampling_rate, persistent=False, window_type="hann"
+            )
+
+    @property
+    def conv_pre(self) -> nn.Conv1d:
+        return self.vocoder.conv_pre
+
+    @property
+    def conv_post(self) -> nn.Conv1d:
+        return self.vocoder.conv_post
+
+    def _compute_mel(self, audio: torch.Tensor) -> torch.Tensor:
+        """Compute log-mel spectrogram from waveform using causal STFT bases.
+        Args:
+            audio: Waveform tensor of shape (B, C, T).
+        Returns:
+            mel: Log-mel spectrogram of shape (B, C, n_mels, T_frames).
+        """
+        batch, n_channels, _ = audio.shape
+        flat = audio.reshape(batch * n_channels, -1)  # (B*C, T)
+        mel, _, _, _ = self.mel_stft.mel_spectrogram(flat)  # (B*C, n_mels, T_frames)
+        return mel.reshape(batch, n_channels, mel.shape[1], mel.shape[2])  # (B, C, n_mels, T_frames)
+
+    def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
+        """Run the full vocoder + BWE forward pass.
+        Args:
+            mel_spec: Mel spectrogram of shape (B, 2, T, mel_bins) for stereo
+                      or (B, T, mel_bins) for mono. Same format as Vocoder.forward.
+        Returns:
+            Waveform tensor of shape (B, out_channels, T_out) clipped to [-1, 1].
+        """
+        x = self.vocoder(mel_spec)
+        _, _, length_low_rate = x.shape
+        output_length = length_low_rate * self.output_sampling_rate // self.input_sampling_rate
+
+        # Pad to multiple of hop_length for exact mel frame count
+        remainder = length_low_rate % self.hop_length
+        if remainder != 0:
+            x = F.pad(x, (0, self.hop_length - remainder))
+
+        # Compute mel spectrogram from vocoder output: (B, C, n_mels, T_frames)
+        mel = self._compute_mel(x)
+
+        # Vocoder.forward expects (B, C, T, mel_bins) — transpose before calling bwe_generator
+        mel_for_bwe = mel.transpose(2, 3)  # (B, C, T_frames, mel_bins)
+        residual = self.bwe_generator(mel_for_bwe)
+        skip = self.resampler(x)
+        assert residual.shape == skip.shape, f"residual {residual.shape} != skip {skip.shape}"
+
+        return torch.clamp(residual + skip, -1, 1)[..., :output_length]

packages/ltx-core/src/ltx_core/model/transformer/adaln.py

@@ -4,6 +4,17 @@ import torch
 
 from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings
 
+# Number of AdaLN modulation parameters per transformer block.
+# Base: 2 params (shift + scale) x 3 norms (self-attn, feed-forward, output).
+ADALN_NUM_BASE_PARAMS = 6
+# Cross-attention AdaLN adds 3 more (scale, shift, gate) for the CA norm.
+ADALN_NUM_CROSS_ATTN_PARAMS = 3
+
+
+def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int:
+    """Total number of AdaLN parameters per block."""
+    return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0)
+
 
 class AdaLayerNormSingle(torch.nn.Module):
     r"""

packages/ltx-core/src/ltx_core/model/transformer/attention.py

@@ -184,21 +184,55 @@ class Attention(torch.nn.Module):
         mask: torch.Tensor | None = None,
         pe: torch.Tensor | None = None,
         k_pe: torch.Tensor | None = None,
+        perturbation_mask: torch.Tensor | None = None,
+        all_perturbed: bool = False,
     ) -> torch.Tensor:
-        q = self.to_q(x)
+        """Multi-head attention with optional RoPE, perturbation masking, and per-head gating.
+        When ``perturbation_mask`` is all zeros, the expensive query/key path
+        (linear projections, RMSNorm, RoPE) is skipped entirely and only the
+        value projection is used as a pass-through.
+        Args:
+            x: Query input tensor of shape ``(B, T, query_dim)``.
+            context: Key/value context tensor of shape ``(B, S, context_dim)``.
+                Falls back to ``x`` (self-attention) when *None*.
+            mask: Optional attention mask. Interpretation depends on the attention
+                backend (additive bias for xformers/PyTorch SDPA).
+            pe: Rotary positional embeddings applied to both ``q`` and ``k``.
+            k_pe: Separate rotary positional embeddings for ``k`` only. When
+                *None*, ``pe`` is reused for keys.
+            perturbation_mask: Optional mask in ``[0, 1]`` that
+                blends the attention output with the raw value projection:
+                ``out = attn_out * mask + v * (1 - mask)``.
+                **1** keeps the full attention output, **0** bypasses attention
+                and passes the value projection through unchanged.
+                *None* or all-ones means standard attention; all-zeros skips
+                the query/key path entirely for efficiency.
+            all_perturbed: Whether all perturbations are active for this block.
+        Returns:
+            Output tensor of shape ``(B, T, query_dim)``.
+        """
         context = x if context is None else context
-        k = self.to_k(context)
+        use_attention = not all_perturbed
+
         v = self.to_v(context)
 
-        q = self.q_norm(q)
-        k = self.k_norm(k)
+        if not use_attention:
+            out = v
+        else:
+            q = self.to_q(x)
+            k = self.to_k(context)
 
-        if pe is not None:
-            q = apply_rotary_emb(q, pe, self.rope_type)
-            k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
+            q = self.q_norm(q)
+            k = self.k_norm(k)
 
-        # attention_function can be an enum *or* a custom callable
-        out = self.attention_function(q, k, v, self.heads, mask)  # (B, T, H*D)
+            if pe is not None:
+                q = apply_rotary_emb(q, pe, self.rope_type)
+                k = apply_rotary_emb(k, pe if k_pe is None else k_pe, self.rope_type)
+
+            out = self.attention_function(q, k, v, self.heads, mask)  # (B, T, H*D)
+
+            if perturbation_mask is not None:
+                out = out * perturbation_mask + v * (1 - perturbation_mask)
 
         # Apply per-head gating if enabled
         if self.to_gate_logits is not None:

packages/ltx-core/src/ltx_core/model/transformer/modality.py

@@ -9,11 +9,27 @@ class Modality:
     Input data for a single modality (video or audio) in the transformer.
     Bundles the latent tokens, timestep embeddings, positional information,
     and text conditioning context for processing by the diffusion transformer.
+    Attributes:
+        latent: Patchified latent tokens, shape ``(B, T, D)`` where *B* is
+            the batch size, *T* is the total number of tokens (noisy +
+            conditioning), and *D* is the input dimension.
+        timesteps: Per-token timestep embeddings, shape ``(B, T)``.
+        positions: Positional coordinates, shape ``(B, 3, T)`` for video
+            (time, height, width) or ``(B, 1, T)`` for audio.
+        context: Text conditioning embeddings from the prompt encoder.
+        enabled: Whether this modality is active in the current forward pass.
+        context_mask: Optional mask for the text context tokens.
+        attention_mask: Optional 2-D self-attention mask, shape ``(B, T, T)``.
+            Values in ``[0, 1]`` where ``1`` = full attention and ``0`` = no
+            attention. ``None`` means unrestricted (full) attention between
+            all tokens. Built incrementally by conditioning items; see
+            :class:`~ltx_core.conditioning.types.attention_strength_wrapper.ConditioningItemAttentionStrengthWrapper`.
     """
 
     latent: (
         torch.Tensor
     )  # Shape: (B, T, D) where B is the batch size, T is the number of tokens, and D is input dimension
+    sigma: torch.Tensor  # Shape: (B,). Current sigma value, used for cross-attention timestep calculation.
     timesteps: torch.Tensor  # Shape: (B, T) where T is the number of timesteps
     positions: (
         torch.Tensor
@@ -21,3 +37,4 @@ class Modality:
     context: torch.Tensor
     enabled: bool = True
     context_mask: torch.Tensor | None = None
+    attention_mask: torch.Tensor | None = None

packages/ltx-core/src/ltx_core/model/transformer/model.py

@@ -3,11 +3,10 @@ from enum import Enum
 import torch
 
 from ltx_core.guidance.perturbations import BatchedPerturbationConfig
-from ltx_core.model.transformer.adaln import AdaLayerNormSingle
+from ltx_core.model.transformer.adaln import AdaLayerNormSingle, adaln_embedding_coefficient
 from ltx_core.model.transformer.attention import AttentionCallable, AttentionFunction
 from ltx_core.model.transformer.modality import Modality
 from ltx_core.model.transformer.rope import LTXRopeType
-from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
 from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig
 from ltx_core.model.transformer.transformer_args import (
     MultiModalTransformerArgsPreprocessor,
@@ -47,7 +46,6 @@ class LTXModel(torch.nn.Module):
         cross_attention_dim: int = 4096,
         norm_eps: float = 1e-06,
         attention_type: AttentionFunction | AttentionCallable = AttentionFunction.DEFAULT,
-        caption_channels: int = 3840,
         positional_embedding_theta: float = 10000.0,
         positional_embedding_max_pos: list[int] | None = None,
         timestep_scale_multiplier: int = 1000,
@@ -62,9 +60,13 @@ class LTXModel(torch.nn.Module):
         rope_type: LTXRopeType = LTXRopeType.INTERLEAVED,
         double_precision_rope: bool = False,
         apply_gated_attention: bool = False,
+        caption_projection: torch.nn.Module | None = None,
+        audio_caption_projection: torch.nn.Module | None = None,
+        cross_attention_adaln: bool = False,
     ):
         super().__init__()
         self._enable_gradient_checkpointing = False
+        self.cross_attention_adaln = cross_attention_adaln
         self.use_middle_indices_grid = use_middle_indices_grid
         self.rope_type = rope_type
         self.double_precision_rope = double_precision_rope
@@ -81,8 +83,8 @@ class LTXModel(torch.nn.Module):
             self._init_video(
                 in_channels=in_channels,
                 out_channels=out_channels,
-                caption_channels=caption_channels,
                 norm_eps=norm_eps,
+                caption_projection=caption_projection,
             )
 
         if model_type.is_audio_enabled():
@@ -94,8 +96,8 @@ class LTXModel(torch.nn.Module):
             self._init_audio(
                 in_channels=audio_in_channels,
                 out_channels=audio_out_channels,
-                caption_channels=caption_channels,
                 norm_eps=norm_eps,
+                caption_projection=audio_caption_projection,
             )
 
         if model_type.is_video_enabled() and model_type.is_audio_enabled():
@@ -117,23 +119,27 @@ class LTXModel(torch.nn.Module):
             apply_gated_attention=apply_gated_attention,
         )
 
+    @property
+    def _adaln_embedding_coefficient(self) -> int:
+        return adaln_embedding_coefficient(self.cross_attention_adaln)
+
     def _init_video(
         self,
         in_channels: int,
         out_channels: int,
-        caption_channels: int,
         norm_eps: float,
+        caption_projection: torch.nn.Module | None = None,
     ) -> None:
         """Initialize video-specific components."""
         # Video input components
         self.patchify_proj = torch.nn.Linear(in_channels, self.inner_dim, bias=True)
+        if caption_projection is not None:
+            self.caption_projection = caption_projection
 
-        self.adaln_single = AdaLayerNormSingle(self.inner_dim)
+        self.adaln_single = AdaLayerNormSingle(self.inner_dim, embedding_coefficient=self._adaln_embedding_coefficient)
 
-        # Video caption projection
-        self.caption_projection = PixArtAlphaTextProjection(
-            in_features=caption_channels,
-            hidden_size=self.inner_dim,
+        self.prompt_adaln_single = (
+            AdaLayerNormSingle(self.inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None
         )
 
         # Video output components
@@ -145,22 +151,23 @@ class LTXModel(torch.nn.Module):
         self,
         in_channels: int,
         out_channels: int,
-        caption_channels: int,
         norm_eps: float,
+        caption_projection: torch.nn.Module | None = None,
     ) -> None:
         """Initialize audio-specific components."""
 
         # Audio input components
         self.audio_patchify_proj = torch.nn.Linear(in_channels, self.audio_inner_dim, bias=True)
+        if caption_projection is not None:
+            self.audio_caption_projection = caption_projection
 
         self.audio_adaln_single = AdaLayerNormSingle(
             self.audio_inner_dim,
+            embedding_coefficient=self._adaln_embedding_coefficient,
         )
 
-        # Audio caption projection
-        self.audio_caption_projection = PixArtAlphaTextProjection(
-            in_features=caption_channels,
-            hidden_size=self.audio_inner_dim,
+        self.audio_prompt_adaln_single = (
+            AdaLayerNormSingle(self.audio_inner_dim, embedding_coefficient=2) if self.cross_attention_adaln else None
         )
 
         # Audio output components
@@ -203,7 +210,6 @@ class LTXModel(torch.nn.Module):
             self.video_args_preprocessor = MultiModalTransformerArgsPreprocessor(
                 patchify_proj=self.patchify_proj,
                 adaln=self.adaln_single,
-                caption_projection=self.caption_projection,
                 cross_scale_shift_adaln=self.av_ca_video_scale_shift_adaln_single,
                 cross_gate_adaln=self.av_ca_a2v_gate_adaln_single,
                 inner_dim=self.inner_dim,
@@ -217,11 +223,12 @@ class LTXModel(torch.nn.Module):
                 positional_embedding_theta=self.positional_embedding_theta,
                 rope_type=self.rope_type,
                 av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
+                caption_projection=getattr(self, "caption_projection", None),
+                prompt_adaln=getattr(self, "prompt_adaln_single", None),
             )
             self.audio_args_preprocessor = MultiModalTransformerArgsPreprocessor(
                 patchify_proj=self.audio_patchify_proj,
                 adaln=self.audio_adaln_single,
-                caption_projection=self.audio_caption_projection,
                 cross_scale_shift_adaln=self.av_ca_audio_scale_shift_adaln_single,
                 cross_gate_adaln=self.av_ca_v2a_gate_adaln_single,
                 inner_dim=self.audio_inner_dim,
@@ -235,12 +242,13 @@ class LTXModel(torch.nn.Module):
                 positional_embedding_theta=self.positional_embedding_theta,
                 rope_type=self.rope_type,
                 av_ca_timestep_scale_multiplier=self.av_ca_timestep_scale_multiplier,
+                caption_projection=getattr(self, "audio_caption_projection", None),
+                prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
             )
         elif self.model_type.is_video_enabled():
             self.video_args_preprocessor = TransformerArgsPreprocessor(
                 patchify_proj=self.patchify_proj,
                 adaln=self.adaln_single,
-                caption_projection=self.caption_projection,
                 inner_dim=self.inner_dim,
                 max_pos=self.positional_embedding_max_pos,
                 num_attention_heads=self.num_attention_heads,
@@ -249,12 +257,13 @@ class LTXModel(torch.nn.Module):
                 double_precision_rope=self.double_precision_rope,
                 positional_embedding_theta=self.positional_embedding_theta,
                 rope_type=self.rope_type,
+                caption_projection=getattr(self, "caption_projection", None),
+                prompt_adaln=getattr(self, "prompt_adaln_single", None),
             )
         elif self.model_type.is_audio_enabled():
             self.audio_args_preprocessor = TransformerArgsPreprocessor(
                 patchify_proj=self.audio_patchify_proj,
                 adaln=self.audio_adaln_single,
-                caption_projection=self.audio_caption_projection,
                 inner_dim=self.audio_inner_dim,
                 max_pos=self.audio_positional_embedding_max_pos,
                 num_attention_heads=self.audio_num_attention_heads,
@@ -263,6 +272,8 @@ class LTXModel(torch.nn.Module):
                 double_precision_rope=self.double_precision_rope,
                 positional_embedding_theta=self.positional_embedding_theta,
                 rope_type=self.rope_type,
+                caption_projection=getattr(self, "audio_caption_projection", None),
+                prompt_adaln=getattr(self, "audio_prompt_adaln_single", None),
             )
 
     def _init_transformer_blocks(
@@ -284,6 +295,7 @@ class LTXModel(torch.nn.Module):
                 d_head=attention_head_dim,
                 context_dim=cross_attention_dim,
                 apply_gated_attention=apply_gated_attention,
+                cross_attention_adaln=self.cross_attention_adaln,
             )
             if self.model_type.is_video_enabled()
             else None
@@ -295,6 +307,7 @@ class LTXModel(torch.nn.Module):
                 d_head=audio_attention_head_dim,
                 context_dim=audio_cross_attention_dim,
                 apply_gated_attention=apply_gated_attention,
+                cross_attention_adaln=self.cross_attention_adaln,
             )
             if self.model_type.is_audio_enabled()
             else None
@@ -386,8 +399,8 @@ class LTXModel(torch.nn.Module):
         if not self.model_type.is_audio_enabled() and audio is not None:
             raise ValueError("Audio is not enabled for this model")
 
-        video_args = self.video_args_preprocessor.prepare(video) if video is not None else None
-        audio_args = self.audio_args_preprocessor.prepare(audio) if audio is not None else None
+        video_args = self.video_args_preprocessor.prepare(video, audio) if video is not None else None
+        audio_args = self.audio_args_preprocessor.prepare(audio, video) if audio is not None else None
         # Process transformer blocks
         video_out, audio_out = self._process_transformer_blocks(
             video=video_args,

packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py

@@ -1,8 +1,11 @@
+import torch
+
 from ltx_core.loader.sd_ops import SDOps
 from ltx_core.model.model_protocol import ModelConfigurator
 from ltx_core.model.transformer.attention import AttentionFunction
 from ltx_core.model.transformer.model import LTXModel, LTXModelType
 from ltx_core.model.transformer.rope import LTXRopeType
+from ltx_core.model.transformer.text_projection import create_caption_projection
 from ltx_core.utils import check_config_value
 
 
@@ -14,6 +17,9 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
 
     @classmethod
     def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+        # Build caption projections for 19B models (projection handled in transformer).
+        caption_projection, audio_caption_projection = _build_caption_projections(config, is_av=True)
+
         config = config.get("transformer", {})
 
         check_config_value(config, "dropout", 0.0)
@@ -45,7 +51,6 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             cross_attention_dim=config.get("cross_attention_dim", 4096),
             norm_eps=config.get("norm_eps", 1e-06),
             attention_type=AttentionFunction(config.get("attention_type", "default")),
-            caption_channels=config.get("caption_channels", 3840),
             positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
             positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
             timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
@@ -60,6 +65,9 @@ class LTXModelConfigurator(ModelConfigurator[LTXModel]):
             rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
             apply_gated_attention=config.get("apply_gated_attention", False),
+            caption_projection=caption_projection,
+            audio_caption_projection=audio_caption_projection,
+            cross_attention_adaln=config.get("cross_attention_adaln", False),
         )
 
 
@@ -71,6 +79,9 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
 
     @classmethod
     def from_config(cls: type[LTXModel], config: dict) -> LTXModel:
+        # Build caption projection for 19B model (projection handled in transformer).
+        caption_projection, _ = _build_caption_projections(config, is_av=False)
+
         config = config.get("transformer", {})
 
         check_config_value(config, "dropout", 0.0)
@@ -99,7 +110,6 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
             cross_attention_dim=config.get("cross_attention_dim", 4096),
             norm_eps=config.get("norm_eps", 1e-06),
             attention_type=AttentionFunction(config.get("attention_type", "default")),
-            caption_channels=config.get("caption_channels", 3840),
             positional_embedding_theta=config.get("positional_embedding_theta", 10000.0),
             positional_embedding_max_pos=config.get("positional_embedding_max_pos", [20, 2048, 2048]),
             timestep_scale_multiplier=config.get("timestep_scale_multiplier", 1000),
@@ -107,9 +117,34 @@ class LTXVideoOnlyModelConfigurator(ModelConfigurator[LTXModel]):
             rope_type=LTXRopeType(config.get("rope_type", "interleaved")),
             double_precision_rope=config.get("frequencies_precision", False) == "float64",
             apply_gated_attention=config.get("apply_gated_attention", False),
+            caption_projection=caption_projection,
+            cross_attention_adaln=config.get("cross_attention_adaln", False),
         )
 
 
+def _build_caption_projections(
+    config: dict,
+    is_av: bool,
+) -> tuple[torch.nn.Module | None, torch.nn.Module | None]:
+    """Build caption projections for the transformer when projection is NOT in the text encoder.
+    19B models: projection is in the transformer (caption_proj_before_connector=False).
+    20B models: projection is in the text encoder, so no projections are created here.
+    Args:
+        config: Full model config dict (must contain "transformer" key).
+        is_av: Whether this is an audio-video model. When False, audio projection is skipped.
+    Returns:
+        Tuple of (video_caption_projection, audio_caption_projection), both None for 20B models.
+    """
+    transformer_config = config.get("transformer", {})
+    if transformer_config.get("caption_proj_before_connector", False):
+        return None, None
+
+    with torch.device("meta"):
+        caption_projection = create_caption_projection(transformer_config)
+        audio_caption_projection = create_caption_projection(transformer_config, audio=True) if is_av else None
+    return caption_projection, audio_caption_projection
+
+
 LTXV_MODEL_COMFY_RENAMING_MAP = (
     SDOps("LTXV_MODEL_COMFY_PREFIX_MAP")
     .with_matching(prefix="model.diffusion_model.")

packages/ltx-core/src/ltx_core/model/transformer/text_projection.py

@@ -3,7 +3,8 @@ import torch
 
 class PixArtAlphaTextProjection(torch.nn.Module):
     """
-    Projects caption embeddings. Also handles dropout for classifier-free guidance.
+    Projects caption embeddings using dual linear layers.
+    Flow: linear_1 → activation → linear_2
     Adapted from https://github.com/PixArt-alpha/PixArt-alpha/blob/master/diffusion/model/nets/PixArt_blocks.py
     """
 
@@ -25,3 +26,13 @@ class PixArtAlphaTextProjection(torch.nn.Module):
         hidden_states = self.act_1(hidden_states)
         hidden_states = self.linear_2(hidden_states)
         return hidden_states
+
+
+def create_caption_projection(transformer_config: dict, audio: bool = False) -> PixArtAlphaTextProjection:
+    """Create a caption projection for the transformer (V1/19B only)."""
+    caption_channels = transformer_config["caption_channels"]
+    if audio:
+        inner_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"]
+    else:
+        inner_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"]
+    return PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim)

packages/ltx-core/src/ltx_core/model/transformer/transformer.py

@@ -3,6 +3,7 @@ from dataclasses import dataclass, replace
 import torch
 
 from ltx_core.guidance.perturbations import BatchedPerturbationConfig, PerturbationType
+from ltx_core.model.transformer.adaln import adaln_embedding_coefficient
 from ltx_core.model.transformer.attention import Attention, AttentionCallable, AttentionFunction
 from ltx_core.model.transformer.feed_forward import FeedForward
 from ltx_core.model.transformer.rope import LTXRopeType
@@ -17,6 +18,7 @@ class TransformerConfig:
     d_head: int
     context_dim: int
     apply_gated_attention: bool = False
+    cross_attention_adaln: bool = False
 
 
 class BasicAVTransformerBlock(torch.nn.Module):
@@ -54,7 +56,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 apply_gated_attention=video.apply_gated_attention,
             )
             self.ff = FeedForward(video.dim, dim_out=video.dim)
-            self.scale_shift_table = torch.nn.Parameter(torch.empty(6, video.dim))
+            video_sst_size = adaln_embedding_coefficient(video.cross_attention_adaln)
+            self.scale_shift_table = torch.nn.Parameter(torch.empty(video_sst_size, video.dim))
 
         if audio is not None:
             self.audio_attn1 = Attention(
@@ -78,7 +81,8 @@ class BasicAVTransformerBlock(torch.nn.Module):
                 apply_gated_attention=audio.apply_gated_attention,
             )
             self.audio_ff = FeedForward(audio.dim, dim_out=audio.dim)
-            self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(6, audio.dim))
+            audio_sst_size = adaln_embedding_coefficient(audio.cross_attention_adaln)
+            self.audio_scale_shift_table = torch.nn.Parameter(torch.empty(audio_sst_size, audio.dim))
 
         if audio is not None and video is not None:
             # Q: Video, K,V: Audio
@@ -108,6 +112,15 @@ class BasicAVTransformerBlock(torch.nn.Module):
             self.scale_shift_table_a2v_ca_audio = torch.nn.Parameter(torch.empty(5, audio.dim))
             self.scale_shift_table_a2v_ca_video = torch.nn.Parameter(torch.empty(5, video.dim))
 
+        self.cross_attention_adaln = (video is not None and video.cross_attention_adaln) or (
+            audio is not None and audio.cross_attention_adaln
+        )
+
+        if self.cross_attention_adaln and video is not None:
+            self.prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, video.dim))
+        if self.cross_attention_adaln and audio is not None:
+            self.audio_prompt_scale_shift_table = torch.nn.Parameter(torch.empty(2, audio.dim))
+
         self.norm_eps = norm_eps
 
     def get_ada_values(
@@ -141,6 +154,35 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
         return (*scale_shift_chunks, *gate_ada_values)
 
+    def _apply_text_cross_attention(
+        self,
+        x: torch.Tensor,
+        context: torch.Tensor,
+        attn: AttentionCallable,
+        scale_shift_table: torch.Tensor,
+        prompt_scale_shift_table: torch.Tensor | None,
+        timestep: torch.Tensor,
+        prompt_timestep: torch.Tensor | None,
+        context_mask: torch.Tensor | None,
+        cross_attention_adaln: bool = False,
+    ) -> torch.Tensor:
+        """Apply text cross-attention, with optional AdaLN modulation."""
+        if cross_attention_adaln:
+            shift_q, scale_q, gate = self.get_ada_values(scale_shift_table, x.shape[0], timestep, slice(6, 9))
+            return apply_cross_attention_adaln(
+                x,
+                context,
+                attn,
+                shift_q,
+                scale_q,
+                gate,
+                prompt_scale_shift_table,
+                prompt_timestep,
+                context_mask,
+                self.norm_eps,
+            )
+        return attn(rms_norm(x, eps=self.norm_eps), context=context, mask=context_mask)
+
     def forward(  # noqa: PLR0915
         self,
         video: TransformerArgs | None,
@@ -168,28 +210,77 @@ class BasicAVTransformerBlock(torch.nn.Module):
             vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
                 self.scale_shift_table, vx.shape[0], video.timesteps, slice(0, 3)
             )
-            if not perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx):
-                norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
-                v_mask = perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
-                vx = vx + self.attn1(norm_vx, pe=video.positional_embeddings) * vgate_msa * v_mask
+            norm_vx = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_msa) + vshift_msa
+            del vshift_msa, vscale_msa
 
-            vx = vx + self.attn2(rms_norm(vx, eps=self.norm_eps), context=video.context, mask=video.context_mask)
-
-            del vshift_msa, vscale_msa, vgate_msa
+            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
+            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx)
+            v_mask = (
+                perturbations.mask_like(PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx, vx)
+                if not all_perturbed and not none_perturbed
+                else None
+            )
+            vx = (
+                vx
+                + self.attn1(
+                    norm_vx,
+                    pe=video.positional_embeddings,
+                    mask=video.self_attention_mask,
+                    perturbation_mask=v_mask,
+                    all_perturbed=all_perturbed,
+                )
+                * vgate_msa
+            )
+            del vgate_msa, norm_vx, v_mask
+            vx = vx + self._apply_text_cross_attention(
+                vx,
+                video.context,
+                self.attn2,
+                self.scale_shift_table,
+                getattr(self, "prompt_scale_shift_table", None),
+                video.timesteps,
+                video.prompt_timestep,
+                video.context_mask,
+                cross_attention_adaln=self.cross_attention_adaln,
+            )
 
         if run_ax:
             ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
                 self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3)
             )
 
-            if not perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx):
-                norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
-                a_mask = perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
-                ax = ax + self.audio_attn1(norm_ax, pe=audio.positional_embeddings) * agate_msa * a_mask
-
-            ax = ax + self.audio_attn2(rms_norm(ax, eps=self.norm_eps), context=audio.context, mask=audio.context_mask)
-
-            del ashift_msa, ascale_msa, agate_msa
+            norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa
+            del ashift_msa, ascale_msa
+            all_perturbed = perturbations.all_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
+            none_perturbed = not perturbations.any_in_batch(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx)
+            a_mask = (
+                perturbations.mask_like(PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx, ax)
+                if not all_perturbed and not none_perturbed
+                else None
+            )
+            ax = (
+                ax
+                + self.audio_attn1(
+                    norm_ax,
+                    pe=audio.positional_embeddings,
+                    mask=audio.self_attention_mask,
+                    perturbation_mask=a_mask,
+                    all_perturbed=all_perturbed,
+                )
+                * agate_msa
+            )
+            del agate_msa, norm_ax, a_mask
+            ax = ax + self._apply_text_cross_attention(
+                ax,
+                audio.context,
+                self.audio_attn2,
+                self.audio_scale_shift_table,
+                getattr(self, "audio_prompt_scale_shift_table", None),
+                audio.timesteps,
+                audio.prompt_timestep,
+                audio.context_mask,
+                cross_attention_adaln=self.cross_attention_adaln,
+            )
 
         # Audio - Video cross attention.
         if run_a2v or run_v2a:
@@ -224,7 +315,9 @@ class BasicAVTransformerBlock(torch.nn.Module):
 
             if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
                 vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
+                del scale_ca_video_hidden_states_a2v, shift_ca_video_hidden_states_a2v
                 ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
+                del scale_ca_audio_hidden_states_a2v, shift_ca_audio_hidden_states_a2v
                 a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
                 vx = vx + (
                     self.audio_to_video_attn(
@@ -236,10 +329,12 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     * gate_out_a2v
                     * a2v_mask
                 )
-
+                del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled
             if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
                 ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
+                del scale_ca_audio_hidden_states_v2a, shift_ca_audio_hidden_states_v2a
                 vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
+                del scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a
                 v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
                 ax = ax + (
                     self.video_to_audio_attn(
@@ -251,35 +346,47 @@ class BasicAVTransformerBlock(torch.nn.Module):
                     * gate_out_v2a
                     * v2a_mask
                 )
-
-            del gate_out_a2v, gate_out_v2a
-            del (
-                scale_ca_video_hidden_states_a2v,
-                shift_ca_video_hidden_states_a2v,
-                scale_ca_audio_hidden_states_a2v,
-                shift_ca_audio_hidden_states_a2v,
-                scale_ca_video_hidden_states_v2a,
-                shift_ca_video_hidden_states_v2a,
-                scale_ca_audio_hidden_states_v2a,
-                shift_ca_audio_hidden_states_v2a,
-            )
+                del gate_out_v2a, v2a_mask, ax_scaled, vx_scaled
+            del vx_norm3, ax_norm3
 
         if run_vx:
             vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
-                self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, None)
+                self.scale_shift_table, vx.shape[0], video.timesteps, slice(3, 6)
             )
             vx_scaled = rms_norm(vx, eps=self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
             vx = vx + self.ff(vx_scaled) * vgate_mlp
 
-            del vshift_mlp, vscale_mlp, vgate_mlp
+            del vshift_mlp, vscale_mlp, vgate_mlp, vx_scaled
 
         if run_ax:
             ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
-                self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, None)
+                self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6)
             )
             ax_scaled = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
             ax = ax + self.audio_ff(ax_scaled) * agate_mlp
 
-            del ashift_mlp, ascale_mlp, agate_mlp
+            del ashift_mlp, ascale_mlp, agate_mlp, ax_scaled
 
         return replace(video, x=vx) if video is not None else None, replace(audio, x=ax) if audio is not None else None
+
+
+def apply_cross_attention_adaln(
+    x: torch.Tensor,
+    context: torch.Tensor,
+    attn: AttentionCallable,
+    q_shift: torch.Tensor,
+    q_scale: torch.Tensor,
+    q_gate: torch.Tensor,
+    prompt_scale_shift_table: torch.Tensor,
+    prompt_timestep: torch.Tensor,
+    context_mask: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+) -> torch.Tensor:
+    batch_size = x.shape[0]
+    shift_kv, scale_kv = (
+        prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
+        + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
+    ).unbind(dim=2)
+    attn_input = rms_norm(x, eps=norm_eps) * (1 + q_scale) + q_shift
+    encoder_hidden_states = context * (1 + scale_kv) + shift_kv
+    return attn(attn_input, context=encoder_hidden_states, mask=context_mask) * q_gate

packages/ltx-core/src/ltx_core/model/transformer/transformer_args.py

@@ -10,7 +10,6 @@ from ltx_core.model.transformer.rope import (
     generate_freq_grid_pytorch,
     precompute_freqs_cis,
 )
-from ltx_core.model.transformer.text_projection import PixArtAlphaTextProjection
 
 
 @dataclass(frozen=True)
@@ -25,6 +24,10 @@ class TransformerArgs:
     cross_scale_shift_timestep: torch.Tensor | None
     cross_gate_timestep: torch.Tensor | None
     enabled: bool
+    prompt_timestep: torch.Tensor | None = None
+    self_attention_mask: torch.Tensor | None = (
+        None  # Additive log-space self-attention bias (B, 1, T, T), None = full attention
+    )
 
 
 class TransformerArgsPreprocessor:
@@ -32,7 +35,6 @@ class TransformerArgsPreprocessor:
         self,
         patchify_proj: torch.nn.Linear,
         adaln: AdaLayerNormSingle,
-        caption_projection: PixArtAlphaTextProjection,
         inner_dim: int,
         max_pos: list[int],
         num_attention_heads: int,
@@ -41,10 +43,11 @@ class TransformerArgsPreprocessor:
         double_precision_rope: bool,
         positional_embedding_theta: float,
         rope_type: LTXRopeType,
+        caption_projection: torch.nn.Module | None = None,
+        prompt_adaln: AdaLayerNormSingle | None = None,
     ) -> None:
         self.patchify_proj = patchify_proj
         self.adaln = adaln
-        self.caption_projection = caption_projection
         self.inner_dim = inner_dim
         self.max_pos = max_pos
         self.num_attention_heads = num_attention_heads
@@ -53,35 +56,34 @@ class TransformerArgsPreprocessor:
         self.double_precision_rope = double_precision_rope
         self.positional_embedding_theta = positional_embedding_theta
         self.rope_type = rope_type
+        self.caption_projection = caption_projection
+        self.prompt_adaln = prompt_adaln
 
     def _prepare_timestep(
-        self, timestep: torch.Tensor, batch_size: int, hidden_dtype: torch.dtype
+        self, timestep: torch.Tensor, adaln: AdaLayerNormSingle, batch_size: int, hidden_dtype: torch.dtype
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Prepare timestep embeddings."""
-
-        timestep = timestep * self.timestep_scale_multiplier
-        timestep, embedded_timestep = self.adaln(
-            timestep.flatten(),
+        timestep_scaled = timestep * self.timestep_scale_multiplier
+        timestep, embedded_timestep = adaln(
+            timestep_scaled.flatten(),
             hidden_dtype=hidden_dtype,
         )
-
         # Second dimension is 1 or number of tokens (if timestep_per_token)
         timestep = timestep.view(batch_size, -1, timestep.shape[-1])
         embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.shape[-1])
+
         return timestep, embedded_timestep
 
     def _prepare_context(
         self,
         context: torch.Tensor,
         x: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+    ) -> torch.Tensor:
         """Prepare context for transformer blocks."""
+        if self.caption_projection is not None:
+            context = self.caption_projection(context)
         batch_size = x.shape[0]
-        context = self.caption_projection(context)
-        context = context.view(batch_size, -1, x.shape[-1])
-
-        return context, attention_mask
+        return context.view(batch_size, -1, x.shape[-1])
 
     def _prepare_attention_mask(self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype) -> torch.Tensor | None:
         """Prepare attention mask."""
@@ -92,6 +94,34 @@ class TransformerArgsPreprocessor:
             (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
         ) * torch.finfo(x_dtype).max
 
+    def _prepare_self_attention_mask(
+        self, attention_mask: torch.Tensor | None, x_dtype: torch.dtype
+    ) -> torch.Tensor | None:
+        """Prepare self-attention mask by converting [0,1] values to additive log-space bias.
+        Input shape: (B, T, T) with values in [0, 1].
+        Output shape: (B, 1, T, T) with 0.0 for full attention and a large negative value
+        for masked positions.
+        Positions with attention_mask <= 0 are fully masked (mapped to the dtype's minimum
+        representable value). Strictly positive entries are converted via log-space for
+        smooth attenuation, with small values clamped for numerical stability.
+        Returns None if input is None (no masking).
+        """
+        if attention_mask is None:
+            return None
+
+        # Convert [0, 1] attention mask to additive log-space bias:
+        #   1.0 -> log(1.0) = 0.0  (no bias, full attention)
+        #   0.0 -> finfo.min        (fully masked)
+        finfo = torch.finfo(x_dtype)
+        eps = finfo.tiny
+
+        bias = torch.full_like(attention_mask, finfo.min, dtype=x_dtype)
+        positive = attention_mask > 0
+        if positive.any():
+            bias[positive] = torch.log(attention_mask[positive].clamp(min=eps)).to(x_dtype)
+
+        return bias.unsqueeze(1)  # (B, 1, T, T) for head broadcast
+
     def _prepare_positional_embeddings(
         self,
         positions: torch.Tensor,
@@ -119,11 +149,20 @@ class TransformerArgsPreprocessor:
     def prepare(
         self,
         modality: Modality,
+        cross_modality: Modality | None = None,  # noqa: ARG002
     ) -> TransformerArgs:
         x = self.patchify_proj(modality.latent)
-        timestep, embedded_timestep = self._prepare_timestep(modality.timesteps, x.shape[0], modality.latent.dtype)
-        context, attention_mask = self._prepare_context(modality.context, x, modality.context_mask)
-        attention_mask = self._prepare_attention_mask(attention_mask, modality.latent.dtype)
+        batch_size = x.shape[0]
+        timestep, embedded_timestep = self._prepare_timestep(
+            modality.timesteps, self.adaln, batch_size, modality.latent.dtype
+        )
+        prompt_timestep = None
+        if self.prompt_adaln is not None:
+            prompt_timestep, _ = self._prepare_timestep(
+                modality.sigma, self.prompt_adaln, batch_size, modality.latent.dtype
+            )
+        context = self._prepare_context(modality.context, x)
+        attention_mask = self._prepare_attention_mask(modality.context_mask, modality.latent.dtype)
         pe = self._prepare_positional_embeddings(
             positions=modality.positions,
             inner_dim=self.inner_dim,
@@ -132,6 +171,7 @@ class TransformerArgsPreprocessor:
             num_attention_heads=self.num_attention_heads,
             x_dtype=modality.latent.dtype,
         )
+        self_attention_mask = self._prepare_self_attention_mask(modality.attention_mask, modality.latent.dtype)
         return TransformerArgs(
             x=x,
             context=context,
@@ -143,6 +183,8 @@ class TransformerArgsPreprocessor:
             cross_scale_shift_timestep=None,
             cross_gate_timestep=None,
             enabled=modality.enabled,
+            prompt_timestep=prompt_timestep,
+            self_attention_mask=self_attention_mask,
         )
 
 
@@ -151,7 +193,6 @@ class MultiModalTransformerArgsPreprocessor:
         self,
         patchify_proj: torch.nn.Linear,
         adaln: AdaLayerNormSingle,
-        caption_projection: PixArtAlphaTextProjection,
         cross_scale_shift_adaln: AdaLayerNormSingle,
         cross_gate_adaln: AdaLayerNormSingle,
         inner_dim: int,
@@ -165,11 +206,12 @@ class MultiModalTransformerArgsPreprocessor:
         positional_embedding_theta: float,
         rope_type: LTXRopeType,
         av_ca_timestep_scale_multiplier: int,
+        caption_projection: torch.nn.Module | None = None,
+        prompt_adaln: AdaLayerNormSingle | None = None,
     ) -> None:
         self.simple_preprocessor = TransformerArgsPreprocessor(
             patchify_proj=patchify_proj,
             adaln=adaln,
-            caption_projection=caption_projection,
             inner_dim=inner_dim,
             max_pos=max_pos,
             num_attention_heads=num_attention_heads,
@@ -178,6 +220,8 @@ class MultiModalTransformerArgsPreprocessor:
             double_precision_rope=double_precision_rope,
             positional_embedding_theta=positional_embedding_theta,
             rope_type=rope_type,
+            caption_projection=caption_projection,
+            prompt_adaln=prompt_adaln,
         )
         self.cross_scale_shift_adaln = cross_scale_shift_adaln
         self.cross_gate_adaln = cross_gate_adaln
@@ -188,8 +232,22 @@ class MultiModalTransformerArgsPreprocessor:
     def prepare(
         self,
         modality: Modality,
+        cross_modality: Modality | None = None,
     ) -> TransformerArgs:
         transformer_args = self.simple_preprocessor.prepare(modality)
+        if cross_modality is None:
+            return transformer_args
+
+        if cross_modality.sigma.numel() > 1:
+            if cross_modality.sigma.shape[0] != modality.timesteps.shape[0]:
+                raise ValueError("Cross modality sigma must have the same batch size as the modality")
+            if cross_modality.sigma.ndim != 1:
+                raise ValueError("Cross modality sigma must be a 1D tensor")
+
+        cross_timestep = cross_modality.sigma.view(
+            modality.timesteps.shape[0], 1, *[1] * len(modality.timesteps.shape[2:])
+        )
+
         cross_pe = self.simple_preprocessor._prepare_positional_embeddings(
             positions=modality.positions[:, 0:1, :],
             inner_dim=self.audio_cross_attention_dim,
@@ -200,7 +258,7 @@ class MultiModalTransformerArgsPreprocessor:
         )
 
         cross_scale_shift_timestep, cross_gate_timestep = self._prepare_cross_attention_timestep(
-            timestep=modality.timesteps,
+            timestep=cross_timestep,
             timestep_scale_multiplier=self.simple_preprocessor.timestep_scale_multiplier,
             batch_size=transformer_args.x.shape[0],
             hidden_dtype=modality.latent.dtype,
@@ -215,7 +273,7 @@ class MultiModalTransformerArgsPreprocessor:
 
     def _prepare_cross_attention_timestep(
         self,
-        timestep: torch.Tensor,
+        timestep: torch.Tensor | None,
         timestep_scale_multiplier: int,
         batch_size: int,
         hidden_dtype: torch.dtype,

packages/ltx-core/src/ltx_core/model/video_vae/model_configurator.py

@@ -13,7 +13,7 @@ class VideoEncoderConfigurator(ModelConfigurator[VideoEncoder]):
         convolution_dimensions = config.get("dims", 3)
         in_channels = config.get("in_channels", 3)
         latent_channels = config.get("latent_channels", 128)
-        encoder_spatial_padding_mode = PaddingModeType(config.get("encoder_spatial_padding_mode", "zeros"))
+        spatial_padding_mode = PaddingModeType(config.get("spatial_padding_mode", "zeros"))
         encoder_blocks = config.get("encoder_blocks", [])
         patch_size = config.get("patch_size", 4)
         norm_layer_str = config.get("norm_layer", "pixel_norm")
@@ -27,7 +27,7 @@ class VideoEncoderConfigurator(ModelConfigurator[VideoEncoder]):
             patch_size=patch_size,
             norm_layer=NormLayerType(norm_layer_str),
             latent_log_var=LogVarianceType(latent_log_var_str),
-            encoder_spatial_padding_mode=encoder_spatial_padding_mode,
+            encoder_spatial_padding_mode=spatial_padding_mode,
         )
 
 
@@ -39,13 +39,14 @@ class VideoDecoderConfigurator(ModelConfigurator[VideoDecoder]):
         config = config.get("vae", {})
         convolution_dimensions = config.get("dims", 3)
         latent_channels = config.get("latent_channels", 128)
-        decoder_spatial_padding_mode = PaddingModeType(config.get("decoder_spatial_padding_mode", "reflect"))
+        spatial_padding_mode = PaddingModeType(config.get("spatial_padding_mode", "reflect"))
         out_channels = config.get("out_channels", 3)
         decoder_blocks = config.get("decoder_blocks", [])
         patch_size = config.get("patch_size", 4)
         norm_layer_str = config.get("norm_layer", "pixel_norm")
         causal = config.get("causal_decoder", False)
         timestep_conditioning = config.get("timestep_conditioning", True)
+        base_channels = config.get("decoder_base_channels", 128)
 
         return VideoDecoder(
             convolution_dimensions=convolution_dimensions,
@@ -56,7 +57,8 @@ class VideoDecoderConfigurator(ModelConfigurator[VideoDecoder]):
             norm_layer=NormLayerType(norm_layer_str),
             causal=causal,
             timestep_conditioning=timestep_conditioning,
-            decoder_spatial_padding_mode=decoder_spatial_padding_mode,
+            decoder_spatial_padding_mode=spatial_padding_mode,
+            base_channels=base_channels,
         )
 
 

packages/ltx-core/src/ltx_core/model/video_vae/ops.py

@@ -70,9 +70,6 @@ class PerChannelStatistics(nn.Module):
         super().__init__()
         self.register_buffer("std-of-means", torch.empty(latent_channels))
         self.register_buffer("mean-of-means", torch.empty(latent_channels))
-        self.register_buffer("mean-of-stds", torch.empty(latent_channels))
-        self.register_buffer("mean-of-stds_over_std-of-means", torch.empty(latent_channels))
-        self.register_buffer("channel", torch.empty(latent_channels))
 
     def un_normalize(self, x: torch.Tensor) -> torch.Tensor:
         return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view(

packages/ltx-core/src/ltx_core/model/video_vae/video_vae.py

@@ -515,17 +515,21 @@ def _make_decoder_block(
             spatial_padding_mode=spatial_padding_mode,
         )
     elif block_name == "compress_time":
+        out_channels = in_channels // block_config.get("multiplier", 1)
         block = DepthToSpaceUpsample(
             dims=convolution_dimensions,
             in_channels=in_channels,
             stride=(2, 1, 1),
+            out_channels_reduction_factor=block_config.get("multiplier", 1),
             spatial_padding_mode=spatial_padding_mode,
         )
     elif block_name == "compress_space":
+        out_channels = in_channels // block_config.get("multiplier", 1)
         block = DepthToSpaceUpsample(
             dims=convolution_dimensions,
             in_channels=in_channels,
             stride=(1, 2, 2),
+            out_channels_reduction_factor=block_config.get("multiplier", 1),
             spatial_padding_mode=spatial_padding_mode,
         )
     elif block_name == "compress_all":
@@ -585,6 +589,7 @@ class VideoDecoder(nn.Module):
         causal: bool = False,
         timestep_conditioning: bool = False,
         decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT,
+        base_channels: int = 128,
     ):
         super().__init__()
 
@@ -612,15 +617,9 @@ class VideoDecoder(nn.Module):
         self.decode_noise_scale = 0.025
         self.decode_timestep = 0.05
 
-        # Compute initial feature_channels by going through blocks in reverse
-        # This determines the channel width at the start of the decoder
-        feature_channels = in_channels
-        for block_name, block_params in list(reversed(decoder_blocks)):
-            block_config = block_params if isinstance(block_params, dict) else {}
-            if block_name == "res_x_y":
-                feature_channels = feature_channels * block_config.get("multiplier", 2)
-            if block_name == "compress_all":
-                feature_channels = feature_channels * block_config.get("multiplier", 1)
+        # LTX VAE decoder architecture uses 3 upsampler blocks with multiplier equals to 2.
+        # Hence the total feature_channels is multiplied by 8 (2^3).
+        feature_channels = base_channels * 8
 
         self.conv_in = make_conv_nd(
             dims=convolution_dimensions,

packages/ltx-core/src/ltx_core/quantization/fp8_cast.py

@@ -128,16 +128,16 @@ TRANSFORMER_LINEAR_DOWNCAST_MAP = (
         key_prefix="transformer_blocks.", key_suffix=".to_out.0.bias", operation=_naive_weight_or_bias_downcast
     )
     .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
+        key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.weight", operation=_naive_weight_or_bias_downcast
     )
     .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
+        key_prefix="transformer_blocks.", key_suffix="ff.net.0.proj.bias", operation=_naive_weight_or_bias_downcast
     )
     .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.2.weight", operation=_naive_weight_or_bias_downcast
+        key_prefix="transformer_blocks.", key_suffix="ff.net.2.weight", operation=_naive_weight_or_bias_downcast
     )
     .with_kv_operation(
-        key_prefix="transformer_blocks.", key_suffix=".ff.net.2.bias", operation=_naive_weight_or_bias_downcast
+        key_prefix="transformer_blocks.", key_suffix="ff.net.2.bias", operation=_naive_weight_or_bias_downcast
     )
 )
 

packages/ltx-core/src/ltx_core/text_encoders/gemma/__init__.py

@@ -1,31 +1,25 @@
 """Gemma text encoder components."""
 
-from ltx_core.text_encoders.gemma.encoders.av_encoder import (
-    AV_GEMMA_TEXT_ENCODER_KEY_OPS,
-    AVGemmaEncoderOutput,
-    AVGemmaTextEncoderModel,
-    AVGemmaTextEncoderModelConfigurator,
-)
 from ltx_core.text_encoders.gemma.encoders.base_encoder import (
-    GemmaTextEncoderModelBase,
+    GemmaEncoderOutput,
+    GemmaTextEncoder,
     encode_text,
     module_ops_from_gemma_root,
 )
-from ltx_core.text_encoders.gemma.encoders.video_only_encoder import (
-    VideoGemmaEncoderOutput,
-    VideoGemmaTextEncoderModel,
-    VideoGemmaTextEncoderModelConfigurator,
+from ltx_core.text_encoders.gemma.encoders.encoder_configurator import (
+    AV_GEMMA_TEXT_ENCODER_KEY_OPS,
+    GEMMA_MODEL_OPS,
+    VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS,
+    GemmaTextEncoderConfigurator,
 )
 
 __all__ = [
     "AV_GEMMA_TEXT_ENCODER_KEY_OPS",
-    "AVGemmaEncoderOutput",
-    "AVGemmaTextEncoderModel",
-    "AVGemmaTextEncoderModelConfigurator",
-    "GemmaTextEncoderModelBase",
-    "VideoGemmaEncoderOutput",
-    "VideoGemmaTextEncoderModel",
-    "VideoGemmaTextEncoderModelConfigurator",
+    "GEMMA_MODEL_OPS",
+    "VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS",
+    "GemmaEncoderOutput",
+    "GemmaTextEncoder",
+    "GemmaTextEncoderConfigurator",
     "encode_text",
     "module_ops_from_gemma_root",
 ]

packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_connector.py

@@ -199,17 +199,63 @@ class Embeddings1DConnector(torch.nn.Module):
 
 
 class Embeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
+    """Configurator for video embeddings connector."""
+
     @classmethod
     def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
-        config = config.get("transformer", {})
-        rope_type = LTXRopeType(config.get("rope_type", "interleaved"))
-        double_precision_rope = config.get("frequencies_precision", False) == "float64"
-        pe_max_pos = config.get("connector_positional_embedding_max_pos", [1])
+        transformer_config = config.get("transformer", {})
+        rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
+        double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
+        pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
+
+        # Video connector dimensions
+        num_attention_heads = transformer_config.get("connector_num_attention_heads", 30)
+        attention_head_dim = transformer_config.get("connector_attention_head_dim", 128)
+        num_layers = transformer_config.get("connector_num_layers", 2)
 
         connector = Embeddings1DConnector(
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            num_layers=num_layers,
             positional_embedding_max_pos=pe_max_pos,
             rope_type=rope_type,
             double_precision_rope=double_precision_rope,
-            apply_gated_attention=config.get("connector_apply_gated_attention", False),
+            apply_gated_attention=transformer_config.get("connector_apply_gated_attention", False),
+        )
+        return connector
+
+
+class AudioEmbeddings1DConnectorConfigurator(ModelConfigurator[Embeddings1DConnector]):
+    """Configurator for audio embeddings connector with separate dimension settings."""
+
+    @classmethod
+    def from_config(cls: type[Embeddings1DConnector], config: dict) -> Embeddings1DConnector:
+        transformer_config = config.get("transformer", {})
+        rope_type = LTXRopeType(transformer_config.get("rope_type", "interleaved"))
+        double_precision_rope = transformer_config.get("frequencies_precision", False) == "float64"
+        pe_max_pos = transformer_config.get("connector_positional_embedding_max_pos", [1])
+
+        # Audio connector dimensions - fall back to video connector config for backwards compatibility
+        num_attention_heads = transformer_config.get(
+            "audio_connector_num_attention_heads",
+            transformer_config.get("connector_num_attention_heads", 30),
+        )
+        attention_head_dim = transformer_config.get(
+            "audio_connector_attention_head_dim",
+            transformer_config.get("connector_attention_head_dim", 128),
+        )
+        num_layers = transformer_config.get(
+            "audio_connector_num_layers",
+            transformer_config.get("connector_num_layers", 2),
+        )
+
+        connector = Embeddings1DConnector(
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            num_layers=num_layers,
+            positional_embedding_max_pos=pe_max_pos,
+            rope_type=rope_type,
+            double_precision_rope=double_precision_rope,
+            apply_gated_attention=transformer_config.get("connector_apply_gated_attention", False),
         )
         return connector

packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_processor.py

@@ -0,0 +1,43 @@
+import torch
+from torch import nn
+
+from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnector
+
+
+def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Convert connector output mask to binary mask and apply to encoded tensor."""
+    binary_mask = (encoded_mask < 0.000001).to(torch.int64)
+    binary_mask = binary_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
+    encoded = encoded * binary_mask
+    return encoded, binary_mask
+
+
+class EmbeddingsProcessor(nn.Module):
+    """Wraps video connector + optional audio connector.
+    Returns (video_encoded, audio_encoded | None, binary_mask).
+    """
+
+    def __init__(self, video_connector: Embeddings1DConnector, audio_connector: Embeddings1DConnector | None = None):
+        super().__init__()
+        self.video_connector = video_connector
+        self.audio_connector = audio_connector
+
+    def create_embeddings(
+        self,
+        video_features: torch.Tensor,
+        audio_features: torch.Tensor | None,
+        additive_attention_mask: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
+        if self.audio_connector is not None and audio_features is None:
+            raise ValueError("Audio connector is configured but no audio features were provided.")
+        if self.audio_connector is None and audio_features is not None:
+            raise ValueError("Audio features were provided but no audio connector is configured.")
+
+        video_encoded, video_mask = self.video_connector(video_features, additive_attention_mask)
+        video_encoded, binary_mask = _to_binary_mask(video_encoded, video_mask)
+
+        audio_encoded = None
+        if self.audio_connector is not None:
+            audio_encoded, _ = self.audio_connector(audio_features, additive_attention_mask)
+
+        return video_encoded, audio_encoded, binary_mask.squeeze(-1)

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/av_encoder.py

@@ -1,152 +0,0 @@
-from typing import NamedTuple
-
-import torch
-from transformers import Gemma3Config
-from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
-from transformers.models.gemma3 import Gemma3ForConditionalGeneration
-
-from ltx_core.loader import KeyValueOperationResult
-from ltx_core.loader.module_ops import ModuleOps
-from ltx_core.loader.sd_ops import SDOps
-from ltx_core.model.model_protocol import ModelConfigurator
-from ltx_core.text_encoders.gemma.config import GEMMA3_CONFIG_FOR_LTX
-from ltx_core.text_encoders.gemma.embeddings_connector import (
-    Embeddings1DConnector,
-    Embeddings1DConnectorConfigurator,
-)
-from ltx_core.text_encoders.gemma.encoders.base_encoder import (
-    GemmaTextEncoderModelBase,
-)
-from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
-from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
-
-
-class AVGemmaEncoderOutput(NamedTuple):
-    video_encoding: torch.Tensor
-    audio_encoding: torch.Tensor
-    attention_mask: torch.Tensor
-
-
-class AVGemmaTextEncoderModel(GemmaTextEncoderModelBase):
-    """
-    AVGemma Text Encoder Model.
-    This class combines the tokenizer, Gemma model, feature extractor from base class and a
-    video and audio embeddings connectors to provide a preprocessing for audio-visual pipeline.
-    """
-
-    def __init__(
-        self,
-        feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
-        embeddings_connector: Embeddings1DConnector,
-        audio_embeddings_connector: Embeddings1DConnector,
-        tokenizer: LTXVGemmaTokenizer | None = None,
-        model: Gemma3ForConditionalGeneration | None = None,
-        dtype: torch.dtype = torch.bfloat16,
-    ) -> None:
-        super().__init__(
-            feature_extractor_linear=feature_extractor_linear,
-            tokenizer=tokenizer,
-            model=model,
-            dtype=dtype,
-        )
-        self.embeddings_connector = embeddings_connector.to(dtype=dtype)
-        self.audio_embeddings_connector = audio_embeddings_connector.to(dtype=dtype)
-
-    def _run_connectors(
-        self, encoded_input: torch.Tensor, attention_mask: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        connector_attention_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
-
-        encoded, encoded_connector_attention_mask = self.embeddings_connector(
-            encoded_input,
-            connector_attention_mask,
-        )
-
-        # restore the mask values to int64
-        attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
-        attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
-        encoded = encoded * attention_mask
-
-        encoded_for_audio, _ = self.audio_embeddings_connector(encoded_input, connector_attention_mask)
-
-        return encoded, encoded_for_audio, attention_mask.squeeze(-1)
-
-    def forward(self, text: str, padding_side: str = "left") -> AVGemmaEncoderOutput:
-        encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
-        video_encoding, audio_encoding, attention_mask = self._run_connectors(encoded_inputs, attention_mask)
-        return AVGemmaEncoderOutput(video_encoding, audio_encoding, attention_mask)
-
-
-class AVGemmaTextEncoderModelConfigurator(ModelConfigurator[AVGemmaTextEncoderModel]):
-    @classmethod
-    def from_config(cls: type["AVGemmaTextEncoderModel"], config: dict) -> "AVGemmaTextEncoderModel":
-        feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
-        embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
-        audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
-        gemma_config = Gemma3Config.from_dict(GEMMA3_CONFIG_FOR_LTX.to_dict())
-        with torch.device("meta"):
-            model = Gemma3ForConditionalGeneration(gemma_config)
-        return AVGemmaTextEncoderModel(
-            model=model,
-            feature_extractor_linear=feature_extractor_linear,
-            embeddings_connector=embeddings_connector,
-            audio_embeddings_connector=audio_embeddings_connector,
-        )
-
-
-AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
-    SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS")
-    # 1. Map the feature extractor
-    .with_matching(prefix="text_embedding_projection.")
-    .with_replacement("text_embedding_projection.", "feature_extractor_linear.")
-    # 2. Map the connectors (fixing the swapped prefixes from before)
-    .with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
-    .with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_connector.")
-    .with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
-    .with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_embeddings_connector.")
-    # 3. Map language model layers (note the double .model prefix)
-    .with_matching(prefix="language_model.model.")
-    .with_replacement("language_model.model.", "model.model.language_model.")
-    # 4. Map the Vision Tower
-    .with_matching(prefix="vision_tower.")
-    .with_replacement("vision_tower.", "model.model.vision_tower.")
-    # 5. Map the Multi-Modal Projector
-    .with_matching(prefix="multi_modal_projector.")
-    .with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
-    .with_kv_operation(
-        operation=lambda key, value: [
-            KeyValueOperationResult(key, value),
-            KeyValueOperationResult("model.lm_head.weight", value),
-        ],
-        key_prefix="model.model.language_model.embed_tokens.weight",
-    )
-)
-
-
-def create_and_populate(module: AVGemmaTextEncoderModel) -> AVGemmaTextEncoderModel:
-    model = module.model
-    v_model = model.model.vision_tower.vision_model
-    l_model = model.model.language_model
-
-    config = model.config.text_config
-    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
-    base = config.rope_local_base_freq
-    local_rope_freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
-    inv_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](config)
-
-    positions_length = len(v_model.embeddings.position_ids[0])
-    position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0)
-    v_model.embeddings.register_buffer("position_ids", position_ids)
-    embed_scale = torch.tensor(model.config.text_config.hidden_size**0.5, device="cpu")
-    l_model.embed_tokens.register_buffer("embed_scale", embed_scale)
-    l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
-    l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
-
-    return module
-
-
-GEMMA_MODEL_OPS = ModuleOps(
-    name="GemmaModel",
-    matcher=lambda module: hasattr(module, "model") and isinstance(module.model, Gemma3ForConditionalGeneration),
-    mutator=create_and_populate,
-)

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py

@@ -1,77 +1,77 @@
 import functools
 from pathlib import Path
+from typing import NamedTuple
 
 import torch
-from einops import rearrange
 from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor
 
 from ltx_core.loader.module_ops import ModuleOps
-from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
+from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
 from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
 from ltx_core.utils import find_matching_file
 
 
-class GemmaTextEncoderModelBase(torch.nn.Module):
-    """
-    Gemma Text Encoder Model.
-    This base class combines the tokenizer, Gemma model and feature extractor to provide a preprocessing
-    for implementation classes for multimodal pipelines. It processes input text through tokenization,
-    obtains hidden states from the base language model, applies a linear feature extractor.
-    Args:
-        tokenizer (LTXVGemmaTokenizer): The tokenizer used for text preprocessing.
-        model (Gemma3ForConditionalGeneration): The base Gemma LLM.
-        feature_extractor_linear (GemmaFeaturesExtractorProjLinear): Linear projection for hidden state aggregation.
-        dtype (torch.dtype, optional): The data type for model parameters (default: torch.bfloat16).
+class GemmaEncoderOutput(NamedTuple):
+    video_encoding: torch.Tensor
+    audio_encoding: torch.Tensor | None
+    attention_mask: torch.Tensor
+
+
+class GemmaTextEncoder(torch.nn.Module):
+    """Unified Gemma text encoder with 3-block pipeline.
+    Block 1: Gemma model (runs LLM, gets hidden states)
+    Block 2: Feature extractor
+    Block 3: Embeddings processor (connector with optional audio)
     """
 
     def __init__(
         self,
-        feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
-        tokenizer: LTXVGemmaTokenizer | None = None,
+        feature_extractor: torch.nn.Module,
+        embeddings_processor: EmbeddingsProcessor,
         model: Gemma3ForConditionalGeneration | None = None,
-        img_processor: Gemma3Processor | None = None,
+        tokenizer: LTXVGemmaTokenizer | None = None,
+        processor: Gemma3Processor | None = None,
         dtype: torch.dtype = torch.bfloat16,
-    ) -> None:
+    ):
         super().__init__()
-        self.tokenizer = tokenizer
         self.model = model
-        self.processor = img_processor
-        self.feature_extractor_linear = feature_extractor_linear.to(dtype=dtype)
-
-    def _run_feature_extractor(
-        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "right"
-    ) -> torch.Tensor:
-        encoded_text_features = torch.stack(hidden_states, dim=-1)
-        encoded_text_features_dtype = encoded_text_features.dtype
-
-        sequence_lengths = attention_mask.sum(dim=-1)
-        normed_concated_encoded_text_features = _norm_and_concat_padded_batch(
-            encoded_text_features, sequence_lengths, padding_side=padding_side
-        )
-
-        return self.feature_extractor_linear(normed_concated_encoded_text_features.to(encoded_text_features_dtype))
+        self.tokenizer = tokenizer
+        self.processor = processor
+        self.feature_extractor = feature_extractor.to(dtype=dtype)
+        self.embeddings_processor = embeddings_processor.to(dtype=dtype)
 
     def _convert_to_additive_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
-        return (attention_mask - 1).to(dtype).reshape(
+        return (attention_mask.to(torch.int64) - 1).to(dtype).reshape(
             (attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
         ) * torch.finfo(dtype).max
 
-    def _preprocess_text(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
-        """
-        Encode a given string into feature tensors suitable for downstream tasks.
-        Args:
-            text (str): Input string to encode.
-        Returns:
-            tuple[torch.Tensor, dict[str, torch.Tensor]]: Encoded features and a dictionary with attention mask.
+    def precompute(
+        self, text: str, padding_side: str = "left"
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
+        """Blocks 1+2: Gemma model -> feature extraction.
+        Used by process_captions.py for offline precomputation.
+        Returns (video_features, audio_features | None, attention_mask).
         """
+        # Block 1: Run Gemma
         token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
         input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
         attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
         outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
-        projected = self._run_feature_extractor(
-            hidden_states=outputs.hidden_states, attention_mask=attention_mask, padding_side=padding_side
+
+        # Block 2: Feature extraction
+        video_feats, audio_feats = self.feature_extractor(outputs.hidden_states, attention_mask, padding_side)
+        return video_feats, audio_feats, attention_mask
+
+    def forward(self, text: str, padding_side: str = "left") -> GemmaEncoderOutput:
+        """Full pipeline: precompute -> embeddings processor."""
+        video_feats, audio_feats, attention_mask = self.precompute(text, padding_side)
+        additive_mask = self._convert_to_additive_mask(attention_mask, video_feats.dtype)
+        video_enc, audio_enc, binary_mask = self.embeddings_processor.create_embeddings(
+            video_feats, audio_feats, additive_mask
         )
-        return projected, attention_mask
+        return GemmaEncoderOutput(video_enc, audio_enc, binary_mask)
+
+    # --- Prompt enhancement methods ---
 
     def _enhance(
         self,
@@ -111,7 +111,6 @@ class GemmaTextEncoderModelBase(torch.nn.Module):
         seed: int = 10,
     ) -> str:
         """Enhance a text prompt for T2V generation."""
-
         system_prompt = system_prompt or self.default_gemma_t2v_system_prompt
 
         messages = [
@@ -151,67 +150,8 @@ class GemmaTextEncoderModelBase(torch.nn.Module):
     def default_gemma_t2v_system_prompt(self) -> str:
         return _load_system_prompt("gemma_t2v_system_prompt.txt")
 
-    def forward(self, text: str, padding_side: str = "left") -> tuple[torch.Tensor, torch.Tensor]:
-        raise NotImplementedError("This method is not implemented for the base class")
 
-
-def _norm_and_concat_padded_batch(
-    encoded_text: torch.Tensor,
-    sequence_lengths: torch.Tensor,
-    padding_side: str = "right",
-) -> torch.Tensor:
-    """Normalize and flatten multi-layer hidden states, respecting padding.
-    Performs per-batch, per-layer normalization using masked mean and range,
-    then concatenates across the layer dimension.
-    Args:
-        encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
-        sequence_lengths: Number of valid (non-padded) tokens per batch item.
-        padding_side: Whether padding is on "left" or "right".
-    Returns:
-        Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
-        with padded positions zeroed out.
-    """
-    b, t, d, l = encoded_text.shape  # noqa: E741
-    device = encoded_text.device
-
-    # Build mask: [B, T, 1, 1]
-    token_indices = torch.arange(t, device=device)[None, :]  # [1, T]
-
-    if padding_side == "right":
-        # For right padding, valid tokens are from 0 to sequence_length-1
-        mask = token_indices < sequence_lengths[:, None]  # [B, T]
-    elif padding_side == "left":
-        # For left padding, valid tokens are from (T - sequence_length) to T-1
-        start_indices = t - sequence_lengths[:, None]  # [B, 1]
-        mask = token_indices >= start_indices  # [B, T]
-    else:
-        raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
-
-    mask = rearrange(mask, "b t -> b t 1 1")
-
-    eps = 1e-6
-
-    # Compute masked mean: [B, 1, 1, L]
-    masked = encoded_text.masked_fill(~mask, 0.0)
-    denom = (sequence_lengths * d).view(b, 1, 1, 1)
-    mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
-
-    # Compute masked min/max: [B, 1, 1, L]
-    x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
-    x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
-    range_ = x_max - x_min
-
-    # Normalize only the valid tokens
-    normed = 8 * (encoded_text - mean) / (range_ + eps)
-
-    # concat to be [Batch, T,  D * L] - this preserves the original structure
-    normed = normed.reshape(b, t, -1)  # [B, T, D * L]
-
-    # Apply mask to preserve original padding (set padded positions to 0)
-    mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
-    normed = normed.masked_fill(~mask_flattened, 0.0)
-
-    return normed
+# --- Standalone utility functions ---
 
 
 @functools.lru_cache(maxsize=2)
@@ -220,50 +160,6 @@ def _load_system_prompt(prompt_name: str) -> str:
         return f.read()
 
 
-def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
-    tokenizer_root = str(find_matching_file(gemma_root, "tokenizer.model").parent)
-    processor_root = str(find_matching_file(gemma_root, "preprocessor_config.json").parent)
-
-    def load_tokenizer(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
-        module.tokenizer = LTXVGemmaTokenizer(tokenizer_root, 1024)
-        return module
-
-    def load_processor(module: GemmaTextEncoderModelBase) -> GemmaTextEncoderModelBase:
-        image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True)
-        if not module.tokenizer:
-            raise ValueError("Tokenizer model operation must be performed before processor model operation")
-        module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer)
-        return module
-
-    tokenizer_load_ops = ModuleOps(
-        "TokenizerLoad",
-        matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.tokenizer is None,
-        mutator=load_tokenizer,
-    )
-    processor_load_ops = ModuleOps(
-        "ProcessorLoad",
-        matcher=lambda module: isinstance(module, GemmaTextEncoderModelBase) and module.processor is None,
-        mutator=load_processor,
-    )
-    return (tokenizer_load_ops, processor_load_ops)
-
-
-def encode_text(text_encoder: GemmaTextEncoderModelBase, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]:
-    """
-    Encode a list of prompts using the provided Gemma text encoder.
-    Args:
-        text_encoder: The Gemma text encoder instance.
-        prompts: List of prompt strings to encode.
-    Returns:
-        List of tuples, each containing (v_context, a_context) tensors for each prompt.
-    """
-    result = []
-    for prompt in prompts:
-        v_context, a_context, _ = text_encoder(prompt)
-        result.append((v_context, a_context))
-    return result
-
-
 def _cat_with_padding(
     tensor: torch.Tensor,
     padding_length: int,
@@ -289,21 +185,55 @@ def _pad_inputs_for_attention_alignment(
     pad_token_id: int = 0,
     alignment: int = 8,
 ) -> dict[str, torch.Tensor]:
-    """Pad sequence length to multiple of alignment for Flash Attention compatibility.
-    Flash Attention within SDPA requires sequence lengths aligned to 8 bytes.
-    This pads input_ids, attention_mask, and token_type_ids (if present) to prevent
-    'p.attn_bias_ptr is not correctly aligned' errors.
-    """
+    """Pad sequence length to multiple of alignment for Flash Attention compatibility."""
     seq_len = model_inputs.input_ids.shape[1]
     padded_len = ((seq_len + alignment - 1) // alignment) * alignment
     padding_length = padded_len - seq_len
 
     if padding_length > 0:
         model_inputs["input_ids"] = _cat_with_padding(model_inputs.input_ids, padding_length, pad_token_id)
-
         model_inputs["attention_mask"] = _cat_with_padding(model_inputs.attention_mask, padding_length, 0)
-
         if "token_type_ids" in model_inputs and model_inputs["token_type_ids"] is not None:
             model_inputs["token_type_ids"] = _cat_with_padding(model_inputs["token_type_ids"], padding_length, 0)
 
     return model_inputs
+
+
+def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
+    tokenizer_root = str(find_matching_file(gemma_root, "tokenizer.model").parent)
+    processor_root = str(find_matching_file(gemma_root, "preprocessor_config.json").parent)
+
+    def load_tokenizer(module: GemmaTextEncoder) -> GemmaTextEncoder:
+        module.tokenizer = LTXVGemmaTokenizer(tokenizer_root, 1024)
+        return module
+
+    def load_processor(module: GemmaTextEncoder) -> GemmaTextEncoder:
+        image_processor = AutoImageProcessor.from_pretrained(processor_root, local_files_only=True)
+        if not module.tokenizer:
+            raise ValueError("Tokenizer model operation must be performed before processor model operation")
+        module.processor = Gemma3Processor(image_processor=image_processor, tokenizer=module.tokenizer.tokenizer)
+        return module
+
+    tokenizer_load_ops = ModuleOps(
+        "TokenizerLoad",
+        matcher=lambda module: isinstance(module, GemmaTextEncoder) and module.tokenizer is None,
+        mutator=load_tokenizer,
+    )
+    processor_load_ops = ModuleOps(
+        "ProcessorLoad",
+        matcher=lambda module: isinstance(module, GemmaTextEncoder) and module.processor is None,
+        mutator=load_processor,
+    )
+    return (tokenizer_load_ops, processor_load_ops)
+
+
+def encode_text(text_encoder: GemmaTextEncoder, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]:
+    """Encode a list of prompts using the provided Gemma text encoder.
+    Returns:
+        List of tuples, each containing (v_context, a_context) tensors for each prompt.
+    """
+    result = []
+    for prompt in prompts:
+        v_context, a_context, _ = text_encoder(prompt)
+        result.append((v_context, a_context))
+    return result

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py

@@ -0,0 +1,174 @@
+import torch
+from transformers import Gemma3Config
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from transformers.models.gemma3 import Gemma3ForConditionalGeneration
+
+from ltx_core.loader import KeyValueOperationResult
+from ltx_core.loader.module_ops import ModuleOps
+from ltx_core.loader.sd_ops import SDOps
+from ltx_core.model.model_protocol import ModelConfigurator
+from ltx_core.text_encoders.gemma.config import GEMMA3_CONFIG_FOR_LTX
+from ltx_core.text_encoders.gemma.embeddings_connector import (
+    AudioEmbeddings1DConnectorConfigurator,
+    Embeddings1DConnectorConfigurator,
+)
+from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_core.text_encoders.gemma.feature_extractor import (
+    FeatureExtractorV1,
+    FeatureExtractorV2,
+)
+
+
+class GemmaTextEncoderConfigurator(ModelConfigurator[GemmaTextEncoder]):
+    @classmethod
+    def from_config(cls, config: dict) -> GemmaTextEncoder:
+        transformer_config = config.get("transformer", {})
+
+        gemma_config = Gemma3Config.from_dict(GEMMA3_CONFIG_FOR_LTX.to_dict())
+        with torch.device("meta"):
+            model = Gemma3ForConditionalGeneration(gemma_config)
+
+        # Create video embeddings connector (always needed)
+        video_connector = Embeddings1DConnectorConfigurator.from_config(config)
+
+        # Create audio embeddings connector
+        audio_connector = AudioEmbeddings1DConnectorConfigurator.from_config(config)
+
+        # Create embeddings processor with both connectors
+        embeddings_processor = EmbeddingsProcessor(
+            video_connector=video_connector,
+            audio_connector=audio_connector,
+        )
+
+        feature_extractor = _create_feature_extractor(transformer_config)
+
+        return GemmaTextEncoder(
+            feature_extractor=feature_extractor,
+            embeddings_processor=embeddings_processor,
+            model=model,
+        )
+
+
+_V2_EXPECTED_CONFIG = {
+    "caption_proj_before_connector": True,
+    "caption_projection_first_linear": False,
+    "caption_proj_input_norm": False,
+    "caption_projection_second_linear": False,
+}
+
+
+def _create_feature_extractor(transformer_config: dict) -> torch.nn.Module:
+    """Select and create the appropriate feature extractor based on config.
+    Detection logic:
+    - V1: V2 config keys absent → projection lives in transformer
+    - V2: V2 config keys present with exact expected values → per-token RMS norm with dual aggregate embeds
+    - Anything else: NotImplementedError (config drift)
+    """
+    gemma_text_config = GEMMA3_CONFIG_FOR_LTX.text_config
+    embedding_dim = gemma_text_config.hidden_size
+    num_layers = gemma_text_config.num_hidden_layers + 1  # +1 for the embedding layer
+    flat_dim = embedding_dim * num_layers
+
+    overlapping_keys = transformer_config.keys() & _V2_EXPECTED_CONFIG.keys()
+    if not overlapping_keys:
+        aggregate_embed = torch.nn.Linear(flat_dim, embedding_dim, bias=False)
+        return FeatureExtractorV1(aggregate_embed=aggregate_embed, is_av=True)
+
+    missing_keys = _V2_EXPECTED_CONFIG.keys() - overlapping_keys
+    if missing_keys:
+        raise NotImplementedError("Partial V2 config — missing keys: " + ", ".join(sorted(missing_keys)))
+
+    unexpected_value_keys = {k for k in overlapping_keys if transformer_config[k] != _V2_EXPECTED_CONFIG[k]}
+    if unexpected_value_keys:
+        raise NotImplementedError(
+            "Unknown config: "
+            + ", ".join(
+                f"{k}={transformer_config[k]!r} (expected {_V2_EXPECTED_CONFIG[k]!r})" for k in unexpected_value_keys
+            )
+        )
+
+    video_inner_dim = transformer_config["num_attention_heads"] * transformer_config["attention_head_dim"]
+    audio_inner_dim = transformer_config["audio_num_attention_heads"] * transformer_config["audio_attention_head_dim"]
+    return FeatureExtractorV2(
+        video_aggregate_embed=torch.nn.Linear(flat_dim, video_inner_dim, bias=True),
+        embedding_dim=embedding_dim,
+        audio_aggregate_embed=torch.nn.Linear(flat_dim, audio_inner_dim, bias=True),
+    )
+
+
+AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
+    SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS")
+    # 1. Map the feature extractor (V1: aggregate_embed inside feature_extractor)
+    .with_matching(prefix="text_embedding_projection.aggregate_embed.")
+    .with_replacement("text_embedding_projection.aggregate_embed.", "feature_extractor.aggregate_embed.")
+    # V2 dual aggregate embeds
+    .with_matching(prefix="text_embedding_projection.video_aggregate_embed.")
+    .with_replacement("text_embedding_projection.video_aggregate_embed.", "feature_extractor.video_aggregate_embed.")
+    .with_matching(prefix="text_embedding_projection.audio_aggregate_embed.")
+    .with_replacement("text_embedding_projection.audio_aggregate_embed.", "feature_extractor.audio_aggregate_embed.")
+    # 2. Map the connectors
+    .with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
+    .with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_processor.video_connector.")
+    .with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
+    .with_replacement("model.diffusion_model.audio_embeddings_connector.", "embeddings_processor.audio_connector.")
+    # 3. Map language model layers (note the double .model prefix)
+    .with_matching(prefix="language_model.model.")
+    .with_replacement("language_model.model.", "model.model.language_model.")
+    # 4. Map the Vision Tower
+    .with_matching(prefix="vision_tower.")
+    .with_replacement("vision_tower.", "model.model.vision_tower.")
+    # 5. Map the Multi-Modal Projector
+    .with_matching(prefix="multi_modal_projector.")
+    .with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
+    .with_kv_operation(
+        operation=lambda key, value: [
+            KeyValueOperationResult(key, value),
+            KeyValueOperationResult("model.lm_head.weight", value),
+        ],
+        key_prefix="model.model.language_model.embed_tokens.weight",
+    )
+)
+
+
+VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS = (
+    SDOps("VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS")
+    # 1. Map the feature extractor (V1: aggregate_embed inside feature_extractor)
+    .with_matching(prefix="text_embedding_projection.aggregate_embed.")
+    .with_replacement("text_embedding_projection.aggregate_embed.", "feature_extractor.aggregate_embed.")
+    # V2 video aggregate embed
+    .with_matching(prefix="text_embedding_projection.video_aggregate_embed.")
+    .with_replacement("text_embedding_projection.video_aggregate_embed.", "feature_extractor.video_aggregate_embed.")
+    # 2. Map the connectors
+    .with_matching(prefix="model.diffusion_model.embeddings_connector.")
+    .with_replacement("model.diffusion_model.embeddings_connector.", "embeddings_processor.video_connector.")
+)
+
+
+def create_and_populate(module: GemmaTextEncoder) -> GemmaTextEncoder:
+    model = module.model
+    v_model = model.model.vision_tower.vision_model
+    l_model = model.model.language_model
+
+    config = model.config.text_config
+    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+    base = config.rope_local_base_freq
+    local_rope_freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim))
+    inv_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](config)
+
+    positions_length = len(v_model.embeddings.position_ids[0])
+    position_ids = torch.arange(positions_length, dtype=torch.long, device="cpu").unsqueeze(0)
+    v_model.embeddings.register_buffer("position_ids", position_ids)
+    embed_scale = torch.tensor(model.config.text_config.hidden_size**0.5, device="cpu")
+    l_model.embed_tokens.register_buffer("embed_scale", embed_scale)
+    l_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
+    l_model.rotary_emb.register_buffer("inv_freq", inv_freqs)
+
+    return module
+
+
+GEMMA_MODEL_OPS = ModuleOps(
+    name="GemmaModel",
+    matcher=lambda module: hasattr(module, "model") and isinstance(module.model, Gemma3ForConditionalGeneration),
+    mutator=create_and_populate,
+)

packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/video_only_encoder.py

@@ -1,85 +0,0 @@
-from typing import NamedTuple
-
-import torch
-from transformers import Gemma3ForConditionalGeneration
-
-from ltx_core.loader.sd_ops import SDOps
-from ltx_core.model.model_protocol import ModelConfigurator
-from ltx_core.text_encoders.gemma.embeddings_connector import (
-    Embeddings1DConnector,
-    Embeddings1DConnectorConfigurator,
-)
-from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoderModelBase
-from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
-from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
-
-
-class VideoGemmaEncoderOutput(NamedTuple):
-    video_encoding: torch.Tensor
-    attention_mask: torch.Tensor
-
-
-class VideoGemmaTextEncoderModel(GemmaTextEncoderModelBase):
-    """
-    Video Gemma Text Encoder Model.
-    This class combines the tokenizer, Gemma model, feature extractor from base class and a
-    video embeddings connector to provide a preprocessing for video only pipeline.
-    """
-
-    def __init__(
-        self,
-        feature_extractor_linear: GemmaFeaturesExtractorProjLinear,
-        embeddings_connector: Embeddings1DConnector,
-        tokenizer: LTXVGemmaTokenizer | None = None,
-        model: Gemma3ForConditionalGeneration | None = None,
-        dtype: torch.dtype = torch.bfloat16,
-    ) -> None:
-        super().__init__(
-            feature_extractor_linear=feature_extractor_linear,
-            tokenizer=tokenizer,
-            model=model,
-            dtype=dtype,
-        )
-        self.embeddings_connector = embeddings_connector.to(dtype=dtype)
-
-    def _run_connector(
-        self, encoded_input: torch.Tensor, attention_mask: torch.Tensor
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        connector_attention_mask = self._convert_to_additive_mask(attention_mask, encoded_input.dtype)
-
-        encoded, encoded_connector_attention_mask = self.embeddings_connector(
-            encoded_input,
-            connector_attention_mask,
-        )
-
-        # restore the mask values to int64
-        attention_mask = (encoded_connector_attention_mask < 0.000001).to(torch.int64)
-        attention_mask = attention_mask.reshape([encoded.shape[0], encoded.shape[1], 1])
-        encoded = encoded * attention_mask
-
-        return encoded, attention_mask.squeeze(-1)
-
-    def forward(self, text: str, padding_side: str = "left") -> VideoGemmaEncoderOutput:
-        encoded_inputs, attention_mask = self._preprocess_text(text, padding_side)
-        video_encoding, attention_mask = self._run_connector(encoded_inputs, attention_mask)
-        return VideoGemmaEncoderOutput(video_encoding, attention_mask)
-
-
-class VideoGemmaTextEncoderModelConfigurator(ModelConfigurator[VideoGemmaTextEncoderModel]):
-    @classmethod
-    def from_config(cls: type["VideoGemmaTextEncoderModel"], config: dict) -> "VideoGemmaTextEncoderModel":
-        feature_extractor_linear = GemmaFeaturesExtractorProjLinear.from_config(config)
-        embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
-        return VideoGemmaTextEncoderModel(
-            feature_extractor_linear=feature_extractor_linear,
-            embeddings_connector=embeddings_connector,
-        )
-
-
-VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS = (
-    SDOps("VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS")
-    .with_matching(prefix="text_embedding_projection.")
-    .with_matching(prefix="model.diffusion_model.embeddings_connector.")
-    .with_replacement("text_embedding_projection.", "feature_extractor_linear.")
-    .with_replacement("model.diffusion_model.embeddings_connector.", "embeddings_connector.")
-)

packages/ltx-core/src/ltx_core/text_encoders/gemma/feature_extractor.py

@@ -1,36 +1,141 @@
+import math
+
 import torch
+from einops import rearrange
+from torch import nn
 
-from ltx_core.model.model_protocol import ModelConfigurator
+# ---------------------------------------------------------------------------
+# Normalization functions
+# ---------------------------------------------------------------------------
 
 
-class GemmaFeaturesExtractorProjLinear(torch.nn.Module, ModelConfigurator["GemmaFeaturesExtractorProjLinear"]):
-    """
-    Feature extractor module for Gemma models.
-    This module applies a single linear projection to the input tensor.
-    It expects a flattened feature tensor of shape (batch_size, 3840*49).
-    The linear layer maps this to a (batch_size, 3840) embedding.
-    Attributes:
-        aggregate_embed (torch.nn.Linear): Linear projection layer.
+def _norm_and_concat_padded_batch(
+    encoded_text: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    padding_side: str = "right",
+) -> torch.Tensor:
+    """Normalize and flatten multi-layer hidden states, respecting padding.
+    Performs per-batch, per-layer normalization using masked mean and range,
+    then concatenates across the layer dimension.
+    Args:
+        encoded_text: Hidden states of shape [batch, seq_len, hidden_dim, num_layers].
+        sequence_lengths: Number of valid (non-padded) tokens per batch item.
+        padding_side: Whether padding is on "left" or "right".
+    Returns:
+        Normalized tensor of shape [batch, seq_len, hidden_dim * num_layers],
+        with padded positions zeroed out.
     """
+    b, t, d, l = encoded_text.shape  # noqa: E741
+    device = encoded_text.device
 
-    def __init__(self) -> None:
-        """
-        Initialize the GemmaFeaturesExtractorProjLinear module.
-        The input dimension is expected to be 3840 * 49, and the output is 3840.
-        """
+    token_indices = torch.arange(t, device=device)[None, :]
+
+    if padding_side == "right":
+        mask = token_indices < sequence_lengths[:, None]
+    elif padding_side == "left":
+        start_indices = t - sequence_lengths[:, None]
+        mask = token_indices >= start_indices
+    else:
+        raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
+
+    mask = rearrange(mask, "b t -> b t 1 1")
+
+    eps = 1e-6
+
+    masked = encoded_text.masked_fill(~mask, 0.0)
+    denom = (sequence_lengths * d).view(b, 1, 1, 1)
+    mean = masked.sum(dim=(1, 2), keepdim=True) / (denom + eps)
+
+    x_min = encoded_text.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True)
+    x_max = encoded_text.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True)
+    range_ = x_max - x_min
+
+    normed = 8 * (encoded_text - mean) / (range_ + eps)
+    normed = normed.reshape(b, t, -1)
+
+    mask_flattened = rearrange(mask, "b t 1 1 -> b t 1").expand(-1, -1, d * l)
+    normed = normed.masked_fill(~mask_flattened, 0.0)
+
+    return normed
+
+
+def norm_and_concat_per_token_rms(
+    encoded_text: torch.Tensor,
+    attention_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Per-token RMSNorm normalization for V2 models.
+    Args:
+        encoded_text: [B, T, D, L]
+        attention_mask: [B, T] binary mask
+    Returns:
+        [B, T, D*L] normalized tensor with padding zeroed out.
+    """
+    B, T, D, L = encoded_text.shape  # noqa: N806
+    variance = torch.mean(encoded_text**2, dim=2, keepdim=True)  # [B,T,1,L]
+    normed = encoded_text * torch.rsqrt(variance + 1e-6)
+    normed = normed.reshape(B, T, D * L)
+    mask_3d = attention_mask.bool().unsqueeze(-1)  # [B, T, 1]
+    return torch.where(mask_3d, normed, torch.zeros_like(normed))
+
+
+def _rescale_norm(x: torch.Tensor, target_dim: int, source_dim: int) -> torch.Tensor:
+    """Rescale normalization: x * sqrt(target_dim / source_dim)."""
+    return x * math.sqrt(target_dim / source_dim)
+
+
+# ---------------------------------------------------------------------------
+# Feature extractor variants
+# ---------------------------------------------------------------------------
+
+
+class FeatureExtractorV1(nn.Module):
+    """19B: per-segment norm -> aggregate_embed -> 3840"""
+
+    def __init__(self, aggregate_embed: nn.Module, is_av: bool = False):
         super().__init__()
-        self.aggregate_embed = torch.nn.Linear(3840 * 49, 3840, bias=False)
+        self.aggregate_embed = aggregate_embed
+        self.is_av = is_av
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass for the feature extractor.
-        Args:
-            x (torch.Tensor): Input tensor of shape (batch_size, 3840 * 49).
-        Returns:
-            torch.Tensor: Output tensor of shape (batch_size, 3840).
-        """
-        return self.aggregate_embed(x)
+    def forward(
+        self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, padding_side: str = "left"
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states
+        dtype = encoded.dtype
+        sequence_lengths = attention_mask.sum(dim=-1)
+        normed = _norm_and_concat_padded_batch(encoded, sequence_lengths, padding_side)
+        features = self.aggregate_embed(normed.to(dtype))
+        if self.is_av:
+            return features, features
+        return features, None
 
-    @classmethod
-    def from_config(cls: type["GemmaFeaturesExtractorProjLinear"], _config: dict) -> "GemmaFeaturesExtractorProjLinear":
-        return cls()
+
+class FeatureExtractorV2(nn.Module):
+    """20B: per-token RMS norm → rescale → dual aggregate embeds"""
+
+    def __init__(
+        self,
+        video_aggregate_embed: nn.Linear,
+        embedding_dim: int,
+        audio_aggregate_embed: nn.Linear | None = None,
+    ):
+        super().__init__()
+        self.video_aggregate_embed = video_aggregate_embed
+        self.audio_aggregate_embed = audio_aggregate_embed
+        self.embedding_dim = embedding_dim
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        padding_side: str = "left",  # noqa: ARG002
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        encoded = torch.stack(hidden_states, dim=-1) if isinstance(hidden_states, (list, tuple)) else hidden_states
+        normed = norm_and_concat_per_token_rms(encoded, attention_mask)
+        normed = normed.to(encoded.dtype)
+        v_dim = self.video_aggregate_embed.out_features
+        video = self.video_aggregate_embed(_rescale_norm(normed, v_dim, self.embedding_dim))
+        audio = None
+        if self.audio_aggregate_embed is not None:
+            a_dim = self.audio_aggregate_embed.out_features
+            audio = self.audio_aggregate_embed(_rescale_norm(normed, a_dim, self.embedding_dim))
+        return video, audio

packages/ltx-core/src/ltx_core/tools.py

@@ -76,7 +76,13 @@ class LatentTools(Protocol):
         denoise_mask = torch.ones_like(latent_state.denoise_mask)[:, :num_tokens]
         positions = latent_state.positions[:, :, :num_tokens]
 
-        return LatentState(latent=latent, denoise_mask=denoise_mask, positions=positions, clean_latent=clean_latent)
+        return LatentState(
+            latent=latent,
+            denoise_mask=denoise_mask,
+            positions=positions,
+            clean_latent=clean_latent,
+            attention_mask=None,
+        )
 
 
 @dataclass(frozen=True)

packages/ltx-core/src/ltx_core/types.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import NamedTuple
 
 import torch
@@ -61,6 +61,10 @@ class VideoLatentShape(NamedTuple):
             width=shape[4],
         )
 
+    def token_count(self) -> int:
+        """Number of tokens after patchification with the default patch size of 1."""
+        return self.frames * self.height * self.width
+
     def mask_shape(self) -> "VideoLatentShape":
         return self._replace(channels=1)
 
@@ -105,6 +109,10 @@ class AudioLatentShape(NamedTuple):
     def to_torch_shape(self) -> torch.Size:
         return torch.Size([self.batch, self.channels, self.frames, self.mel_bins])
 
+    def token_count(self) -> int:
+        """Number of tokens after patchification."""
+        return self.frames
+
     def mask_shape(self) -> "AudioLatentShape":
         return self._replace(channels=1, mel_bins=1)
 
@@ -156,6 +164,22 @@ class AudioLatentShape(NamedTuple):
         )
 
 
+@dataclass(frozen=True)
+class Audio:
+    """
+    Container for decoded audio samples and metadata.
+    Attributes:
+        waveform: Audio waveform tensor.
+        sampling_rate: Sampling rate (Hz) of the waveform.
+    """
+
+    waveform: torch.Tensor
+    sampling_rate: int
+
+    def to(self, **kwargs: object) -> "Audio":
+        return replace(self, waveform=self.waveform.to(**kwargs))
+
+
 @dataclass(frozen=True)
 class LatentState:
     """
@@ -165,12 +189,15 @@ class LatentState:
         denoise_mask: Mask encoding the denoising strength for each token (1 = full denoising, 0 = no denoising).
         positions: Positional indices for each latent element, used for positional embeddings.
         clean_latent: Initial state of the latent before denoising, may include conditioning latents.
+        attention_mask: Optional 2D self-attention mask of shape (B, T, T). Values in [0, 1] where 1 = full attention,
+            0 = no attention. None means full attention everywhere. Built incrementally by conditioning items.
     """
 
     latent: torch.Tensor
     denoise_mask: torch.Tensor
     positions: torch.Tensor
     clean_latent: torch.Tensor
+    attention_mask: torch.Tensor | None = None
 
     def clone(self) -> "LatentState":
         return LatentState(
@@ -178,4 +205,5 @@ class LatentState:
             denoise_mask=self.denoise_mask.clone(),
             positions=self.positions.clone(),
             clean_latent=self.clean_latent.clone(),
+            attention_mask=self.attention_mask.clone() if self.attention_mask is not None else None,
         )

packages/ltx-pipelines/README.md

@@ -12,7 +12,7 @@ LTX-2 Pipelines provides production-ready implementations that abstract away the
 
 **Key Features:**
 
-- 🎬 **Multiple Pipeline Types**: Text-to-video, image-to-video, video-to-video, and keyframe interpolation
+- 🎬 **Multiple Pipeline Types**: Text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake
 - ⚡ **Optimized Performance**: Support for FP8 transformers, gradient estimation, and memory optimization
 - 🎯 **Production Ready**: Two-stage pipelines for best quality output
 - 🔧 **LoRA Support**: Easy integration with trained LoRA adapters
@@ -23,7 +23,7 @@ LTX-2 Pipelines provides production-ready implementations that abstract away the
 
 ## 🚀 Quick Start
 
-`ltx-pipelines` provides ready-made inference pipelines for text-to-video, image-to-video, video-to-video, and keyframe interpolation. Built using building blocks from [`ltx-core`](../ltx-core/), these pipelines handle the complete inference flow including model loading, encoding, decoding, and file I/O.
+`ltx-pipelines` provides ready-made inference pipelines for text-to-video, image-to-video, video-to-video, audio-to-video, keyframe interpolation, and retake. Built using building blocks from [`ltx-core`](../ltx-core/), these pipelines handle the complete inference flow including model loading, encoding, decoding, and file I/O.
 
 ## 🔧 Installation
 
@@ -56,10 +56,13 @@ python -m ltx_pipelines.ti2vid_two_stages --help
 Available pipeline modules:
 
 - `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
+- `ltx_pipelines.ti2vid_two_stages_res2s` - Two-stage text/image-to-video (use 2 times less steps).
 - `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
 - `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.
+- `ltx_pipelines.a2vid_two_stage` - Audio-to-video generation conditioned on an input audio.
+- `ltx_pipelines.retake` - Regenerate a time region of an existing video.
 
 Use `--help` with any pipeline module to see all available options and parameters.
 
@@ -70,6 +73,12 @@ Use `--help` with any pipeline module to see all available options and parameter
 ### Quick Decision Tree
 
 ```text
+Do you have an existing video to modify?
+├─ YES → Use RetakePipeline (regenerate a specific time region)
+│
+Do you have an audio file to drive generation?
+├─ YES → Use A2VidPipelineTwoStage (audio-to-video)
+│
 Do you need to condition on existing images/videos?
 ├─ YES → Do you have reference videos for video-to-video?
 │  ├─ YES → Use ICLoraPipeline
@@ -85,17 +94,20 @@ Do you need to condition on existing images/videos?
       └─ YES → Use DistilledPipeline (with 8 predefined sigmas)
 ```
 
-> **Note:** [`TI2VidOneStagePipeline`](src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](src/ltx_pipelines/ti2vid_two_stages.py), [`ICLoraPipeline`](src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](src/ltx_pipelines/keyframe_interpolation.py), or [`DistilledPipeline`](src/ltx_pipelines/distilled.py)).
+> **Note:** [`TI2VidOneStagePipeline`](src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](src/ltx_pipelines/ti2vid_two_stages.py), [`TI2VidTwoStagesRes2sPipeline`](src/ltx_pipelines/ti2vid_two_stages_res2s.py), [`ICLoraPipeline`](src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](src/ltx_pipelines/keyframe_interpolation.py), [`A2VidPipelineTwoStage`](src/ltx_pipelines/a2vid_two_stage.py), or [`DistilledPipeline`](src/ltx_pipelines/distilled.py)). For editing existing videos, use [`RetakePipeline`](src/ltx_pipelines/retake.py).
 
 ### Features Comparison
 
 | Pipeline | Stages | [Multimodal Guidance](#%EF%B8%8F-multimodal-guidance) | Upsampling | Conditioning | Best For |
 | -------- | ------ | --- | ---------- | ------------- | -------- |
 | **TI2VidTwoStagesPipeline** | 2 | ✅ | ✅ | Image | **Production quality** (recommended) |
+| **TI2VidTwoStagesRes2sPipeline** | 2 | ✅ | ✅ | Image | Same as above, res_2s sampler (fewer steps) |
 | **TI2VidOneStagePipeline** | 1 | ✅ | ❌ | Image | Educational, prototyping |
 | **DistilledPipeline** | 2 | ❌ | ✅ | Image | Fastest inference (8 sigmas) |
 | **ICLoraPipeline** | 2 | ✅ | ✅ | Image + Video | Video-to-video transformations |
 | **KeyframeInterpolationPipeline** | 2 | ✅ | ✅ | Keyframes | Animation, interpolation |
+| **A2VidPipelineTwoStage** | 2 | ✅ | ✅ | Audio + Image | Audio-driven video generation |
+| **RetakePipeline** | 1 | ✅ | ❌ | Source Video | Regenerating a time region of a video |
 
 ---
 
@@ -113,7 +125,19 @@ Two-stage generation: Stage 1 generates low-resolution video with [multimodal gu
 
 ---
 
-### 2. TI2VidOneStagePipeline
+### 2. TI2VidTwoStagesRes2sPipeline
+
+**Best for:** Same two-stage text/image-to-video as TI2VidTwoStagesPipeline but with a different sampler and step count.
+
+**Source**: [`src/ltx_pipelines/ti2vid_two_stages_res2s.py`](src/ltx_pipelines/ti2vid_two_stages_res2s.py)
+
+Uses the **res_2s** second-order sampler instead of Euler. Same stage structure (stage 1 at target resolution with CFG, stage 2 upsampling with distilled LoRA) and image conditioning support. Typically allows fewer steps for comparable quality; trade-offs differ from the default Euler-based pipeline.
+
+**Use when:** You want the same two-stage workflow with fewer steps or prefer the res_2s sampling behavior.
+
+---
+
+### 3. TI2VidOneStagePipeline
 
 **Best for:** Educational purposes and quick prototyping.
 
@@ -127,7 +151,7 @@ Single-stage generation (no upsampling) with [multimodal guidance](#%EF%B8%8F-mu
 
 ---
 
-### 3. DistilledPipeline
+### 4. DistilledPipeline
 
 **Best for:** Fastest inference with good quality using a distilled model with predefined sigma schedule.
 
@@ -139,7 +163,7 @@ Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in st
 
 ---
 
-### 4. ICLoraPipeline
+### 5. ICLoraPipeline
 
 **Best for:** Video-to-video and image-to-video transformations using IC-LoRA.
 
@@ -147,11 +171,13 @@ Two-stage generation with 8 predefined sigmas (8 steps in stage 1, 4 steps in st
 
 Two-stage generation with IC-LoRA support. Can condition on reference videos (video-to-video) or images at specific frames. CFG guidance in stage 1, upsampling in stage 2. Requires IC-LoRA trained model.
 
+**Note:** ICLoraPipeline can only be used with a distilled model.
+
 **Use when:** Video-to-video transformations, image-to-video with strong control, or when you have reference videos to guide generation.
 
 ---
 
-### 5. KeyframeInterpolationPipeline
+### 6. KeyframeInterpolationPipeline
 
 **Best for:** Generating videos by interpolating between keyframe images.
 
@@ -163,6 +189,36 @@ Two-stage generation with keyframe interpolation. Uses guiding latents (additive
 
 ---
 
+### 7. A2VidPipelineTwoStage
+
+**Best for:** Generating video driven by an input audio.
+
+**Source**: [`src/ltx_pipelines/a2vid_two_stage.py`](src/ltx_pipelines/a2vid_two_stage.py)
+
+Two-stage audio-to-video generation. Stage 1 generates video at half resolution with audio conditioning (video-only denoising with the audio frozen), then Stage 2 upsamples by 2x and refines the video while keeping the audio fixed, using a distilled LoRA. The input audio is encoded via the audio VAE and used as the initial audio latent, but the original audio waveform is passed through and returned in the output to preserve fidelity. Supports image conditioning and prompt enhancement.
+
+**Extra CLI arguments:** `--audio-path` (required), `--audio-start-time`, `--audio-max-duration`.
+
+**Use when:** You have an audio clip and want to generate a matching video, audio-reactive video generation, or music visualization.
+
+---
+
+### 8. RetakePipeline
+
+**Best for:** Regenerating a specific time region of an existing video while keeping the rest unchanged.
+
+**Source**: [`src/ltx_pipelines/retake.py`](src/ltx_pipelines/retake.py)
+
+Single-stage generation that encodes the source video and audio into latents, applies a temporal region mask to mark `[start_time, end_time]` for regeneration, and denoises only the masked region from a text prompt. Content outside the time window is preserved. Supports independent control over video and audio regeneration (`regenerate_video`, `regenerate_audio` flags), and can use either the full model with CFG guidance or the distilled model with a fixed sigma schedule.
+
+**Extra CLI arguments:** `--video-path` (required), `--start-time` (required), `--end-time` (required).
+
+**Constraints:** Source video frame count must satisfy the 8k+1 format (e.g. 97, 193) and resolution must be multiples of 32.
+
+**Use when:** You want to re-do a specific section of a generated video (e.g. fix a bad segment), selectively regenerate audio or video in a time window, or iterate on part of a result without re-generating the entire clip.
+
+---
+
 ## 🎨 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.
@@ -250,28 +306,43 @@ audio_guider_params = MultiModalGuiderParams(
 
 ### Memory Optimization
 
-**FP8 Transformer (Lower Memory Footprint):**
+**FP8 Quantization (Lower Memory Footprint):**
 
-For smaller GPU memory footprint, use the `enable-fp8` flag and use the `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` environment variable.
+For smaller GPU memory footprint, use the `--quantization` flag and set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.
+
+Two quantization policies are available:
+
+| Policy | CLI Flag | Description |
+| ------ | -------- | ----------- |
+| **FP8 Cast** | `--quantization fp8-cast` | Downcasts transformer linear weights to FP8 during loading; upcasts on the fly during inference. No extra dependencies. |
+| **FP8 Scaled MM** | `--quantization fp8-scaled-mm` | Uses FP8 scaled matrix multiplication via TensorRT-LLM (`tensorrt_llm` must be installed). Best performance on Hopper GPUs. |
 
 **CLI:**
 
 ```bash
-PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_one_stage --enable-fp8 --checkpoint-path=...
+# FP8 Cast (works on any GPU with FP8 support)
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
+    --quantization fp8-cast --checkpoint-path=...
+
+# FP8 Scaled MM (requires tensorrt_llm, best on Hopper GPUs)
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -m ltx_pipelines.ti2vid_two_stages \
+    --quantization fp8-scaled-mm --checkpoint-path=...
 ```
 
 **Programmatically:**
 
-When authoring custom scripts, pass the `fp8transformer` flag to pipeline classes or construct your own by analogy:
+When authoring custom scripts, pass a `QuantizationPolicy` to pipeline classes:
 
 ```python
+from ltx_core.quantization import QuantizationPolicy
+
 pipeline = TI2VidTwoStagesPipeline(
     checkpoint_path=ltx_model_path,
     distilled_lora=distilled_lora,
     spatial_upsampler_path=upsampler_path,
     gemma_root=gemma_root_path,
     loras=[],
-    fp8transformer=True,
+    quantization=QuantizationPolicy.fp8_cast(),  # or QuantizationPolicy.fp8_scaled_mm()
 )
 pipeline(...)
 ```
@@ -299,7 +370,7 @@ By default, pipelines clean GPU memory (especially transformer weights) between
 Instead of the standard Euler denoising loop, you can use gradient estimation for fewer steps (~20-30 instead of 40):
 
 ```python
-from ltx_pipelines.utils.helpers import gradient_estimating_euler_denoising_loop
+from ltx_pipelines.utils import gradient_estimating_euler_denoising_loop
 
 # Use gradient estimation denoising loop
 def denoising_loop(sigmas, video_state, audio_state, stepper):
@@ -379,7 +450,7 @@ pipeline(
     num_inference_steps=40,
     video_guider_params=video_guider_params,
     audio_guider_params=audio_guider_params,
-    images=[("input_image.jpg", 0, 1.0)],  # Image at frame 0, strength 1.0
+    images=[ImageConditioningInput("input_image.jpg", 0, 1.0, 33)],  # Image at frame 0, strength 1.0, CRF 33
 )
 ```
 

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

@@ -6,21 +6,26 @@ This package provides ready-to-use pipelines for video generation:
 - DistilledPipeline: Fast distilled two-stage generation
 - ICLoraPipeline: Image/video conditioning with distilled LoRA
 - KeyframeInterpolationPipeline: Keyframe-based video interpolation
+- RetakePipeline: Regenerate a time region (retake) of an existing video
 - ModelLedger: Central coordinator for loading and building models
 For more detailed components and utilities, import from specific submodules
 like `ltx_pipelines.utils.media_io` or `ltx_pipelines.utils.constants`.
 """
 
+from ltx_pipelines.a2vid_two_stage import A2VidPipelineTwoStage
 from ltx_pipelines.distilled import DistilledPipeline
 from ltx_pipelines.ic_lora import ICLoraPipeline
 from ltx_pipelines.keyframe_interpolation import KeyframeInterpolationPipeline
+from ltx_pipelines.retake import RetakePipeline
 from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
 from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline
 
 __all__ = [
+    "A2VidPipelineTwoStage",
     "DistilledPipeline",
     "ICLoraPipeline",
     "KeyframeInterpolationPipeline",
+    "RetakePipeline",
     "TI2VidOneStagePipeline",
     "TI2VidTwoStagesPipeline",
 ]

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

@@ -0,0 +1,323 @@
+import logging
+from collections.abc import Iterator
+
+import torch
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+from ltx_core.model.upsampler import upsample_video
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoPixelShape
+from ltx_pipelines.utils import ModelLedger
+from ltx_pipelines.utils.args import default_2_stage_arg_parser
+from ltx_pipelines.utils.constants import (
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+)
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    cleanup_memory,
+    denoise_video_only,
+    generate_enhanced_prompt,
+    get_device,
+    image_conditionings_by_replacing_latent,
+    multi_modal_guider_denoising_func,
+    simple_denoising_func,
+)
+from ltx_pipelines.utils.media_io import decode_audio_from_file, encode_video
+from ltx_pipelines.utils.samplers import euler_denoising_loop
+from ltx_pipelines.utils.types import PipelineComponents
+
+device = get_device()
+
+
+class A2VidPipelineTwoStage:
+    """
+    Two-stage audio to video generation pipeline.
+    Stage 1 generates video at half the target resolution with audio conditioning
+    (video-only denoising, audio frozen), then Stage 2 upsamples by 2x and refines
+    both video and audio using a distilled LoRA for higher quality output.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        distilled_lora: list[LoraPathStrengthAndSDOps],
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device = device,
+        quantization: QuantizationPolicy | None = None,
+    ):
+        self.device = device
+        self.dtype = torch.bfloat16
+
+        self.stage_1_model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            gemma_root_path=gemma_root,
+            spatial_upsampler_path=spatial_upsampler_path,
+            loras=loras,
+            quantization=quantization,
+        )
+
+        self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
+            loras=distilled_lora,
+        )
+
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        video_guider_params: MultiModalGuiderParams,
+        images: list[tuple[str, int, float]],
+        audio_path: str,
+        audio_start_time: float = 0.0,
+        audio_max_duration: float | None = None,
+        tiling_config: TilingConfig | None = None,
+        enhance_prompt: bool = False,
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
+        assert_resolution(height=height, width=width, is_two_stage=True)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+        stepper = EulerDiffusionStep()
+        dtype = torch.bfloat16
+
+        text_encoder = self.stage_1_model_ledger.text_encoder()
+        if enhance_prompt:
+            prompt = generate_enhanced_prompt(text_encoder, prompt, images[0][0] if len(images) > 0 else None)
+        context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
+        v_context_p, a_context_p = context_p
+        v_context_n, _ = context_n
+
+        torch.cuda.synchronize()
+        del text_encoder
+        cleanup_memory()
+
+        # Encode audio.
+        decoded_audio = decode_audio_from_file(audio_path, self.device, audio_start_time, audio_max_duration)
+        encoded_audio_latent = vae_encode_audio(decoded_audio, self.stage_1_model_ledger.audio_encoder())
+        audio_shape = AudioLatentShape.from_duration(batch=1, duration=num_frames / frame_rate, channels=8, mel_bins=16)
+        encoded_audio_latent = encoded_audio_latent[:, :, : audio_shape.frames]
+
+        cleanup_memory()
+        # Stage 1: Initial low resolution video generation with audio conditioning.
+        transformer = self.stage_1_model_ledger.transformer()
+        sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
+
+        def first_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=multi_modal_guider_denoising_func(
+                    video_guider=MultiModalGuider(
+                        params=video_guider_params,
+                        negative_context=v_context_n,
+                    ),
+                    audio_guider=MultiModalGuider(
+                        params=MultiModalGuiderParams(),
+                    ),
+                    v_context=v_context_p,
+                    a_context=a_context_p,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+
+        video_encoder = self.stage_1_model_ledger.video_encoder()
+        stage_1_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_1_output_shape.height,
+            width=stage_1_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        video_state = denoise_video_only(
+            output_shape=stage_1_output_shape,
+            conditionings=stage_1_conditionings,
+            noiser=noiser,
+            sigmas=sigmas,
+            stepper=stepper,
+            denoising_loop_fn=first_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            initial_audio_latent=encoded_audio_latent,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        cleanup_memory()
+
+        # Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
+        upscaled_video_latent = upsample_video(
+            latent=video_state.latent[:1],
+            video_encoder=video_encoder,
+            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
+        )
+
+        torch.cuda.synchronize()
+        cleanup_memory()
+
+        transformer = self.stage_2_model_ledger.transformer()
+        distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
+
+        def second_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=simple_denoising_func(
+                    video_context=v_context_p,
+                    audio_context=a_context_p,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_2_output_shape.height,
+            width=stage_2_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_state = denoise_video_only(
+            output_shape=stage_2_output_shape,
+            conditionings=stage_2_conditionings,
+            noiser=noiser,
+            sigmas=distilled_sigmas,
+            stepper=stepper,
+            denoising_loop_fn=second_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            noise_scale=distilled_sigmas[0],
+            initial_video_latent=upscaled_video_latent,
+            initial_audio_latent=encoded_audio_latent,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        del video_encoder
+        cleanup_memory()
+
+        decoded_video = vae_decode_video(
+            video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
+        )
+
+        # Return the original input audio instead of VAE-decoded audio to preserve fidelity.
+        # decode_audio_from_file already returns normalised [-1, 1] float values.
+        original_audio = Audio(waveform=decoded_audio.waveform.squeeze(0), sampling_rate=decoded_audio.sampling_rate)
+
+        return decoded_video, original_audio
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    parser = default_2_stage_arg_parser()
+    parser.add_argument(
+        "--audio-path",
+        type=str,
+        required=True,
+        help="Path to the audio file to condition the video generation.",
+    )
+    parser.add_argument(
+        "--audio-start-time",
+        type=float,
+        default=0.0,
+        help="Start time in seconds to read audio from (default: 0.0).",
+    )
+    parser.add_argument(
+        "--audio-max-duration",
+        type=float,
+        default=None,
+        help="Maximum audio duration in seconds. Defaults to video duration (num_frames / frame_rate).",
+    )
+    args = parser.parse_args()
+    pipeline = A2VidPipelineTwoStage(
+        checkpoint_path=args.checkpoint_path,
+        distilled_lora=args.distilled_lora,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        loras=args.lora,
+        quantization=args.quantization,
+    )
+    tiling_config = TilingConfig.default()
+    video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        video_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.video_cfg_guidance_scale,
+            stg_scale=args.video_stg_guidance_scale,
+            rescale_scale=args.video_rescale_scale,
+            modality_scale=args.a2v_guidance_scale,
+            skip_step=args.video_skip_step,
+            stg_blocks=args.video_stg_blocks,
+        ),
+        images=args.images,
+        tiling_config=tiling_config,
+        enhance_prompt=args.enhance_prompt,
+        audio_path=args.audio_path,
+        audio_start_time=args.audio_start_time,
+        audio_max_duration=args.audio_max_duration
+        if args.audio_max_duration is not None
+        else args.num_frames / args.frame_rate,
+    )
+
+    encode_video(
+        video=video,
+        fps=args.frame_rate,
+        audio=audio,
+        output_path=args.output_path,
+        video_chunks_number=video_chunks_number,
+    )
+
+
+if __name__ == "__main__":
+    main()

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

@@ -13,19 +13,22 @@ from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
-from ltx_core.types import LatentState, VideoPixelShape
-from ltx_pipelines.utils import ModelLedger
-from ltx_pipelines.utils.args import default_2_stage_distilled_arg_parser
+from ltx_core.types import Audio, LatentState, VideoPixelShape
+from ltx_pipelines.utils import ModelLedger, euler_denoising_loop
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    default_2_stage_distilled_arg_parser,
+    detect_checkpoint_path,
+)
 from ltx_pipelines.utils.constants import (
-    AUDIO_SAMPLE_RATE,
     DISTILLED_SIGMA_VALUES,
     STAGE_2_DISTILLED_SIGMA_VALUES,
+    detect_params,
 )
 from ltx_pipelines.utils.helpers import (
     assert_resolution,
     cleanup_memory,
     denoise_audio_video,
-    euler_denoising_loop,
     generate_enhanced_prompt,
     get_device,
     image_conditionings_by_replacing_latent,
@@ -40,13 +43,13 @@ device = get_device()
 class DistilledPipeline:
     """
     Two-stage distilled video generation pipeline.
-    Stage 1 generates video at the target resolution, then Stage 2 upsamples
+    Stage 1 generates video at half of the target resolution, then Stage 2 upsamples
     by 2x and refines with additional denoising steps for higher quality output.
     """
 
     def __init__(
         self,
-        checkpoint_path: str,
+        distilled_checkpoint_path: str,
         gemma_root: str,
         spatial_upsampler_path: str,
         loras: list[LoraPathStrengthAndSDOps],
@@ -59,7 +62,7 @@ class DistilledPipeline:
         self.model_ledger = ModelLedger(
             dtype=self.dtype,
             device=device,
-            checkpoint_path=checkpoint_path,
+            checkpoint_path=distilled_checkpoint_path,
             spatial_upsampler_path=spatial_upsampler_path,
             gemma_root_path=gemma_root,
             loras=loras,
@@ -79,10 +82,10 @@ class DistilledPipeline:
         width: int,
         num_frames: int,
         frame_rate: float,
-        images: list[tuple[str, int, float]],
+        images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -198,10 +201,12 @@ class DistilledPipeline:
 @torch.inference_mode()
 def main() -> None:
     logging.getLogger().setLevel(logging.INFO)
-    parser = default_2_stage_distilled_arg_parser()
+    checkpoint_path = detect_checkpoint_path(distilled=True)
+    params = detect_params(checkpoint_path)
+    parser = default_2_stage_distilled_arg_parser(params=params)
     args = parser.parse_args()
     pipeline = DistilledPipeline(
-        checkpoint_path=args.checkpoint_path,
+        distilled_checkpoint_path=args.distilled_checkpoint_path,
         spatial_upsampler_path=args.spatial_upsampler_path,
         gemma_root=args.gemma_root,
         loras=args.lora,
@@ -225,7 +230,6 @@ def main() -> None:
         video=video,
         fps=args.frame_rate,
         audio=audio,
-        audio_sample_rate=AUDIO_SAMPLE_RATE,
         output_path=args.output_path,
         video_chunks_number=video_chunks_number,
     )

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

@@ -2,12 +2,17 @@ import logging
 from collections.abc import Iterator
 
 import torch
+from einops import rearrange
 from safetensors import safe_open
 
 from ltx_core.components.diffusion_steps import EulerDiffusionStep
 from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.protocols import DiffusionStepProtocol
-from ltx_core.conditioning import ConditioningItem, VideoConditionByReferenceLatent
+from ltx_core.conditioning import (
+    ConditioningItem,
+    ConditioningItemAttentionStrengthWrapper,
+    VideoConditionByReferenceLatent,
+)
 from ltx_core.loader import LoraPathStrengthAndSDOps
 from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
 from ltx_core.model.upsampler import upsample_video
@@ -15,15 +20,9 @@ from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunk
 from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
-from ltx_core.types import LatentState, VideoPixelShape
-from ltx_pipelines.utils import ModelLedger
-from ltx_pipelines.utils.args import VideoConditioningAction, default_2_stage_distilled_arg_parser
-from ltx_pipelines.utils.constants import (
-    AUDIO_SAMPLE_RATE,
-    DISTILLED_SIGMA_VALUES,
-    STAGE_2_DISTILLED_SIGMA_VALUES,
-)
-from ltx_pipelines.utils.helpers import (
+from ltx_core.types import Audio, LatentState, VideoLatentShape, VideoPixelShape
+from ltx_pipelines.utils import (
+    ModelLedger,
     assert_resolution,
     cleanup_memory,
     denoise_audio_video,
@@ -33,6 +32,18 @@ from ltx_pipelines.utils.helpers import (
     image_conditionings_by_replacing_latent,
     simple_denoising_func,
 )
+from ltx_pipelines.utils.args import (
+    ImageConditioningInput,
+    VideoConditioningAction,
+    VideoMaskConditioningAction,
+    default_2_stage_distilled_arg_parser,
+    detect_checkpoint_path,
+)
+from ltx_pipelines.utils.constants import (
+    DISTILLED_SIGMA_VALUES,
+    STAGE_2_DISTILLED_SIGMA_VALUES,
+    detect_params,
+)
 from ltx_pipelines.utils.media_io import encode_video, load_video_conditioning
 from ltx_pipelines.utils.types import PipelineComponents
 
@@ -45,13 +56,14 @@ class ICLoraPipeline:
     Allows conditioning the generated video on control signals such as depth maps,
     human pose, or image edges via the video_conditioning parameter.
     The specific IC-LoRA model should be provided via the loras parameter.
-    Stage 1 generates video at the target resolution, then Stage 2 upsamples
+    Stage 1 generates video at half of the target resolution, then Stage 2 upsamples
     by 2x and refines with additional denoising steps for higher quality output.
+    Both stages use distilled models for efficiency.
     """
 
     def __init__(
         self,
-        checkpoint_path: str,
+        distilled_checkpoint_path: str,
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
@@ -62,7 +74,7 @@ class ICLoraPipeline:
         self.stage_1_model_ledger = ModelLedger(
             dtype=self.dtype,
             device=device,
-            checkpoint_path=checkpoint_path,
+            checkpoint_path=distilled_checkpoint_path,
             spatial_upsampler_path=spatial_upsampler_path,
             gemma_root_path=gemma_root,
             loras=loras,
@@ -71,7 +83,7 @@ class ICLoraPipeline:
         self.stage_2_model_ledger = ModelLedger(
             dtype=self.dtype,
             device=device,
-            checkpoint_path=checkpoint_path,
+            checkpoint_path=distilled_checkpoint_path,
             spatial_upsampler_path=spatial_upsampler_path,
             gemma_root_path=gemma_root,
             loras=[],
@@ -98,8 +110,7 @@ class ICLoraPipeline:
                     )
                 self.reference_downscale_factor = scale
 
-    @torch.inference_mode()
-    def __call__(
+    def __call__(  # noqa: PLR0913
         self,
         prompt: str,
         seed: int,
@@ -107,12 +118,51 @@ class ICLoraPipeline:
         width: int,
         num_frames: int,
         frame_rate: float,
-        images: list[tuple[str, int, float]],
+        images: list[ImageConditioningInput],
         video_conditioning: list[tuple[str, float]],
         enhance_prompt: bool = False,
         tiling_config: TilingConfig | None = None,
-    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+        conditioning_attention_strength: float = 1.0,
+        skip_stage_2: bool = False,
+        conditioning_attention_mask: torch.Tensor | None = None,
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
+        """
+        Generate video with IC-LoRA conditioning.
+        Args:
+            prompt: Text prompt for video generation.
+            seed: Random seed for reproducibility.
+            height: Output video height in pixels (must be divisible by 64).
+            width: Output video width in pixels (must be divisible by 64).
+            num_frames: Number of frames to generate.
+            frame_rate: Output video frame rate.
+            images: List of (path, frame_idx, strength) tuples for image conditioning.
+            video_conditioning: List of (path, strength) tuples for IC-LoRA video conditioning.
+            enhance_prompt: Whether to enhance the prompt using the text encoder.
+            tiling_config: Optional tiling configuration for VAE decoding.
+            conditioning_attention_strength: Scale factor for IC-LoRA conditioning attention.
+                Controls how strongly the conditioning video influences the output.
+                0.0 = ignore conditioning, 1.0 = full conditioning influence. Default 1.0.
+                When conditioning_attention_mask is provided, the mask is multiplied by
+                this strength before being passed to the conditioning items.
+            skip_stage_2: If True, skip Stage 2 upsampling and refinement. Output will be
+                at half resolution (height//2, width//2). Default is False.
+            conditioning_attention_mask: Optional pixel-space attention mask with the same
+                spatial-temporal dimensions as the input reference video. Shape should be
+                (B, 1, F, H, W) or (1, 1, F, H, W) where F, H, W match the reference
+                video's pixel dimensions. Values in [0, 1].
+                The mask is downsampled to latent space using VAE scale factors (with
+                causal temporal handling for the first frame), then multiplied by
+                conditioning_attention_strength.
+                When None (default): scalar conditioning_attention_strength is used
+                directly.
+        Returns:
+            Tuple of (video_iterator, audio_tensor).
+        """
         assert_resolution(height=height, width=width, is_two_stage=True)
+        if not (0.0 <= conditioning_attention_strength <= 1.0):
+            raise ValueError(
+                f"conditioning_attention_strength must be in [0.0, 1.0], got {conditioning_attention_strength}"
+            )
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
@@ -158,6 +208,7 @@ class ICLoraPipeline:
             height=height // 2,
             fps=frame_rate,
         )
+
         stage_1_conditionings = self._create_conditionings(
             images=images,
             video_conditioning=video_conditioning,
@@ -165,7 +216,10 @@ class ICLoraPipeline:
             width=stage_1_output_shape.width,
             video_encoder=video_encoder,
             num_frames=num_frames,
+            conditioning_attention_strength=conditioning_attention_strength,
+            conditioning_attention_mask=conditioning_attention_mask,
         )
+
         video_state, audio_state = denoise_audio_video(
             output_shape=stage_1_output_shape,
             conditionings=stage_1_conditionings,
@@ -182,6 +236,19 @@ class ICLoraPipeline:
         del transformer
         cleanup_memory()
 
+        if skip_stage_2:
+            # Skip Stage 2: Decode directly from Stage 1 output at half resolution
+            logging.info("[IC-LoRA] Skipping Stage 2 (--skip-stage-2 enabled)")
+            decoded_video = vae_decode_video(
+                video_state.latent, self.stage_1_model_ledger.video_decoder(), tiling_config, generator
+            )
+            decoded_audio = vae_decode_audio(
+                audio_state.latent, self.stage_1_model_ledger.audio_decoder(), self.stage_1_model_ledger.vocoder()
+            )
+            del video_encoder
+            cleanup_memory()
+            return decoded_video, decoded_audio
+
         # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
         upscaled_video_latent = upsample_video(
             latent=video_state.latent[:1],
@@ -250,13 +317,29 @@ class ICLoraPipeline:
 
     def _create_conditionings(
         self,
-        images: list[tuple[str, int, float]],
+        images: list[ImageConditioningInput],
         video_conditioning: list[tuple[str, float]],
         height: int,
         width: int,
         num_frames: int,
         video_encoder: VideoEncoder,
+        conditioning_attention_strength: float = 1.0,
+        conditioning_attention_mask: torch.Tensor | None = None,
     ) -> list[ConditioningItem]:
+        """
+        Create conditioning items for video generation.
+        Args:
+            conditioning_attention_strength: Scalar attention weight in [0, 1].
+                If conditioning_attention_mask is also provided, the downsampled mask
+                is multiplied by this strength. Otherwise this scalar is passed
+                directly as the attention mask.
+            conditioning_attention_mask: Optional pixel-space attention mask with shape
+                (B, 1, F_pixel, H_pixel, W_pixel) matching the reference video's
+                pixel dimensions. Downsampled to latent space with causal temporal
+                handling, then multiplied by conditioning_attention_strength.
+        Returns:
+            List of conditioning items. IC-LoRA conditionings are appended last.
+        """
         conditionings = image_conditionings_by_replacing_latent(
             images=images,
             height=height,
@@ -287,21 +370,96 @@ class ICLoraPipeline:
                 device=self.device,
             )
             encoded_video = video_encoder(video)
-            conditionings.append(
-                VideoConditionByReferenceLatent(
-                    latent=encoded_video,
-                    downscale_factor=scale,
-                    strength=strength,
+            reference_video_shape = VideoLatentShape.from_torch_shape(encoded_video.shape)
+
+            # Build attention_mask for ConditioningItemAttentionStrengthWrapper
+            if conditioning_attention_mask is not None:
+                # Downsample pixel-space mask to latent space, then scale by strength
+                latent_mask = self._downsample_mask_to_latent(
+                    mask=conditioning_attention_mask,
+                    target_latent_shape=reference_video_shape,
                 )
+                attn_mask = latent_mask * conditioning_attention_strength
+            elif conditioning_attention_strength < 1.0:
+                # Use scalar strength only
+                attn_mask = conditioning_attention_strength
+            else:
+                attn_mask = None
+
+            cond = VideoConditionByReferenceLatent(
+                latent=encoded_video,
+                downscale_factor=scale,
+                strength=strength,
             )
+            if attn_mask is not None:
+                cond = ConditioningItemAttentionStrengthWrapper(cond, attention_mask=attn_mask)
+            conditionings.append(cond)
+
+        if video_conditioning:
+            logging.info(f"[IC-LoRA] Added {len(video_conditioning)} video conditioning(s)")
 
         return conditionings
 
+    @staticmethod
+    def _downsample_mask_to_latent(
+        mask: torch.Tensor,
+        target_latent_shape: VideoLatentShape,
+    ) -> torch.Tensor:
+        """
+        Downsample a pixel-space mask to latent space using VAE scale factors.
+        Handles causal temporal downsampling: the first frame is kept separately
+        (temporal scale factor = 1 for the first frame), while the remaining
+        frames are downsampled by the VAE's temporal scale factor.
+        Args:
+            mask: Pixel-space mask of shape (B, 1, F_pixel, H_pixel, W_pixel).
+                Values in [0, 1].
+            target_latent_shape: Expected latent shape after VAE encoding.
+                Used to determine the target (F_latent, H_latent, W_latent).
+        Returns:
+            Flattened latent-space mask of shape (B, F_lat * H_lat * W_lat),
+            matching the patchifier's token ordering (f, h, w).
+        """
+        b = mask.shape[0]
+        f_lat = target_latent_shape.frames
+        h_lat = target_latent_shape.height
+        w_lat = target_latent_shape.width
+
+        # Step 1: Spatial downsampling (area interpolation per frame)
+        f_pix = mask.shape[2]
+        spatial_down = torch.nn.functional.interpolate(
+            rearrange(mask, "b 1 f h w -> (b f) 1 h w"),
+            size=(h_lat, w_lat),
+            mode="area",
+        )
+        spatial_down = rearrange(spatial_down, "(b f) 1 h w -> b 1 f h w", b=b)
+
+        # Step 2: Causal temporal downsampling
+        # First frame: kept as-is (causal VAE encodes first frame independently)
+        first_frame = spatial_down[:, :, :1, :, :]  # (B, 1, 1, H_lat, W_lat)
+
+        if f_pix > 1 and f_lat > 1:
+            # Remaining frames: downsample by temporal factor via group-mean
+            t = (f_pix - 1) // (f_lat - 1)  # temporal downscale factor
+            assert (f_pix - 1) % (f_lat - 1) == 0, (
+                f"Pixel frames ({f_pix}) not compatible with latent frames ({f_lat}): "
+                f"(f_pix - 1) must be divisible by (f_lat - 1)"
+            )
+            rest = rearrange(spatial_down[:, :, 1:, :, :], "b 1 (f t) h w -> b 1 f t h w", t=t)
+            rest = rest.mean(dim=3)  # (B, 1, F_lat-1, H_lat, W_lat)
+            latent_mask = torch.cat([first_frame, rest], dim=2)  # (B, 1, F_lat, H_lat, W_lat)
+        else:
+            latent_mask = first_frame
+
+        # Flatten to (B, F_lat * H_lat * W_lat) matching patchifier token order (f, h, w)
+        return rearrange(latent_mask, "b 1 f h w -> b (f h w)")
+
 
 @torch.inference_mode()
 def main() -> None:
     logging.getLogger().setLevel(logging.INFO)
-    parser = default_2_stage_distilled_arg_parser()
+    checkpoint_path = detect_checkpoint_path(distilled=True)
+    params = detect_params(checkpoint_path)
+    parser = default_2_stage_distilled_arg_parser(params=params)
     parser.add_argument(
         "--video-conditioning",
         action=VideoConditioningAction,
@@ -309,9 +467,47 @@ def main() -> None:
         metavar=("PATH", "STRENGTH"),
         required=True,
     )
+    parser.add_argument(
+        "--conditioning-attention-mask",
+        action=VideoMaskConditioningAction,
+        nargs=2,
+        metavar=("MASK_PATH", "STRENGTH"),
+        default=None,
+        help=(
+            "Optional spatial attention mask: path to a grayscale mask video and "
+            "attention strength. The mask video pixel values in [0,1] control "
+            "per-region conditioning attention strength. The strength scalar is "
+            "multiplied with the spatial mask. "
+            "0.0 = ignore IC-LoRA conditioning, 1.0 = full conditioning influence. "
+            "When not provided, full conditioning strength (1.0) is used. "
+            "Example: --conditioning-attention-mask path/to/mask.mp4 0.5"
+        ),
+    )
+    parser.add_argument(
+        "--skip-stage-2",
+        action="store_true",
+        help=(
+            "Skip Stage 2 upsampling and refinement. Output will be at half resolution "
+            "(height//2, width//2). Useful for faster iteration or when GPU memory is limited."
+        ),
+    )
     args = parser.parse_args()
+
+    # Load mask video if provided via --conditioning-attention-mask
+    conditioning_attention_mask = None
+    conditioning_attention_strength = 1.0
+    if args.conditioning_attention_mask is not None:
+        mask_path, mask_strength = args.conditioning_attention_mask
+        conditioning_attention_strength = mask_strength
+        conditioning_attention_mask = _load_mask_video(
+            mask_path=mask_path,
+            height=args.height // 2,  # Stage 1 operates at half resolution
+            width=args.width // 2,
+            num_frames=args.num_frames,
+        )
+
     pipeline = ICLoraPipeline(
-        checkpoint_path=args.checkpoint_path,
+        distilled_checkpoint_path=args.distilled_checkpoint_path,
         spatial_upsampler_path=args.spatial_upsampler_path,
         gemma_root=args.gemma_root,
         loras=args.lora,
@@ -329,18 +525,53 @@ def main() -> None:
         images=args.images,
         video_conditioning=args.video_conditioning,
         tiling_config=tiling_config,
+        conditioning_attention_strength=conditioning_attention_strength,
+        skip_stage_2=args.skip_stage_2,
+        conditioning_attention_mask=conditioning_attention_mask,
     )
 
     encode_video(
         video=video,
         fps=args.frame_rate,
         audio=audio,
-        audio_sample_rate=AUDIO_SAMPLE_RATE,
         output_path=args.output_path,
         video_chunks_number=video_chunks_number,
     )
 
 
+def _load_mask_video(
+    mask_path: str,
+    height: int,
+    width: int,
+    num_frames: int,
+) -> torch.Tensor:
+    """Load a mask video and return a pixel-space tensor of shape (1, 1, F, H, W).
+    The mask video is loaded, resized to (height, width), converted to
+    grayscale, and normalised to [0, 1].
+    Args:
+        mask_path: Path to the mask video file.
+        height: Target height in pixels.
+        width: Target width in pixels.
+        num_frames: Maximum number of frames to load.
+    Returns:
+        Tensor of shape ``(1, 1, F, H, W)`` with values in ``[0, 1]``.
+    """
+    mask_video = load_video_conditioning(
+        video_path=mask_path,
+        height=height,
+        width=width,
+        frame_cap=num_frames,
+        dtype=torch.bfloat16,
+        device=device,
+    )
+    # mask_video shape: (1, C, F, H, W) — take mean over channels for grayscale
+    mask = mask_video.mean(dim=1, keepdim=True)  # (1, 1, F, H, W)
+    # Normalise to [0, 1] — load_video_conditioning applies normalize_latent,
+    # so undo that: values are in [-1, 1], remap to [0, 1]
+    mask = (mask + 1.0) / 2.0
+    return mask.clamp(0.0, 1.0)
+
+
 def _read_lora_reference_downscale_factor(lora_path: str) -> int:
     """Read reference_downscale_factor from LoRA safetensors metadata.
     Some IC-LoRA models are trained with reference videos at lower resolution than

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

@@ -4,7 +4,11 @@ from collections.abc import Iterator
 import torch
 
 from ltx_core.components.diffusion_steps import EulerDiffusionStep
-from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+from ltx_core.components.guiders import (
+    MultiModalGuiderFactory,
+    MultiModalGuiderParams,
+    create_multimodal_guider_factory,
+)
 from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
@@ -15,25 +19,22 @@ from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
-from ltx_core.types import LatentState, VideoPixelShape
+from ltx_core.types import Audio, LatentState, VideoPixelShape
 from ltx_pipelines.utils import ModelLedger
-from ltx_pipelines.utils.args import default_2_stage_arg_parser
-from ltx_pipelines.utils.constants import (
-    AUDIO_SAMPLE_RATE,
-    STAGE_2_DISTILLED_SIGMA_VALUES,
-)
+from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
 from ltx_pipelines.utils.helpers import (
     assert_resolution,
     cleanup_memory,
     denoise_audio_video,
-    euler_denoising_loop,
     generate_enhanced_prompt,
     get_device,
     image_conditionings_by_adding_guiding_latent,
-    multi_modal_guider_denoising_func,
+    multi_modal_guider_factory_denoising_func,
     simple_denoising_func,
 )
 from ltx_pipelines.utils.media_io import encode_video
+from ltx_pipelines.utils.samplers import euler_denoising_loop
 from ltx_pipelines.utils.types import PipelineComponents
 
 device = get_device()
@@ -43,8 +44,10 @@ class KeyframeInterpolationPipeline:
     """
     Keyframe-based Two-stage video interpolation pipeline.
     Interpolates between keyframes to generate a video with smoother transitions.
-    Stage 1 generates video at the target resolution, then Stage 2 upsamples
+    Stage 1 generates video at half of the target resolution, then Stage 2 upsamples
     by 2x and refines with additional denoising steps for higher quality output.
+    Stage 1 uses full model while Stage 2 uses distilled LORA for efficiency,
+    as the upsampled video already has good quality and just needs refinement.
     """
 
     def __init__(
@@ -76,7 +79,6 @@ class KeyframeInterpolationPipeline:
             device=device,
         )
 
-    @torch.inference_mode()
     def __call__(  # noqa: PLR0913
         self,
         prompt: str,
@@ -87,12 +89,12 @@ class KeyframeInterpolationPipeline:
         num_frames: int,
         frame_rate: float,
         num_inference_steps: int,
-        video_guider_params: MultiModalGuiderParams,
-        audio_guider_params: MultiModalGuiderParams,
-        images: list[tuple[str, int, float]],
+        video_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -126,12 +128,12 @@ class KeyframeInterpolationPipeline:
                 video_state=video_state,
                 audio_state=audio_state,
                 stepper=stepper,
-                denoise_fn=multi_modal_guider_denoising_func(
-                    video_guider=MultiModalGuider(
+                denoise_fn=multi_modal_guider_factory_denoising_func(
+                    video_guider_factory=create_multimodal_guider_factory(
                         params=video_guider_params,
                         negative_context=v_context_n,
                     ),
-                    audio_guider=MultiModalGuider(
+                    audio_guider_factory=create_multimodal_guider_factory(
                         params=audio_guider_params,
                         negative_context=a_context_n,
                     ),
@@ -241,7 +243,9 @@ class KeyframeInterpolationPipeline:
 @torch.inference_mode()
 def main() -> None:
     logging.getLogger().setLevel(logging.INFO)
-    parser = default_2_stage_arg_parser()
+    checkpoint_path = detect_checkpoint_path()
+    params = detect_params(checkpoint_path)
+    parser = default_2_stage_arg_parser(params=params)
     args = parser.parse_args()
     pipeline = KeyframeInterpolationPipeline(
         checkpoint_path=args.checkpoint_path,
@@ -286,7 +290,6 @@ def main() -> None:
         video=video,
         fps=args.frame_rate,
         audio=audio,
-        audio_sample_rate=AUDIO_SAMPLE_RATE,
         output_path=args.output_path,
         video_chunks_number=video_chunks_number,
     )

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

@@ -0,0 +1,476 @@
+from __future__ import annotations
+
+import argparse
+import logging
+from collections.abc import Iterator
+from dataclasses import dataclass
+
+import torch
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import get_pixel_coords
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.conditioning import ConditioningItem
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.tools import LatentTools
+from ltx_core.types import (
+    Audio,
+    AudioLatentShape,
+    LatentState,
+    SpatioTemporalScaleFactors,
+    VideoPixelShape,
+)
+from ltx_pipelines.utils import ModelLedger
+from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES
+from ltx_pipelines.utils.helpers import (
+    cleanup_memory,
+    generate_enhanced_prompt,
+    get_device,
+    multi_modal_guider_denoising_func,
+    noise_audio_state,
+    noise_video_state,
+    simple_denoising_func,
+)
+from ltx_pipelines.utils.media_io import (
+    decode_audio_from_file,
+    encode_video,
+    get_videostream_metadata,
+    load_video_conditioning,
+)
+from ltx_pipelines.utils.samplers import euler_denoising_loop
+from ltx_pipelines.utils.types import PipelineComponents
+
+device = get_device()
+
+
+def _encode_video_for_retake(
+    video_encoder: torch.nn.Module,
+    video_path: str,
+    output_shape: VideoPixelShape,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor:
+    """Load video and encode to latents."""
+    pixel_video = load_video_conditioning(
+        video_path=video_path,
+        height=output_shape.height,
+        width=output_shape.width,
+        frame_cap=output_shape.frames,
+        dtype=dtype,
+        device=device,
+    )  # (1, C, F, H, W)
+    return video_encoder(pixel_video)
+
+
+def _encode_audio_for_retake(
+    audio_encoder: torch.nn.Module,
+    waveform: torch.Tensor,
+    waveform_sr: int,
+    output_shape: VideoPixelShape,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Encode audio to latents and trim/pad to match output_shape."""
+    waveform_batch = waveform.unsqueeze(0) if waveform.dim() == 2 else waveform
+    initial_audio_latent = vae_encode_audio(
+        Audio(waveform=waveform_batch.to(dtype), sampling_rate=waveform_sr), audio_encoder, None
+    )
+    expected_audio_shape = AudioLatentShape.from_video_pixel_shape(output_shape)
+    expected_frames = expected_audio_shape.frames
+    actual_frames = initial_audio_latent.shape[2]
+    if actual_frames > expected_frames:
+        initial_audio_latent = initial_audio_latent[:, :, :expected_frames, :]
+    elif actual_frames < expected_frames:
+        pad = torch.zeros(
+            initial_audio_latent.shape[0],
+            initial_audio_latent.shape[1],
+            expected_frames - actual_frames,
+            initial_audio_latent.shape[3],
+            device=initial_audio_latent.device,
+            dtype=initial_audio_latent.dtype,
+        )
+        initial_audio_latent = torch.cat([initial_audio_latent, pad], dim=2)
+    return initial_audio_latent
+
+
+# ---------------------------------------------------------------------------
+# Custom conditioning item: temporal region mask
+# ---------------------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class TemporalRegionMask:
+    """Conditioning item that sets ``denoise_mask = 0`` outside a time range
+    and ``1`` inside, so only the specified temporal region is regenerated.
+    Uses ``start_time`` and ``end_time`` in seconds. Works in *patchified*
+    (token) space using the patchifier's ``get_patch_grid_bounds``: for video
+    coords are latent frame indices (converted from seconds via ``fps``), for
+    audio coords are already in seconds.
+    """
+
+    start_time: float  # seconds, inclusive
+    end_time: float  # seconds, exclusive
+    fps: float
+
+    def apply_to(self, latent_state: LatentState, latent_tools: LatentTools) -> LatentState:
+        coords = latent_tools.patchifier.get_patch_grid_bounds(
+            latent_tools.target_shape, device=latent_state.denoise_mask.device
+        )
+        # coords: [B, 3, N, 2] (video) or [B, 1, N, 2] (audio); temporal dim is index 0
+        if coords.shape[1] == 1:
+            # Audio: patchifier returns seconds
+            t_start = coords[:, 0, :, 0]  # [B, N]
+            t_end = coords[:, 0, :, 1]  # [B, N]
+            in_region = (t_end > self.start_time) & (t_start < self.end_time)
+        else:
+            # Video: get pixel bounds per patch, find patches for start/end frame, read latent from coords.
+            scale_factors = getattr(latent_tools, "scale_factors", SpatioTemporalScaleFactors.default())
+            pixel_bounds = get_pixel_coords(coords, scale_factors, causal_fix=getattr(latent_tools, "causal_fix", True))
+            timestamp_bounds = pixel_bounds[0, 0] / self.fps
+            t_start, t_end = timestamp_bounds.unbind(dim=-1)
+            in_region = (t_end > self.start_time) & (t_start < self.end_time)
+        state = latent_state.clone()
+        mask_val = in_region.to(state.denoise_mask.dtype)
+        if state.denoise_mask.dim() == 3:
+            mask_val = mask_val.unsqueeze(-1)
+        state.denoise_mask.copy_(mask_val)
+        return state
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+
+class RetakePipeline:
+    """Regenerate a time region (retake) of an existing video.
+    Given a source video file and a time window ``[start_time, end_time]``
+    (in seconds), this pipeline keeps the video/audio outside that window
+    unchanged and *regenerates* the content inside the window from a text
+    prompt using the LTX-2 diffusion model.
+    Parameters
+    ----------
+    checkpoint_path : str
+        Path to the LTX-2 model checkpoint.
+    gemma_root : str
+        Root directory containing Gemma text-encoder weights.
+    loras : list[LoraPathStrengthAndSDOps]
+        Optional LoRA configs applied to the transformer.
+    device : torch.device
+        Target device (default: CUDA if available).
+    quantization : QuantizationPolicy | None
+        Optional quantization policy for the transformer.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: torch.device = device,
+        quantization: QuantizationPolicy | None = None,
+    ):
+        self.device = device
+        self.dtype = torch.bfloat16
+        self.model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            gemma_root_path=gemma_root,
+            loras=loras,
+            quantization=quantization,
+        )
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+
+    # --------------------------------------------------------------------- #
+    #  Public entry point                                                     #
+    # --------------------------------------------------------------------- #
+
+    @torch.inference_mode()
+    def __call__(  # noqa: PLR0913, PLR0915
+        self,
+        video_path: str,
+        prompt: str,
+        start_time: float,
+        end_time: float,
+        seed: int,
+        *,
+        negative_prompt: str = "",
+        num_inference_steps: int = 40,
+        video_guider_params: MultiModalGuiderParams | None = None,
+        audio_guider_params: MultiModalGuiderParams | None = None,
+        regenerate_video: bool = True,
+        regenerate_audio: bool = True,
+        enhance_prompt: bool = False,
+        distilled: bool = False,
+    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+        """Regenerate ``[start_time, end_time]`` of the source video (retake).
+        Parameters
+        ----------
+        video_path : str
+            Path to the source video file (must contain video; audio is optional).
+        prompt : str
+            Text prompt describing the *regenerated* section.
+        start_time, end_time : float
+            Time window (in seconds) of the section to regenerate.
+        seed : int
+            Random seed for reproducibility.
+        negative_prompt : str
+            Negative prompt for CFG guidance (ignored in distilled mode).
+        num_inference_steps : int
+            Number of Euler denoising steps (ignored in distilled mode which
+            uses a fixed 8-step schedule).
+        video_guider_params, audio_guider_params : MultiModalGuiderParams | None
+            Guidance parameters for video and audio modalities.  Ignored in
+            distilled mode.
+        regenerate_video : bool
+            If ``True`` (default), preserve video outside ``[start_time, end_time]``
+            and only regenerate the masked region.  If ``False``, fully regenerate
+            all video frames (the encoded video is still used as the initial latent
+            but with ``denoise_mask = 1`` everywhere).
+        regenerate_audio : bool
+            If True, regenerate audio in the [start_time, end_time] window; if False,
+            audio is preserved as-is (no regeneration).
+        enhance_prompt : bool
+            Whether to enhance the prompt via the text encoder.
+        distilled : bool
+            If ``True``, use the distilled sigma schedule
+            (``DISTILLED_SIGMA_VALUES``) and a simple (non-guided) denoising
+            function.  The model checkpoint must be the distilled variant.
+        Returns
+        -------
+        tuple[Iterator[torch.Tensor], torch.Tensor]
+            ``(video_frames_iterator, audio_waveform)``
+        """
+        if start_time >= end_time:
+            raise ValueError(f"start_time ({start_time}) must be less than end_time ({end_time})")
+
+        effective_seed = torch.randint(0, 2**31, (1,), device=self.device).item() if seed < 0 else seed
+        generator = torch.Generator(device=self.device).manual_seed(effective_seed)
+        noiser = GaussianNoiser(generator=generator)
+        stepper = EulerDiffusionStep()
+        dtype = self.dtype
+
+        video_encoder = self.model_ledger.video_encoder()
+
+        # Use av to get metadata
+        fps, num_pixel_frames, src_width, src_height = get_videostream_metadata(video_path)
+
+        output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_pixel_frames,
+            width=src_width,
+            height=src_height,
+            fps=fps,
+        )
+        initial_video_latent = _encode_video_for_retake(
+            video_encoder=video_encoder,
+            video_path=video_path,
+            output_shape=output_shape,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_conditionings: list[ConditioningItem] = [
+            TemporalRegionMask(
+                start_time=start_time if regenerate_video else 0.0,
+                end_time=end_time if regenerate_video else 0.0,
+                fps=fps,
+            )
+        ]
+        del video_encoder
+        cleanup_memory()
+
+        initial_audio_latent: torch.Tensor | None = None
+        audio_conditionings: list[ConditioningItem] = []
+
+        audio_in = decode_audio_from_file(video_path, self.device)
+        audio_encoder = self.model_ledger.audio_encoder()
+
+        if audio_in is not None:
+            waveform = audio_in.waveform.squeeze(0)
+            waveform_sr = audio_in.sampling_rate
+        else:
+            waveform, waveform_sr = None, None
+        if waveform is not None:
+            initial_audio_latent = _encode_audio_for_retake(
+                audio_encoder=audio_encoder,
+                waveform=waveform,
+                waveform_sr=waveform_sr,
+                output_shape=output_shape,
+                dtype=dtype,
+            )
+            audio_conditionings = [
+                TemporalRegionMask(
+                    start_time=start_time if regenerate_audio else 0.0,
+                    end_time=end_time if regenerate_audio else 0.0,
+                    fps=fps,
+                )
+            ]
+
+        del audio_encoder
+        cleanup_memory()
+
+        text_encoder = self.model_ledger.text_encoder()
+        if enhance_prompt:
+            prompt = generate_enhanced_prompt(text_encoder, prompt, None, seed=effective_seed)
+
+        if distilled:
+            # Distilled mode: single prompt, no negative
+            context_p = encode_text(text_encoder, prompts=[prompt])[0]
+            v_context_p, a_context_p = context_p
+        else:
+            context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
+            v_context_p, a_context_p = context_p
+            v_context_n, a_context_n = context_n
+
+        torch.cuda.synchronize()
+        del text_encoder
+        cleanup_memory()
+
+        transformer = self.model_ledger.transformer()
+
+        sigmas = (
+            torch.tensor(DISTILLED_SIGMA_VALUES) if distilled else LTX2Scheduler().execute(steps=num_inference_steps)
+        ).to(dtype=torch.float32, device=self.device)
+        if distilled:
+            denoise_fn = simple_denoising_func(
+                video_context=v_context_p,
+                audio_context=a_context_p,
+                transformer=transformer,
+            )
+        else:
+            video_guider = MultiModalGuider(
+                params=video_guider_params,
+                negative_context=v_context_n,
+            )
+            audio_guider = MultiModalGuider(
+                params=audio_guider_params,
+                negative_context=a_context_n,
+            )
+            denoise_fn = multi_modal_guider_denoising_func(
+                video_guider,
+                audio_guider,
+                v_context=v_context_p,
+                a_context=a_context_p,
+                transformer=transformer,
+            )
+
+        def denoising_loop(
+            sigmas: torch.Tensor,
+            video_state: LatentState,
+            audio_state: LatentState,
+            stepper: DiffusionStepProtocol,
+        ) -> tuple[LatentState, LatentState]:
+            return euler_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=denoise_fn,
+            )
+
+        # Build noised states with the encoded latents as initial values and
+        # the temporal masks applied via conditionings.
+        video_state, video_tools = noise_video_state(
+            output_shape=output_shape,
+            noiser=noiser,
+            conditionings=video_conditionings,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            initial_latent=initial_video_latent,
+        )
+        audio_state, audio_tools = noise_audio_state(
+            output_shape=output_shape,
+            noiser=noiser,
+            conditionings=audio_conditionings,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            initial_latent=initial_audio_latent,
+        )
+
+        video_state, audio_state = denoising_loop(sigmas, video_state, audio_state, stepper)
+
+        video_state = video_tools.clear_conditioning(video_state)
+        video_state = video_tools.unpatchify(video_state)
+        audio_state = audio_tools.clear_conditioning(audio_state)
+        audio_state = audio_tools.unpatchify(audio_state)
+
+        torch.cuda.synchronize()
+        del transformer
+        cleanup_memory()
+
+        decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder(), generator=generator)
+        decoded_audio = vae_decode_audio(
+            audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
+        )
+
+        return decoded_video, decoded_audio
+
+
+def main() -> None:
+    """CLI entry point for retake (regenerate a time region)."""
+    logging.getLogger().setLevel(logging.INFO)
+    parser = argparse.ArgumentParser(description="Retake: regenerate a time region of a video with LTX-2.")
+    parser.add_argument("--video-path", type=str, required=True, help="Path to the source video.")
+    parser.add_argument("--prompt", type=str, required=True, help="Text prompt for the regenerated region.")
+    parser.add_argument("--start-time", type=float, required=True, help="Start time of the region to regenerate (s).")
+    parser.add_argument("--end-time", type=float, required=True, help="End time of the region to regenerate (s).")
+    parser.add_argument("--output-path", type=str, required=True, help="Path for the output video.")
+    parser.add_argument("--checkpoint-path", type=str, required=True, help="Path to the LTX-2 checkpoint.")
+    parser.add_argument("--gemma-root", type=str, required=True, help="Path to Gemma text encoder weights.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed. Use -1 for a random seed.")
+    parser.add_argument("--loras", nargs="*", default=[], help="LoRA paths (optional).")
+    args = parser.parse_args()
+
+    if args.start_time >= args.end_time:
+        raise ValueError("start_time must be less than end_time")
+
+    # Validate frame count (8k+1) and resolution (multiples of 32) at CLI stage
+    video_scale = SpatioTemporalScaleFactors.default()
+    fps, num_frames, width, height = get_videostream_metadata(args.video_path)
+    if (num_frames - 1) % video_scale.time != 0:
+        snapped = ((num_frames - 1) // video_scale.time) * video_scale.time + 1
+        raise ValueError(
+            f"Video frame count must satisfy 8k+1 (e.g. 97, 193). Got {num_frames}; use a video with {snapped} frames."
+        )
+    if width % 32 != 0 or height % 32 != 0:
+        raise ValueError(f"Video width and height must be multiples of 32. Got {width}x{height}.")
+
+    pipeline = RetakePipeline(
+        checkpoint_path=args.checkpoint_path,
+        gemma_root=args.gemma_root,
+        loras=args.loras or [],
+    )
+    video_iter, audio = pipeline(
+        video_path=args.video_path,
+        prompt=args.prompt,
+        start_time=args.start_time,
+        end_time=args.end_time,
+        seed=args.seed,
+    )
+    tiling_config = TilingConfig.default()
+    video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
+    encode_video(
+        video=video_iter,
+        fps=int(fps),
+        audio=audio,
+        output_path=args.output_path,
+        video_chunks_number=video_chunks_number,
+    )
+
+
+if __name__ == "__main__":
+    main()

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

@@ -4,7 +4,11 @@ from collections.abc import Iterator
 import torch
 
 from ltx_core.components.diffusion_steps import EulerDiffusionStep
-from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+from ltx_core.components.guiders import (
+    MultiModalGuiderFactory,
+    MultiModalGuiderParams,
+    create_multimodal_guider_factory,
+)
 from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
@@ -13,11 +17,9 @@ from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
 from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
-from ltx_core.types import LatentState, VideoPixelShape
-from ltx_pipelines.utils import ModelLedger
-from ltx_pipelines.utils.args import default_1_stage_arg_parser
-from ltx_pipelines.utils.constants import AUDIO_SAMPLE_RATE
-from ltx_pipelines.utils.helpers import (
+from ltx_core.types import Audio, LatentState, VideoPixelShape
+from ltx_pipelines.utils import (
+    ModelLedger,
     assert_resolution,
     cleanup_memory,
     denoise_audio_video,
@@ -25,8 +27,10 @@ from ltx_pipelines.utils.helpers import (
     generate_enhanced_prompt,
     get_device,
     image_conditionings_by_replacing_latent,
-    multi_modal_guider_denoising_func,
+    multi_modal_guider_factory_denoising_func,
 )
+from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.constants import detect_params
 from ltx_pipelines.utils.media_io import encode_video
 from ltx_pipelines.utils.types import PipelineComponents
 
@@ -39,6 +43,7 @@ class TI2VidOneStagePipeline:
     Generates video at the target resolution in a single diffusion pass with
     classifier-free guidance (CFG). Supports optional image conditioning via
     the images parameter.
+    Assumes full non distilled model is provided in the checkpoint_path.
     """
 
     def __init__(
@@ -74,11 +79,11 @@ class TI2VidOneStagePipeline:
         num_frames: int,
         frame_rate: float,
         num_inference_steps: int,
-        video_guider_params: MultiModalGuiderParams,
-        audio_guider_params: MultiModalGuiderParams,
-        images: list[tuple[str, int, float]],
+        video_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        images: list[ImageConditioningInput],
         enhance_prompt: bool = False,
-    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=False)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -104,6 +109,15 @@ class TI2VidOneStagePipeline:
         transformer = self.model_ledger.transformer()
         sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
 
+        video_guider_factory = create_multimodal_guider_factory(
+            params=video_guider_params,
+            negative_context=v_context_n,
+        )
+        audio_guider_factory = create_multimodal_guider_factory(
+            params=audio_guider_params,
+            negative_context=a_context_n,
+        )
+
         def first_stage_denoising_loop(
             sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
         ) -> tuple[LatentState, LatentState]:
@@ -112,15 +126,9 @@ class TI2VidOneStagePipeline:
                 video_state=video_state,
                 audio_state=audio_state,
                 stepper=stepper,
-                denoise_fn=multi_modal_guider_denoising_func(
-                    video_guider=MultiModalGuider(
-                        params=video_guider_params,
-                        negative_context=v_context_n,
-                    ),
-                    audio_guider=MultiModalGuider(
-                        params=audio_guider_params,
-                        negative_context=a_context_n,
-                    ),
+                denoise_fn=multi_modal_guider_factory_denoising_func(
+                    video_guider_factory=video_guider_factory,
+                    audio_guider_factory=audio_guider_factory,
                     v_context=v_context_p,
                     a_context=a_context_p,
                     transformer=transformer,  # noqa: F821
@@ -157,14 +165,15 @@ class TI2VidOneStagePipeline:
         decoded_audio = vae_decode_audio(
             audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
         )
-
         return decoded_video, decoded_audio
 
 
 @torch.inference_mode()
 def main() -> None:
     logging.getLogger().setLevel(logging.INFO)
-    parser = default_1_stage_arg_parser()
+    checkpoint_path = detect_checkpoint_path()
+    params = detect_params(checkpoint_path)
+    parser = default_1_stage_arg_parser(params=params)
     args = parser.parse_args()
     pipeline = TI2VidOneStagePipeline(
         checkpoint_path=args.checkpoint_path,
@@ -204,7 +213,6 @@ def main() -> None:
         video=video,
         fps=args.frame_rate,
         audio=audio,
-        audio_sample_rate=AUDIO_SAMPLE_RATE,
         output_path=args.output_path,
         video_chunks_number=1,
     )

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

@@ -4,7 +4,11 @@ from collections.abc import Iterator
 import torch
 
 from ltx_core.components.diffusion_steps import EulerDiffusionStep
-from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+from ltx_core.components.guiders import (
+    MultiModalGuiderFactory,
+    MultiModalGuiderParams,
+    create_multimodal_guider_factory,
+)
 from ltx_core.components.noisers import GaussianNoiser
 from ltx_core.components.protocols import DiffusionStepProtocol
 from ltx_core.components.schedulers import LTX2Scheduler
@@ -15,14 +19,9 @@ from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
 from ltx_core.model.video_vae import decode_video as vae_decode_video
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import encode_text
-from ltx_core.types import LatentState, VideoPixelShape
-from ltx_pipelines.utils import ModelLedger
-from ltx_pipelines.utils.args import default_2_stage_arg_parser
-from ltx_pipelines.utils.constants import (
-    AUDIO_SAMPLE_RATE,
-    STAGE_2_DISTILLED_SIGMA_VALUES,
-)
-from ltx_pipelines.utils.helpers import (
+from ltx_core.types import Audio, LatentState, VideoPixelShape
+from ltx_pipelines.utils import (
+    ModelLedger,
     assert_resolution,
     cleanup_memory,
     denoise_audio_video,
@@ -30,9 +29,11 @@ from ltx_pipelines.utils.helpers import (
     generate_enhanced_prompt,
     get_device,
     image_conditionings_by_replacing_latent,
-    multi_modal_guider_denoising_func,
+    multi_modal_guider_factory_denoising_func,
     simple_denoising_func,
 )
+from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
 from ltx_pipelines.utils.media_io import encode_video
 from ltx_pipelines.utils.types import PipelineComponents
 
@@ -42,9 +43,10 @@ device = get_device()
 class TI2VidTwoStagesPipeline:
     """
     Two-stage text/image-to-video generation pipeline.
-    Stage 1 generates video at the target resolution with CFG guidance, then
-    Stage 2 upsamples by 2x and refines using a distilled LoRA for higher
-    quality output. Supports optional image conditioning via the images parameter.
+    Stage 1 generates video at half of the target resolution with CFG guidance (assuming
+    full model is used), then Stage 2 upsamples by 2x and refines using a distilled
+    LoRA for higher quality output. Supports optional image conditioning via the
+    images parameter.
     """
 
     def __init__(
@@ -54,7 +56,7 @@ class TI2VidTwoStagesPipeline:
         spatial_upsampler_path: str,
         gemma_root: str,
         loras: list[LoraPathStrengthAndSDOps],
-        device: str = device,
+        device: torch.device = device,
         quantization: QuantizationPolicy | None = None,
     ):
         self.device = device
@@ -78,7 +80,6 @@ class TI2VidTwoStagesPipeline:
             device=device,
         )
 
-    @torch.inference_mode()
     def __call__(  # noqa: PLR0913
         self,
         prompt: str,
@@ -89,12 +90,12 @@ class TI2VidTwoStagesPipeline:
         num_frames: int,
         frame_rate: float,
         num_inference_steps: int,
-        video_guider_params: MultiModalGuiderParams,
-        audio_guider_params: MultiModalGuiderParams,
-        images: list[tuple[str, int, float]],
+        video_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        audio_guider_params: MultiModalGuiderParams | MultiModalGuiderFactory,
+        images: list[ImageConditioningInput],
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
-    ) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
         assert_resolution(height=height, width=width, is_two_stage=True)
 
         generator = torch.Generator(device=self.device).manual_seed(seed)
@@ -128,12 +129,12 @@ class TI2VidTwoStagesPipeline:
                 video_state=video_state,
                 audio_state=audio_state,
                 stepper=stepper,
-                denoise_fn=multi_modal_guider_denoising_func(
-                    video_guider=MultiModalGuider(
+                denoise_fn=multi_modal_guider_factory_denoising_func(
+                    video_guider_factory=create_multimodal_guider_factory(
                         params=video_guider_params,
                         negative_context=v_context_n,
                     ),
-                    audio_guider=MultiModalGuider(
+                    audio_guider_factory=create_multimodal_guider_factory(
                         params=audio_guider_params,
                         negative_context=a_context_n,
                     ),
@@ -237,14 +238,15 @@ class TI2VidTwoStagesPipeline:
         decoded_audio = vae_decode_audio(
             audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
         )
-
         return decoded_video, decoded_audio
 
 
 @torch.inference_mode()
 def main() -> None:
     logging.getLogger().setLevel(logging.INFO)
-    parser = default_2_stage_arg_parser()
+    checkpoint_path = detect_checkpoint_path()
+    params = detect_params(checkpoint_path)
+    parser = default_2_stage_arg_parser(params=params)
     args = parser.parse_args()
     pipeline = TI2VidTwoStagesPipeline(
         checkpoint_path=args.checkpoint_path,
@@ -289,7 +291,6 @@ def main() -> None:
         video=video,
         fps=args.frame_rate,
         audio=audio,
-        audio_sample_rate=AUDIO_SAMPLE_RATE,
         output_path=args.output_path,
         video_chunks_number=video_chunks_number,
     )

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

@@ -0,0 +1,307 @@
+import logging
+from collections.abc import Iterator
+
+import torch
+
+from ltx_core.components.diffusion_steps import Res2sDiffusionStep
+from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
+from ltx_core.model.upsampler import upsample_video
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_core.model.video_vae import decode_video as vae_decode_video
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.text_encoders.gemma import encode_text
+from ltx_core.tools import VideoLatentShape
+from ltx_core.types import Audio, LatentState, VideoPixelShape
+from ltx_pipelines.utils import (
+    ModelLedger,
+    assert_resolution,
+    cleanup_memory,
+    denoise_audio_video,
+    generate_enhanced_prompt,
+    get_device,
+    image_conditionings_by_replacing_latent,
+    multi_modal_guider_denoising_func,
+    res2s_audio_video_denoising_loop,
+    simple_denoising_func,
+)
+from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
+from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
+from ltx_pipelines.utils.media_io import encode_video
+from ltx_pipelines.utils.types import PipelineComponents
+
+device = get_device()
+
+
+class TI2VidTwoStagesRes2sPipeline:
+    """
+    Two-stage text/image-to-video generation pipeline using the res_2s sampler.
+    Same structure as :class:`TI2VidTwoStagesPipeline`: stage 1 generates video at
+    half of the target resolution with CFG guidance (assuming  full model is used),
+    then Stage 2 upsamples by 2x and refines using a distilled LoRA for higher
+    quality output.
+    Uses the res_2s second-order sampler instead of Euler, allowing fewer
+    steps for comparable quality. Supports optional image conditioning via
+    the images parameter.
+    """
+
+    def __init__(
+        self,
+        checkpoint_path: str,
+        distilled_lora: list[LoraPathStrengthAndSDOps],
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        loras: list[LoraPathStrengthAndSDOps],
+        device: str = device,
+        quantization: QuantizationPolicy | None = None,
+    ):
+        self.device = device
+        self.dtype = torch.bfloat16
+        self.stage_1_model_ledger = ModelLedger(
+            dtype=self.dtype,
+            device=device,
+            checkpoint_path=checkpoint_path,
+            gemma_root_path=gemma_root,
+            spatial_upsampler_path=spatial_upsampler_path,
+            loras=loras,
+            quantization=quantization,
+        )
+
+        self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
+            loras=distilled_lora,
+        )
+
+        self.pipeline_components = PipelineComponents(
+            dtype=self.dtype,
+            device=device,
+        )
+
+    @torch.inference_mode()
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        video_guider_params: MultiModalGuiderParams,
+        audio_guider_params: MultiModalGuiderParams,
+        images: list[ImageConditioningInput],
+        tiling_config: TilingConfig | None = None,
+        enhance_prompt: bool = False,
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
+        assert_resolution(height=height, width=width, is_two_stage=True)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+        dtype = torch.bfloat16
+
+        text_encoder = self.stage_1_model_ledger.text_encoder()
+        if enhance_prompt:
+            prompt = generate_enhanced_prompt(
+                text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
+            )
+        context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
+        v_context_p, a_context_p = context_p
+        v_context_n, a_context_n = context_n
+
+        torch.cuda.synchronize()
+        del text_encoder
+        cleanup_memory()
+
+        # Stage 1: Initial low resolution video generation.
+        video_encoder = self.stage_1_model_ledger.video_encoder()
+        transformer = self.stage_1_model_ledger.transformer()
+
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+        empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
+        stepper = Res2sDiffusionStep()
+        sigmas = (
+            LTX2Scheduler()
+            .execute(latent=empty_latent, steps=num_inference_steps)
+            .to(dtype=torch.float32, device=self.device)
+        )
+
+        def first_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return res2s_audio_video_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=multi_modal_guider_denoising_func(
+                    video_guider=MultiModalGuider(
+                        params=video_guider_params,
+                        negative_context=v_context_n,
+                    ),
+                    audio_guider=MultiModalGuider(
+                        params=audio_guider_params,
+                        negative_context=a_context_n,
+                    ),
+                    v_context=v_context_p,
+                    a_context=a_context_p,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_1_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_1_output_shape.height,
+            width=stage_1_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_1_output_shape,
+            conditionings=stage_1_conditionings,
+            noiser=noiser,
+            sigmas=sigmas,
+            stepper=stepper,
+            denoising_loop_fn=first_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        cleanup_memory()
+
+        # Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
+        upscaled_video_latent = upsample_video(
+            latent=video_state.latent[:1],
+            video_encoder=video_encoder,
+            upsampler=self.stage_2_model_ledger.spatial_upsampler(),
+        )
+
+        torch.cuda.synchronize()
+        cleanup_memory()
+
+        transformer = self.stage_2_model_ledger.transformer()
+        distilled_sigmas = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=self.device)
+
+        def second_stage_denoising_loop(
+            sigmas: torch.Tensor, video_state: LatentState, audio_state: LatentState, stepper: DiffusionStepProtocol
+        ) -> tuple[LatentState, LatentState]:
+            return res2s_audio_video_denoising_loop(
+                sigmas=sigmas,
+                video_state=video_state,
+                audio_state=audio_state,
+                stepper=stepper,
+                denoise_fn=simple_denoising_func(
+                    video_context=v_context_p,
+                    audio_context=a_context_p,
+                    transformer=transformer,  # noqa: F821
+                ),
+            )
+
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = image_conditionings_by_replacing_latent(
+            images=images,
+            height=stage_2_output_shape.height,
+            width=stage_2_output_shape.width,
+            video_encoder=video_encoder,
+            dtype=dtype,
+            device=self.device,
+        )
+        video_state, audio_state = denoise_audio_video(
+            output_shape=stage_2_output_shape,
+            conditionings=stage_2_conditionings,
+            noiser=noiser,
+            sigmas=distilled_sigmas,
+            stepper=stepper,
+            denoising_loop_fn=second_stage_denoising_loop,
+            components=self.pipeline_components,
+            dtype=dtype,
+            device=self.device,
+            noise_scale=distilled_sigmas[0],
+            initial_video_latent=upscaled_video_latent,
+            initial_audio_latent=audio_state.latent,
+        )
+
+        torch.cuda.synchronize()
+        del transformer
+        del video_encoder
+        cleanup_memory()
+
+        decoded_video = vae_decode_video(
+            video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
+        )
+        decoded_audio = vae_decode_audio(
+            audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
+        )
+        return decoded_video, decoded_audio
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    checkpoint_path = detect_checkpoint_path()
+    params = detect_params(checkpoint_path)
+    parser = default_2_stage_arg_parser(params=params)
+    args = parser.parse_args()
+    pipeline = TI2VidTwoStagesRes2sPipeline(
+        checkpoint_path=args.checkpoint_path,
+        distilled_lora=args.distilled_lora,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        loras=args.lora,
+        quantization=args.quantization,
+    )
+    tiling_config = TilingConfig.default()
+    video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        video_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.video_cfg_guidance_scale,
+            stg_scale=args.video_stg_guidance_scale,
+            rescale_scale=args.video_rescale_scale,
+            modality_scale=args.a2v_guidance_scale,
+            skip_step=args.video_skip_step,
+            stg_blocks=args.video_stg_blocks,
+        ),
+        audio_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.audio_cfg_guidance_scale,
+            stg_scale=args.audio_stg_guidance_scale,
+            rescale_scale=args.audio_rescale_scale,
+            modality_scale=args.v2a_guidance_scale,
+            skip_step=args.audio_skip_step,
+            stg_blocks=args.audio_stg_blocks,
+        ),
+        images=args.images,
+        tiling_config=tiling_config,
+    )
+
+    encode_video(
+        video=video,
+        fps=args.frame_rate,
+        audio=audio,
+        output_path=args.output_path,
+        video_chunks_number=video_chunks_number,
+    )
+
+
+if __name__ == "__main__":
+    main()

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

@@ -1,5 +1,33 @@
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    cleanup_memory,
+    denoise_audio_video,
+    generate_enhanced_prompt,
+    get_device,
+    image_conditionings_by_replacing_latent,
+    multi_modal_guider_denoising_func,
+    multi_modal_guider_factory_denoising_func,
+    simple_denoising_func,
+)
 from ltx_pipelines.utils.model_ledger import ModelLedger
+from ltx_pipelines.utils.samplers import (
+    euler_denoising_loop,
+    gradient_estimating_euler_denoising_loop,
+    res2s_audio_video_denoising_loop,
+)
 
 __all__ = [
     "ModelLedger",
+    "assert_resolution",
+    "cleanup_memory",
+    "denoise_audio_video",
+    "euler_denoising_loop",
+    "generate_enhanced_prompt",
+    "get_device",
+    "gradient_estimating_euler_denoising_loop",
+    "image_conditionings_by_replacing_latent",
+    "multi_modal_guider_denoising_func",
+    "multi_modal_guider_factory_denoising_func",
+    "res2s_audio_video_denoising_loop",
+    "simple_denoising_func",
 ]

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

@@ -1,24 +1,25 @@
 import argparse
 from pathlib import Path
+from typing import NamedTuple
 
 from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps
 from ltx_core.quantization import QuantizationPolicy
 from ltx_pipelines.utils.constants import (
-    DEFAULT_1_STAGE_HEIGHT,
-    DEFAULT_1_STAGE_WIDTH,
-    DEFAULT_2_STAGE_HEIGHT,
-    DEFAULT_2_STAGE_WIDTH,
-    DEFAULT_AUDIO_GUIDER_PARAMS,
-    DEFAULT_FRAME_RATE,
+    DEFAULT_IMAGE_CRF,
     DEFAULT_LORA_STRENGTH,
     DEFAULT_NEGATIVE_PROMPT,
-    DEFAULT_NUM_FRAMES,
-    DEFAULT_NUM_INFERENCE_STEPS,
-    DEFAULT_SEED,
-    DEFAULT_VIDEO_GUIDER_PARAMS,
+    LTX_2_3_PARAMS,
+    PipelineParams,
 )
 
 
+class ImageConditioningInput(NamedTuple):
+    path: str
+    frame_idx: int
+    strength: float
+    crf: int = DEFAULT_IMAGE_CRF
+
+
 class VideoConditioningAction(argparse.Action):
     def __call__(
         self,
@@ -35,20 +36,50 @@ class VideoConditioningAction(argparse.Action):
         setattr(namespace, self.dest, current)
 
 
+class VideoMaskConditioningAction(argparse.Action):
+    """Parse ``--conditioning-attention-mask PATH STRENGTH``.
+    Stores a ``(mask_path, strength)`` tuple on the namespace.  The mask video
+    should be grayscale with pixel values in [0, 1] controlling per-region
+    conditioning attention strength.  The scalar *STRENGTH* is multiplied with
+    the spatial mask before it is applied.
+    """
+
+    def __call__(
+        self,
+        parser: argparse.ArgumentParser,  # noqa: ARG002
+        namespace: argparse.Namespace,
+        values: list[str],
+        option_string: str | None = None,
+    ) -> None:
+        if len(values) != 2:
+            msg = f"{option_string} requires exactly 2 arguments (MASK_PATH STRENGTH), got {len(values)}"
+            raise argparse.ArgumentError(self, msg)
+
+        mask_path = resolve_path(values[0])
+        strength = float(values[1])
+        setattr(namespace, self.dest, (mask_path, strength))
+
+
 class ImageAction(argparse.Action):
     def __call__(
         self,
         parser: argparse.ArgumentParser,  # noqa: ARG002
         namespace: argparse.Namespace,
         values: list[str],
-        option_string: str | None = None,  # noqa: ARG002
+        option_string: str | None = None,
     ) -> None:
-        path, frame_idx, strength_str = values
-        resolved_path = resolve_path(path)
-        frame_idx = int(frame_idx)
-        strength = float(strength_str)
+        if len(values) not in (3, 4):
+            msg = f"{option_string} requires 3 or 4 arguments (PATH FRAME_IDX STRENGTH [CRF]), got {len(values)}"
+            raise argparse.ArgumentError(self, msg)
+
+        conditioning = ImageConditioningInput(
+            path=resolve_path(values[0]),
+            frame_idx=int(values[1]),
+            strength=float(values[2]),
+            crf=int(values[3]) if len(values) > 3 else DEFAULT_IMAGE_CRF,
+        )
         current = getattr(namespace, self.dest) or []
-        current.append((resolved_path, frame_idx, strength))
+        current.append(conditioning)
         setattr(namespace, self.dest, current)
 
 
@@ -113,14 +144,34 @@ class QuantizationAction(argparse.Action):
         setattr(namespace, self.dest, policy)
 
 
-def basic_arg_parser() -> argparse.ArgumentParser:
+def detect_checkpoint_path(distilled: bool = False) -> str:
+    """Pre-parse argv to extract the checkpoint path before building the full parser."""
+    pre = argparse.ArgumentParser(add_help=False)
+    flag = "--distilled-checkpoint-path" if distilled else "--checkpoint-path"
+    pre.add_argument(flag, type=resolve_path, required=True)
+    known, _ = pre.parse_known_args()
+    return known.distilled_checkpoint_path if distilled else known.checkpoint_path
+
+
+def basic_arg_parser(
+    params: PipelineParams = LTX_2_3_PARAMS,
+    distilled: bool = False,
+) -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--checkpoint-path",
-        type=resolve_path,
-        required=True,
-        help="Path to LTX-2 model checkpoint (.safetensors file).",
-    )
+    if distilled:
+        parser.add_argument(
+            "--distilled-checkpoint-path",
+            type=resolve_path,
+            required=True,
+            help="Path to LTX-2 distilled model checkpoint (.safetensors file).",
+        )
+    else:
+        parser.add_argument(
+            "--checkpoint-path",
+            type=resolve_path,
+            required=True,
+            help="Path to LTX-2 model checkpoint (.safetensors file).",
+        )
     parser.add_argument(
         "--gemma-root",
         type=resolve_path,
@@ -142,58 +193,57 @@ def basic_arg_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--seed",
         type=int,
-        default=DEFAULT_SEED,
-        help=(
-            f"Random seed value used to initialize the noise tensor for "
-            f"reproducible generation (default: {DEFAULT_SEED})."
-        ),
+        default=params.seed,
+        help=f"Random seed for reproducible generation (default: {params.seed}).",
     )
     parser.add_argument(
         "--height",
         type=int,
-        default=DEFAULT_1_STAGE_HEIGHT,
-        help=f"Height of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_HEIGHT}).",
+        default=params.stage_1_height,
+        help=f"Video height in pixels, divisible by 32 (default: {params.stage_1_height}).",
     )
     parser.add_argument(
         "--width",
         type=int,
-        default=DEFAULT_1_STAGE_WIDTH,
-        help=f"Width of the generated video in pixels, should be divisible by 32 (default: {DEFAULT_1_STAGE_WIDTH}).",
+        default=params.stage_1_width,
+        help=f"Width of the generated video in pixels, should be divisible by 32 (default: {params.stage_1_width}).",
     )
     parser.add_argument(
         "--num-frames",
         type=int,
-        default=DEFAULT_NUM_FRAMES,
+        default=params.num_frames,
         help=f"Number of frames to generate in the output video sequence, num-frames = (8 x K) + 1, "
-        f"where k is a non-negative integer (default: {DEFAULT_NUM_FRAMES}).",
+        f"where k is a non-negative integer (default: {params.num_frames}).",
     )
     parser.add_argument(
         "--frame-rate",
         type=float,
-        default=DEFAULT_FRAME_RATE,
-        help=f"Frame rate of the generated video (fps) (default: {DEFAULT_FRAME_RATE}).",
+        default=params.frame_rate,
+        help=f"Frame rate of the generated video (fps) (default: {params.frame_rate}).",
     )
     parser.add_argument(
         "--num-inference-steps",
         type=int,
-        default=DEFAULT_NUM_INFERENCE_STEPS,
+        default=params.num_inference_steps,
         help=(
             f"Number of denoising steps in the diffusion sampling process. "
-            f"Higher values improve quality but increase generation time (default: {DEFAULT_NUM_INFERENCE_STEPS})."
+            f"Higher values improve quality but increase generation time (default: {params.num_inference_steps})."
         ),
     )
     parser.add_argument(
         "--image",
         dest="images",
         action=ImageAction,
-        nargs=3,
-        metavar=("PATH", "FRAME_IDX", "STRENGTH"),
+        nargs="+",
+        metavar="ARG",
         default=[],
         help=(
-            "Image conditioning input: path to image file, target frame index, "
-            "and conditioning strength (all three required). Default: empty list [] (no image conditioning). "
+            "Image conditioning input: PATH FRAME_IDX STRENGTH [CRF]. "
+            "PATH is the image file, FRAME_IDX is the target frame index, "
+            "STRENGTH is the conditioning strength (all three required). "
+            f"CRF is the optional H.264 compression quality (0=lossless, default: {DEFAULT_IMAGE_CRF}). "
             "Can be specified multiple times. Example: --image path/to/image1.jpg 0 0.8 "
-            "--image path/to/image2.jpg 160 0.9"
+            "--image path/to/image2.jpg 160 0.9 0"
         ),
     )
     parser.add_argument(
@@ -228,8 +278,10 @@ def basic_arg_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def default_1_stage_arg_parser() -> argparse.ArgumentParser:
-    parser = basic_arg_parser()
+def default_1_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
+    video_guider = params.video_guider_params
+    audio_guider = params.audio_guider_params
+    parser = basic_arg_parser(params=params)
     parser.add_argument(
         "--negative-prompt",
         type=str,
@@ -243,139 +295,139 @@ def default_1_stage_arg_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--video-cfg-guidance-scale",
         type=float,
-        default=DEFAULT_VIDEO_GUIDER_PARAMS.cfg_scale,
+        default=video_guider.cfg_scale,
         help=(
             f"Classifier-free guidance (CFG) scale controlling how strongly "
             f"the model adheres to the video prompt. Higher values increase prompt "
-            "adherence but may reduce diversity. 1.0 means no effect "
-            f"(default: {DEFAULT_VIDEO_GUIDER_PARAMS.cfg_scale})."
+            f"adherence but may reduce diversity. 1.0 means no effect "
+            f"(default: {video_guider.cfg_scale})."
         ),
     )
     parser.add_argument(
         "--video-stg-guidance-scale",
         type=float,
-        default=DEFAULT_VIDEO_GUIDER_PARAMS.stg_scale,
+        default=video_guider.stg_scale,
         help=(
             f"STG (Spatio-Temporal Guidance) scale controlling how strongly "
             f"the model reacts to the perturbation of the video modality. Higher values increase "
             f"the effect but may reduce quality. 0.0 means no effect "
-            f"(default: {DEFAULT_VIDEO_GUIDER_PARAMS.stg_scale})."
+            f"(default: {video_guider.stg_scale})."
         ),
     )
     parser.add_argument(
         "--video-rescale-scale",
         type=float,
-        default=DEFAULT_VIDEO_GUIDER_PARAMS.rescale_scale,
+        default=video_guider.rescale_scale,
         help=(
             f"Rescale scale controlling how strongly "
             f"the model rescales the video modality after applying other guidance. Higher values tend to decrease "
-            f"oversaturation effects. 0.0 means no effect (default: {DEFAULT_VIDEO_GUIDER_PARAMS.rescale_scale})."
+            f"oversaturation effects. 0.0 means no effect (default: {video_guider.rescale_scale})."
         ),
     )
     parser.add_argument(
         "--video-stg-blocks",
         type=int,
         nargs="*",
-        default=DEFAULT_VIDEO_GUIDER_PARAMS.stg_blocks,
-        help=(f"Which transformer blocks to perturb for STG. Default: {DEFAULT_VIDEO_GUIDER_PARAMS.stg_blocks}."),
+        default=video_guider.stg_blocks,
+        help=(f"Which transformer blocks to perturb for STG. Default: {video_guider.stg_blocks}."),
     )
     parser.add_argument(
         "--a2v-guidance-scale",
         type=float,
-        default=DEFAULT_VIDEO_GUIDER_PARAMS.modality_scale,
+        default=video_guider.modality_scale,
         help=(
             f"A2V (Audio-to-Video) guidance scale controlling how strongly "
             f"the model reacts to the perturbation of the audio-to-video cross-attention. Higher values may increase "
-            f"lipsync quality. 1.0 means no effect (default: {DEFAULT_VIDEO_GUIDER_PARAMS.modality_scale})."
+            f"lipsync quality. 1.0 means no effect (default: {video_guider.modality_scale})."
         ),
     )
     parser.add_argument(
         "--video-skip-step",
         type=int,
-        default=DEFAULT_VIDEO_GUIDER_PARAMS.skip_step,
+        default=video_guider.skip_step,
         help=(
             "Video skip step N controls periodic skipping during the video diffusion process: "
             "only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
             f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
-            f"default: {DEFAULT_VIDEO_GUIDER_PARAMS.skip_step})."
+            f"default: {video_guider.skip_step})."
         ),
     )
     parser.add_argument(
         "--audio-cfg-guidance-scale",
         type=float,
-        default=DEFAULT_AUDIO_GUIDER_PARAMS.cfg_scale,
+        default=audio_guider.cfg_scale,
         help=(
             f"Audio CFG (Classifier-free guidance) scale controlling how strongly "
             f"the model adheres to the audio prompt. Higher values increase prompt "
             f"adherence but may reduce diversity. 1.0 means no effect "
-            f"(default: {DEFAULT_AUDIO_GUIDER_PARAMS.cfg_scale})."
+            f"(default: {audio_guider.cfg_scale})."
         ),
     )
     parser.add_argument(
         "--audio-stg-guidance-scale",
         type=float,
-        default=DEFAULT_AUDIO_GUIDER_PARAMS.stg_scale,
+        default=audio_guider.stg_scale,
         help=(
             f"Audio STG (Spatio-Temporal Guidance) scale controlling how strongly "
             f"the model reacts to the perturbation of the audio modality. Higher values increase "
             f"the effect but may reduce quality. 0.0 means no effect "
-            f"(default: {DEFAULT_AUDIO_GUIDER_PARAMS.stg_scale})."
+            f"(default: {audio_guider.stg_scale})."
         ),
     )
     parser.add_argument(
         "--audio-rescale-scale",
         type=float,
-        default=DEFAULT_AUDIO_GUIDER_PARAMS.rescale_scale,
+        default=audio_guider.rescale_scale,
         help=(
             f"Audio rescale scale controlling how strongly "
             f"the model rescales the audio modality after applying other guidance. "
-            f"Experimental. 0.0 means no effect (default: {DEFAULT_AUDIO_GUIDER_PARAMS.rescale_scale})."
+            f"Experimental. 0.0 means no effect (default: {audio_guider.rescale_scale})."
         ),
     )
     parser.add_argument(
         "--audio-stg-blocks",
         type=int,
         nargs="*",
-        default=DEFAULT_AUDIO_GUIDER_PARAMS.stg_blocks,
-        help=(f"Which transformer blocks to perturb for Audio STG. Default: {DEFAULT_AUDIO_GUIDER_PARAMS.stg_blocks}."),
+        default=audio_guider.stg_blocks,
+        help=(f"Which transformer blocks to perturb for Audio STG. Default: {audio_guider.stg_blocks}."),
     )
     parser.add_argument(
         "--v2a-guidance-scale",
         type=float,
-        default=DEFAULT_AUDIO_GUIDER_PARAMS.modality_scale,
+        default=audio_guider.modality_scale,
         help=(
             f"V2A (Video-to-Audio) guidance scale controlling how strongly "
             f"the model reacts to the perturbation of the video-to-audio cross-attention. Higher values may increase "
-            f"lipsync quality. 1.0 means no effect (default: {DEFAULT_AUDIO_GUIDER_PARAMS.modality_scale})."
+            f"lipsync quality. 1.0 means no effect (default: {audio_guider.modality_scale})."
         ),
     )
     parser.add_argument(
         "--audio-skip-step",
         type=int,
-        default=DEFAULT_AUDIO_GUIDER_PARAMS.skip_step,
+        default=audio_guider.skip_step,
         help=(
             "Audio skip step N controls periodic skipping during the audio diffusion process: "
             "only steps where step_index % (N + 1) == 0 are processed, all others are skipped "
             f"(e.g., 0 = no skipping; 1 = skip every other step; 2 = skip 2 of every 3 steps; "
-            f"default: {DEFAULT_AUDIO_GUIDER_PARAMS.skip_step})."
+            f"default: {audio_guider.skip_step})."
         ),
     )
     return parser
 
 
-def default_2_stage_arg_parser() -> argparse.ArgumentParser:
-    parser = default_1_stage_arg_parser()
-    parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
+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)
     # Update help text to reflect 2-stage defaults
     for action in parser._actions:
         if "--height" in action.option_strings:
             action.help = (
                 f"Height of the generated video in pixels, should be divisible by 64 "
-                f"(default: {DEFAULT_2_STAGE_HEIGHT})."
+                f"(default: {params.stage_2_height})."
             )
         if "--width" in action.option_strings:
             action.help = (
-                f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
+                f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width})."
             )
     parser.add_argument(
         "--distilled-lora",
@@ -405,19 +457,19 @@ def default_2_stage_arg_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def default_2_stage_distilled_arg_parser() -> argparse.ArgumentParser:
-    parser = basic_arg_parser()
-    parser.set_defaults(height=DEFAULT_2_STAGE_HEIGHT, width=DEFAULT_2_STAGE_WIDTH)
+def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
+    parser = basic_arg_parser(params=params, distilled=True)
+    parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
     # Update help text to reflect 2-stage defaults
     for action in parser._actions:
         if "--height" in action.option_strings:
             action.help = (
                 f"Height of the generated video in pixels, should be divisible by 64 "
-                f"(default: {DEFAULT_2_STAGE_HEIGHT})."
+                f"(default: {params.stage_2_height})."
             )
         if "--width" in action.option_strings:
             action.help = (
-                f"Width of the generated video in pixels, should be divisible by 64 (default: {DEFAULT_2_STAGE_WIDTH})."
+                f"Width of the generated video in pixels, should be divisible by 64 (default: {params.stage_2_width})."
             )
     parser.add_argument(
         "--spatial-upsampler-path",

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

@@ -1,12 +1,17 @@
+import logging
+from dataclasses import dataclass, field, replace
+
+from safetensors import safe_open
+
+from ltx_core.components.guiders import MultiModalGuiderParams
+from ltx_core.types import SpatioTemporalScaleFactors
+
 # =============================================================================
 # Diffusion Schedule
 # =============================================================================
 
 # Noise schedule for the distilled pipeline. These sigma values control noise
 # levels at each denoising step and were tuned to match the distillation process.
-from ltx_core.components.guiders import MultiModalGuiderParams
-from ltx_core.types import SpatioTemporalScaleFactors
-
 DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]
 
 # Reduced schedule for super-resolution stage 2 (subset of distilled values)
@@ -14,64 +19,88 @@ STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
 
 
 # =============================================================================
-# Video Generation Defaults
+# Pipeline Parameters
 # =============================================================================
 
-DEFAULT_SEED = 10
-DEFAULT_1_STAGE_HEIGHT = 512
-DEFAULT_1_STAGE_WIDTH = 768
-DEFAULT_2_STAGE_HEIGHT = DEFAULT_1_STAGE_HEIGHT * 2
-DEFAULT_2_STAGE_WIDTH = DEFAULT_1_STAGE_WIDTH * 2
-DEFAULT_NUM_FRAMES = 121
-DEFAULT_FRAME_RATE = 24.0
-DEFAULT_NUM_INFERENCE_STEPS = 40
-DEFAULT_VIDEO_GUIDER_PARAMS = MultiModalGuiderParams(
-    cfg_scale=3.0,
-    stg_scale=1.0,
-    rescale_scale=0.7,
-    modality_scale=3.0,
-    skip_step=0,
-    stg_blocks=[29],
+
+@dataclass(frozen=True)
+class PipelineParams:
+    seed: int = 10
+    stage_1_height: int = 512
+    stage_1_width: int = 768
+    num_frames: int = 121
+    frame_rate: float = 24.0
+    num_inference_steps: int = 40
+    video_guider_params: MultiModalGuiderParams = field(
+        default_factory=lambda: MultiModalGuiderParams(
+            cfg_scale=3.0,
+            stg_scale=1.0,
+            rescale_scale=0.7,
+            modality_scale=3.0,
+            skip_step=0,
+            stg_blocks=[29],
+        )
+    )
+    audio_guider_params: MultiModalGuiderParams = field(
+        default_factory=lambda: MultiModalGuiderParams(
+            cfg_scale=7.0,
+            stg_scale=1.0,
+            rescale_scale=0.7,
+            modality_scale=3.0,
+            skip_step=0,
+            stg_blocks=[29],
+        )
+    )
+
+    @property
+    def stage_2_height(self) -> int:
+        return int(self.stage_1_height * 2)
+
+    @property
+    def stage_2_width(self) -> int:
+        return int(self.stage_1_width * 2)
+
+
+# Default params for LTX-2.0 non-distilled models. These can be overridden by detecting from checkpoint metadata.
+LTX_2_PARAMS = PipelineParams()
+
+# Default params for LTX-2.3 non-distilled models. These override some of the LTX-2.0 defaults.
+LTX_2_3_PARAMS = replace(
+    LTX_2_PARAMS,
+    num_inference_steps=30,
+    video_guider_params=replace(LTX_2_PARAMS.video_guider_params, stg_blocks=[28]),
+    audio_guider_params=replace(LTX_2_PARAMS.audio_guider_params, stg_blocks=[28]),
 )
 
-
-# =============================================================================
-# Audio
-# =============================================================================
-
-DEFAULT_AUDIO_GUIDER_PARAMS = MultiModalGuiderParams(
-    cfg_scale=7.0,
-    stg_scale=1.0,
-    rescale_scale=0.7,
-    modality_scale=3.0,
-    skip_step=0,
-    stg_blocks=[29],
-)
-AUDIO_SAMPLE_RATE = 24000
-
-
-# =============================================================================
-# LoRA
-# =============================================================================
-
 DEFAULT_LORA_STRENGTH = 1.0
-
-
-# =============================================================================
-# Video VAE Architecture
-# =============================================================================
-
+DEFAULT_IMAGE_CRF = 33
 VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
 VIDEO_LATENT_CHANNELS = 128
 
+_LTX_2_3_MODEL_VERSION_PREFIX = "2.3"
 
-# =============================================================================
-# Image Preprocessing
-# =============================================================================
 
-# CRF (Constant Rate Factor) for H.264 encoding used in image conditioning.
-# Lower = higher quality, 0 = lossless. This mimics compression artifacts.
-DEFAULT_IMAGE_CRF = 33
+def detect_params(checkpoint_path: str) -> PipelineParams:
+    """Detect pipeline params from checkpoint metadata.
+    Reads the ``model_version`` field from the safetensors config metadata.
+    Returns ``LTX_2_3_PARAMS`` when the version starts with "2.3",
+    otherwise falls back to ``LTX_2_PARAMS``.
+    """
+    logger = logging.getLogger(__name__)
+
+    try:
+        with safe_open(checkpoint_path, framework="pt") as f:
+            metadata = f.metadata() or {}
+        version = metadata.get("model_version", "")
+    except Exception:
+        logger.warning("Could not read checkpoint metadata from %s, using LTX-2 defaults", checkpoint_path)
+        return LTX_2_PARAMS
+
+    if version.startswith(_LTX_2_3_MODEL_VERSION_PREFIX):
+        return LTX_2_3_PARAMS
+
+    logger.info("Using LTX_2_PARAMS for checkpoint (version=%s)", version or "unknown")
+    return LTX_2_PARAMS
 
 
 # =============================================================================

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

@@ -3,9 +3,8 @@ import logging
 from dataclasses import replace
 
 import torch
-from tqdm import tqdm
 
-from ltx_core.components.guiders import MultiModalGuider
+from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderFactory
 from ltx_core.components.noisers import Noiser
 from ltx_core.components.protocols import DiffusionStepProtocol, GuiderProtocol
 from ltx_core.conditioning import (
@@ -21,10 +20,10 @@ from ltx_core.guidance.perturbations import (
 )
 from ltx_core.model.transformer import Modality, X0Model
 from ltx_core.model.video_vae import VideoEncoder
-from ltx_core.text_encoders.gemma import GemmaTextEncoderModelBase
+from ltx_core.text_encoders.gemma import GemmaTextEncoder
 from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
 from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
-from ltx_core.utils import to_denoised, to_velocity
+from ltx_pipelines.utils.args import ImageConditioningInput
 from ltx_pipelines.utils.media_io import decode_image, load_image_conditioning, resize_aspect_ratio_preserving
 from ltx_pipelines.utils.types import (
     DenoisingFunc,
@@ -46,7 +45,7 @@ def cleanup_memory() -> None:
 
 
 def image_conditionings_by_replacing_latent(
-    images: list[tuple[str, int, float]],
+    images: list[ImageConditioningInput],
     height: int,
     width: int,
     video_encoder: VideoEncoder,
@@ -54,20 +53,21 @@ def image_conditionings_by_replacing_latent(
     device: torch.device,
 ) -> list[ConditioningItem]:
     conditionings = []
-    for image_path, frame_idx, strength in images:
+    for img in images:
         image = load_image_conditioning(
-            image_path=image_path,
+            image_path=img.path,
             height=height,
             width=width,
             dtype=dtype,
             device=device,
+            crf=img.crf,
         )
         encoded_image = video_encoder(image)
         conditionings.append(
             VideoConditionByLatentIndex(
                 latent=encoded_image,
-                strength=strength,
-                latent_idx=frame_idx,
+                strength=img.strength,
+                latent_idx=img.frame_idx,
             )
         )
 
@@ -75,7 +75,7 @@ def image_conditionings_by_replacing_latent(
 
 
 def image_conditionings_by_adding_guiding_latent(
-    images: list[tuple[str, int, float]],
+    images: list[ImageConditioningInput],
     height: int,
     width: int,
     video_encoder: VideoEncoder,
@@ -83,131 +83,22 @@ def image_conditionings_by_adding_guiding_latent(
     device: torch.device,
 ) -> list[ConditioningItem]:
     conditionings = []
-    for image_path, frame_idx, strength in images:
+    for img in images:
         image = load_image_conditioning(
-            image_path=image_path,
+            image_path=img.path,
             height=height,
             width=width,
             dtype=dtype,
             device=device,
+            crf=img.crf,
         )
         encoded_image = video_encoder(image)
         conditionings.append(
-            VideoConditionByKeyframeIndex(keyframes=encoded_image, frame_idx=frame_idx, strength=strength)
+            VideoConditionByKeyframeIndex(keyframes=encoded_image, frame_idx=img.frame_idx, strength=img.strength)
         )
     return conditionings
 
 
-def euler_denoising_loop(
-    sigmas: torch.Tensor,
-    video_state: LatentState,
-    audio_state: LatentState,
-    stepper: DiffusionStepProtocol,
-    denoise_fn: DenoisingFunc,
-) -> tuple[LatentState, LatentState]:
-    """
-    Perform the joint audio-video denoising loop over a diffusion schedule.
-    This function iterates over all but the final value in ``sigmas`` and, at
-    each diffusion step, calls ``denoise_fn`` to obtain denoised video and
-    audio latents. The denoised latents are post-processed with their
-    respective denoise masks and clean latents, then passed to ``stepper`` to
-    advance the noisy latents one step along the diffusion schedule.
-    ### Parameters
-    sigmas:
-        A 1D tensor of noise levels (diffusion sigmas) defining the sampling
-        schedule. All steps except the last element are iterated over.
-    video_state:
-        The current video :class:`LatentState`, containing the noisy latent,
-        its clean reference latent, and the denoising mask.
-    audio_state:
-        The current audio :class:`LatentState`, analogous to ``video_state``
-        but for the audio modality.
-    stepper:
-        An implementation of :class:`DiffusionStepProtocol` that updates a
-        latent given the current latent, its denoised estimate, the full
-        ``sigmas`` schedule, and the current step index.
-    denoise_fn:
-        A callable implementing :class:`DenoisingFunc`. It is invoked as
-        ``denoise_fn(video_state, audio_state, sigmas, step_index)`` and must
-        return a tuple ``(denoised_video, denoised_audio)``, where each element
-        is a tensor with the same shape as the corresponding latent.
-    ### Returns
-    tuple[LatentState, LatentState]
-        A pair ``(video_state, audio_state)`` containing the final video and
-        audio latent states after completing the denoising loop.
-    """
-    for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
-        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
-
-        denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
-        denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
-
-        video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
-        audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
-
-    return (video_state, audio_state)
-
-
-def gradient_estimating_euler_denoising_loop(
-    sigmas: torch.Tensor,
-    video_state: LatentState,
-    audio_state: LatentState,
-    stepper: DiffusionStepProtocol,
-    denoise_fn: DenoisingFunc,
-    ge_gamma: float = 2.0,
-) -> tuple[LatentState, LatentState]:
-    """
-    Perform the joint audio-video denoising loop using gradient-estimation sampling.
-    This function is similar to :func:`euler_denoising_loop`, but applies
-    gradient estimation to improve the denoised estimates by tracking velocity
-    changes across steps. See the referenced function for detailed parameter
-    documentation.
-    ### Parameters
-    ge_gamma:
-        Gradient estimation coefficient controlling the velocity correction term.
-        Default is 2.0. Paper: https://openreview.net/pdf?id=o2ND9v0CeK
-    sigmas, video_state, audio_state, stepper, denoise_fn:
-        See :func:`euler_denoising_loop` for parameter descriptions.
-    ### Returns
-    tuple[LatentState, LatentState]
-        See :func:`euler_denoising_loop` for return value description.
-    """
-
-    previous_audio_velocity = None
-    previous_video_velocity = None
-
-    def update_velocity_and_sample(
-        noisy_sample: torch.Tensor, denoised_sample: torch.Tensor, sigma: float, previous_velocity: torch.Tensor | None
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        current_velocity = to_velocity(noisy_sample, sigma, denoised_sample)
-        if previous_velocity is not None:
-            delta_v = current_velocity - previous_velocity
-            total_velocity = ge_gamma * delta_v + previous_velocity
-            denoised_sample = to_denoised(noisy_sample, total_velocity, sigma)
-        return current_velocity, denoised_sample
-
-    for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
-        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
-
-        denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
-        denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
-
-        if sigmas[step_idx + 1] == 0:
-            return replace(video_state, latent=denoised_video), replace(audio_state, latent=denoised_audio)
-
-        previous_video_velocity, denoised_video = update_velocity_and_sample(
-            video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
-        )
-        previous_audio_velocity, denoised_audio = update_velocity_and_sample(
-            audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
-        )
-
-        video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
-        audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
-
-    return (video_state, audio_state)
-
-
 def noise_video_state(
     output_shape: VideoPixelShape,
     noiser: Noiser,
@@ -313,7 +204,10 @@ def post_process_latent(denoised: torch.Tensor, denoise_mask: torch.Tensor, clea
 
 
 def modality_from_latent_state(
-    state: LatentState, context: torch.Tensor, sigma: float | torch.Tensor, enabled: bool = True
+    state: LatentState,
+    context: torch.Tensor,
+    sigma: torch.Tensor,
+    enabled: bool = True,
 ) -> Modality:
     """Create a Modality from a latent state.
     Constructs a Modality object with the latent state's data, timesteps derived
@@ -322,10 +216,12 @@ def modality_from_latent_state(
     return Modality(
         enabled=enabled,
         latent=state.latent,
+        sigma=sigma,
         timesteps=timesteps_from_mask(state.denoise_mask, sigma),
         positions=state.positions,
         context=context,
         context_mask=None,
+        attention_mask=state.attention_mask,
     )
 
 
@@ -389,10 +285,10 @@ def multi_modal_guider_denoising_func(
     v_context: torch.Tensor,
     a_context: torch.Tensor,
     transformer: X0Model,
+    *,
+    last_denoised_video: torch.Tensor | None = None,
+    last_denoised_audio: torch.Tensor | None = None,
 ) -> DenoisingFunc:
-    last_denoised_video = None
-    last_denoised_audio = None
-
     def guider_denoising_step(
         video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
     ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -490,6 +386,43 @@ def multi_modal_guider_denoising_func(
     return guider_denoising_step
 
 
+def multi_modal_guider_factory_denoising_func(
+    video_guider_factory: MultiModalGuiderFactory,
+    audio_guider_factory: MultiModalGuiderFactory | None,
+    v_context: torch.Tensor,
+    a_context: torch.Tensor,
+    transformer: X0Model,
+) -> DenoisingFunc:
+    """Resolve guiders per step via factory.build_from_sigma, then multi_modal_guider_denoising_func."""
+    last_denoised_video: torch.Tensor | None = None
+    last_denoised_audio: torch.Tensor | None = None
+    sigma_vals_cached: list[float] | None = None
+
+    def guider_denoising_step(
+        video_state: LatentState, audio_state: LatentState, sigmas: torch.Tensor, step_index: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        nonlocal last_denoised_video, last_denoised_audio, sigma_vals_cached
+        if sigma_vals_cached is None:
+            sigma_vals_cached = sigmas.detach().cpu().tolist()
+        sigma_val = sigma_vals_cached[step_index]
+        video_guider = video_guider_factory.build_from_sigma(sigma_val)
+        audio_guider = (audio_guider_factory or video_guider_factory).build_from_sigma(sigma_val)
+        denoise_fn = multi_modal_guider_denoising_func(
+            video_guider,
+            audio_guider,
+            v_context,
+            a_context,
+            transformer,
+            last_denoised_video=last_denoised_video,
+            last_denoised_audio=last_denoised_audio,
+        )
+        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_index)
+        last_denoised_video, last_denoised_audio = denoised_video, denoised_audio
+        return denoised_video, denoised_audio
+
+    return guider_denoising_step
+
+
 def denoise_audio_video(  # noqa: PLR0913
     output_shape: VideoPixelShape,
     conditionings: list[ConditioningItem],
@@ -540,6 +473,57 @@ def denoise_audio_video(  # noqa: PLR0913
     return video_state, audio_state
 
 
+def denoise_video_only(  # noqa: PLR0913
+    output_shape: VideoPixelShape,
+    conditionings: list[ConditioningItem],
+    noiser: Noiser,
+    sigmas: torch.Tensor,
+    stepper: DiffusionStepProtocol,
+    denoising_loop_fn: DenoisingLoopFunc,
+    components: PipelineComponents,
+    dtype: torch.dtype,
+    device: torch.device,
+    noise_scale: float = 1.0,
+    initial_video_latent: torch.Tensor | None = None,
+    initial_audio_latent: torch.Tensor | None = None,
+) -> LatentState:
+    video_state, video_tools = noise_video_state(
+        output_shape=output_shape,
+        noiser=noiser,
+        conditionings=conditionings,
+        components=components,
+        dtype=dtype,
+        device=device,
+        noise_scale=noise_scale,
+        initial_latent=initial_video_latent,
+    )
+
+    audio_state, _ = noise_audio_state(
+        output_shape=output_shape,
+        noiser=noiser,
+        conditionings=[],
+        components=components,
+        dtype=dtype,
+        device=device,
+        noise_scale=0.0,
+        initial_latent=initial_audio_latent,
+    )
+
+    audio_state = replace(audio_state, denoise_mask=torch.zeros_like(audio_state.denoise_mask))
+
+    video_state, audio_state = denoising_loop_fn(
+        sigmas,
+        video_state,
+        audio_state,
+        stepper,
+    )
+
+    video_state = video_tools.clear_conditioning(video_state)
+    video_state = video_tools.unpatchify(video_state)
+
+    return video_state
+
+
 _UNICODE_REPLACEMENTS = str.maketrans("\u2018\u2019\u201c\u201d\u2014\u2013\u00a0\u2032\u2212", "''\"\"-- '-")
 
 
@@ -555,7 +539,7 @@ def clean_response(text: str) -> str:
 
 
 def generate_enhanced_prompt(
-    text_encoder: GemmaTextEncoderModelBase,
+    text_encoder: GemmaTextEncoder,
     prompt: str,
     image_path: str | None = None,
     image_long_side: int = 896,

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

@@ -12,6 +12,7 @@ from PIL import Image
 from torch._prims_common import DeviceLikeType
 from tqdm import tqdm
 
+from ltx_core.types import Audio
 from ltx_pipelines.utils.constants import DEFAULT_IMAGE_CRF
 
 logger = logging.getLogger(__name__)
@@ -79,14 +80,19 @@ def normalize_latent(latent: torch.Tensor, device: torch.device, dtype: torch.dt
 
 
 def load_image_conditioning(
-    image_path: str, height: int, width: int, dtype: torch.dtype, device: torch.device
+    image_path: str,
+    height: int,
+    width: int,
+    dtype: torch.dtype,
+    device: torch.device,
+    crf: int = DEFAULT_IMAGE_CRF,
 ) -> torch.Tensor:
     """
     Loads an image from a path and preprocesses it for conditioning.
     Note: The image is resized to the nearest multiple of 2 for compatibility with video codecs.
     """
     image = decode_image(image_path=image_path)
-    image = preprocess(image=image)
+    image = preprocess(image=image, crf=crf)
     image = torch.tensor(image, dtype=torch.float32, device=device)
     image = resize_and_center_crop(image, height, width)
     image = normalize_latent(image, device, dtype)
@@ -115,9 +121,8 @@ def decode_image(image_path: str) -> np.ndarray:
     return np_array
 
 
-def _write_audio(
-    container: av.container.Container, audio_stream: av.audio.AudioStream, samples: torch.Tensor, audio_sample_rate: int
-) -> None:
+def _write_audio(container: av.container.Container, audio_stream: av.audio.AudioStream, audio: Audio) -> None:
+    samples = audio.waveform
     if samples.ndim == 1:
         samples = samples[:, None]
 
@@ -137,7 +142,7 @@ def _write_audio(
         format="s16",
         layout="stereo",
     )
-    frame_in.sample_rate = audio_sample_rate
+    frame_in.sample_rate = audio.sampling_rate
 
     _resample_audio(container, audio_stream, frame_in)
 
@@ -185,8 +190,7 @@ def _resample_audio(
 def encode_video(
     video: torch.Tensor | Iterator[torch.Tensor],
     fps: int,
-    audio: torch.Tensor | None,
-    audio_sample_rate: int | None,
+    audio: Audio | None,
     output_path: str,
     video_chunks_number: int,
 ) -> None:
@@ -204,10 +208,7 @@ def encode_video(
     stream.pix_fmt = "yuv420p"
 
     if audio is not None:
-        if audio_sample_rate is None:
-            raise ValueError("audio_sample_rate is required when audio is provided")
-
-        audio_stream = _prepare_audio_stream(container, audio_sample_rate)
+        audio_stream = _prepare_audio_stream(container, audio.sampling_rate)
 
     def all_tiles(
         first_chunk: torch.Tensor, tiles_generator: Generator[tuple[torch.Tensor, int], None, None]
@@ -227,27 +228,114 @@ def encode_video(
         container.mux(packet)
 
     if audio is not None:
-        _write_audio(container, audio_stream, audio, audio_sample_rate)
+        _write_audio(container, audio_stream, audio)
 
     container.close()
     logger.info(f"Video saved to {output_path}")
 
 
-def decode_audio_from_file(path: str, device: torch.device) -> torch.Tensor | None:
+_INT_FORMAT_MAX: dict[str, float] = {
+    "u8": 128.0,
+    "u8p": 128.0,
+    "s16": 32768.0,
+    "s16p": 32768.0,
+    "s32": 2147483648.0,
+    "s32p": 2147483648.0,
+}
+
+
+def _audio_frame_to_float(frame: av.AudioFrame) -> np.ndarray:
+    """Convert an audio frame to a float32 ndarray with values in [-1, 1] and shape (channels, samples)."""
+    fmt = frame.format.name
+    arr = frame.to_ndarray().astype(np.float32)
+    if fmt in _INT_FORMAT_MAX:
+        arr = arr / _INT_FORMAT_MAX[fmt]
+    if not frame.format.is_planar:
+        # Interleaved formats have shape (1, samples * channels) — reshape to (channels, samples).
+        channels = len(frame.layout.channels)
+        arr = arr.reshape(-1, channels).T
+    return arr
+
+
+def get_videostream_metadata(path: str) -> tuple[float, int, int, int]:
+    """Read video stream metadata: (fps, num_frames, width, height).
+    If frame count is missing in the container, decodes the stream to count frames.
+    """
     container = av.open(path)
     try:
-        audio = []
-        audio_stream = next(s for s in container.streams if s.type == "audio")
-        for frame in container.decode(audio_stream):
-            audio.append(torch.tensor(frame.to_ndarray(), dtype=torch.float32, device=device).unsqueeze(0))
-        container.close()
-        audio = torch.cat(audio)
-    except StopIteration:
-        audio = None
+        video_stream = next(s for s in container.streams if s.type == "video")
+        fps = float(video_stream.average_rate)
+        num_frames = video_stream.frames or 0
+        if num_frames == 0:
+            num_frames = sum(1 for _ in container.decode(video_stream))
+        width = video_stream.codec_context.width
+        height = video_stream.codec_context.height
+        return fps, num_frames, width, height
     finally:
         container.close()
 
-    return audio
+
+def decode_audio_from_file(
+    path: str, device: torch.device, start_time: float = 0.0, max_duration: float | None = None
+) -> Audio | None:
+    """Decodes audio from a file, optionally seeking to a start time and limiting duration.
+    Args:
+        path: Path to the audio/video file containing an audio stream.
+        device: Device to place the resulting tensor on.
+        start_time: Start time in seconds to begin reading audio from.
+        max_duration: Maximum audio duration in seconds. If None, reads to end of stream.
+    Returns:
+        An Audio object with waveform of shape (1, channels, samples), or None if no audio stream.
+    """
+    container = av.open(path)
+    try:
+        audio_stream = next(s for s in container.streams if s.type == "audio")
+    except StopIteration:
+        container.close()
+        return None
+
+    sample_rate = audio_stream.rate
+    start_pts = int(start_time / audio_stream.time_base)
+    end_time = start_time + max_duration if max_duration else audio_stream.duration * audio_stream.time_base
+    container.seek(start_pts, stream=audio_stream)
+
+    samples = []
+    first_frame_time = None
+    for frame in container.decode(audio=0):
+        if frame.pts is None:
+            continue
+        frame_time = float(frame.pts * audio_stream.time_base)
+        frame_end = frame_time + frame.samples / frame.sample_rate
+        if frame_end < start_time:
+            continue
+        if frame_time > end_time:
+            break
+        if first_frame_time is None:
+            first_frame_time = frame_time
+        samples.append(_audio_frame_to_float(frame))
+
+    container.close()
+
+    if not samples:
+        return None
+
+    audio = np.concatenate(samples, axis=-1)
+
+    # Trim samples that fall outside the requested [start_time, start_time + max_duration] window.
+    # Audio codecs decode in fixed-size frames whose boundaries may not align with the requested
+    # time range, so the first frame can start before start_time and the last frame can end after
+    # start_time + max_duration.
+    skip_samples = round((start_time - first_frame_time) * sample_rate)
+    if skip_samples > 0:
+        audio = audio[..., skip_samples:]
+
+    if max_duration is not None:
+        max_samples = round(max_duration * sample_rate)
+        audio = audio[..., :max_samples]
+
+    waveform = torch.from_numpy(audio).to(device).unsqueeze(0)
+
+    return Audio(waveform=waveform, sampling_rate=sample_rate)
 
 
 def decode_video_from_file(path: str, frame_cap: int, device: DeviceLikeType) -> Generator[torch.Tensor]:

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

@@ -8,9 +8,12 @@ from ltx_core.loader.registry import DummyRegistry, Registry
 from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder
 from ltx_core.model.audio_vae import (
     AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
     VOCODER_COMFY_KEYS_FILTER,
     AudioDecoder,
     AudioDecoderConfigurator,
+    AudioEncoder,
+    AudioEncoderConfigurator,
     Vocoder,
     VocoderConfigurator,
 )
@@ -31,11 +34,11 @@ from ltx_core.model.video_vae import (
 from ltx_core.quantization import QuantizationPolicy
 from ltx_core.text_encoders.gemma import (
     AV_GEMMA_TEXT_ENCODER_KEY_OPS,
-    AVGemmaTextEncoderModel,
-    AVGemmaTextEncoderModelConfigurator,
+    GEMMA_MODEL_OPS,
+    GemmaTextEncoder,
+    GemmaTextEncoderConfigurator,
     module_ops_from_gemma_root,
 )
-from ltx_core.text_encoders.gemma.encoders.av_encoder import GEMMA_MODEL_OPS
 from ltx_core.utils import find_matching_file
 
 
@@ -131,6 +134,13 @@ class ModelLedger:
                 registry=self.registry,
             )
 
+            self.audio_encoder_builder = Builder[AudioEncoder](
+                model_path=self.checkpoint_path,
+                model_class_configurator=AudioEncoderConfigurator,
+                model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+                registry=self.registry,
+            )
+
             self.audio_decoder_builder = Builder(
                 model_path=self.checkpoint_path,
                 model_class_configurator=AudioDecoderConfigurator,
@@ -152,7 +162,7 @@ class ModelLedger:
 
                 self.text_encoder_builder = Builder(
                     model_path=(str(self.checkpoint_path), *weight_paths),
-                    model_class_configurator=AVGemmaTextEncoderModelConfigurator,
+                    model_class_configurator=GemmaTextEncoderConfigurator,
                     model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
                     registry=self.registry,
                     module_ops=(GEMMA_MODEL_OPS, *module_ops),
@@ -225,7 +235,7 @@ class ModelLedger:
 
         return self.vae_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
 
-    def text_encoder(self) -> AVGemmaTextEncoderModel:
+    def text_encoder(self) -> GemmaTextEncoder:
         if not hasattr(self, "text_encoder_builder"):
             raise ValueError(
                 "Text encoder not initialized. Please provide a checkpoint path and gemma root path to the "
@@ -234,6 +244,14 @@ class ModelLedger:
 
         return self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
 
+    def audio_encoder(self) -> AudioEncoder:
+        if not hasattr(self, "audio_encoder_builder"):
+            raise ValueError(
+                "Audio encoder not initialized. Please provide a checkpoint path to the ModelLedger constructor."
+            )
+
+        return self.audio_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
+
     def audio_decoder(self) -> AudioDecoder:
         if not hasattr(self, "audio_decoder_builder"):
             raise ValueError(

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

@@ -0,0 +1,62 @@
+import math
+
+
+def phi(j: int, neg_h: float) -> float:
+    """
+    Compute φⱼ(z) where z = -h (negative step size in log-space)
+    φ₁(z) = (e^z - 1) / z
+    φ₂(z) = (e^z - 1 - z) / z²
+    φⱼ(z) = (e^z - Σₖ₌₀^(j-1) zᵏ/k!) / zʲ
+    These functions naturally appear when solving:
+        dx/dt = A*x + g(x,t)  (linear drift + nonlinear part)
+    """
+    if abs(neg_h) < 1e-10:
+        # Taylor series for small h to avoid division by zero
+        # φⱼ(0) = 1/j!
+        return 1.0 / math.factorial(j)
+
+    # Compute the "remainder" sum: Σₖ₌₀^(j-1) z^k/k!
+    remainder = sum(neg_h**k / math.factorial(k) for k in range(j))
+
+    # φⱼ(z) = (e^z - remainder) / z^j
+    return (math.exp(neg_h) - remainder) / (neg_h**j)
+
+
+def get_res2s_coefficients(h: float, phi_cache: dict, c2: float = 0.5) -> tuple[float, float, float]:
+    """
+    Compute res_2s Runge-Kutta coefficients for a given step size.
+    Args:
+        h: Step size in log-space = log(sigma / sigma_next)
+        phi_cache: Dictionary to cache phi function results. Cache key: (j, neg_h)
+        c2: Substep position (default 0.5 = midpoint)
+    Returns:
+        a21: Coefficient for computing intermediate x
+        b1, b2: Coefficients for final combination
+    """
+
+    def get_phi(j: int, neg_h: float) -> float:
+        """Get phi value with caching."""
+        cache_key = (j, neg_h)
+        if cache_key in phi_cache:
+            return phi_cache[cache_key]
+        result = phi(j, neg_h)
+        phi_cache[cache_key] = result
+        return result
+
+    # Substep coefficient: how much of ε₁ to use for intermediate point
+    # a21 = c2 * φ₁(-h*c2)
+    neg_h_c2 = -h * c2
+    phi_1_c2 = get_phi(1, neg_h_c2)
+    a21 = c2 * phi_1_c2
+
+    # Final combination weights
+    # b2 = φ₂(-h) / c2
+    neg_h_full = -h
+    phi_2_full = get_phi(2, neg_h_full)
+    b2 = phi_2_full / c2
+
+    # b1 = φ₁(-h) - b2
+    phi_1_full = get_phi(1, neg_h_full)
+    b1 = phi_1_full - b2
+
+    return a21, b1, b2

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

@@ -0,0 +1,363 @@
+import logging
+from dataclasses import replace
+from functools import partial
+from typing import Callable
+
+import torch
+from tqdm import tqdm
+
+from ltx_core.components.diffusion_steps import Res2sDiffusionStep
+from ltx_core.components.protocols import DiffusionStepProtocol
+from ltx_core.utils import to_denoised, to_velocity
+from ltx_pipelines.utils.helpers import post_process_latent, timesteps_from_mask
+from ltx_pipelines.utils.res2s import get_res2s_coefficients
+from ltx_pipelines.utils.types import DenoisingFunc, LatentState
+
+logger = logging.getLogger(__name__)
+
+
+def euler_denoising_loop(
+    sigmas: torch.Tensor,
+    video_state: LatentState,
+    audio_state: LatentState,
+    stepper: DiffusionStepProtocol,
+    denoise_fn: DenoisingFunc,
+) -> tuple[LatentState, LatentState]:
+    """
+    Perform the joint audio-video denoising loop over a diffusion schedule.
+    This function iterates over all but the final value in ``sigmas`` and, at
+    each diffusion step, calls ``denoise_fn`` to obtain denoised video and
+    audio latents. The denoised latents are post-processed with their
+    respective denoise masks and clean latents, then passed to ``stepper`` to
+    advance the noisy latents one step along the diffusion schedule.
+    ### Parameters
+    sigmas:
+        A 1D tensor of noise levels (diffusion sigmas) defining the sampling
+        schedule. All steps except the last element are iterated over.
+    video_state:
+        The current video :class:`LatentState`, containing the noisy latent,
+        its clean reference latent, and the denoising mask.
+    audio_state:
+        The current audio :class:`LatentState`, analogous to ``video_state``
+        but for the audio modality.
+    stepper:
+        An implementation of :class:`DiffusionStepProtocol` that updates a
+        latent given the current latent, its denoised estimate, the full
+        ``sigmas`` schedule, and the current step index.
+    denoise_fn:
+        A callable implementing :class:`DenoisingFunc`. It is invoked as
+        ``denoise_fn(video_state, audio_state, sigmas, step_index)`` and must
+        return a tuple ``(denoised_video, denoised_audio)``, where each element
+        is a tensor with the same shape as the corresponding latent.
+    ### Returns
+    tuple[LatentState, LatentState]
+        A pair ``(video_state, audio_state)`` containing the final video and
+        audio latent states after completing the denoising loop.
+    """
+    for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
+        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
+
+        denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
+        denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
+
+        video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
+        audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
+
+    return (video_state, audio_state)
+
+
+def gradient_estimating_euler_denoising_loop(
+    sigmas: torch.Tensor,
+    video_state: LatentState,
+    audio_state: LatentState,
+    stepper: DiffusionStepProtocol,
+    denoise_fn: DenoisingFunc,
+    ge_gamma: float = 2.0,
+) -> tuple[LatentState, LatentState]:
+    """
+    Perform the joint audio-video denoising loop using gradient-estimation sampling.
+    This function is similar to :func:`euler_denoising_loop`, but applies
+    gradient estimation to improve the denoised estimates by tracking velocity
+    changes across steps. See the referenced function for detailed parameter
+    documentation.
+    ### Parameters
+    ge_gamma:
+        Gradient estimation coefficient controlling the velocity correction term.
+        Default is 2.0. Paper: https://openreview.net/pdf?id=o2ND9v0CeK
+    sigmas, video_state, audio_state, stepper, denoise_fn:
+        See :func:`euler_denoising_loop` for parameter descriptions.
+    ### Returns
+    tuple[LatentState, LatentState]
+        See :func:`euler_denoising_loop` for return value description.
+    """
+
+    previous_audio_velocity = None
+    previous_video_velocity = None
+
+    def update_velocity_and_sample(
+        noisy_sample: torch.Tensor, denoised_sample: torch.Tensor, sigma: float, previous_velocity: torch.Tensor | None
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        current_velocity = to_velocity(noisy_sample, sigma, denoised_sample)
+        if previous_velocity is not None:
+            delta_v = current_velocity - previous_velocity
+            total_velocity = ge_gamma * delta_v + previous_velocity
+            denoised_sample = to_denoised(noisy_sample, total_velocity, sigma)
+        return current_velocity, denoised_sample
+
+    for step_idx, _ in enumerate(tqdm(sigmas[:-1])):
+        denoised_video, denoised_audio = denoise_fn(video_state, audio_state, sigmas, step_idx)
+
+        denoised_video = post_process_latent(denoised_video, video_state.denoise_mask, video_state.clean_latent)
+        denoised_audio = post_process_latent(denoised_audio, audio_state.denoise_mask, audio_state.clean_latent)
+
+        if sigmas[step_idx + 1] == 0:
+            return replace(video_state, latent=denoised_video), replace(audio_state, latent=denoised_audio)
+
+        previous_video_velocity, denoised_video = update_velocity_and_sample(
+            video_state.latent, denoised_video, sigmas[step_idx], previous_video_velocity
+        )
+        previous_audio_velocity, denoised_audio = update_velocity_and_sample(
+            audio_state.latent, denoised_audio, sigmas[step_idx], previous_audio_velocity
+        )
+
+        video_state = replace(video_state, latent=stepper.step(video_state.latent, denoised_video, sigmas, step_idx))
+        audio_state = replace(audio_state, latent=stepper.step(audio_state.latent, denoised_audio, sigmas, step_idx))
+
+    return (video_state, audio_state)
+
+
+def _channelwise_normalize(x: torch.Tensor) -> torch.Tensor:
+    return x.sub_(x.mean(dim=(-2, -1), keepdim=True)).div_(x.std(dim=(-2, -1), keepdim=True))
+
+
+def _get_new_noise(x: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
+    noise = torch.randn(x.shape, generator=generator, dtype=torch.float64, device=generator.device)
+    noise = (noise - noise.mean()) / noise.std()
+    return _channelwise_normalize(noise)
+
+
+def _inject_sde_noise(
+    state: LatentState,
+    sample: torch.Tensor,
+    denoised_sample: torch.Tensor,
+    step_noise_generator: torch.Generator,
+    new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor],
+    stepper: DiffusionStepProtocol,
+    sigmas: torch.Tensor,
+    step_idx: int,
+    legacy_mode: bool = False,
+) -> torch.Tensor:
+    sigmas_copy = sigmas.clone()
+    new_noise = new_noise_fn(state.latent, step_noise_generator)
+    if not legacy_mode:
+        timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx].double())
+        next_timesteps = timesteps_from_mask(state.denoise_mask.double(), sigmas_copy[step_idx + 1].double())
+        sigmas = torch.stack([timesteps, next_timesteps])
+        step_idx = 0
+    x_next = stepper.step(
+        sample=sample,
+        denoised_sample=denoised_sample,
+        sigmas=sigmas,
+        step_index=step_idx,
+        noise=new_noise,
+    )
+
+    if legacy_mode:
+        x_next = post_process_latent(x_next, state.denoise_mask, state.clean_latent)
+
+    return x_next
+
+
+def res2s_audio_video_denoising_loop(  # noqa: PLR0913,PLR0915
+    sigmas: torch.Tensor,
+    video_state: LatentState,
+    audio_state: LatentState,
+    stepper: DiffusionStepProtocol,
+    denoise_fn: DenoisingFunc,
+    noise_seed: int = -1,
+    noise_seed_substep: int | None = None,
+    bongmath: bool = True,
+    bongmath_max_iter: int = 100,
+    new_noise_fn: Callable[[torch.Tensor, torch.Generator], torch.Tensor] = _get_new_noise,
+    model_dtype: torch.dtype = torch.bfloat16,
+    legacy_mode: bool = True,
+) -> tuple[LatentState, LatentState]:
+    """
+    Joint audio-video denoising loop using the res_2s second-order sampler.
+    Iterates over the diffusion schedule with a two-stage Runge-Kutta step:
+    evaluates the denoiser at the current point and at a midpoint (with SDE
+    noise), then combines both with RK coefficients. Supports anchor-point
+    refinement (bong iteration) and optional SDE noise injection. Requires
+    :class:`Res2sDiffusionStep` as ``stepper``.
+    ### Parameters
+    sigmas:
+        A 1D tensor of noise levels defining the sampling schedule.
+    video_state:
+        Current video :class:`LatentState` (noisy latent, clean reference, mask).
+    audio_state:
+        Current audio :class:`LatentState`, same structure as ``video_state``.
+    stepper:
+        Must be an instance of :class:`Res2sDiffusionStep`; performs SDE step
+        with noise injection.
+    denoise_fn:
+        Callable ``(video_state, audio_state, sigmas, step_index)`` returning
+        ``(denoised_video, denoised_audio)``.
+    noise_seed:
+        Seed for step-level SDE noise; substep seed defaults to ``noise_seed + 10000``.
+    noise_seed_substep:
+        Optional seed for substep SDE noise; if None, derived from ``noise_seed``.
+    bongmath:
+        Whether to run iterative anchor refinement (bong iteration) when step size is small.
+    bongmath_max_iter:
+        Max iterations for bong refinement when enabled.
+    new_noise_fn:
+        Callable ``(latent, generator) -> noise`` for SDE injection; default
+        uses normalized channel-wise Gaussian noise.
+    model_dtype:
+        Dtype for latent state updates (e.g. bfloat16).
+    ### Returns
+    tuple[LatentState, LatentState]
+        Final ``(video_state, audio_state)`` after the denoising loop.
+    """
+    # Initialize noise generators with different seeds
+    if noise_seed_substep is None:
+        noise_seed_substep = noise_seed + 10000  # Offset to ensure different seeds
+    step_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed)
+    substep_noise_generator = torch.Generator(device=video_state.latent.device).manual_seed(noise_seed_substep)
+    sde_noise_injecting_fn = partial(
+        _inject_sde_noise, stepper=stepper, new_noise_fn=new_noise_fn, legacy_mode=legacy_mode
+    )
+    step_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=step_noise_generator)
+    substep_noise_injecting_fn = partial(sde_noise_injecting_fn, step_noise_generator=substep_noise_generator)
+
+    if not isinstance(stepper, Res2sDiffusionStep):
+        raise ValueError("stepper must be an instance of Res2sDiffusionStep")
+
+    n_full_steps = len(sigmas) - 1
+    # inject minimal sigma value to avoid division by zero
+    if sigmas[-1] == 0:
+        sigmas = torch.cat([sigmas[:-1], torch.tensor([0.0011, 0.0], device=sigmas.device)], dim=0)
+    # Compute step sizes in hyperbolic space
+    hs = -torch.log(sigmas[1:].double().cpu() / (sigmas[:-1].double().cpu()))
+
+    # Initialize phi cache for reuse across loop iterations
+    # Cache key: (j, neg_h) where j is phi order and neg_h is negative step value
+    phi_cache = {}
+    c2 = 0.5  # Midpoint for res_2s
+
+    # Progress bar shows only full two-stage steps; final (sigma_next==0) step is done silently
+
+    for step_idx in tqdm(range(n_full_steps)):
+        sigma = sigmas[step_idx].double()
+        sigma_next = sigmas[step_idx + 1].double()
+
+        # Initialize anchor point
+        x_anchor_video = video_state.latent.clone().double()
+        x_anchor_audio = audio_state.latent.clone().double()
+
+        # ====================================================================
+        # STAGE 1: Evaluate at current point
+        # ====================================================================
+        denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, step_idx)
+        denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
+        denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
+
+        h = hs[step_idx].item()
+
+        # Compute RK coefficients (pass phi_cache for caching)
+        a21, b1, b2 = get_res2s_coefficients(h, phi_cache, c2)
+
+        # Compute substep sigma, sqrt is a hardcode for c2 = 0.5
+        sub_sigma = torch.sqrt(sigma * sigma_next)
+
+        # ====================================================================
+        # Compute substep x using RK coefficient a21
+        # ====================================================================
+        eps_1_video = denoised_video_1.double() - x_anchor_video
+        eps_1_audio = denoised_audio_1.double() - x_anchor_audio
+
+        x_mid_video = x_anchor_video.double() + h * a21 * eps_1_video
+        x_mid_audio = x_anchor_audio.double() + h * a21 * eps_1_audio
+
+        # ====================================================================
+        # SDE noise injection at substep
+        # ====================================================================
+        x_mid_video = substep_noise_injecting_fn(
+            state=video_state,
+            sample=x_anchor_video,
+            denoised_sample=x_mid_video,
+            sigmas=torch.stack([sigma, sub_sigma]),
+            step_idx=0,
+        )
+        x_mid_audio = substep_noise_injecting_fn(
+            state=audio_state,
+            sample=x_anchor_audio,
+            denoised_sample=x_mid_audio,
+            sigmas=torch.stack([sigma, sub_sigma]),
+            step_idx=0,
+        )
+        # ====================================================================
+        # ITERATIVE REFINEMENT (Bong Iteration) - Stabilize anchor point
+        # ====================================================================
+        if bongmath and h < 0.5 and sigma > 0.03:
+            for _ in range(bongmath_max_iter):
+                x_anchor_video = x_mid_video - h * a21 * eps_1_video
+                eps_1_video = denoised_video_1.double() - x_anchor_video
+                x_anchor_audio = x_mid_audio - h * a21 * eps_1_audio
+                eps_1_audio = denoised_audio_1.double() - x_anchor_audio
+
+        # ====================================================================
+        # STAGE 2: Evaluate at substep point (WITH NOISE)
+        # ====================================================================
+        mid_video_state = replace(video_state, latent=x_mid_video.to(model_dtype))
+        mid_audio_state = replace(audio_state, latent=x_mid_audio.to(model_dtype))
+
+        denoised_video_2, denoised_audio_2 = denoise_fn(
+            video_state=mid_video_state,
+            audio_state=mid_audio_state,
+            sigmas=torch.stack([sub_sigma]).to(sigmas.device),
+            step_index=0,
+        )
+        denoised_video_2 = post_process_latent(denoised_video_2, video_state.denoise_mask, video_state.clean_latent)
+        denoised_audio_2 = post_process_latent(denoised_audio_2, audio_state.denoise_mask, audio_state.clean_latent)
+
+        # ====================================================================
+        # FINAL COMBINATION: Compute x_next using RK coefficients
+        # ====================================================================
+        eps_2_video = denoised_video_2.double() - x_anchor_video
+        eps_2_audio = denoised_audio_2.double() - x_anchor_audio
+
+        x_next_video = x_anchor_video + h * (b1 * eps_1_video + b2 * eps_2_video)
+        x_next_audio = x_anchor_audio + h * (b1 * eps_1_audio + b2 * eps_2_audio)
+
+        # ====================================================================
+        # SDE NOISE INJECTION AT STEP LEVEL
+        # ====================================================================
+        x_next_video = step_noise_injecting_fn(
+            state=video_state,
+            sample=x_anchor_video,
+            denoised_sample=x_next_video,
+            sigmas=sigmas,
+            step_idx=step_idx,
+        )
+        x_next_audio = step_noise_injecting_fn(
+            state=audio_state,
+            sample=x_anchor_audio,
+            denoised_sample=x_next_audio,
+            sigmas=sigmas,
+            step_idx=step_idx,
+        )
+
+        # Update states
+        video_state = replace(video_state, latent=x_next_video.to(model_dtype))
+        audio_state = replace(audio_state, latent=x_next_audio.to(model_dtype))
+
+    # Final step if we need to fully remove the noise
+    if sigmas[-1] == 0:
+        denoised_video_1, denoised_audio_1 = denoise_fn(video_state, audio_state, sigmas, n_full_steps)
+        denoised_video_1 = post_process_latent(denoised_video_1, video_state.denoise_mask, video_state.clean_latent)
+        denoised_audio_1 = post_process_latent(denoised_audio_1, audio_state.denoise_mask, audio_state.clean_latent)
+        video_state = replace(video_state, latent=denoised_video_1.to(model_dtype))
+        audio_state = replace(audio_state, latent=denoised_audio_1.to(model_dtype))
+
+    return video_state, audio_state

packages/ltx-trainer/AGENTS.md

@@ -4,19 +4,28 @@ This file provides guidance to AI coding assistants (Claude, Cursor, etc.) when
 
 ## Project Overview
 
-**LTX-2 Trainer** is a training toolkit for fine-tuning the Lightricks LTX-2 audio-video generation model. It supports:
+**LTX Trainer** is a training toolkit for fine-tuning the Lightricks LTX audio-video generation models. It supports:
 
 - **LoRA training** - Efficient fine-tuning with adapters
 - **Full fine-tuning** - Complete model training
 - **Audio-video training** - Joint audio and video generation
 - **IC-LoRA training** - In-context control adapters for video-to-video transformations
 
+**Supported model versions:**
+
+- **LTX-2** (19B, initial audio-video model)
+- **LTX-2.3** (20B, improved text conditioning and audio quality)
+
+Version detection is fully automatic — ltx-core reads the checkpoint config and selects the correct architecture
+components. The trainer does not need version-specific code paths.
+
 **Key Dependencies:**
 
-- **[`ltx-core`](../ltx-core/)** - Core model implementations (transformer, VAE, text encoder)
+- **[`ltx-core`](../ltx-core/)** - Core model implementations (transformer, VAE, text encoder, scheduler)
 - **[`ltx-pipelines`](../ltx-pipelines/)** - Inference pipeline components
 
-> **Important:** This trainer only supports **LTX-2** (the audio-video model). The older LTXV models are not supported.
+> **Important:** This trainer only supports **LTX-2 and later** (audio-video models). The older LTXV (video-only) models
+> are not supported.
 
 ## Architecture Overview
 
@@ -24,36 +33,45 @@ This file provides guidance to AI coding assistants (Claude, Cursor, etc.) when
 
 ```
 packages/ltx-trainer/
-├── src/ltx_trainer/           # Main training module
-│   ├── config.py              # Pydantic configuration models
-│   ├── trainer.py             # Main training orchestration with Accelerate
-│   ├── model_loader.py        # Model loading using ltx-core
-│   ├── validation_sampler.py  # Inference for validation samples
-│   ├── datasets.py            # PrecomputedDataset for latent-based training
-│   ├── training_strategies/   # Strategy pattern for different training modes
-│   │   ├── __init__.py        # Factory function: get_training_strategy()
-│   │   ├── base_strategy.py   # TrainingStrategy ABC, ModelInputs, TrainingStrategyConfigBase
-│   │   ├── text_to_video.py   # TextToVideoStrategy, TextToVideoConfig
-│   │   └── video_to_video.py  # VideoToVideoStrategy, VideoToVideoConfig
-│   ├── timestep_samplers.py   # Flow matching timestep sampling
-│   ├── captioning.py          # Video captioning utilities
-│   ├── video_utils.py         # Video processing utilities
-│   └── hf_hub_utils.py        # HuggingFace Hub integration
-├── scripts/                   # User-facing CLI tools
-│   ├── train.py               # Main training script
-│   ├── process_dataset.py     # Dataset preprocessing
-│   ├── process_videos.py      # Video latent encoding
-│   ├── process_captions.py    # Text embedding computation
-│   ├── caption_videos.py      # Automatic video captioning
-│   ├── decode_latents.py      # Latent decoding for debugging
-│   ├── inference.py           # Inference with trained models
-│   ├── compute_reference.py   # Generate IC-LoRA reference videos
-│   └── split_scenes.py        # Scene detection and splitting
-├── configs/                   # Example training configurations
-│   ├── ltx2_av_lora.yaml      # Audio-video LoRA training
-│   ├── ltx2_v2v_ic_lora.yaml  # IC-LoRA video-to-video
-│   └── accelerate/            # Accelerate configs for distributed training
-└── docs/                      # Documentation
+├── src/ltx_trainer/              # Main training module
+│   ├── __init__.py               # Logger setup, path config
+│   ├── config.py                 # Pydantic configuration models
+│   ├── config_display.py         # Config pretty-printing
+│   ├── trainer.py                # Main training orchestration with Accelerate
+│   ├── model_loader.py           # Model loading using ltx-core
+│   ├── validation_sampler.py     # Inference for validation samples
+│   ├── datasets.py               # PrecomputedDataset, DummyDataset
+│   ├── training_strategies/      # Strategy pattern for different training modes
+│   │   ├── __init__.py           # Factory function: get_training_strategy()
+│   │   ├── base_strategy.py      # TrainingStrategy ABC, ModelInputs, TrainingStrategyConfigBase
+│   │   ├── text_to_video.py      # TextToVideoStrategy, TextToVideoConfig
+│   │   └── video_to_video.py     # VideoToVideoStrategy, VideoToVideoConfig
+│   ├── timestep_samplers.py      # Flow matching timestep sampling
+│   ├── gemma_8bit.py             # 8-bit Gemma text encoder loading (bitsandbytes)
+│   ├── quantization.py           # Transformer INT8/INT4/FP8 quantization
+│   ├── captioning.py             # Video captioning utilities
+│   ├── video_utils.py            # Video I/O and processing
+│   ├── gpu_utils.py              # GPU memory helpers
+│   ├── hf_hub_utils.py           # HuggingFace Hub integration
+│   ├── progress.py               # Training progress display
+│   └── utils.py                  # Image I/O helpers
+├── scripts/                      # User-facing CLI tools
+│   ├── train.py                  # Main training script
+│   ├── process_dataset.py        # Dataset preprocessing (latents + captions)
+│   ├── process_videos.py         # Video latent encoding
+│   ├── process_captions.py       # Text embedding computation
+│   ├── caption_videos.py         # Automatic video captioning
+│   ├── decode_latents.py         # Latent decoding for debugging
+│   ├── inference.py              # Inference with trained models
+│   ├── compute_reference.py      # Generate IC-LoRA reference videos
+│   └── split_scenes.py           # Scene detection and splitting
+├── configs/                      # Example training configurations
+│   ├── ltx2_av_lora.yaml         # Audio-video LoRA training
+│   ├── ltx2_av_lora_low_vram.yaml
+│   ├── ltx2_v2v_ic_lora.yaml     # IC-LoRA video-to-video
+│   └── accelerate/               # FSDP, DDP configs
+├── tests/                        # Pytest tests
+└── docs/                         # Documentation
 ```
 
 ### Key Architectural Patterns
@@ -61,35 +79,42 @@ packages/ltx-trainer/
 **Model Loading:**
 
 - `ltx_trainer.model_loader` provides component loaders using `ltx-core`
-- Individual loaders: `load_transformer()`, `load_video_vae_encoder()`, `load_video_vae_decoder()`, `load_text_encoder()`, etc.
+- Individual loaders: `load_transformer()`, `load_video_vae_encoder()`, `load_video_vae_decoder()`,
+  `load_text_encoder()`, etc.
 - Combined loader: `load_model()` returns `LtxModelComponents` dataclass
 - Uses `SingleGPUModelBuilder` from ltx-core internally
+- 8-bit text encoder loading via `gemma_8bit.py` (bitsandbytes)
 
 **Training Flow:**
 
 1. Configuration loaded via Pydantic models in `config.py`
-2. `Trainer` class orchestrates the training loop
-3. Training strategies (`TextToVideoStrategy`, `VideoToVideoStrategy`) prepare inputs and compute loss
-4. Accelerate handles distributed training and device placement
-5. Data flows as precomputed latents through `PrecomputedDataset`
+2. `LtxvTrainer` class orchestrates the training loop
+3. Text encoder loaded on GPU → validation embeddings cached → heavy components unloaded (only `embeddings_processor`
+   kept)
+4. Each training step: embedding connectors applied → strategy prepares `ModelInputs` → transformer forward pass →
+   strategy computes loss
+5. Training strategies (`TextToVideoStrategy`, `VideoToVideoStrategy`) handle mode-specific logic
+6. Accelerate handles distributed training, mixed precision, and device placement
+7. Data flows as precomputed latents through `PrecomputedDataset`
 
 **Model Interface (Modality-based):**
 
 ```python
 from ltx_core.model.transformer.modality import Modality
 
-# Create modality objects for video and audio
 video = Modality(
     enabled=True,
-    latent=video_latents,      # [B, seq_len, 128]
-    timesteps=video_timesteps,  # [B, seq_len] per-token
-    positions=video_positions,  # [B, 3, seq_len, 2]
-    context=video_embeds,
-    context_mask=None,
+    latent=video_latents,  # [B, seq_len, 128] patchified latent tokens
+    sigma=sigma,  # [B,] current noise level (per-batch)
+    timesteps=video_timesteps,  # [B, seq_len] per-token timestep embeddings
+    positions=video_positions,  # [B, 3, seq_len, 2] positional coordinates
+    context=video_embeds,  # text conditioning embeddings
+    context_mask=None,  # optional attention mask for text context
 )
 audio = Modality(
     enabled=True,
     latent=audio_latents,
+    sigma=sigma,
     timesteps=audio_timesteps,
     positions=audio_positions,  # [B, 1, seq_len, 2]
     context=audio_embeds,
@@ -102,14 +127,91 @@ video_pred, audio_pred = model(video=video, audio=audio, perturbations=None)
 
 > **Note:** `Modality` is immutable (frozen dataclass). Use `dataclasses.replace()` to modify.
 
+**`sigma` vs `timesteps`:** These serve different roles. `timesteps` is per-token (e.g. `sigma * denoise_mask` —
+conditioning tokens get 0, noisy tokens get sigma). `sigma` is per-batch and is used for prompt AdaLN conditioning (
+LTX-2.3) and cross-modality (video↔audio) attention conditioning (both versions).
+
 **Configuration System:**
 
 - All config in `src/ltx_trainer/config.py`
 - Main class: `LtxTrainerConfig`
 - Training strategy configs: `TextToVideoConfig`, `VideoToVideoConfig`
 - Uses Pydantic field validators and model validators
+- Config uses `extra="forbid"` — unknown fields cause validation errors
 - Config files in `configs/` directory
 
+## LTX-2 vs LTX-2.3: Differences
+
+Both model versions share the same latent space interface (see [Latent Space Constants](#latent-space-constants)).
+The differences lie in how text conditioning and audio generation work. Version detection is automatic via checkpoint
+config — the trainer uses a unified API.
+
+| Component             | LTX-2 (19B)                                                                     | LTX-2.3 (20B)                                                                                       |
+|-----------------------|---------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------|
+| Feature extractor     | `FeatureExtractorV1`: single `aggregate_embed`, same output for video and audio | `FeatureExtractorV2`: separate `video_aggregate_embed` + `audio_aggregate_embed`, per-token RMSNorm |
+| Caption projection    | Inside the transformer (`caption_projection`)                                   | Inside the feature extractor (before connector)                                                     |
+| Embeddings connectors | Same dimensions for video and audio                                             | Separate dimensions (`AudioEmbeddings1DConnectorConfigurator`)                                      |
+| Prompt AdaLN          | Not present (`cross_attention_adaln=False`)                                     | Active — modulates cross-attention to text using `sigma`                                            |
+| Vocoder               | HiFi-GAN (`Vocoder`)                                                            | BigVGAN v2 + bandwidth extension (`VocoderWithBWE`)                                                 |
+
+**How version detection works in ltx-core:**
+
+- **Feature extractor:** `_create_feature_extractor()` checks for V2 config keys (`caption_proj_before_connector`,
+  etc.). Present → V2; absent → V1.
+- **Vocoder:** `VocoderConfigurator` checks for `config["vocoder"]["bwe"]`. Present → `VocoderWithBWE`; absent →
+  `Vocoder`.
+- **Transformer:** `_build_caption_projections()` checks `caption_proj_before_connector`. True (V2) → no caption
+  projection in transformer; False (V1) → caption projection created in transformer.
+- **Embeddings connectors:** `AudioEmbeddings1DConnectorConfigurator` reads `audio_connector_*` keys, falling back to
+  video connector keys for V1 backward compatibility.
+
+## Text Encoder Pipeline
+
+The `GemmaTextEncoder` implements a 3-block pipeline:
+
+1. **Block 1 — Gemma LLM:** Tokenizes text → runs through Gemma → extracts hidden states
+2. **Block 2 — Feature extractor:** Hidden states → normalized features (V1: single stream duplicated for video/audio;
+   V2: separate video and audio projections)
+3. **Block 3 — Embeddings processor:** Features → embeddings connectors → final context embeddings for the transformer
+
+**Precomputed embeddings (offline):** `process_captions.py` runs Blocks 1+2 via `text_encoder.precompute()` and saves
+the results. Block 3 (connectors) is applied during training via
+`text_encoder.embeddings_processor.create_embeddings()`.
+
+**Precomputed embeddings formats:**
+
+- **New format** (from `precompute()`): saves `video_prompt_embeds`, `audio_prompt_embeds` (optional),
+  `prompt_attention_mask`
+- **Legacy format** (from old `_preprocess_text()`): saves `prompt_embeds`, `prompt_attention_mask`
+
+The trainer handles both formats in `_training_step()`: if `video_prompt_embeds` is present, it uses the new format;
+otherwise, it duplicates `prompt_embeds` for both modalities (mirroring V1 behavior).
+
+**After caching validation embeddings**, the trainer unloads heavy components to free VRAM:
+
+```python
+self._text_encoder.model = None
+self._text_encoder.tokenizer = None
+self._text_encoder.feature_extractor = None
+# Only embeddings_processor (connectors) remains — used during training
+```
+
+## Latent Space Constants
+
+These values are shared across all supported model versions:
+
+| Constant                     | Value                            | Where used                                                |
+|------------------------------|----------------------------------|-----------------------------------------------------------|
+| Video latent channels        | 128                              | VAE encoder/decoder, patchifier, `VideoLatentShape`       |
+| Spatial compression          | 32× (H and W)                    | `SpatioTemporalScaleFactors.default()`, config validators |
+| Temporal compression         | 8×                               | `SpatioTemporalScaleFactors.default()`, config validators |
+| Frame constraint             | `frames % 8 == 1`                | Config validators, validation sampler                     |
+| Resolution constraint        | Width and height divisible by 32 | Config validators, validation sampler                     |
+| Audio latent channels        | 8                                | `AudioLatentShape`, audio patchifier                      |
+| Audio mel bins               | 16                               | `AudioLatentShape`, audio patchifier                      |
+| Patchified token dim (video) | 128 (`128 × 1 × 1 × 1`)          | Transformer `in_channels`                                 |
+| Patchified token dim (audio) | 128 (`8 × 16`)                   | Transformer `audio_in_channels`                           |
+
 ## Development Commands
 
 ### Setup and Installation
@@ -180,25 +282,34 @@ uv run accelerate launch scripts/train.py configs/ltx2_av_lora.yaml
 **`src/ltx_trainer/config.py`** - Master config definitions
 
 Key classes:
+
 - `LtxTrainerConfig` - Main configuration container
-- `ModelConfig` - Model paths and training mode
-- `TrainingStrategyConfig` - Union of `TextToVideoConfig` | `VideoToVideoConfig`
-- `LoraConfig` - LoRA hyperparameters
-- `OptimizationConfig` - Learning rate, batch size, etc.
-- `ValidationConfig` - Validation settings
-- `WandbConfig` - W&B logging settings
+- `ModelConfig` - Model paths, training mode (`lora` | `full`), checkpoint loading
+- `TrainingStrategyConfig` - Union of `TextToVideoConfig` | `VideoToVideoConfig` (discriminated by `name`)
+- `LoraConfig` - Rank, alpha, dropout, target modules
+- `OptimizationConfig` - Learning rate, batch size, gradient accumulation, scheduler, gradient checkpointing
+- `AccelerationConfig` - Mixed precision, quantization, 8-bit text encoder
+- `DataConfig` - Preprocessed data root, dataloader workers
+- `ValidationConfig` - Prompts, video dimensions, CFG/STG guidance, audio generation, inference steps
+- `CheckpointsConfig` - Save interval, retention, precision
+- `FlowMatchingConfig` - Timestep sampling mode and parameters
+- `HubConfig` - HuggingFace Hub push settings
+- `WandbConfig` - Weights & Biases logging
 
 **⚠️ When modifying config.py:**
+
 1. Update ALL config files in `configs/`
 2. Update `docs/configuration-reference.md`
 3. Test that all configs remain valid
 
 ### Training Core
 
-**`src/ltx_trainer/trainer.py`** - Main training loop
+**`src/ltx_trainer/trainer.py`** - Main training loop (`LtxvTrainer`)
 
 - Implements distributed training with Accelerate
 - Handles mixed precision, gradient accumulation, checkpointing
+- `_training_step()` applies embedding connectors then delegates to strategy
+- `_load_text_encoder_and_cache_embeddings()` caches validation embeddings and unloads heavy components
 - Uses training strategies for mode-specific logic
 
 **`src/ltx_trainer/training_strategies/`** - Strategy pattern
@@ -208,35 +319,55 @@ Key classes:
 - `video_to_video.py`: IC-LoRA video-to-video transformations
 
 Key methods each strategy implements:
+
 - `get_data_sources()` - Required data directories
-- `prepare_training_inputs()` - Convert batch to `ModelInputs`
-- `compute_loss()` - Calculate training loss
+- `prepare_training_inputs()` - Convert batch to `ModelInputs` with `Modality` objects
+- `compute_loss()` - Calculate training loss (velocity prediction, MSE with masking)
 - `requires_audio` property - Whether audio components needed
 
 **`src/ltx_trainer/model_loader.py`** - Model loading
 
 Component loaders:
+
 - `load_transformer()` → `LTXModel`
-- `load_video_vae_encoder()` → `VideoVAEEncoder`
-- `load_video_vae_decoder()` → `VideoVAEDecoder`
-- `load_audio_vae_decoder()` → `AudioVAEDecoder`
-- `load_vocoder()` → `Vocoder`
-- `load_text_encoder()` → `AVGemmaTextEncoderModel`
+- `load_video_vae_encoder()` → `VideoEncoder`
+- `load_video_vae_decoder()` → `VideoDecoder`
+- `load_audio_vae_decoder()` → `AudioDecoder`
+- `load_vocoder()` → `Vocoder` or `VocoderWithBWE` (auto-detected)
+- `load_text_encoder()` → `GemmaTextEncoder` (unified, handles V1/V2 automatically)
 - `load_model()` → `LtxModelComponents` (convenience wrapper)
 
 **`src/ltx_trainer/validation_sampler.py`** - Inference for validation
 
 Uses ltx-core components for denoising:
+
 - `LTX2Scheduler` for sigma scheduling
 - `EulerDiffusionStep` for diffusion steps
 - `CFGGuider` for classifier-free guidance
+- `STGGuider` for spatio-temporal guidance
+
+**`src/ltx_trainer/timestep_samplers.py`** - Flow matching timestep sampling
+
+- `UniformTimestepSampler` - Uniform sampling in `[min, max]`
+- `ShiftedLogitNormalTimestepSampler` - Stretched shifted logit-normal distribution with:
+    - Shift determined by sequence length (more noise at higher token counts)
+    - Percentile stretching for better `[0, 1]` coverage
+    - Uniform fallback (10% of samples) to prevent distribution collapse
+    - Reflection around `eps` for numerical stability near zero
+
+**`src/ltx_trainer/gemma_8bit.py`** - 8-bit text encoder loading
+
+Bypasses ltx-core's standard loading path to enable bitsandbytes 8-bit quantization of the Gemma backbone. Manually
+constructs the `GemmaTextEncoder` with quantized model, feature extractor, and embeddings processor.
 
 ### Data
 
 **`src/ltx_trainer/datasets.py`** - Dataset handling
 
-- `PrecomputedDataset` loads pre-computed VAE latents
-- Supports video latents, audio latents, text embeddings, reference latents
+- `PrecomputedDataset` loads pre-computed VAE latents and text embeddings
+- Supports video latents, audio latents, text embeddings, reference latents (for IC-LoRA)
+- Handles legacy patchified format `[seq_len, C]` → automatically unpatchifies to `[C, F, H, W]`
+- `DummyDataset` for benchmarking and minimal testing
 
 ## Common Development Tasks
 
@@ -263,23 +394,40 @@ Uses ltx-core components for denoising:
 from dataclasses import replace
 from ltx_core.model.transformer.modality import Modality
 
-# Create modality
+# Create modality — all fields except enabled and masks are required
 video = Modality(
     enabled=True,
-    latent=latents,
-    timesteps=timesteps,
-    positions=positions,
-    context=context,
+    latent=latents,  # [B, seq_len, 128]
+    sigma=sigma,  # [B,] — the per-batch noise level
+    timesteps=timesteps,  # [B, seq_len] — per-token (sigma * denoise_mask)
+    positions=positions,  # [B, 3, seq_len, 2]
+    context=context,  # text embeddings from embeddings_processor
     context_mask=None,
 )
 
-# Update (immutable - must use replace)
-video = replace(video, latent=new_latent, timesteps=new_timesteps)
+# Update (immutable — must use replace)
+video = replace(video, latent=new_latent, sigma=new_sigma, timesteps=new_timesteps)
 
 # Disable a modality
 audio = replace(audio, enabled=False)
 ```
 
+### Working with the Text Encoder
+
+```python
+# Full forward pass (used for validation — runs all 3 blocks)
+video_embeds, audio_embeds, attention_mask = text_encoder(prompt)
+
+# Precompute features (used in process_captions.py — runs blocks 1+2 only)
+video_features, audio_features, attention_mask = text_encoder.precompute(prompt, padding_side="left")
+
+# Apply connectors during training (block 3 only)
+additive_mask = text_encoder._convert_to_additive_mask(attention_mask, video_features.dtype)
+video_embeds, audio_embeds, binary_mask = text_encoder.embeddings_processor.create_embeddings(
+    video_features, audio_features, additive_mask
+)
+```
+
 ## Debugging Tips
 
 **Training Issues:**
@@ -293,18 +441,27 @@ audio = replace(audio, enabled=False)
 - Ensure `model_path` points to a local `.safetensors` file
 - Ensure `text_encoder_path` points to a Gemma model directory
 - URLs are NOT supported for model paths
+- For 8-bit loading: ensure `bitsandbytes` is installed
 
 **Configuration:**
 
 - Validation errors: Check validators in `config.py`
-- Unknown fields: Config uses `extra="forbid"` - all fields must be defined
+- Unknown fields: Config uses `extra="forbid"` — all fields must be defined
 - Strategy validation: IC-LoRA requires `reference_videos` in validation config
+- Video-to-video strategy requires `training_mode: "lora"`
+
+**Precomputed Data:**
+
+- Legacy data (`prompt_embeds`) works via backward-compat in `_training_step()`
+- New data (`video_prompt_embeds` + `audio_prompt_embeds`) is the expected format
+- Latents must be in `[C, F, H, W]` format (legacy `[seq_len, C]` is auto-converted)
 
 ## Key Constraints
 
-### LTX-2 Frame Requirements
+### Frame Requirements
 
 Frames must satisfy `frames % 8 == 1`:
+
 - ✅ Valid: 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, 97, 121
 - ❌ Invalid: 24, 32, 48, 64, 100
 
@@ -321,7 +478,7 @@ Width and height must be divisible by 32.
 ### Platform Requirements
 
 - Linux required (uses `triton` which is Linux-only)
-- CUDA GPU with 24GB+ VRAM recommended
+- CUDA GPU with 24GB+ VRAM recommended (80GB+ for full fine-tuning)
 
 ## Reference: ltx-core Key Components
 
@@ -329,24 +486,41 @@ Width and height must be divisible by 32.
 packages/ltx-core/src/ltx_core/
 ├── model/
 │   ├── transformer/
-│   │   ├── model.py              # LTXModel
-│   │   ├── modality.py           # Modality dataclass
-│   │   └── transformer.py        # BasicAVTransformerBlock
+│   │   ├── model.py                # LTXModel (diffusion transformer)
+│   │   ├── modality.py             # Modality dataclass
+│   │   ├── transformer.py          # BasicAVTransformerBlock
+│   │   ├── transformer_args.py     # TransformerArgsPreprocessor (sigma → prompt AdaLN)
+│   │   ├── model_configurator.py   # LTXModelConfigurator (version-aware)
+│   │   └── timestep_embedding.py   # Timestep/sigma embedding
 │   ├── video_vae/
-│   │   └── video_vae.py          # Encoder, Decoder
+│   │   ├── video_vae.py            # VideoEncoder, VideoDecoder
+│   │   └── model_configurator.py   # VideoEncoderConfigurator, VideoDecoderConfigurator
 │   ├── audio_vae/
-│   │   ├── audio_vae.py          # Decoder
-│   │   └── vocoder.py            # Vocoder
-│   └── clip/gemma/
-│       └── encoders/av_encoder.py  # AVGemmaTextEncoderModel
-├── pipeline/
-│   ├── components/
-│   │   ├── schedulers.py         # LTX2Scheduler
-│   │   ├── diffusion_steps.py    # EulerDiffusionStep
-│   │   ├── guiders.py            # CFGGuider
-│   │   └── patchifiers.py        # VideoLatentPatchifier, AudioPatchifier
-│   └── conditioning/             # VideoLatentTools, AudioLatentTools
-└── loader/
-    ├── single_gpu_model_builder.py  # SingleGPUModelBuilder
-    └── sd_ops.py                    # Key remapping (SDOps)
+│   │   ├── audio_vae.py            # AudioEncoder, AudioDecoder
+│   │   └── vocoder.py              # Vocoder, VocoderWithBWE (output_sampling_rate)
+│   └── common/                     # Shared model components
+├── text_encoders/gemma/
+│   ├── __init__.py                 # Exports: GemmaTextEncoder, GemmaTextEncoderConfigurator,
+│   │                               #   AV_GEMMA_TEXT_ENCODER_KEY_OPS, GEMMA_MODEL_OPS,
+│   │                               #   module_ops_from_gemma_root
+│   ├── encoders/
+│   │   ├── base_encoder.py         # GemmaTextEncoder (unified 3-block pipeline)
+│   │   └── encoder_configurator.py # GemmaTextEncoderConfigurator, _create_feature_extractor
+│   ├── feature_extractor.py        # FeatureExtractorV1 (19B), FeatureExtractorV2 (20B)
+│   ├── embeddings_connector.py     # Embeddings1DConnector, Embeddings1DConnectorConfigurator,
+│   │                               #   AudioEmbeddings1DConnectorConfigurator
+│   ├── embeddings_processor.py     # EmbeddingsProcessor (wraps video + audio connectors)
+│   └── tokenizer.py               # LTXVGemmaTokenizer
+├── components/
+│   ├── schedulers.py               # LTX2Scheduler
+│   ├── diffusion_steps.py          # EulerDiffusionStep
+│   ├── guiders.py                  # CFGGuider, STGGuider
+│   └── patchifiers.py              # VideoLatentPatchifier, AudioPatchifier
+├── conditioning/                   # ConditioningItem, mask_utils, types
+├── tools.py                        # VideoLatentTools, AudioLatentTools
+├── loader/
+│   ├── single_gpu_model_builder.py # SingleGPUModelBuilder
+│   ├── sft_loader.py              # SafetensorsModelStateDictLoader
+│   └── sd_ops.py                  # Key remapping (SDOps)
+└── types.py                       # SpatioTemporalScaleFactors, VideoLatentShape, AudioLatentShape
 ```

packages/ltx-trainer/docs/custom-training-strategies.md

@@ -282,6 +282,7 @@ class InpaintingStrategy(TrainingStrategy):
         video_modality = Modality(
             enabled=True,
             latent=noisy_latents,
+            sigma=sigmas,
             timesteps=timesteps,
             positions=positions,
             context=video_prompt_embeds,

packages/ltx-trainer/scripts/decode_latents.py

@@ -287,7 +287,7 @@ class LatentsDecoder:
 
         # Save as WAV
         output_path = output_dir / f"{latent_file.stem}.wav"
-        sample_rate = self.vocoder.output_sample_rate
+        sample_rate = self.vocoder.output_sampling_rate
         torchaudio.save(str(output_path), waveform[0].cpu(), sample_rate)
 
 

packages/ltx-trainer/scripts/inference.py

@@ -410,7 +410,7 @@ def main() -> None:  # noqa: PLR0912, PLR0915
     # Get audio sample rate from vocoder if audio was generated
     audio_sample_rate = None
     if audio is not None and components.vocoder is not None:
-        audio_sample_rate = components.vocoder.output_sample_rate
+        audio_sample_rate = components.vocoder.output_sampling_rate
 
     save_video(
         video_tensor=video,

packages/ltx-trainer/scripts/process_captions.py

@@ -303,14 +303,13 @@ def compute_captions_embeddings(  # noqa: PLR0913
     ) as progress:
         task = progress.add_task("Processing captions", total=len(dataloader))
         for batch in dataloader:
-            # Encode prompts using _preprocess_text (returns embeddings before connector)
-            # This is what we want to save - the connector is applied during training
+            # Encode prompts using precompute() (returns video/audio features before connector)
+            # The connector is applied during training via embeddings_processor
             with torch.inference_mode():
-                # TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once:
-                #   prompt_embeds, prompt_attention_mask = text_encoder._preprocess_text(batch["prompt"]) # noqa: ERA001
+                # TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once.
                 # For now, process one at a time:
                 for i in range(len(batch["prompt"])):
-                    prompt_embeds, prompt_attention_mask = text_encoder._preprocess_text(
+                    video_prompt_embeds, audio_prompt_embeds, prompt_attention_mask = text_encoder.precompute(
                         batch["prompt"][i], padding_side="left"
                     )
 
@@ -321,9 +320,11 @@ def compute_captions_embeddings(  # noqa: PLR0913
                     output_dir_path.mkdir(parents=True, exist_ok=True)
 
                     embedding_data = {
-                        "prompt_embeds": prompt_embeds[0].cpu().contiguous(),
+                        "video_prompt_embeds": video_prompt_embeds[0].cpu().contiguous(),
                         "prompt_attention_mask": prompt_attention_mask[0].cpu().contiguous(),
                     }
+                    if audio_prompt_embeds is not None:
+                        embedding_data["audio_prompt_embeds"] = audio_prompt_embeds[0].cpu().contiguous()
 
                     output_file = output_path / output_rel_path
                     torch.save(embedding_data, output_file)

packages/ltx-trainer/scripts/process_videos.py

@@ -41,6 +41,7 @@ from torchvision.transforms.functional import crop, resize, to_tensor
 from transformers.utils.logging import disable_progress_bar
 
 from ltx_core.model.audio_vae import AudioProcessor
+from ltx_core.types import Audio
 from ltx_trainer import logger
 from ltx_trainer.model_loader import load_audio_vae_encoder, load_video_vae_encoder
 from ltx_trainer.utils import open_image_as_srgb
@@ -503,7 +504,7 @@ def compute_latents(  # noqa: PLR0913, PLR0915
             )
             # Create audio processor for waveform-to-spectrogram conversion
             audio_processor = AudioProcessor(
-                sample_rate=audio_vae_encoder.sample_rate,
+                target_sample_rate=audio_vae_encoder.sample_rate,
                 mel_bins=audio_vae_encoder.mel_bins,
                 mel_hop_length=audio_vae_encoder.mel_hop_length,
                 n_fft=audio_vae_encoder.n_fft,
@@ -567,10 +568,10 @@ def compute_latents(  # noqa: PLR0913, PLR0915
                     if audio_batch is not None:
                         # Extract the i-th item from batched audio data
                         # DataLoader collates [channels, samples] -> [batch, channels, samples]
-                        audio_data = {
-                            "waveform": audio_batch["waveform"][i],
-                            "sample_rate": audio_batch["sample_rate"][i].item(),
-                        }
+                        audio_data = Audio(
+                            waveform=audio_batch["waveform"][i],
+                            sampling_rate=audio_batch["sample_rate"][i].item(),
+                        )
 
                         # Encode audio
                         with torch.inference_mode():
@@ -822,13 +823,13 @@ def tiled_encode_video(  # noqa: PLR0912, PLR0915
 def encode_audio(
     audio_vae_encoder: torch.nn.Module,
     audio_processor: torch.nn.Module,
-    audio_data: dict[str, torch.Tensor | int],
+    audio: Audio,
 ) -> dict[str, torch.Tensor | int | float]:
     """Encode audio waveform into latent representation.
     Args:
         audio_vae_encoder: Audio VAE encoder model from ltx-core
         audio_processor: AudioProcessor for waveform-to-spectrogram conversion
-        audio_data: Dict with {"waveform": Tensor[channels, samples], "sample_rate": int}
+        audio: Audio container with waveform tensor and sampling rate.
     Returns:
         Dict containing audio latents and shape information:
         {
@@ -841,18 +842,17 @@ def encode_audio(
     device = next(audio_vae_encoder.parameters()).device
     dtype = next(audio_vae_encoder.parameters()).dtype
 
-    waveform = audio_data["waveform"].to(device=device, dtype=dtype)
-    sample_rate = audio_data["sample_rate"]
+    waveform = audio.waveform.to(device=device, dtype=dtype)
 
     # Add batch dimension if needed: [channels, samples] -> [batch, channels, samples]
     if waveform.dim() == 2:
         waveform = waveform.unsqueeze(0)
 
     # Calculate duration
-    duration = waveform.shape[-1] / sample_rate
+    duration = waveform.shape[-1] / audio.sampling_rate
 
     # Convert waveform to mel spectrogram using AudioProcessor
-    mel_spectrogram = audio_processor.waveform_to_mel(waveform, waveform_sample_rate=sample_rate)
+    mel_spectrogram = audio_processor.waveform_to_mel(Audio(waveform=waveform, sampling_rate=audio.sampling_rate))
     mel_spectrogram = mel_spectrogram.to(dtype=dtype)
 
     # Encode mel spectrogram to latents

packages/ltx-trainer/src/ltx_trainer/datasets.py

@@ -67,14 +67,18 @@ class DummyDataset(Dataset):
                 "fps": self.fps,
             },
             "text_conditions": {
-                "prompt_embeds": torch.randn(
+                "video_prompt_embeds": torch.randn(
                     self.prompt_sequence_length,
                     self.prompt_embed_dim,
-                ),  # random text embeddings
+                ),
+                "audio_prompt_embeds": torch.randn(
+                    self.prompt_sequence_length,
+                    self.prompt_embed_dim,
+                ),
                 "prompt_attention_mask": torch.ones(
                     self.prompt_sequence_length,
                     dtype=torch.bool,
-                ),  # random attention mask
+                ),
             },
         }
 

packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py

@@ -18,28 +18,26 @@ import logging
 from collections.abc import Generator
 from contextlib import contextmanager
 from pathlib import Path
-from typing import TYPE_CHECKING
 
 import torch
 
 from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
-from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnectorConfigurator
-from ltx_core.text_encoders.gemma.encoders.av_encoder import (
-    AV_GEMMA_TEXT_ENCODER_KEY_OPS,
-    AVGemmaTextEncoderModel,
+from ltx_core.text_encoders.gemma import AV_GEMMA_TEXT_ENCODER_KEY_OPS
+from ltx_core.text_encoders.gemma.embeddings_connector import (
+    AudioEmbeddings1DConnectorConfigurator,
+    Embeddings1DConnectorConfigurator,
 )
-from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
+from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_core.text_encoders.gemma.encoders.encoder_configurator import _create_feature_extractor
 from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
 
-if TYPE_CHECKING:
-    from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
-
 
 def load_8bit_gemma(
     checkpoint_path: str | Path,
     gemma_model_path: str | Path,
     dtype: torch.dtype = torch.bfloat16,
-) -> "AVGemmaTextEncoderModel":
+) -> GemmaTextEncoder:
     """Load the Gemma text encoder in 8-bit precision using bitsandbytes.
     This function bypasses ltx-core's standard loading path to enable 8-bit quantization
     via the bitsandbytes library. The Gemma model is loaded with load_in_8bit=True and
@@ -50,7 +48,7 @@ def load_8bit_gemma(
         gemma_model_path: Path to Gemma model directory
         dtype: Data type for non-quantized model weights (feature extractor, connectors)
     Returns:
-        Loaded AVGemmaTextEncoderModel with 8-bit quantized Gemma backbone
+        Loaded GemmaTextEncoder with 8-bit quantized Gemma backbone
     Raises:
         ImportError: If bitsandbytes is not installed
         FileNotFoundError: If required model files are not found
@@ -88,28 +86,35 @@ def load_8bit_gemma(
     def extract_state_dict(prefix: str) -> dict[str, torch.Tensor]:
         return {k.replace(prefix, ""): v for k, v in sd.sd.items() if k.startswith(prefix)}
 
-    # Create and load feature extractor
-    feature_extractor = GemmaFeaturesExtractorProjLinear()
-    feature_extractor.load_state_dict(extract_state_dict("feature_extractor_linear."))
-    feature_extractor = feature_extractor.to(device=gemma_model.device, dtype=dtype)
-
     # Create and load video embeddings connector
     embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
-    embeddings_connector.load_state_dict(extract_state_dict("embeddings_connector."))
+    embeddings_connector.load_state_dict(extract_state_dict("embeddings_processor.video_connector."))
     embeddings_connector = embeddings_connector.to(device=gemma_model.device, dtype=dtype)
 
     # Create and load audio embeddings connector
-    audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
-    audio_embeddings_connector.load_state_dict(extract_state_dict("audio_embeddings_connector."))
+    audio_embeddings_connector = AudioEmbeddings1DConnectorConfigurator.from_config(config)
+    audio_embeddings_connector.load_state_dict(extract_state_dict("embeddings_processor.audio_connector."))
     audio_embeddings_connector = audio_embeddings_connector.to(device=gemma_model.device, dtype=dtype)
 
-    # Construct the text encoder
-    text_encoder = AVGemmaTextEncoderModel(
-        feature_extractor_linear=feature_extractor,
-        embeddings_connector=embeddings_connector,
-        audio_embeddings_connector=audio_embeddings_connector,
+    # Create embeddings processor
+    embeddings_processor = EmbeddingsProcessor(
+        video_connector=embeddings_connector,
+        audio_connector=audio_embeddings_connector,
+    )
+
+    transformer_config = config.get("transformer", {})
+    feature_extractor = _create_feature_extractor(transformer_config)
+    feature_extractor.load_state_dict(
+        {k.removeprefix("feature_extractor."): v for k, v in sd.sd.items() if k.startswith("feature_extractor.")},
+    )
+    feature_extractor = feature_extractor.to(device=gemma_model.device, dtype=dtype)
+
+    text_encoder = GemmaTextEncoder(
+        feature_extractor=feature_extractor,
+        embeddings_processor=embeddings_processor,
         tokenizer=tokenizer,
         model=gemma_model,
+        dtype=dtype,
     )
 
     return text_encoder

packages/ltx-trainer/src/ltx_trainer/model_loader.py

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
     from ltx_core.model.audio_vae import AudioDecoder, AudioEncoder, Vocoder
     from ltx_core.model.transformer import LTXModel
     from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
-    from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
+    from ltx_core.text_encoders.gemma import GemmaTextEncoder
 
 
 def _to_torch_device(device: Device) -> torch.device:
@@ -192,7 +192,7 @@ def load_text_encoder(
     device: Device = "cpu",
     dtype: torch.dtype = torch.bfloat16,
     load_in_8bit: bool = False,
-) -> "AVGemmaTextEncoderModel":
+) -> "GemmaTextEncoder":
     """Load the Gemma text encoder.
     Args:
         checkpoint_path: Path to the LTX-2 safetensors checkpoint file
@@ -203,7 +203,7 @@ def load_text_encoder(
             When True, the model is loaded with device_map="auto" and the device argument
             is ignored for the Gemma backbone (feature extractor still uses dtype).
     Returns:
-        Loaded AVGemmaTextEncoderModel
+        Loaded GemmaTextEncoder (unified encoder handling V1/V2/V3)
     """
     if not Path(gemma_model_path).is_dir():
         raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
@@ -216,12 +216,12 @@ def load_text_encoder(
 
     # Standard loading path
     from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
-    from ltx_core.text_encoders.gemma.encoders.av_encoder import (
+    from ltx_core.text_encoders.gemma import (
         AV_GEMMA_TEXT_ENCODER_KEY_OPS,
         GEMMA_MODEL_OPS,
-        AVGemmaTextEncoderModelConfigurator,
+        GemmaTextEncoderConfigurator,
+        module_ops_from_gemma_root,
     )
-    from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
     from ltx_core.utils import find_matching_file
 
     torch_device = _to_torch_device(device)
@@ -231,7 +231,7 @@ def load_text_encoder(
 
     text_encoder = SingleGPUModelBuilder(
         model_path=(str(checkpoint_path), *gemma_weight_paths),
-        model_class_configurator=AVGemmaTextEncoderModelConfigurator,
+        model_class_configurator=GemmaTextEncoderConfigurator,
         model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
         module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
     ).build(device=torch_device, dtype=dtype)
@@ -253,7 +253,7 @@ class LtxModelComponents:
     video_vae_decoder: "VideoDecoder | None" = None
     audio_vae_decoder: "AudioDecoder | None" = None
     vocoder: "Vocoder | None" = None
-    text_encoder: "AVGemmaTextEncoderModel | None" = None
+    text_encoder: "GemmaTextEncoder | None" = None
     scheduler: "LTX2Scheduler | None" = None
 
 

packages/ltx-trainer/src/ltx_trainer/timestep_samplers.py

@@ -45,29 +45,61 @@ class UniformTimestepSampler(TimestepSampler):
         return self.sample(batch.shape[0], device=batch.device)
 
 
-class ShiftedLogitNormalTimestepSampler:
+class ShiftedLogitNormalTimestepSampler(TimestepSampler):
     """
-    Samples timesteps from a shifted logit-normal distribution,
+    Samples timesteps from a stretched shifted logit-normal distribution,
     where the shift is determined by the sequence length.
+    The stretching normalizes samples between percentile bounds to ensure
+    the distribution covers [0, 1] more evenly. A uniform fallback prevents
+    collapse at high token counts.
     """
 
-    def __init__(self, std: float = 1.0):
+    def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1):
         self.std = std
+        self.eps = eps
+        self.uniform_prob = uniform_prob
+        # Percentile values for stretching (scaled by std)
+        # 99.9th percentile of standard normal ≈ 3.0902
+        # 0.5th percentile of standard normal ≈ -2.5758
+        self.normal_999_percentile = 3.0902 * std
+        self.normal_005_percentile = -2.5758 * std
 
     def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:
-        """Sample timesteps for a batch from a shifted logit-normal distribution.
+        """Sample timesteps for a batch from a stretched shifted logit-normal distribution.
         Args:
             batch_size: Number of timesteps to sample
             seq_length: Length of the sequence being processed, used to determine the shift
             device: Device to place the samples on
         Returns:
-            Tensor of shape (batch_size,) containing timesteps sampled from a shifted
-            logit-normal distribution, where the shift is determined by seq_length
+            Tensor of shape (batch_size,) containing timesteps sampled from a stretched
+            shifted logit-normal distribution, where the shift is determined by seq_length
         """
-        shift = self._get_shift_for_sequence_length(seq_length)
-        normal_samples = torch.randn((batch_size,), device=device) * self.std + shift
-        timesteps = torch.sigmoid(normal_samples)
-        return timesteps
+        mu = self._get_shift_for_sequence_length(seq_length)
+
+        # Sample from shifted logit-normal
+        normal_samples = torch.randn((batch_size,), device=device) * self.std + mu
+        logitnormal_samples = torch.sigmoid(normal_samples)
+
+        # Compute percentile bounds for stretching
+        percentile_999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device))
+        percentile_005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device))
+
+        # Stretch to [0, 1] range by normalizing between percentiles
+        zero_terminal_raw = (logitnormal_samples - percentile_005) / (percentile_999 - percentile_005)
+
+        # Reflect small values around eps for numerical stability
+        stretched_logit = torch.where(
+            zero_terminal_raw >= self.eps,
+            zero_terminal_raw,
+            2 * self.eps - zero_terminal_raw,
+        )
+        stretched_logit = torch.clamp(stretched_logit, 0, 1)
+
+        # Mix with uniform samples (uniform_prob of the time)
+        uniform = (1 - self.eps) * torch.rand((batch_size,), device=device) + self.eps
+        prob = torch.rand((batch_size,), device=device)
+
+        return torch.where(prob > self.uniform_prob, stretched_logit, uniform)
 
     def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
         """Sample timesteps for a specific batch tensor.

packages/ltx-trainer/src/ltx_trainer/trainer.py

@@ -309,9 +309,22 @@ class LtxvTrainer:
         """Perform a single training step using the configured strategy."""
         # Apply embedding connectors to transform pre-computed text embeddings
         conditions = batch["conditions"]
-        video_embeds, audio_embeds, attention_mask = self._text_encoder._run_connectors(
-            conditions["prompt_embeds"], conditions["prompt_attention_mask"]
+
+        if "video_prompt_embeds" in conditions:
+            # New format: separate video/audio features from precompute()
+            video_features = conditions["video_prompt_embeds"]
+            audio_features = conditions.get("audio_prompt_embeds")
+        else:
+            # Legacy format: single prompt_embeds tensor — duplicate for both modalities
+            video_features = conditions["prompt_embeds"]
+            audio_features = conditions["prompt_embeds"]
+
+        mask = conditions["prompt_attention_mask"]
+        additive_mask = self._text_encoder._convert_to_additive_mask(mask, video_features.dtype)
+        video_embeds, audio_embeds, attention_mask = self._text_encoder.embeddings_processor.create_embeddings(
+            video_features, audio_features, additive_mask
         )
+
         conditions["video_prompt_embeds"] = video_embeds
         conditions["audio_prompt_embeds"] = audio_embeds
         conditions["prompt_attention_mask"] = attention_mask
@@ -375,7 +388,7 @@ class LtxvTrainer:
         # Unload heavy components to free VRAM, keeping only the embedding connectors
         self._text_encoder.model = None
         self._text_encoder.tokenizer = None
-        self._text_encoder.feature_extractor_linear = None
+        self._text_encoder.feature_extractor = None
 
         logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
         return cached_embeddings
@@ -822,7 +835,7 @@ class LtxvTrainer:
                         output_path=output_path,
                         fps=self._config.validation.frame_rate,
                         audio=audio,
-                        audio_sample_rate=self._vocoder.output_sample_rate if audio is not None else None,
+                        audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
                     )
                 video_paths.append(output_path)
 

packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py

@@ -163,6 +163,7 @@ class TextToVideoStrategy(TrainingStrategy):
         # Create video Modality
         video_modality = Modality(
             enabled=True,
+            sigma=sigmas,
             latent=noisy_video,
             timesteps=video_timesteps,
             positions=video_positions,
@@ -254,6 +255,7 @@ class TextToVideoStrategy(TrainingStrategy):
         audio_modality = Modality(
             enabled=True,
             latent=noisy_audio,
+            sigma=sigmas,
             timesteps=audio_timesteps,
             positions=audio_positions,
             context=audio_prompt_embeds,

packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py

@@ -210,6 +210,7 @@ class VideoToVideoStrategy(TrainingStrategy):
         video_modality = Modality(
             enabled=True,
             latent=combined_latents,
+            sigma=sigmas,
             timesteps=timesteps,
             positions=positions,
             context=prompt_embeds,

packages/ltx-trainer/src/ltx_trainer/validation_sampler.py

@@ -36,7 +36,7 @@ if TYPE_CHECKING:
     from ltx_core.model.audio_vae import AudioDecoder, Vocoder
     from ltx_core.model.transformer import LTXModel
     from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
-    from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
+    from ltx_core.text_encoders.gemma import GemmaTextEncoder
 
 VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
 
@@ -124,7 +124,7 @@ class ValidationSampler:
         transformer: "LTXModel",
         vae_decoder: "VideoDecoder",
         vae_encoder: "VideoEncoder | None",
-        text_encoder: "AVGemmaTextEncoderModel | None" = None,
+        text_encoder: "GemmaTextEncoder | None" = None,
         audio_decoder: "AudioDecoder | None" = None,
         vocoder: "Vocoder | None" = None,
         sampling_context: SamplingContext | None = None,
@@ -497,6 +497,7 @@ class ValidationSampler:
         video = Modality(
             enabled=True,
             latent=video_state.latent,
+            sigma=sigmas[0].repeat(video_state.latent.shape[0]),
             timesteps=video_state.denoise_mask,
             positions=video_state.positions,
             context=v_ctx_pos,
@@ -509,6 +510,7 @@ class ValidationSampler:
             audio = Modality(
                 enabled=True,
                 latent=audio_state.latent,
+                sigma=sigmas[0].repeat(audio_state.latent.shape[0]),
                 timesteps=audio_state.denoise_mask,
                 positions=audio_state.positions,
                 context=a_ctx_pos,
@@ -525,6 +527,7 @@ class ValidationSampler:
                 video = replace(
                     video,
                     latent=video_state.latent,
+                    sigma=sigma.repeat(video_state.latent.shape[0]),
                     timesteps=sigma * video_state.denoise_mask,
                     positions=video_state.positions,
                 )
@@ -533,6 +536,7 @@ class ValidationSampler:
                     audio = replace(
                         audio,
                         latent=audio_state.latent,
+                        sigma=sigma.repeat(audio_state.latent.shape[0]),
                         timesteps=sigma * audio_state.denoise_mask,
                         positions=audio_state.positions,
                     )
@@ -703,7 +707,8 @@ class ValidationSampler:
         # Move the base Gemma model to CPU but keep embeddings connectors on GPU
         # as this module is also used during training
         self._text_encoder.model.to("cpu")
-        self._text_encoder.feature_extractor_linear.to("cpu")
+        if self._text_encoder.feature_extractor is not None:
+            self._text_encoder.feature_extractor.to("cpu")
 
         return v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg