[diffusion] model: LTX-2 Support (2/2) (#17496)

Co-authored-by: Fan Yin <1106310035@qq.com>
Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>

python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py

@@ -0,0 +1,594 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
+
+from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import (
+    LTX2ConnectorConfig,
+)
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+def apply_interleaved_rotary_emb(
+    x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
+) -> torch.Tensor:
+    cos, sin = freqs
+    x_real, x_imag = x.unflatten(2, (-1, 2)).unbind(-1)  # [B, S, C // 2]
+    x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+    out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+    return out
+
+
+def apply_split_rotary_emb(
+    x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor]
+) -> torch.Tensor:
+    cos, sin = freqs
+
+    x_dtype = x.dtype
+    needs_reshape = False
+    if x.ndim != 4 and cos.ndim == 4:
+        # cos is (#b, h, t, r) -> reshape x to (b, h, t, dim_per_head)
+        # The cos/sin batch dim may only be broadcastable, so take batch size from x
+        b = x.shape[0]
+        _, h, t, _ = cos.shape
+        x = x.reshape(b, t, h, -1).swapaxes(1, 2)
+        needs_reshape = True
+
+    # Split last dim (2*r) into (d=2, r)
+    last = x.shape[-1]
+    if last % 2 != 0:
+        raise ValueError(
+            f"Expected x.shape[-1] to be even for split rotary, got {last}."
+        )
+    r = last // 2
+
+    # (..., 2, r)
+    split_x = x.reshape(*x.shape[:-1], 2, r).float()  # Explicitly upcast to float
+    first_x = split_x[..., :1, :]  # (..., 1, r)
+    second_x = split_x[..., 1:, :]  # (..., 1, r)
+
+    cos_u = cos.unsqueeze(-2)  # broadcast to (..., 1, r) against (..., 2, r)
+    sin_u = sin.unsqueeze(-2)
+
+    out = split_x * cos_u
+    first_out = out[..., :1, :]
+    second_out = out[..., 1:, :]
+
+    first_out.addcmul_(-sin_u, second_x)
+    second_out.addcmul_(sin_u, first_x)
+
+    out = out.reshape(*out.shape[:-2], last)
+
+    if needs_reshape:
+        out = out.swapaxes(1, 2).reshape(b, t, -1)
+
+    out = out.to(dtype=x_dtype)
+    return out
+
+
+class LTX2Attention(torch.nn.Module):
+    r"""
+    Attention class for all LTX-2.0 attention layers. Compared to LTX-1.0, this supports specifying the query and key
+    RoPE embeddings separately for audio-to-video (a2v) and video-to-audio (v2a) cross-attention.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        kv_heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = True,
+        cross_attention_dim: Optional[int] = None,
+        out_bias: bool = True,
+        qk_norm: str = "rms_norm_across_heads",
+        norm_eps: float = 1e-6,
+        norm_elementwise_affine: bool = True,
+        rope_type: str = "interleaved",
+        processor=None,
+    ):
+        super().__init__()
+        if qk_norm != "rms_norm_across_heads":
+            raise NotImplementedError(
+                "Only 'rms_norm_across_heads' is supported as a valid value for `qk_norm`."
+            )
+
+        self.head_dim = dim_head
+        self.inner_dim = dim_head * heads
+        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+        self.query_dim = query_dim
+        self.cross_attention_dim = (
+            cross_attention_dim if cross_attention_dim is not None else query_dim
+        )
+        self.use_bias = bias
+        self.dropout = dropout
+        self.out_dim = query_dim
+        self.heads = heads
+        self.rope_type = rope_type
+
+        self.norm_q = torch.nn.RMSNorm(
+            dim_head * heads, eps=norm_eps, elementwise_affine=norm_elementwise_affine
+        )
+        self.norm_k = torch.nn.RMSNorm(
+            dim_head * kv_heads,
+            eps=norm_eps,
+            elementwise_affine=norm_elementwise_affine,
+        )
+        self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = torch.nn.Linear(
+            self.cross_attention_dim, self.inner_kv_dim, bias=bias
+        )
+        self.to_v = torch.nn.Linear(
+            self.cross_attention_dim, self.inner_kv_dim, bias=bias
+        )
+        self.to_out = torch.nn.ModuleList([])
+        self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+        self.to_out.append(torch.nn.Dropout(dropout))
+
+        # Scaled dot product attention
+        self.attn = USPAttention(
+            num_heads=heads,
+            head_size=self.head_dim,
+            dropout_rate=0,
+            softmax_scale=None,
+            causal=False,
+            supported_attention_backends={
+                AttentionBackendEnum.FA,
+                AttentionBackendEnum.AITER,
+                AttentionBackendEnum.TORCH_SDPA,
+                AttentionBackendEnum.SAGE_ATTN,
+                AttentionBackendEnum.SAGE_ATTN_3,
+            },
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        batch_size, sequence_length, _ = (
+            hidden_states.shape
+            if encoder_hidden_states is None
+            else encoder_hidden_states.shape
+        )
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(encoder_hidden_states)
+        value = self.to_v(encoder_hidden_states)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if query_rotary_emb is not None:
+            if self.rope_type == "interleaved":
+                query = apply_interleaved_rotary_emb(query, query_rotary_emb)
+                key = apply_interleaved_rotary_emb(
+                    key,
+                    key_rotary_emb if key_rotary_emb is not None else query_rotary_emb,
+                )
+            elif self.rope_type == "split":
+                query = apply_split_rotary_emb(query, query_rotary_emb)
+                key = apply_split_rotary_emb(
+                    key,
+                    key_rotary_emb if key_rotary_emb is not None else query_rotary_emb,
+                )
+
+        query = query.unflatten(2, (self.heads, -1))
+        key = key.unflatten(2, (self.heads, -1))
+        value = value.unflatten(2, (self.heads, -1))
+
+        hidden_states = self.attn(
+            query,
+            key,
+            value,
+        )
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states.to(query.dtype)
+
+        hidden_states = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+        return hidden_states
+
+
+class LTX2RotaryPosEmbed1d(nn.Module):
+    """
+    1D rotary positional embeddings (RoPE) for the LTX 2.0 text encoder connectors.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base_seq_len: int = 4096,
+        theta: float = 10000.0,
+        double_precision: bool = True,
+        rope_type: str = "interleaved",
+        num_attention_heads: int = 32,
+    ):
+        super().__init__()
+        if rope_type not in ["interleaved", "split"]:
+            raise ValueError(
+                f"{rope_type=} not supported. Choose between 'interleaved' and 'split'."
+            )
+
+        self.dim = dim
+        self.base_seq_len = base_seq_len
+        self.theta = theta
+        self.double_precision = double_precision
+        self.rope_type = rope_type
+        self.num_attention_heads = num_attention_heads
+
+    def forward(
+        self,
+        batch_size: int,
+        pos: int,
+        device: Union[str, torch.device],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # 1. Get 1D position ids
+        grid_1d = torch.arange(pos, dtype=torch.float32, device=device)
+        # Get fractional indices relative to self.base_seq_len
+        grid_1d = grid_1d / self.base_seq_len
+        grid = grid_1d.unsqueeze(0).repeat(batch_size, 1)  # [batch_size, seq_len]
+
+        # 2. Calculate 1D RoPE frequencies
+        num_rope_elems = 2  # 1 (because 1D) * 2 (for cos, sin) = 2
+        freqs_dtype = torch.float64 if self.double_precision else torch.float32
+        pow_indices = torch.pow(
+            self.theta,
+            torch.linspace(
+                start=0.0,
+                end=1.0,
+                steps=self.dim // num_rope_elems,
+                dtype=freqs_dtype,
+                device=device,
+            ),
+        )
+        freqs = (pow_indices * torch.pi / 2.0).to(dtype=torch.float32)
+
+        # 3. Matrix-vector outer product between pos ids of shape (batch_size, seq_len) and freqs vector of shape
+        # (self.dim // 2,).
+        freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs  # [B, seq_len, self.dim // 2]
+
+        # 4. Get real, interleaved (cos, sin) frequencies, padded to self.dim
+        if self.rope_type == "interleaved":
+            cos_freqs = freqs.cos().repeat_interleave(2, dim=-1)
+            sin_freqs = freqs.sin().repeat_interleave(2, dim=-1)
+
+            if self.dim % num_rope_elems != 0:
+                cos_padding = torch.ones_like(
+                    cos_freqs[:, :, : self.dim % num_rope_elems]
+                )
+                sin_padding = torch.zeros_like(
+                    sin_freqs[:, :, : self.dim % num_rope_elems]
+                )
+                cos_freqs = torch.cat([cos_padding, cos_freqs], dim=-1)
+                sin_freqs = torch.cat([sin_padding, sin_freqs], dim=-1)
+
+        elif self.rope_type == "split":
+            expected_freqs = self.dim // 2
+            current_freqs = freqs.shape[-1]
+            pad_size = expected_freqs - current_freqs
+            cos_freq = freqs.cos()
+            sin_freq = freqs.sin()
+
+            if pad_size != 0:
+                cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
+                sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
+
+                cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
+                sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
+
+            # Reshape freqs to be compatible with multi-head attention
+            b = cos_freq.shape[0]
+            t = cos_freq.shape[1]
+
+            cos_freq = cos_freq.reshape(b, t, self.num_attention_heads, -1)
+            sin_freq = sin_freq.reshape(b, t, self.num_attention_heads, -1)
+
+            cos_freqs = torch.swapaxes(cos_freq, 1, 2)  # (B,H,T,D//2)
+            sin_freqs = torch.swapaxes(sin_freq, 1, 2)  # (B,H,T,D//2)
+
+        return cos_freqs, sin_freqs
+
+
+class LTX2TransformerBlock1d(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        activation_fn: str = "gelu-approximate",
+        eps: float = 1e-6,
+        rope_type: str = "interleaved",
+    ):
+        super().__init__()
+
+        self.norm1 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.attn1 = LTX2Attention(
+            query_dim=dim,
+            heads=num_attention_heads,
+            kv_heads=num_attention_heads,
+            dim_head=attention_head_dim,
+            rope_type=rope_type,
+        )
+
+        self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False)
+        self.ff = FeedForward(dim, activation_fn=activation_fn)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        norm_hidden_states = self.norm1(hidden_states)
+        attn_hidden_states = self.attn1(
+            norm_hidden_states,
+            attention_mask=attention_mask,
+            query_rotary_emb=rotary_emb,
+        )
+        hidden_states = hidden_states + attn_hidden_states
+
+        norm_hidden_states = self.norm2(hidden_states)
+        ff_hidden_states = self.ff(norm_hidden_states)
+        hidden_states = hidden_states + ff_hidden_states
+
+        return hidden_states
+
+
+class LTX2ConnectorTransformer1d(nn.Module):
+    """
+    A 1D sequence transformer for modalities such as text.
+    In LTX 2.0, this is used to process the text encoder hidden states for each of the video and audio streams.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    def __init__(
+        self,
+        num_attention_heads: int = 30,
+        attention_head_dim: int = 128,
+        num_layers: int = 2,
+        num_learnable_registers: int | None = 128,
+        rope_base_seq_len: int = 4096,
+        rope_theta: float = 10000.0,
+        rope_double_precision: bool = True,
+        eps: float = 1e-6,
+        causal_temporal_positioning: bool = False,
+        rope_type: str = "interleaved",
+    ):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        self.inner_dim = num_attention_heads * attention_head_dim
+        self.causal_temporal_positioning = causal_temporal_positioning
+
+        self.num_learnable_registers = num_learnable_registers
+        self.learnable_registers = None
+        if num_learnable_registers is not None:
+            init_registers = (
+                torch.rand(num_learnable_registers, self.inner_dim) * 2.0 - 1.0
+            )
+            self.learnable_registers = torch.nn.Parameter(init_registers)
+
+        self.rope = LTX2RotaryPosEmbed1d(
+            self.inner_dim,
+            base_seq_len=rope_base_seq_len,
+            theta=rope_theta,
+            double_precision=rope_double_precision,
+            rope_type=rope_type,
+            num_attention_heads=num_attention_heads,
+        )
+
+        self.transformer_blocks = torch.nn.ModuleList(
+            [
+                LTX2TransformerBlock1d(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    rope_type=rope_type,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.norm_out = torch.nn.RMSNorm(
+            self.inner_dim, eps=eps, elementwise_affine=False
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        attn_mask_binarize_threshold: float = -9000.0,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # hidden_states shape: [batch_size, seq_len, hidden_dim]
+        # attention_mask shape: [batch_size, seq_len] or [batch_size, 1, 1, seq_len]
+        batch_size, seq_len, _ = hidden_states.shape
+
+        # 1. Replace padding with learned registers, if using
+        if self.learnable_registers is not None:
+            if seq_len % self.num_learnable_registers != 0:
+                raise ValueError(
+                    f"The `hidden_states` sequence length {hidden_states.shape[1]} should be divisible by the number"
+                    f" of learnable registers {self.num_learnable_registers}"
+                )
+
+            num_register_repeats = seq_len // self.num_learnable_registers
+            registers = torch.tile(
+                self.learnable_registers, (num_register_repeats, 1)
+            )  # [seq_len, inner_dim]
+
+            binary_attn_mask = (attention_mask >= attn_mask_binarize_threshold).int()
+            if binary_attn_mask.ndim == 4:
+                binary_attn_mask = binary_attn_mask.squeeze(1).squeeze(
+                    1
+                )  # [B, 1, 1, L] --> [B, L]
+
+            hidden_states_non_padded = [
+                hidden_states[i, binary_attn_mask[i].bool(), :]
+                for i in range(batch_size)
+            ]
+            valid_seq_lens = [x.shape[0] for x in hidden_states_non_padded]
+            pad_lengths = [seq_len - valid_seq_len for valid_seq_len in valid_seq_lens]
+            padded_hidden_states = [
+                F.pad(x, pad=(0, 0, 0, p), value=0)
+                for x, p in zip(hidden_states_non_padded, pad_lengths)
+            ]
+            padded_hidden_states = torch.cat(
+                [x.unsqueeze(0) for x in padded_hidden_states], dim=0
+            )  # [B, L, D]
+
+            flipped_mask = torch.flip(binary_attn_mask, dims=[1]).unsqueeze(
+                -1
+            )  # [B, L, 1]
+            hidden_states = (
+                flipped_mask * padded_hidden_states + (1 - flipped_mask) * registers
+            )
+
+            # Overwrite attention_mask with an all-zeros mask if using registers.
+            attention_mask = torch.zeros_like(attention_mask)
+
+        # 2. Calculate 1D RoPE positional embeddings
+        rotary_emb = self.rope(batch_size, seq_len, device=hidden_states.device)
+
+        # 3. Run 1D transformer blocks
+        for block in self.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, attention_mask, rotary_emb
+                )
+            else:
+                hidden_states = block(
+                    hidden_states, attention_mask=attention_mask, rotary_emb=rotary_emb
+                )
+
+        hidden_states = self.norm_out(hidden_states)
+
+        return hidden_states, attention_mask
+
+
+class LTX2TextConnectors(nn.Module):
+    """
+    Text connector stack used by LTX 2.0 to process the packed text encoder hidden states for both the video and audio
+    streams.
+    """
+
+    def __init__(
+        self,
+        config: LTX2ConnectorConfig,
+    ):
+        super().__init__()
+        caption_channels = config.caption_channels
+        text_proj_in_factor = config.text_proj_in_factor
+        video_connector_num_attention_heads = config.video_connector_num_attention_heads
+        video_connector_attention_head_dim = config.video_connector_attention_head_dim
+        video_connector_num_layers = config.video_connector_num_layers
+        video_connector_num_learnable_registers = (
+            config.video_connector_num_learnable_registers
+        )
+        audio_connector_num_attention_heads = config.audio_connector_num_attention_heads
+        audio_connector_attention_head_dim = config.audio_connector_attention_head_dim
+        audio_connector_num_layers = config.audio_connector_num_layers
+        audio_connector_num_learnable_registers = (
+            config.audio_connector_num_learnable_registers
+        )
+        connector_rope_base_seq_len = config.connector_rope_base_seq_len
+        rope_theta = config.rope_theta
+        rope_double_precision = config.rope_double_precision
+        causal_temporal_positioning = config.causal_temporal_positioning
+        rope_type = config.rope_type
+
+        self.text_proj_in = nn.Linear(
+            caption_channels * text_proj_in_factor, caption_channels, bias=False
+        )
+        self.video_connector = LTX2ConnectorTransformer1d(
+            num_attention_heads=video_connector_num_attention_heads,
+            attention_head_dim=video_connector_attention_head_dim,
+            num_layers=video_connector_num_layers,
+            num_learnable_registers=video_connector_num_learnable_registers,
+            rope_base_seq_len=connector_rope_base_seq_len,
+            rope_theta=rope_theta,
+            rope_double_precision=rope_double_precision,
+            causal_temporal_positioning=causal_temporal_positioning,
+            rope_type=rope_type,
+        )
+        self.audio_connector = LTX2ConnectorTransformer1d(
+            num_attention_heads=audio_connector_num_attention_heads,
+            attention_head_dim=audio_connector_attention_head_dim,
+            num_layers=audio_connector_num_layers,
+            num_learnable_registers=audio_connector_num_learnable_registers,
+            rope_base_seq_len=connector_rope_base_seq_len,
+            rope_theta=rope_theta,
+            rope_double_precision=rope_double_precision,
+            causal_temporal_positioning=causal_temporal_positioning,
+            rope_type=rope_type,
+        )
+
+    def forward(
+        self,
+        text_encoder_hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        additive_mask: bool = False,
+    ):
+        # Convert to additive attention mask, if necessary
+        if not additive_mask:
+            text_dtype = text_encoder_hidden_states.dtype
+            attention_mask = (attention_mask - 1).reshape(
+                attention_mask.shape[0], 1, -1, attention_mask.shape[-1]
+            )
+            attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max
+
+        # Ensure input dtype matches the layer's weight dtype
+        if text_encoder_hidden_states.dtype != self.text_proj_in.weight.dtype:
+            text_encoder_hidden_states = text_encoder_hidden_states.to(
+                self.text_proj_in.weight.dtype
+            )
+
+        # Ensure sequence length is divisible by num_learnable_registers (128)
+        seq_len = text_encoder_hidden_states.shape[1]
+        num_learnable_registers = self.video_connector.num_learnable_registers
+        if (
+            num_learnable_registers is not None
+            and seq_len % num_learnable_registers != 0
+        ):
+            pad_len = num_learnable_registers - (seq_len % num_learnable_registers)
+            text_encoder_hidden_states = F.pad(
+                text_encoder_hidden_states, (0, 0, 0, pad_len), value=0.0
+            )
+
+            if attention_mask.shape[-1] == seq_len:
+                # Pad with a large negative value to mask out the new tokens
+                attention_mask = F.pad(attention_mask, (0, pad_len), value=-1000000.0)
+
+        text_encoder_hidden_states = self.text_proj_in(text_encoder_hidden_states)
+
+        video_text_embedding, new_attn_mask = self.video_connector(
+            text_encoder_hidden_states, attention_mask
+        )
+
+        attn_mask = (new_attn_mask < 1e-6).to(torch.int64)
+        attn_mask = attn_mask.reshape(
+            video_text_embedding.shape[0], video_text_embedding.shape[1], 1
+        )
+        video_text_embedding = video_text_embedding * attn_mask
+        new_attn_mask = attn_mask.squeeze(-1)
+
+        audio_text_embedding, _ = self.audio_connector(
+            text_encoder_hidden_states, attention_mask
+        )
+
+        return video_text_embedding, audio_text_embedding, new_attn_mask
+
+
+EntryClass = LTX2TextConnectors

python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py

@@ -0,0 +1,905 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from diffusers.models.autoencoders.vae import (
+    DecoderOutput,
+    DiagonalGaussianDistribution,
+)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.ltx_audio import LTXAudioVAEConfig
+from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
+
+LATENT_DOWNSAMPLE_FACTOR = 4
+
+
+class LTX2AudioCausalConv2d(nn.Module):
+    """
+    A causal 2D convolution that pads asymmetrically along the causal axis.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Tuple[int, int]],
+        stride: int = 1,
+        dilation: Union[int, Tuple[int, int]] = 1,
+        groups: int = 1,
+        bias: bool = True,
+        causality_axis: str = "height",
+    ) -> None:
+        super().__init__()
+
+        self.causality_axis = causality_axis
+        kernel_size = (
+            (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size
+        )
+        dilation = (dilation, dilation) if isinstance(dilation, int) else dilation
+
+        pad_h = (kernel_size[0] - 1) * dilation[0]
+        pad_w = (kernel_size[1] - 1) * dilation[1]
+
+        if self.causality_axis == "none":
+            padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
+        elif self.causality_axis in {"width", "width-compatibility"}:
+            padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
+        elif self.causality_axis == "height":
+            padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
+        else:
+            raise ValueError(f"Invalid causality_axis: {causality_axis}")
+
+        self.padding = padding
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.pad(x, self.padding)
+        return self.conv(x)
+
+
+class LTX2AudioPixelNorm(nn.Module):
+    """
+    Per-pixel (per-location) RMS normalization layer.
+    """
+
+    def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True)
+        rms = torch.sqrt(mean_sq + self.eps)
+        return x / rms
+
+
+class LTX2AudioAttnBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        norm_type: str = "group",
+    ) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+
+        if norm_type == "group":
+            self.norm = nn.GroupNorm(
+                num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(
+            in_channels, in_channels, kernel_size=1, stride=1, padding=0
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h_ = self.norm(x)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        batch, channels, height, width = q.shape
+        q = q.reshape(batch, channels, height * width).permute(0, 2, 1).contiguous()
+        k = k.reshape(batch, channels, height * width).contiguous()
+        attn = torch.bmm(q, k) * (int(channels) ** (-0.5))
+        attn = torch.nn.functional.softmax(attn, dim=2)
+
+        v = v.reshape(batch, channels, height * width)
+        attn = attn.permute(0, 2, 1).contiguous()
+        h_ = torch.bmm(v, attn).reshape(batch, channels, height, width)
+
+        h_ = self.proj_out(h_)
+        return x + h_
+
+
+class LTX2AudioResnetBlock(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        conv_shortcut: bool = False,
+        dropout: float = 0.0,
+        temb_channels: int = 512,
+        norm_type: str = "group",
+        causality_axis: str = "height",
+    ) -> None:
+        super().__init__()
+        self.causality_axis = causality_axis
+
+        if (
+            self.causality_axis is not None
+            and self.causality_axis != "none"
+            and norm_type == "group"
+        ):
+            raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        if norm_type == "group":
+            self.norm1 = nn.GroupNorm(
+                num_groups=32, num_channels=in_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm1 = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.non_linearity = nn.SiLU()
+        if causality_axis is not None:
+            self.conv1 = LTX2AudioCausalConv2d(
+                in_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=causality_axis,
+            )
+        else:
+            self.conv1 = nn.Conv2d(
+                in_channels, out_channels, kernel_size=3, stride=1, padding=1
+            )
+        if temb_channels > 0:
+            self.temb_proj = nn.Linear(temb_channels, out_channels)
+        if norm_type == "group":
+            self.norm2 = nn.GroupNorm(
+                num_groups=32, num_channels=out_channels, eps=1e-6, affine=True
+            )
+        elif norm_type == "pixel":
+            self.norm2 = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {norm_type}")
+        self.dropout = nn.Dropout(dropout)
+        if causality_axis is not None:
+            self.conv2 = LTX2AudioCausalConv2d(
+                out_channels,
+                out_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=causality_axis,
+            )
+        else:
+            self.conv2 = nn.Conv2d(
+                out_channels, out_channels, kernel_size=3, stride=1, padding=1
+            )
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                if causality_axis is not None:
+                    self.conv_shortcut = LTX2AudioCausalConv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=3,
+                        stride=1,
+                        causality_axis=causality_axis,
+                    )
+                else:
+                    self.conv_shortcut = nn.Conv2d(
+                        in_channels, out_channels, kernel_size=3, stride=1, padding=1
+                    )
+            else:
+                if causality_axis is not None:
+                    self.nin_shortcut = LTX2AudioCausalConv2d(
+                        in_channels,
+                        out_channels,
+                        kernel_size=1,
+                        stride=1,
+                        causality_axis=causality_axis,
+                    )
+                else:
+                    self.nin_shortcut = nn.Conv2d(
+                        in_channels, out_channels, kernel_size=1, stride=1, padding=0
+                    )
+
+    def forward(
+        self, x: torch.Tensor, temb: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = self.norm1(x)
+        h = self.non_linearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(self.non_linearity(temb))[:, :, None, None]
+
+        h = self.norm2(h)
+        h = self.non_linearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = (
+                self.conv_shortcut(x)
+                if self.use_conv_shortcut
+                else self.nin_shortcut(x)
+            )
+
+        return x + h
+
+
+class LTX2AudioDownsample(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: Optional[str] = "height",
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(
+                in_channels, in_channels, kernel_size=3, stride=2, padding=0
+            )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.with_conv:
+            # Padding tuple is in the order: (left, right, top, bottom).
+            if self.causality_axis == "none":
+                pad = (0, 1, 0, 1)
+            elif self.causality_axis == "width":
+                pad = (2, 0, 0, 1)
+            elif self.causality_axis == "height":
+                pad = (0, 1, 2, 0)
+            elif self.causality_axis == "width-compatibility":
+                pad = (1, 0, 0, 1)
+            else:
+                raise ValueError(
+                    f"Invalid `causality_axis` {self.causality_axis}; supported values are `none`, `width`, `height`,"
+                    f" and `width-compatibility`."
+                )
+
+            x = F.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            # with_conv=False implies that causality_axis is "none"
+            x = F.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class LTX2AudioUpsample(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        with_conv: bool,
+        causality_axis: Optional[str] = "height",
+    ) -> None:
+        super().__init__()
+        self.with_conv = with_conv
+        self.causality_axis = causality_axis
+        if self.with_conv:
+            if causality_axis is not None:
+                self.conv = LTX2AudioCausalConv2d(
+                    in_channels,
+                    in_channels,
+                    kernel_size=3,
+                    stride=1,
+                    causality_axis=causality_axis,
+                )
+            else:
+                self.conv = nn.Conv2d(
+                    in_channels, in_channels, kernel_size=3, stride=1, padding=1
+                )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+            if self.causality_axis is None or self.causality_axis == "none":
+                pass
+            elif self.causality_axis == "height":
+                x = x[:, :, 1:, :]
+            elif self.causality_axis == "width":
+                x = x[:, :, :, 1:]
+            elif self.causality_axis == "width-compatibility":
+                pass
+            else:
+                raise ValueError(f"Invalid causality_axis: {self.causality_axis}")
+
+        return x
+
+
+class LTX2AudioAudioPatchifier:
+    """
+    Patchifier for spectrogram/audio latents.
+    """
+
+    def __init__(
+        self,
+        patch_size: int,
+        sample_rate: int = 16000,
+        hop_length: int = 160,
+        audio_latent_downsample_factor: int = 4,
+        is_causal: bool = True,
+    ):
+        self.hop_length = hop_length
+        self.sample_rate = sample_rate
+        self.audio_latent_downsample_factor = audio_latent_downsample_factor
+        self.is_causal = is_causal
+        self._patch_size = (1, patch_size, patch_size)
+
+    def patchify(self, audio_latents: torch.Tensor) -> torch.Tensor:
+        batch, channels, time, freq = audio_latents.shape
+        return audio_latents.permute(0, 2, 1, 3).reshape(batch, time, channels * freq)
+
+    def unpatchify(
+        self, audio_latents: torch.Tensor, channels: int, mel_bins: int
+    ) -> torch.Tensor:
+        batch, time, _ = audio_latents.shape
+        return audio_latents.view(batch, time, channels, mel_bins).permute(0, 2, 1, 3)
+
+    @property
+    def patch_size(self) -> Tuple[int, int, int]:
+        return self._patch_size
+
+
+class LTX2AudioEncoder(nn.Module):
+    def __init__(
+        self,
+        base_channels: int = 128,
+        output_channels: int = 1,
+        num_res_blocks: int = 2,
+        attn_resolutions: Optional[Tuple[int, ...]] = None,
+        in_channels: int = 2,
+        resolution: int = 256,
+        latent_channels: int = 8,
+        ch_mult: Tuple[int, ...] = (1, 2, 4),
+        norm_type: str = "group",
+        causality_axis: Optional[str] = "width",
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = False,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: Optional[int] = 64,
+        double_z: bool = True,
+    ):
+        super().__init__()
+
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+
+        self.base_channels = base_channels
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_ch = output_channels
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.latent_channels = latent_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+
+        base_block_channels = base_channels
+        base_resolution = resolution
+        self.z_shape = (1, latent_channels, base_resolution, base_resolution)
+
+        if self.causality_axis is not None:
+            self.conv_in = LTX2AudioCausalConv2d(
+                in_channels,
+                base_block_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_in = nn.Conv2d(
+                in_channels, base_block_channels, kernel_size=3, stride=1, padding=1
+            )
+
+        self.down = nn.ModuleList()
+        block_in = base_block_channels
+        curr_res = self.resolution
+
+        for level in range(self.num_resolutions):
+            stage = nn.Module()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_out = self.base_channels * self.channel_multipliers[level]
+
+            for _ in range(self.num_res_blocks):
+                stage.block.append(
+                    LTX2AudioResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                        norm_type=self.norm_type,
+                        causality_axis=self.causality_axis,
+                    )
+                )
+                block_in = block_out
+                if self.attn_resolutions:
+                    if curr_res in self.attn_resolutions:
+                        stage.attn.append(
+                            LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+                        )
+
+            if level != self.num_resolutions - 1:
+                stage.downsample = LTX2AudioDownsample(
+                    block_in, True, causality_axis=self.causality_axis
+                )
+                curr_res = curr_res // 2
+
+            self.down.append(stage)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = LTX2AudioResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+        if mid_block_add_attention:
+            self.mid.attn_1 = LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+        else:
+            self.mid.attn_1 = nn.Identity()
+        self.mid.block_2 = LTX2AudioResnetBlock(
+            in_channels=block_in,
+            out_channels=block_in,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+
+        final_block_channels = block_in
+        z_channels = 2 * latent_channels if double_z else latent_channels
+        if self.norm_type == "group":
+            self.norm_out = nn.GroupNorm(
+                num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True
+            )
+        elif self.norm_type == "pixel":
+            self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {self.norm_type}")
+        self.non_linearity = nn.SiLU()
+
+        if self.causality_axis is not None:
+            self.conv_out = LTX2AudioCausalConv2d(
+                final_block_channels,
+                z_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_out = nn.Conv2d(
+                final_block_channels, z_channels, kernel_size=3, stride=1, padding=1
+            )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states expected shape: (batch_size, channels, time, num_mel_bins)
+        hidden_states = self.conv_in(hidden_states)
+
+        for level in range(self.num_resolutions):
+            stage = self.down[level]
+            for block_idx, block in enumerate(stage.block):
+                hidden_states = block(hidden_states, temb=None)
+                if stage.attn:
+                    hidden_states = stage.attn[block_idx](hidden_states)
+
+            if level != self.num_resolutions - 1 and hasattr(stage, "downsample"):
+                hidden_states = stage.downsample(hidden_states)
+
+        hidden_states = self.mid.block_1(hidden_states, temb=None)
+        hidden_states = self.mid.attn_1(hidden_states)
+        hidden_states = self.mid.block_2(hidden_states, temb=None)
+
+        hidden_states = self.norm_out(hidden_states)
+        hidden_states = self.non_linearity(hidden_states)
+        hidden_states = self.conv_out(hidden_states)
+
+        return hidden_states
+
+
+class LTX2AudioDecoder(nn.Module):
+    """
+    Symmetric decoder that reconstructs audio spectrograms from latent features.
+
+    The decoder mirrors the encoder structure with configurable channel multipliers, attention resolutions, and causal
+    convolutions.
+    """
+
+    def __init__(
+        self,
+        base_channels: int = 128,
+        output_channels: int = 1,
+        num_res_blocks: int = 2,
+        attn_resolutions: Optional[Tuple[int, ...]] = None,
+        in_channels: int = 2,
+        resolution: int = 256,
+        latent_channels: int = 8,
+        ch_mult: Tuple[int, ...] = (1, 2, 4),
+        norm_type: str = "group",
+        causality_axis: Optional[str] = "width",
+        dropout: float = 0.0,
+        mid_block_add_attention: bool = False,
+        sample_rate: int = 16000,
+        mel_hop_length: int = 160,
+        is_causal: bool = True,
+        mel_bins: Optional[int] = 64,
+    ) -> None:
+        super().__init__()
+
+        self.sample_rate = sample_rate
+        self.mel_hop_length = mel_hop_length
+        self.is_causal = is_causal
+        self.mel_bins = mel_bins
+        self.patchifier = LTX2AudioAudioPatchifier(
+            patch_size=1,
+            audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
+            sample_rate=sample_rate,
+            hop_length=mel_hop_length,
+            is_causal=is_causal,
+        )
+
+        self.base_channels = base_channels
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.out_ch = output_channels
+        self.give_pre_end = False
+        self.tanh_out = False
+        self.norm_type = norm_type
+        self.latent_channels = latent_channels
+        self.channel_multipliers = ch_mult
+        self.attn_resolutions = attn_resolutions
+        self.causality_axis = causality_axis
+
+        base_block_channels = base_channels * self.channel_multipliers[-1]
+        base_resolution = resolution // (2 ** (self.num_resolutions - 1))
+        self.z_shape = (1, latent_channels, base_resolution, base_resolution)
+
+        if self.causality_axis is not None:
+            self.conv_in = LTX2AudioCausalConv2d(
+                latent_channels,
+                base_block_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_in = nn.Conv2d(
+                latent_channels, base_block_channels, kernel_size=3, stride=1, padding=1
+            )
+        self.non_linearity = nn.SiLU()
+        self.mid = nn.Module()
+        self.mid.block_1 = LTX2AudioResnetBlock(
+            in_channels=base_block_channels,
+            out_channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+        if mid_block_add_attention:
+            self.mid.attn_1 = LTX2AudioAttnBlock(
+                base_block_channels, norm_type=self.norm_type
+            )
+        else:
+            self.mid.attn_1 = nn.Identity()
+        self.mid.block_2 = LTX2AudioResnetBlock(
+            in_channels=base_block_channels,
+            out_channels=base_block_channels,
+            temb_channels=self.temb_ch,
+            dropout=dropout,
+            norm_type=self.norm_type,
+            causality_axis=self.causality_axis,
+        )
+
+        self.up = nn.ModuleList()
+        block_in = base_block_channels
+        curr_res = self.resolution // (2 ** (self.num_resolutions - 1))
+
+        for level in reversed(range(self.num_resolutions)):
+            stage = nn.Module()
+            stage.block = nn.ModuleList()
+            stage.attn = nn.ModuleList()
+            block_out = self.base_channels * self.channel_multipliers[level]
+
+            for _ in range(self.num_res_blocks + 1):
+                stage.block.append(
+                    LTX2AudioResnetBlock(
+                        in_channels=block_in,
+                        out_channels=block_out,
+                        temb_channels=self.temb_ch,
+                        dropout=dropout,
+                        norm_type=self.norm_type,
+                        causality_axis=self.causality_axis,
+                    )
+                )
+                block_in = block_out
+                if self.attn_resolutions:
+                    if curr_res in self.attn_resolutions:
+                        stage.attn.append(
+                            LTX2AudioAttnBlock(block_in, norm_type=self.norm_type)
+                        )
+
+            if level != 0:
+                stage.upsample = LTX2AudioUpsample(
+                    block_in, True, causality_axis=self.causality_axis
+                )
+                curr_res *= 2
+
+            self.up.insert(0, stage)
+
+        final_block_channels = block_in
+
+        if self.norm_type == "group":
+            self.norm_out = nn.GroupNorm(
+                num_groups=32, num_channels=final_block_channels, eps=1e-6, affine=True
+            )
+        elif self.norm_type == "pixel":
+            self.norm_out = LTX2AudioPixelNorm(dim=1, eps=1e-6)
+        else:
+            raise ValueError(f"Invalid normalization type: {self.norm_type}")
+
+        if self.causality_axis is not None:
+            self.conv_out = LTX2AudioCausalConv2d(
+                final_block_channels,
+                output_channels,
+                kernel_size=3,
+                stride=1,
+                causality_axis=self.causality_axis,
+            )
+        else:
+            self.conv_out = nn.Conv2d(
+                final_block_channels,
+                output_channels,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            )
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        _, _, frames, mel_bins = sample.shape
+
+        target_frames = frames * LATENT_DOWNSAMPLE_FACTOR
+
+        if self.causality_axis is not None:
+            target_frames = max(target_frames - (LATENT_DOWNSAMPLE_FACTOR - 1), 1)
+
+        target_channels = self.out_ch
+        target_mel_bins = self.mel_bins if self.mel_bins is not None else mel_bins
+
+        hidden_features = self.conv_in(sample)
+        hidden_features = self.mid.block_1(hidden_features, temb=None)
+        hidden_features = self.mid.attn_1(hidden_features)
+        hidden_features = self.mid.block_2(hidden_features, temb=None)
+
+        for level in reversed(range(self.num_resolutions)):
+            stage = self.up[level]
+            for block_idx, block in enumerate(stage.block):
+                hidden_features = block(hidden_features, temb=None)
+                if stage.attn:
+                    hidden_features = stage.attn[block_idx](hidden_features)
+
+            if level != 0 and hasattr(stage, "upsample"):
+                hidden_features = stage.upsample(hidden_features)
+
+        if self.give_pre_end:
+            return hidden_features
+
+        hidden = self.norm_out(hidden_features)
+        hidden = self.non_linearity(hidden)
+        decoded_output = self.conv_out(hidden)
+        decoded_output = torch.tanh(decoded_output) if self.tanh_out else decoded_output
+
+        _, _, current_time, current_freq = decoded_output.shape
+        target_time = target_frames
+        target_freq = target_mel_bins
+
+        decoded_output = decoded_output[
+            :,
+            :target_channels,
+            : min(current_time, target_time),
+            : min(current_freq, target_freq),
+        ]
+
+        time_padding_needed = target_time - decoded_output.shape[2]
+        freq_padding_needed = target_freq - decoded_output.shape[3]
+
+        if time_padding_needed > 0 or freq_padding_needed > 0:
+            padding = (
+                0,
+                max(freq_padding_needed, 0),
+                0,
+                max(time_padding_needed, 0),
+            )
+            decoded_output = F.pad(decoded_output, padding)
+
+        decoded_output = decoded_output[:, :target_channels, :target_time, :target_freq]
+
+        return decoded_output
+
+
+class AutoencoderKLLTX2Audio(ParallelTiledVAE):
+    r"""
+    LTX2 audio VAE for encoding and decoding audio latent representations.
+    """
+
+    _supports_gradient_checkpointing = False
+
+    def __init__(
+        self,
+        config: LTXAudioVAEConfig,
+    ) -> None:
+        super().__init__(config=config)
+
+        causality_axis = config.arch_config.causality_axis
+        attn_resolutions = config.arch_config.attn_resolutions
+        base_channels = config.arch_config.base_channels
+        output_channels = config.arch_config.output_channels
+        ch_mult = config.arch_config.ch_mult
+        num_res_blocks = config.arch_config.num_res_blocks
+        in_channels = config.arch_config.in_channels
+        resolution = config.arch_config.resolution
+        latent_channels = config.arch_config.latent_channels
+        norm_type = config.arch_config.norm_type
+        dropout = config.arch_config.dropout
+        mid_block_add_attention = config.arch_config.mid_block_add_attention
+        sample_rate = config.arch_config.sample_rate
+        mel_hop_length = config.arch_config.mel_hop_length
+        is_causal = config.arch_config.is_causal
+        mel_bins = config.arch_config.mel_bins
+        double_z = config.arch_config.double_z
+
+        supported_causality_axes = {"none", "width", "height", "width-compatibility"}
+        if causality_axis not in supported_causality_axes:
+            raise ValueError(
+                f"{causality_axis=} is not valid. Supported values: {supported_causality_axes}"
+            )
+
+        attn_resolution_set = (
+            set(attn_resolutions) if attn_resolutions else attn_resolutions
+        )
+
+        self.encoder = LTX2AudioEncoder(
+            base_channels=base_channels,
+            output_channels=output_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolution_set,
+            in_channels=in_channels,
+            resolution=resolution,
+            latent_channels=latent_channels,
+            norm_type=norm_type,
+            causality_axis=causality_axis,
+            dropout=dropout,
+            mid_block_add_attention=mid_block_add_attention,
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+            double_z=double_z,
+        )
+
+        self.decoder = LTX2AudioDecoder(
+            base_channels=base_channels,
+            output_channels=output_channels,
+            ch_mult=ch_mult,
+            num_res_blocks=num_res_blocks,
+            attn_resolutions=attn_resolution_set,
+            in_channels=in_channels,
+            resolution=resolution,
+            latent_channels=latent_channels,
+            norm_type=norm_type,
+            causality_axis=causality_axis,
+            dropout=dropout,
+            mid_block_add_attention=mid_block_add_attention,
+            sample_rate=sample_rate,
+            mel_hop_length=mel_hop_length,
+            is_causal=is_causal,
+            mel_bins=mel_bins,
+        )
+
+        # Per-channel statistics for normalizing and denormalizing the latent representation. This statistics is computed over
+        # the entire dataset and stored in model's checkpoint under AudioVAE state_dict
+        latents_std = torch.zeros((base_channels,))
+        latents_mean = torch.ones((base_channels,))
+        self.register_buffer("latents_mean", latents_mean, persistent=True)
+        self.register_buffer("latents_std", latents_std, persistent=True)
+
+        # TODO: confirm whether the mel compression ratio below is correct
+        self.mel_compression_ratio = LATENT_DOWNSAMPLE_FACTOR
+        self.use_slicing = False
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encoder(x)
+
+    def encode(self, x: torch.Tensor, return_dict: bool = True):
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+        posterior = DiagonalGaussianDistribution(h)
+
+        if not return_dict:
+            return (posterior,)
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(z)
+
+    def decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice) for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z)
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        posterior = self.encode(sample).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z)
+        if not return_dict:
+            return (dec.sample,)
+        return dec
+
+
+EntryClass = AutoencoderKLLTX2Audio

python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py

@@ -0,0 +1,193 @@
+import math
+from abc import ABC
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig
+
+
+class ResBlock(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        dilations: Tuple[int, ...] = (1, 3, 5),
+        leaky_relu_negative_slope: float = 0.1,
+        padding_mode: str = "same",
+    ):
+        super().__init__()
+        self.dilations = dilations
+        self.negative_slope = leaky_relu_negative_slope
+
+        self.convs1 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=dilation,
+                    padding=padding_mode,
+                )
+                for dilation in dilations
+            ]
+        )
+
+        self.convs2 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    stride=stride,
+                    dilation=1,
+                    padding=padding_mode,
+                )
+                for _ in range(len(dilations))
+            ]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv1, conv2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, negative_slope=self.negative_slope)
+            xt = conv1(xt)
+            xt = F.leaky_relu(xt, negative_slope=self.negative_slope)
+            xt = conv2(xt)
+            x = x + xt
+        return x
+
+
+class LTX2Vocoder(ABC, nn.Module):
+    r"""
+    LTX 2.0 vocoder for converting generated mel spectrograms back to audio waveforms.
+    """
+
+    def __init__(
+        self,
+        config: LTXVocoderConfig,
+    ):
+        super().__init__()
+        self.config = config
+        self.sample_rate = (
+            getattr(config.arch_config, "sample_rate", None)
+            or getattr(config.arch_config, "sampling_rate", None)
+            or getattr(config.arch_config, "audio_sample_rate", None)
+        )
+
+        in_channels = config.arch_config.in_channels
+        hidden_channels = config.arch_config.hidden_channels
+        out_channels = config.arch_config.out_channels
+        upsample_kernel_sizes = config.arch_config.upsample_kernel_sizes
+        upsample_factors = config.arch_config.upsample_factors
+        resnet_kernel_sizes = config.arch_config.resnet_kernel_sizes
+        resnet_dilations = config.arch_config.resnet_dilations
+        leaky_relu_negative_slope = config.arch_config.leaky_relu_negative_slope
+
+        self.num_upsample_layers = len(upsample_kernel_sizes)
+        self.resnets_per_upsample = len(resnet_kernel_sizes)
+        self.out_channels = out_channels
+        self.total_upsample_factor = math.prod(upsample_factors)
+        self.negative_slope = leaky_relu_negative_slope
+
+        if self.num_upsample_layers != len(upsample_factors):
+            raise ValueError(
+                f"`upsample_kernel_sizes` and `upsample_factors` should be lists of the same length but are length"
+                f" {self.num_upsample_layers} and {len(upsample_factors)}, respectively."
+            )
+
+        if self.resnets_per_upsample != len(resnet_dilations):
+            raise ValueError(
+                f"`resnet_kernel_sizes` and `resnet_dilations` should be lists of the same length but are length"
+                f" {len(self.resnets_per_upsample)} and {len(resnet_dilations)}, respectively."
+            )
+
+        self.conv_in = nn.Conv1d(
+            in_channels, hidden_channels, kernel_size=7, stride=1, padding=3
+        )
+
+        self.upsamplers = nn.ModuleList()
+        self.resnets = nn.ModuleList()
+        input_channels = hidden_channels
+        for i, (stride, kernel_size) in enumerate(
+            zip(upsample_factors, upsample_kernel_sizes)
+        ):
+            output_channels = input_channels // 2
+            self.upsamplers.append(
+                nn.ConvTranspose1d(
+                    input_channels,  # hidden_channels // (2 ** i)
+                    output_channels,  # hidden_channels // (2 ** (i + 1))
+                    kernel_size,
+                    stride=stride,
+                    padding=(kernel_size - stride) // 2,
+                )
+            )
+
+            for kernel_size, dilations in zip(resnet_kernel_sizes, resnet_dilations):
+                self.resnets.append(
+                    ResBlock(
+                        output_channels,
+                        kernel_size,
+                        dilations=dilations,
+                        leaky_relu_negative_slope=leaky_relu_negative_slope,
+                    )
+                )
+            input_channels = output_channels
+
+        self.conv_out = nn.Conv1d(output_channels, out_channels, 7, stride=1, padding=3)
+
+    def forward(
+        self, hidden_states: torch.Tensor, time_last: bool = False
+    ) -> torch.Tensor:
+        r"""
+        Forward pass of the vocoder.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                Input Mel spectrogram tensor of shape `(batch_size, num_channels, time, num_mel_bins)` if `time_last`
+                is `False` (the default) or shape `(batch_size, num_channels, num_mel_bins, time)` if `time_last` is
+                `True`.
+            time_last (`bool`, *optional*, defaults to `False`):
+                Whether the last dimension of the input is the time/frame dimension or the Mel bins dimension.
+
+        Returns:
+            `torch.Tensor`:
+                Audio waveform tensor of shape (batch_size, out_channels, audio_length)
+        """
+
+        # Ensure that the time/frame dimension is last
+        if not time_last:
+            hidden_states = hidden_states.transpose(2, 3)
+        # Combine channels and frequency (mel bins) dimensions
+        hidden_states = hidden_states.flatten(1, 2)
+
+        hidden_states = self.conv_in(hidden_states)
+
+        for i in range(self.num_upsample_layers):
+            hidden_states = F.leaky_relu(
+                hidden_states, negative_slope=self.negative_slope
+            )
+            hidden_states = self.upsamplers[i](hidden_states)
+
+            # Run all resnets in parallel on hidden_states
+            start = i * self.resnets_per_upsample
+            end = (i + 1) * self.resnets_per_upsample
+            resnet_outputs = torch.stack(
+                [self.resnets[j](hidden_states) for j in range(start, end)], dim=0
+            )
+
+            hidden_states = torch.mean(resnet_outputs, dim=0)
+
+        # NOTE: unlike the first leaky ReLU, this leaky ReLU is set to use the default F.leaky_relu negative slope of
+        # 0.01 (whereas the others usually use a slope of 0.1). Not sure if this is intended
+        hidden_states = F.leaky_relu(hidden_states, negative_slope=0.01)
+        hidden_states = self.conv_out(hidden_states)
+        hidden_states = torch.tanh(hidden_states)
+
+        return hidden_states
+
+
+EntryClass = LTX2Vocoder

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -0,0 +1,191 @@
+import inspect
+import json
+import os
+
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    InputValidationStage,
+    LTX2AVDecodingStage,
+    LTX2AVDenoisingStage,
+    LTX2AVLatentPreparationStage,
+    LTX2TextConnectorStage,
+    TextEncodingStage,
+    TimestepPreparationStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+def prepare_mu(batch: Req, server_args: ServerArgs):
+    height = batch.height
+    width = batch.width
+    num_frames = batch.num_frames
+
+    vae_arch = getattr(
+        getattr(server_args.pipeline_config, "vae_config", None), "arch_config", None
+    )
+    vae_scale_factor = (
+        getattr(vae_arch, "spatial_compression_ratio", None)
+        or getattr(vae_arch, "vae_scale_factor", None)
+        or getattr(server_args.pipeline_config, "vae_scale_factor", None)
+    )
+    vae_temporal_compression = getattr(
+        vae_arch, "temporal_compression_ratio", None
+    ) or getattr(server_args.pipeline_config, "vae_temporal_compression", None)
+
+    latent_num_frames = (int(num_frames) - 1) // int(vae_temporal_compression) + 1
+    latent_height = int(height) // int(vae_scale_factor)
+    latent_width = int(width) // int(vae_scale_factor)
+    video_sequence_length = latent_num_frames * latent_height * latent_width
+
+    # Values from LTX2Pipeline in diffusers
+    mu = calculate_shift(
+        video_sequence_length,
+        base_seq_len=1024,
+        max_seq_len=4096,
+        base_shift=0.95,
+        max_shift=2.05,
+    )
+    return "mu", mu
+
+
+def _load_component_config(model_path: str, component_name: str):
+    """Helper to load component config from model_index.json or config.json"""
+    try:
+        # Try loading model_index.json first
+        index_path = os.path.join(model_path, "model_index.json")
+        if os.path.exists(index_path):
+            with open(index_path, "r") as f:
+                index = json.load(f)
+
+            if component_name in index:
+                # It's a subfolder
+                subfolder = index[component_name][1]
+                config_path = os.path.join(model_path, subfolder, "config.json")
+                if os.path.exists(config_path):
+                    with open(config_path, "r") as f:
+                        return json.load(f)
+
+        # Fallback to direct config.json in subfolder if standard structure
+        config_path = os.path.join(model_path, component_name, "config.json")
+        if os.path.exists(config_path):
+            with open(config_path, "r") as f:
+                return json.load(f)
+
+    except Exception as e:
+        logger.warning(f"Failed to load config for {component_name}: {e}")
+
+    return {}
+
+
+def _filter_kwargs_for_cls(cls, kwargs):
+    """Filter kwargs to only include those accepted by cls.__init__"""
+    sig = inspect.signature(cls.__init__)
+    return {k: v for k, v in kwargs.items() if k in sig.parameters}
+
+
+class LTX2Pipeline(ComposedPipelineBase):
+    # NOTE: must match `model_index.json`'s `_class_name` for native dispatch.
+    pipeline_name = "LTX2Pipeline"
+
+    _required_config_modules = [
+        "transformer",
+        "text_encoder",
+        "tokenizer",
+        "scheduler",
+        "vae",
+        "audio_vae",
+        "vocoder",
+        "connectors",
+    ]
+
+    def create_pipeline_stages(self, server_args: ServerArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        # 1. Input Validation
+        self.add_stage(
+            stage_name="input_validation_stage", stage=InputValidationStage()
+        )
+
+        # 2. Text Encoding
+        self.add_stage(
+            stage_name="text_encoding_stage",
+            stage=TextEncodingStage(
+                # LTX-2 needs two contexts (video/audio). We reuse the same
+                # underlying Gemma encoder/tokenizer twice.
+                text_encoders=[
+                    self.get_module("text_encoder"),
+                ],
+                tokenizers=[
+                    self.get_module("tokenizer"),
+                ],
+            ),
+        )
+
+        # 3. connector stage
+        self.add_stage(
+            stage_name="text_connector_stage",
+            stage=LTX2TextConnectorStage(connectors=self.get_module("connectors")),
+        )
+
+        # 4. Timestep Preparation
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=TimestepPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                prepare_extra_set_timesteps_kwargs=[prepare_mu],
+            ),
+        )
+
+        # 4. Latent Preparation
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=LTX2AVLatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                transformer=self.get_module("transformer"),
+                audio_vae=self.get_module("audio_vae"),
+            ),
+        )
+
+        # 5. Denoising
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=LTX2AVDenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae"),
+                audio_vae=self.get_module("audio_vae"),
+            ),
+        )
+
+        # 6. Decoding
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=LTX2AVDecodingStage(
+                vae=self.get_module("vae"),
+                audio_vae=self.get_module("audio_vae"),
+                vocoder=self.get_module("vocoder"),
+                pipeline=self,
+            ),
+        )
+
+
+EntryClass = LTX2Pipeline

