Automated PR - 2026-03-04

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,
         )