python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py

@@ -19,7 +19,13 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.conditioning import (
     ConditioningStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import (
+    LTX2AVDecodingStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import (
+    LTX2AVDenoisingStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import (
     DmdDenoisingStage,
 )
@@ -34,6 +40,12 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import
 from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
     LatentPreparationStage,
 )
+from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import (
+    LTX2AVLatentPreparationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import (
+    LTX2TextConnectorStage,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
     TextEncodingStage,
 )
@@ -47,13 +59,17 @@ __all__ = [
     "TimestepPreparationStage",
     "LatentPreparationStage",
     "ComfyUILatentPreparationStage",
+    "LTX2AVLatentPreparationStage",
     "ConditioningStage",
     "DenoisingStage",
     "DmdDenoisingStage",
+    "LTX2AVDenoisingStage",
     "CausalDMDDenoisingStage",
     "EncodingStage",
     "DecodingStage",
+    "LTX2AVDecodingStage",
     "ImageEncodingStage",
     "ImageVAEEncodingStage",
     "TextEncodingStage",
+    "LTX2TextConnectorStage",
 ]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py

@@ -0,0 +1,147 @@
+import torch
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class LTX2AVDecodingStage(DecodingStage):
+    """
+    LTX-2 specific decoding stage that handles both video and audio decoding.
+    """
+
+    def __init__(self, vae, audio_vae, vocoder, pipeline=None):
+        super().__init__(vae, pipeline)
+        self.audio_vae = audio_vae
+        self.vocoder = vocoder
+        # Add video processor for postprocessing
+        from diffusers.video_processor import VideoProcessor
+
+        self.video_processor = VideoProcessor(vae_scale_factor=32)
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
+        self.load_model()
+
+        self.vae = self.vae.to(get_local_torch_device())
+        self.vae.eval()
+        latents = batch.latents.to(get_local_torch_device())
+
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32
+        ) and not server_args.disable_autocast
+
+        latents = self.scale_and_shift(latents, server_args)
+        latents = server_args.pipeline_config.preprocess_decoding(
+            latents, server_args, vae=self.vae
+        )
+
+        with torch.autocast(
+            device_type=current_platform.device_type,
+            dtype=vae_dtype,
+            enabled=vae_autocast_enabled,
+        ):
+            try:
+                if server_args.pipeline_config.vae_tiling:
+                    self.vae.enable_tiling()
+            except Exception:
+                pass
+            if not vae_autocast_enabled:
+                latents = latents.to(vae_dtype)
+            decode_output = self.vae.decode(latents)
+            if isinstance(decode_output, tuple):
+                video = decode_output[0]
+            elif hasattr(decode_output, "sample"):
+                video = decode_output.sample
+            else:
+                video = decode_output
+
+        video = self.video_processor.postprocess_video(video, output_type="np")
+
+        output_batch = OutputBatch(
+            output=video,
+            trajectory_timesteps=batch.trajectory_timesteps,
+            trajectory_latents=batch.trajectory_latents,
+            trajectory_decoded=None,
+            timings=batch.timings,
+        )
+
+        # 2. Decode Audio
+        try:
+            audio_latents = batch.audio_latents
+        except AttributeError:
+            audio_latents = None
+        if audio_latents is not None:
+            # Ensure device/dtype
+            device = get_local_torch_device()
+            self.audio_vae = self.audio_vae.to(device)
+            self.vocoder = self.vocoder.to(device)
+            self.audio_vae.eval()
+            self.vocoder.eval()
+            try:
+                dtype = self.audio_vae.dtype
+            except AttributeError:
+                dtype = None
+            if dtype is None:
+                try:
+                    dtype = next(self.audio_vae.parameters()).dtype
+                except StopIteration:
+                    dtype = torch.float32
+            audio_latents = audio_latents.to(device, dtype=dtype)
+            try:
+                latents_std = self.audio_vae.latents_std
+            except AttributeError:
+                latents_std = None
+            if isinstance(latents_std, torch.Tensor) and torch.all(latents_std == 0):
+                logger.warning(
+                    "audio_vae.latents_std is all zeros; audio denorm may be incorrect."
+                )
+
+            with torch.no_grad():
+                # Decode latents to spectrogram
+                spectrogram = self.audio_vae.decode(audio_latents, return_dict=False)[0]
+                if hasattr(self.vocoder, "conv_in") and hasattr(
+                    self.vocoder.conv_in, "in_channels"
+                ):
+                    expected_in = int(self.vocoder.conv_in.in_channels)
+                    actual_in = int(spectrogram.shape[1]) * int(spectrogram.shape[3])
+                    if actual_in != expected_in:
+                        raise ValueError(
+                            f"Vocoder expects channels*mel_bins={expected_in}, got {actual_in} from spectrogram shape {tuple(spectrogram.shape)}"
+                        )
+                # Decode spectrogram to waveform
+                waveform = self.vocoder(spectrogram)
+            output_batch.audio = waveform.cpu().float()
+            try:
+                pipeline_audio_cfg = server_args.pipeline_config.audio_vae_config
+            except AttributeError:
+                pipeline_audio_cfg = None
+            try:
+                pipeline_audio_arch = pipeline_audio_cfg.arch_config  # type: ignore[union-attr]
+            except AttributeError:
+                pipeline_audio_arch = None
+            try:
+                pipeline_audio_sr = pipeline_audio_arch.sample_rate  # type: ignore[union-attr]
+            except AttributeError:
+                pipeline_audio_sr = None
+
+            try:
+                vocoder_sr = self.vocoder.sample_rate
+            except AttributeError:
+                vocoder_sr = None
+            try:
+                audio_vae_sr = self.audio_vae.sample_rate
+            except AttributeError:
+                audio_vae_sr = None
+            output_batch.audio_sample_rate = (
+                vocoder_sr or audio_vae_sr or pipeline_audio_sr
+            )
+
+        self.offload_model()
+        return output_batch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py

@@ -0,0 +1,729 @@
+import copy
+import time
+
+import PIL.Image
+import torch
+from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.models.vision_utils import (
+    load_image,
+    normalize,
+    numpy_to_pt,
+    pil_to_numpy,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    VerificationResult,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class LTX2AVDenoisingStage(DenoisingStage):
+    """
+    LTX-2 specific denoising stage that handles joint video and audio generation.
+    """
+
+    def __init__(self, transformer, scheduler, vae=None, audio_vae=None, **kwargs):
+        super().__init__(
+            transformer=transformer, scheduler=scheduler, vae=vae, **kwargs
+        )
+        self.audio_vae = audio_vae
+
+    @staticmethod
+    def _get_video_latent_num_frames_for_model(
+        batch: Req, server_args: ServerArgs, latents: torch.Tensor
+    ) -> int:
+        """Return the latent-frame length the DiT model should see.
+
+        - If video latents were time-sharded for SP and are packed as token latents
+          ([B, S, D]), the model only sees the local shard and must use the local
+          latent-frame count (stored on the batch during SP sharding).
+        - Otherwise, fall back to the global latent-frame count inferred from the
+          requested output frames and the VAE temporal compression ratio.
+        """
+        did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False))
+        is_token_latents = isinstance(latents, torch.Tensor) and latents.ndim == 3
+
+        if did_sp_shard and is_token_latents:
+            if not hasattr(batch, "sp_video_latent_num_frames"):
+                raise ValueError(
+                    "SP-sharded LTX2 token latents require `batch.sp_video_latent_num_frames` "
+                    "to be set by `LTX2PipelineConfig.shard_latents_for_sp()`."
+                )
+            return int(batch.sp_video_latent_num_frames)
+
+        pc = server_args.pipeline_config
+        return int((batch.num_frames - 1) // int(pc.vae_temporal_compression) + 1)
+
+    @staticmethod
+    def _truncate_sp_padded_token_latents(
+        batch: Req, latents: torch.Tensor
+    ) -> torch.Tensor:
+        """Remove token padding introduced by SP time-sharding (if applicable)."""
+        did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False))
+        if not did_sp_shard or not (
+            isinstance(latents, torch.Tensor) and latents.ndim == 3
+        ):
+            return latents
+
+        raw_shape = getattr(batch, "raw_latent_shape", None)
+        if not (isinstance(raw_shape, tuple) and len(raw_shape) == 3):
+            return latents
+
+        orig_s = int(raw_shape[1])
+        cur_s = int(latents.shape[1])
+        if cur_s == orig_s:
+            return latents
+        if cur_s < orig_s:
+            raise ValueError(
+                f"Unexpected gathered token-latents seq_len {cur_s} < original seq_len {orig_s}."
+            )
+        return latents[:, :orig_s, :].contiguous()
+
+    def _maybe_enable_cache_dit(self, num_inference_steps: int, batch: Req) -> None:
+        """Disable cache-dit for TI2V-style requests (image-conditioned), to avoid stale activations.
+
+        NOTE: base denoising stage calls this hook with (num_inference_steps, batch).
+        """
+        if getattr(self, "_disable_cache_dit_for_request", False):
+            return
+        return super()._maybe_enable_cache_dit(num_inference_steps, batch)
+
+    @staticmethod
+    def _resize_center_crop(
+        img: PIL.Image.Image, *, width: int, height: int
+    ) -> PIL.Image.Image:
+        return img.resize((width, height), resample=PIL.Image.Resampling.BILINEAR)
+
+    @staticmethod
+    def _pil_to_normed_tensor(img: PIL.Image.Image) -> torch.Tensor:
+        # PIL -> numpy [0,1] -> torch [B,C,H,W], then [-1,1]
+        arr = pil_to_numpy(img)
+        t = numpy_to_pt(arr)
+        return normalize(t)
+
+    @staticmethod
+    def _should_apply_ltx2_ti2v(batch: Req) -> bool:
+        """True if we have an image-latent token prefix to condition with.
+
+        SP note: when token latents are time-sharded, only the rank that owns the
+        *global* first latent frame should apply TI2V conditioning (rank with start_frame==0).
+        """
+        if (
+            batch.image_latent is None
+            or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0
+        ):
+            return False
+        did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False))
+        if not did_sp_shard:
+            return True
+        return int(getattr(batch, "sp_video_start_frame", 0)) == 0
+
+    def _prepare_ltx2_image_latent(self, batch: Req, server_args: ServerArgs) -> None:
+        """Encode `batch.image_path` into packed token latents for LTX-2 TI2V."""
+        if (
+            batch.image_latent is not None
+            and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0
+        ):
+            return
+        batch.ltx2_num_image_tokens = 0
+        batch.image_latent = None
+
+        if batch.image_path is None:
+            return
+        if batch.width is None or batch.height is None:
+            raise ValueError("width/height must be provided for LTX-2 TI2V.")
+        if self.vae is None:
+            raise ValueError("VAE must be provided for LTX-2 TI2V.")
+
+        image_path = (
+            batch.image_path[0]
+            if isinstance(batch.image_path, list)
+            else batch.image_path
+        )
+
+        img = load_image(image_path)
+        img = self._resize_center_crop(
+            img, width=int(batch.width), height=int(batch.height)
+        )
+        batch.condition_image = img
+
+        latents_device = (
+            batch.latents.device
+            if isinstance(batch.latents, torch.Tensor)
+            else torch.device("cpu")
+        )
+        image_tensor = self._pil_to_normed_tensor(img).to(
+            latents_device, dtype=torch.float32
+        )
+        # [B, C, H, W] -> [B, C, 1, H, W]
+        video_condition = image_tensor.unsqueeze(2)
+
+        self.vae = self.vae.to(latents_device)
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32
+        ) and not server_args.disable_autocast
+
+        with torch.autocast(
+            device_type=current_platform.device_type,
+            dtype=vae_dtype,
+            enabled=vae_autocast_enabled,
+        ):
+            try:
+                if server_args.pipeline_config.vae_tiling:
+                    self.vae.enable_tiling()
+            except Exception:
+                pass
+            if not vae_autocast_enabled:
+                video_condition = video_condition.to(vae_dtype)
+
+            latent_dist: DiagonalGaussianDistribution = self.vae.encode(video_condition)
+            if isinstance(latent_dist, AutoencoderKLOutput):
+                latent_dist = latent_dist.latent_dist
+
+        mode = server_args.pipeline_config.vae_config.encode_sample_mode()
+        if mode == "argmax":
+            latent = latent_dist.mode()
+        elif mode == "sample":
+            if batch.generator is None:
+                raise ValueError("Generator must be provided for VAE sampling.")
+            latent = latent_dist.sample(batch.generator)
+        else:
+            raise ValueError(f"Unsupported encode_sample_mode: {mode}")
+
+        # Match the normalized latent space used by this pipeline (inverse of DecodingStage.scale_and_shift).
+        scaling_factor, shift_factor = (
+            server_args.pipeline_config.get_decode_scale_and_shift(
+                device=latent.device, dtype=latent.dtype, vae=self.vae
+            )
+        )
+        if isinstance(shift_factor, torch.Tensor):
+            shift_factor = shift_factor.to(latent.device)
+        if isinstance(scaling_factor, torch.Tensor):
+            scaling_factor = scaling_factor.to(latent.device)
+        if shift_factor is not None:
+            latent = latent - shift_factor
+        latent = latent * scaling_factor
+
+        packed = server_args.pipeline_config.maybe_pack_latents(
+            latent, latent.shape[0], batch
+        )
+        if not (isinstance(packed, torch.Tensor) and packed.ndim == 3):
+            raise ValueError("Expected packed image latents [B, S0, D].")
+
+        # Fail-fast token count: must match one latent frame's tokens.
+        vae_sf = int(server_args.pipeline_config.vae_scale_factor)
+        patch = int(server_args.pipeline_config.patch_size)
+        latent_h = int(batch.height) // vae_sf
+        latent_w = int(batch.width) // vae_sf
+        expected_tokens = (latent_h // patch) * (latent_w // patch)
+        if int(packed.shape[1]) != int(expected_tokens):
+            raise ValueError(
+                "LTX-2 conditioning token count mismatch: "
+                f"{int(packed.shape[1])=} {int(expected_tokens)=}."
+            )
+
+        batch.image_latent = packed
+        batch.ltx2_num_image_tokens = int(packed.shape[1])
+
+        if batch.debug:
+            logger.info(
+                "LTX2 TI2V conditioning prepared: %d tokens (shape=%s) for %sx%s",
+                batch.ltx2_num_image_tokens,
+                tuple(batch.image_latent.shape),
+                batch.width,
+                batch.height,
+            )
+
+        if server_args.vae_cpu_offload:
+            self.vae = self.vae.to("cpu")
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        """
+         Run the denoising loop.
+
+        Args:
+            batch: The current batch information.
+            server_args: The inference arguments.
+
+        Returns:
+            The batch with denoised latents.
+        """
+        # Disable cache-dit for image-conditioned requests (TI2V-style) for correctness/debuggability.
+        self._disable_cache_dit_for_request = batch.image_path is not None
+
+        # Prepare variables for the denoising loop
+
+        prepared_vars = self._prepare_denoising_loop(batch, server_args)
+        extra_step_kwargs = prepared_vars["extra_step_kwargs"]
+        target_dtype = prepared_vars["target_dtype"]
+        autocast_enabled = prepared_vars["autocast_enabled"]
+        timesteps = prepared_vars["timesteps"]
+        num_inference_steps = prepared_vars["num_inference_steps"]
+        num_warmup_steps = prepared_vars["num_warmup_steps"]
+        image_kwargs = prepared_vars["image_kwargs"]
+        pos_cond_kwargs = prepared_vars["pos_cond_kwargs"]
+        neg_cond_kwargs = prepared_vars["neg_cond_kwargs"]
+        latents = prepared_vars["latents"]
+        boundary_timestep = prepared_vars["boundary_timestep"]
+        z = prepared_vars["z"]
+        reserved_frames_mask = prepared_vars["reserved_frames_mask"]
+        seq_len = prepared_vars["seq_len"]
+        guidance = prepared_vars["guidance"]
+
+        audio_latents = batch.audio_latents
+        audio_scheduler = copy.deepcopy(self.scheduler)
+
+        # Prepare TI2V conditioning once (encode image -> patchify tokens).
+        self._prepare_ltx2_image_latent(batch, server_args)
+
+        # For LTX-2 packed token latents, SP sharding happens on the time dimension
+        # (frames). The model must see local latent frames (RoPE offset is applied
+        # inside the model using SP rank).
+        latent_num_frames_for_model = self._get_video_latent_num_frames_for_model(
+            batch=batch, server_args=server_args, latents=latents
+        )
+        latent_height = batch.height // server_args.pipeline_config.vae_scale_factor
+        latent_width = batch.width // server_args.pipeline_config.vae_scale_factor
+
+        # Initialize lists for ODE trajectory
+        trajectory_timesteps: list[torch.Tensor] = []
+        trajectory_latents: list[torch.Tensor] = []
+        trajectory_audio_latents: list[torch.Tensor] = []
+
+        # Run denoising loop
+        denoising_start_time = time.time()
+
+        # to avoid device-sync caused by timestep comparison
+        is_warmup = batch.is_warmup
+        self.scheduler.set_begin_index(0)
+        audio_scheduler.set_begin_index(0)
+        timesteps_cpu = timesteps.cpu()
+        num_timesteps = timesteps_cpu.shape[0]
+
+        do_ti2v = self._should_apply_ltx2_ti2v(batch)
+        num_img_tokens = int(getattr(batch, "ltx2_num_image_tokens", 0))
+        denoise_mask = None
+        clean_latent = None
+        if do_ti2v:
+            if not (isinstance(latents, torch.Tensor) and latents.ndim == 3):
+                raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].")
+            latents[:, :num_img_tokens, :] = batch.image_latent[
+                :, :num_img_tokens, :
+            ].to(device=latents.device, dtype=latents.dtype)
+            denoise_mask = torch.ones(
+                (latents.shape[0], latents.shape[1], 1),
+                device=latents.device,
+                dtype=torch.float32,
+            )
+            denoise_mask[:, :num_img_tokens, :] = 0.0
+            clean_latent = latents.detach().clone()
+            clean_latent[:, :num_img_tokens, :] = batch.image_latent[
+                :, :num_img_tokens, :
+            ].to(device=latents.device, dtype=latents.dtype)
+
+        with torch.autocast(
+            device_type=current_platform.device_type,
+            dtype=target_dtype,
+            enabled=autocast_enabled,
+        ):
+            with self.progress_bar(total=num_inference_steps) as progress_bar:
+                for i, t_host in enumerate(timesteps_cpu):
+                    with StageProfiler(
+                        f"denoising_step_{i}",
+                        logger=logger,
+                        timings=batch.timings,
+                        perf_dump_path_provided=batch.perf_dump_path is not None,
+                    ):
+                        t_int = int(t_host.item())
+                        t_device = timesteps[i]
+                        current_model, current_guidance_scale = (
+                            self._select_and_manage_model(
+                                t_int=t_int,
+                                boundary_timestep=boundary_timestep,
+                                server_args=server_args,
+                                batch=batch,
+                            )
+                        )
+
+                        # Predict noise residual
+                        attn_metadata = self._build_attn_metadata(i, batch, server_args)
+
+                        # === LTX-2 sigma-space Euler step (flow matching) ===
+                        # Use scheduler-generated sigmas (includes terminal sigma=0).
+                        sigmas = getattr(self.scheduler, "sigmas", None)
+                        if sigmas is None or not isinstance(sigmas, torch.Tensor):
+                            raise ValueError(
+                                "Expected scheduler.sigmas to be a tensor for LTX-2."
+                            )
+                        sigma = sigmas[i].to(device=latents.device, dtype=torch.float32)
+                        sigma_next = sigmas[i + 1].to(
+                            device=latents.device, dtype=torch.float32
+                        )
+                        dt = sigma_next - sigma
+
+                        latent_model_input = latents.to(target_dtype)
+                        audio_latent_model_input = audio_latents.to(target_dtype)
+
+                        latent_num_frames = latent_num_frames_for_model
+
+                        # Audio latent dims
+                        if audio_latent_model_input.ndim == 3:
+                            audio_num_frames_latent = int(
+                                audio_latent_model_input.shape[1]
+                            )
+                        elif audio_latent_model_input.ndim == 4:
+                            audio_num_frames_latent = int(
+                                audio_latent_model_input.shape[2]
+                            )
+                        else:
+                            raise ValueError(
+                                f"Unexpected audio latents rank: {audio_latent_model_input.ndim}, shape={tuple(audio_latent_model_input.shape)}"
+                            )
+
+                        # LTX-2 model can generate coords internally.
+                        video_coords = None
+                        audio_coords = None
+
+                        timestep = t_device.expand(int(latent_model_input.shape[0]))
+                        if do_ti2v and denoise_mask is not None:
+                            timestep_video = timestep.unsqueeze(
+                                -1
+                            ) * denoise_mask.squeeze(-1)
+                        else:
+                            timestep_video = timestep
+                        timestep_audio = timestep
+
+                        # Conditions
+                        encoder_hidden_states = batch.prompt_embeds[0]
+                        audio_encoder_hidden_states = batch.audio_prompt_embeds[0]
+                        encoder_attention_mask = batch.prompt_attention_mask
+
+                        # Follow ltx-pipelines structure: separate pos/neg forward passes,
+                        # then apply CFG on denoised (x0) predictions.
+                        with set_forward_context(
+                            current_timestep=i, attn_metadata=attn_metadata
+                        ):
+                            v_pos, a_v_pos = current_model(
+                                hidden_states=latent_model_input,
+                                audio_hidden_states=audio_latent_model_input,
+                                encoder_hidden_states=encoder_hidden_states,
+                                audio_encoder_hidden_states=audio_encoder_hidden_states,
+                                timestep=timestep_video,
+                                audio_timestep=timestep_audio,
+                                encoder_attention_mask=encoder_attention_mask,
+                                audio_encoder_attention_mask=encoder_attention_mask,
+                                num_frames=latent_num_frames,
+                                height=latent_height,
+                                width=latent_width,
+                                fps=batch.fps,
+                                audio_num_frames=audio_num_frames_latent,
+                                video_coords=video_coords,
+                                audio_coords=audio_coords,
+                                return_latents=False,
+                                return_dict=False,
+                            )
+
+                            if batch.do_classifier_free_guidance:
+                                neg_encoder_hidden_states = (
+                                    batch.negative_prompt_embeds[0]
+                                )
+                                neg_audio_encoder_hidden_states = (
+                                    batch.negative_audio_prompt_embeds[0]
+                                )
+                                neg_encoder_attention_mask = (
+                                    batch.negative_attention_mask
+                                )
+
+                                v_neg, a_v_neg = current_model(
+                                    hidden_states=latent_model_input,
+                                    audio_hidden_states=audio_latent_model_input,
+                                    encoder_hidden_states=neg_encoder_hidden_states,
+                                    audio_encoder_hidden_states=neg_audio_encoder_hidden_states,
+                                    timestep=timestep_video,
+                                    audio_timestep=timestep_audio,
+                                    encoder_attention_mask=neg_encoder_attention_mask,
+                                    audio_encoder_attention_mask=neg_encoder_attention_mask,
+                                    num_frames=latent_num_frames,
+                                    height=latent_height,
+                                    width=latent_width,
+                                    fps=batch.fps,
+                                    audio_num_frames=audio_num_frames_latent,
+                                    video_coords=video_coords,
+                                    audio_coords=audio_coords,
+                                    return_latents=False,
+                                    return_dict=False,
+                                )
+                            else:
+                                v_neg = None
+                                a_v_neg = None
+
+                        v_pos = v_pos.float()
+                        a_v_pos = a_v_pos.float()
+                        if v_neg is not None:
+                            v_neg = v_neg.float()
+                        if a_v_neg is not None:
+                            a_v_neg = a_v_neg.float()
+
+                        # Velocity -> denoised (x0): x0 = x - sigma * v
+                        sigma_val = float(sigma.item())
+                        denoised_video = latents.float() - sigma_val * v_pos
+                        denoised_audio = audio_latents.float() - sigma_val * a_v_pos
+
+                        if (
+                            batch.do_classifier_free_guidance
+                            and v_neg is not None
+                            and a_v_neg is not None
+                        ):
+                            denoised_video_neg = latents.float() - sigma_val * v_neg
+                            denoised_audio_neg = (
+                                audio_latents.float() - sigma_val * a_v_neg
+                            )
+                            denoised_video = denoised_video + (
+                                batch.guidance_scale - 1.0
+                            ) * (denoised_video - denoised_video_neg)
+                            denoised_audio = denoised_audio + (
+                                batch.guidance_scale - 1.0
+                            ) * (denoised_audio - denoised_audio_neg)
+
+                        # Apply conditioning mask (keep conditioned tokens clean).
+                        if (
+                            do_ti2v
+                            and denoise_mask is not None
+                            and clean_latent is not None
+                        ):
+                            denoised_video = (
+                                denoised_video * denoise_mask
+                                + clean_latent.float() * (1.0 - denoise_mask)
+                            )
+
+                        # Euler step in sigma space: x_next = x + (sigma_next - sigma) * v,
+                        # where v = (x - x0) / sigma.
+                        if sigma_val == 0.0:
+                            v_video = torch.zeros_like(denoised_video)
+                            v_audio = torch.zeros_like(denoised_audio)
+                        else:
+                            v_video = (latents.float() - denoised_video) / sigma_val
+                            v_audio = (
+                                audio_latents.float() - denoised_audio
+                            ) / sigma_val
+
+                        latents = (latents.float() + v_video * dt).to(
+                            dtype=latents.dtype
+                        )
+                        audio_latents = (audio_latents.float() + v_audio * dt).to(
+                            dtype=audio_latents.dtype
+                        )
+
+                        if do_ti2v:
+                            latents[:, :num_img_tokens, :] = batch.image_latent[
+                                :, :num_img_tokens, :
+                            ].to(device=latents.device, dtype=latents.dtype)
+
+                        latents = self.post_forward_for_ti2v_task(
+                            batch, server_args, reserved_frames_mask, latents, z
+                        )
+
+                        # save trajectory latents if needed
+                        if batch.return_trajectory_latents:
+                            trajectory_timesteps.append(t_host)
+                            trajectory_latents.append(latents)
+                            if audio_latents is not None:
+                                trajectory_audio_latents.append(audio_latents)
+
+                        # Update progress bar
+                        if i == num_timesteps - 1 or (
+                            (i + 1) > num_warmup_steps
+                            and (i + 1) % self.scheduler.order == 0
+                            and progress_bar is not None
+                        ):
+                            progress_bar.update()
+
+                        if not is_warmup:
+                            self.step_profile()
+
+        denoising_end_time = time.time()
+
+        if num_timesteps > 0 and not is_warmup:
+            self.log_info(
+                "average time per step: %.4f seconds",
+                (denoising_end_time - denoising_start_time) / len(timesteps),
+            )
+
+        batch.audio_latents = audio_latents
+        self._post_denoising_loop(
+            batch=batch,
+            latents=latents,
+            trajectory_latents=trajectory_latents,
+            trajectory_timesteps=trajectory_timesteps,
+            trajectory_audio_latents=trajectory_audio_latents,
+            server_args=server_args,
+            is_warmup=is_warmup,
+        )
+
+        return batch
+
+    def _post_denoising_loop(
+        self,
+        batch: Req,
+        latents: torch.Tensor,
+        trajectory_latents: list,
+        trajectory_timesteps: list,
+        trajectory_audio_latents: list,
+        server_args: ServerArgs,
+        is_warmup: bool = False,
+    ):
+        # 1. Handle Trajectory (Video) - Copy from base
+        if trajectory_latents:
+            trajectory_tensor = torch.stack(trajectory_latents, dim=1)
+            trajectory_timesteps_tensor = torch.stack(trajectory_timesteps, dim=0)
+        else:
+            trajectory_tensor = None
+            trajectory_timesteps_tensor = None
+
+        latents, trajectory_tensor = self._postprocess_sp_latents(
+            batch, latents, trajectory_tensor
+        )
+
+        # If SP time-sharding padded whole frames worth of tokens, remove padding
+        # after gather and before unpacking.
+        latents = self._truncate_sp_padded_token_latents(batch, latents)
+
+        if trajectory_tensor is not None and trajectory_timesteps_tensor is not None:
+            batch.trajectory_timesteps = trajectory_timesteps_tensor.cpu()
+            batch.trajectory_latents = trajectory_tensor.cpu()
+
+        # 2. Handle Trajectory (Audio) - LTX-2 specific
+        if trajectory_audio_latents:
+            trajectory_audio_tensor = torch.stack(trajectory_audio_latents, dim=1)
+            # We don't have SP support for audio latents yet (or needed?)
+            batch.trajectory_audio_latents = trajectory_audio_tensor.cpu()
+
+        # 3. Unpack and Denormalize
+        # Call pipeline_config._unpad_and_unpack_latents
+        # latents is video latents.
+        # batch.audio_latents is audio latents.
+
+        audio_latents = batch.audio_latents
+
+        # NOTE: self.vae and self.audio_vae should be populated via __init__ or manual setting
+        if self.vae is None or self.audio_vae is None:
+            logger.warning(
+                "VAE or Audio VAE not found in DenoisingStage. Skipping unpack and denormalize."
+            )
+            batch.latents = latents
+            batch.audio_latents = audio_latents
+        else:
+            latents, audio_latents = (
+                server_args.pipeline_config._unpad_and_unpack_latents(
+                    latents, audio_latents, batch, self.vae, self.audio_vae
+                )
+            )
+
+            batch.latents = latents
+            batch.audio_latents = audio_latents
+
+        # 4. Cleanup
+        offload_mgr = getattr(self.transformer, "_layerwise_offload_manager", None)
+        if offload_mgr is not None and getattr(offload_mgr, "enabled", False):
+            offload_mgr.release_all()
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        """Verify denoising stage inputs.
+
+        Note: LTX-2 connector stage converts `prompt_embeds`/`negative_prompt_embeds`
+        from list-of-tensors to a single tensor (video context) and stores audio
+        context separately.
+        """
+
+        result = VerificationResult()
+        result.add_check("timesteps", batch.timesteps, [V.is_tensor, V.min_dims(1)])
+
+        # LTX-2 may carry prompt embeddings as either a tensor (preferred) or legacy list.
+        result.add_check(
+            "prompt_embeds",
+            batch.prompt_embeds,
+            lambda x: V.is_tensor(x) or V.list_not_empty(x),
+        )
+
+        # Keep base expectation: image_embeds is always a list (may be empty).
+        result.add_check("image_embeds", batch.image_embeds, V.is_list)
+
+        result.add_check(
+            "num_inference_steps", batch.num_inference_steps, V.positive_int
+        )
+        result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float)
+        result.add_check("eta", batch.eta, V.non_negative_float)
+        result.add_check("generator", batch.generator, V.generator_or_list_generators)
+        result.add_check(
+            "do_classifier_free_guidance",
+            batch.do_classifier_free_guidance,
+            V.bool_value,
+        )
+
+        # When CFG is enabled, negative prompt embeddings must exist (tensor or legacy list).
+        result.add_check(
+            "negative_prompt_embeds",
+            batch.negative_prompt_embeds,
+            lambda x: (not batch.do_classifier_free_guidance)
+            or V.is_tensor(x)
+            or V.list_not_empty(x),
+        )
+        return result
+
+    def do_classifier_free_guidance(self, batch: Req) -> bool:
+        return batch.guidance_scale > 1.0
+
+
+class LTX2RefinementStage(LTX2AVDenoisingStage):
+    def __init__(
+        self, transformer, scheduler, distilled_sigmas, vae=None, audio_vae=None
+    ):
+        super().__init__(transformer, scheduler, vae, audio_vae)
+        self.distilled_sigmas = torch.tensor(distilled_sigmas)
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        # 1. Add noise to latents
+        noise_scale = self.distilled_sigmas[0].to(batch.latents.device)
+        noise = torch.randn_like(batch.latents)
+        batch.latents = batch.latents + noise * noise_scale
+
+        # 2. Run denoising loop with distilled_sigmas
+        # Save original sigmas
+        original_sigmas = self.scheduler.sigmas
+        original_timesteps = self.scheduler.timesteps
+        original_num_inference_steps = self.scheduler.num_inference_steps
+
+        # Set distilled sigmas
+        self.scheduler.sigmas = self.distilled_sigmas.to(self.scheduler.sigmas.device)
+        # Approximation for timesteps
+        self.scheduler.timesteps = self.scheduler.sigmas * 1000
+        self.scheduler.num_inference_steps = len(self.distilled_sigmas) - 1
+
+        # Call parent forward
+        try:
+            batch = super().forward(batch, server_args)
+        finally:
+            # Restore original sigmas
+            self.scheduler.sigmas = original_sigmas
+            self.scheduler.timesteps = original_timesteps
+            self.scheduler.num_inference_steps = original_num_inference_steps
+
+        return batch
+
+    def do_classifier_free_guidance(self, batch: Req) -> bool:
+        return False  # Stage 2 uses simple denoising (no CFG)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py

@@ -0,0 +1,104 @@
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
+    LatentPreparationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    VerificationResult,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class LTX2AVLatentPreparationStage(LatentPreparationStage):
+    """
+    LTX-2 specific latent preparation stage that handles both video and audio latents.
+    """
+
+    def __init__(self, scheduler, transformer=None, audio_vae=None):
+        super().__init__(scheduler, transformer)
+        self.audio_vae = audio_vae
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        """Verify latent preparation stage inputs."""
+        result = VerificationResult()
+        result.add_check(
+            "prompt_or_embeds",
+            None,
+            lambda _: V.string_or_list_strings(batch.prompt)
+            or V.list_not_empty(batch.prompt_embeds)
+            or V.is_tensor(batch.prompt_embeds),
+        )
+
+        if isinstance(batch.prompt_embeds, list):
+            result.add_check("prompt_embeds", batch.prompt_embeds, V.list_of_tensors)
+        else:
+            result.add_check("prompt_embeds", batch.prompt_embeds, V.is_tensor)
+
+        result.add_check(
+            "num_videos_per_prompt", batch.num_outputs_per_prompt, V.positive_int
+        )
+        result.add_check("generator", batch.generator, V.generator_or_list_generators)
+        result.add_check("num_frames", batch.num_frames, V.positive_int)
+        result.add_check("height", batch.height, V.positive_int)
+        result.add_check("width", batch.width, V.positive_int)
+        result.add_check("latents", batch.latents, V.none_or_tensor)
+        return result
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        # 1. Prepare Video Latents using base class logic
+        # This sets batch.latents and batch.raw_latent_shape
+        batch = super().forward(batch, server_args)
+
+        # 2. Prepare Audio Latents (optional)
+        # Default to True if not specified
+        try:
+            generate_audio = batch.generate_audio
+        except AttributeError:
+            generate_audio = True
+        if not generate_audio:
+            batch.audio_latents = None
+            batch.raw_audio_latent_shape = None
+            return batch
+
+        device = get_local_torch_device()
+        if isinstance(batch.prompt_embeds, list) and batch.prompt_embeds:
+            dtype = batch.prompt_embeds[0].dtype
+        elif isinstance(batch.prompt_embeds, torch.Tensor):
+            dtype = batch.prompt_embeds.dtype
+        else:
+            dtype = torch.float16
+        generator = batch.generator
+
+        audio_latents = batch.audio_latents
+        batch_size = batch.batch_size
+        num_frames = batch.num_frames
+
+        if audio_latents is None:
+            shape = server_args.pipeline_config.prepare_audio_latent_shape(
+                batch, batch_size, num_frames
+            )
+
+            audio_latents = randn_tensor(
+                shape, generator=generator, device=device, dtype=dtype
+            )
+        else:
+            audio_latents = audio_latents.to(device)
+
+        audio_latents = server_args.pipeline_config.maybe_pack_audio_latents(
+            audio_latents, batch_size, batch
+        )
+
+        # Store in batch
+        batch.audio_latents = audio_latents
+        batch.raw_audio_latent_shape = audio_latents.shape
+
+        return batch

python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py

@@ -0,0 +1,92 @@
+import torch
+
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+
+
+class LTX2TextConnectorStage(PipelineStage):
+    """
+    Stage for applying LTX-2 Text Connectors to split/transform text embeddings
+    into video and audio contexts.
+    """
+
+    def __init__(self, connectors):
+        super().__init__()
+        self.connectors = connectors
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        # Input: batch.prompt_embeds (from Gemma, [B, S, D])
+        # Output: batch.prompt_embeds (Video Context), batch.audio_prompt_embeds (Audio Context)
+
+        prompt_embeds = batch.prompt_embeds
+        prompt_attention_mask = batch.prompt_attention_mask
+        neg_prompt_embeds = batch.negative_prompt_embeds
+        neg_prompt_attention_mask = batch.negative_attention_mask
+
+        if isinstance(prompt_embeds, list):
+            prompt_embeds = prompt_embeds[0] if len(prompt_embeds) > 0 else None
+
+        if isinstance(prompt_attention_mask, list):
+            prompt_attention_mask = (
+                prompt_attention_mask[0] if len(prompt_attention_mask) > 0 else None
+            )
+
+        if isinstance(neg_prompt_embeds, list):
+            neg_prompt_embeds = (
+                neg_prompt_embeds[0] if len(neg_prompt_embeds) > 0 else None
+            )
+
+        if isinstance(neg_prompt_attention_mask, list):
+            neg_prompt_attention_mask = (
+                neg_prompt_attention_mask[0]
+                if len(neg_prompt_attention_mask) > 0
+                else None
+            )
+
+        # Handle CFG: Concatenate negative and positive inputs
+        if batch.do_classifier_free_guidance:
+
+            # Concatenate: [Negative, Positive]
+            prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat(
+                [neg_prompt_attention_mask, prompt_attention_mask], dim=0
+            )
+
+        # Prepare additive mask for connectors (as per Diffusers implementation)
+        dtype = prompt_embeds.dtype
+
+        additive_attention_mask = (1 - prompt_attention_mask.to(dtype)) * -1000000.0
+
+        # Call connectors
+        # Expects: prompt_embeds, attention_mask, additive_mask=True
+        with set_forward_context(current_timestep=None, attn_metadata=None):
+            connector_prompt_embeds, connector_audio_prompt_embeds, connector_mask = (
+                self.connectors(
+                    prompt_embeds, additive_attention_mask, additive_mask=True
+                )
+            )
+
+        # Split results if CFG was enabled
+        if batch.do_classifier_free_guidance:
+            neg_embeds, pos_embeds = connector_prompt_embeds.chunk(2, dim=0)
+            neg_audio_embeds, pos_audio_embeds = connector_audio_prompt_embeds.chunk(
+                2, dim=0
+            )
+            neg_mask, pos_mask = connector_mask.chunk(2, dim=0)
+
+            batch.prompt_embeds = [pos_embeds]
+            batch.audio_prompt_embeds = [pos_audio_embeds]
+            batch.prompt_attention_mask = pos_mask
+
+            batch.negative_prompt_embeds = [neg_embeds]
+            batch.negative_audio_prompt_embeds = [neg_audio_embeds]
+            batch.negative_attention_mask = neg_mask
+        else:
+            # Update positive fields
+            batch.prompt_embeds = [connector_prompt_embeds]
+            batch.audio_prompt_embeds = [connector_audio_prompt_embeds]
+            batch.prompt_attention_mask = connector_mask
+
+        return batch