[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/configs/models/adapter/base.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Any
+
+from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+
+
+@dataclass
+class AdapterArchConfig(ArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=list)
+    _compile_conditions: list = field(default_factory=list)
+
+    # convert weights name from HF-format to SGLang-dit-format
+    param_names_mapping: dict = field(default_factory=dict)
+
+    # Reverse mapping for saving checkpoints: custom -> hf
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    _supported_attention_backends: set[AttentionBackendEnum] = field(
+        default_factory=lambda: {
+            AttentionBackendEnum.SLIDING_TILE_ATTN,
+            AttentionBackendEnum.SAGE_ATTN,
+            AttentionBackendEnum.FA,
+            AttentionBackendEnum.AITER,
+            AttentionBackendEnum.TORCH_SDPA,
+            AttentionBackendEnum.VIDEO_SPARSE_ATTN,
+            AttentionBackendEnum.VMOBA_ATTN,
+            AttentionBackendEnum.SAGE_ATTN_3,
+        }
+    )
+
+    hidden_size: int = 0
+    num_attention_heads: int = 0
+    num_channels_latents: int = 0
+    exclude_lora_layers: list[str] = field(default_factory=list)
+    boundary_ratio: float | None = None
+
+    def __post_init__(self) -> None:
+        if not self._compile_conditions:
+            self._compile_conditions = self._fsdp_shard_conditions.copy()
+
+
+@dataclass
+class AdapterConfig(ModelConfig):
+    arch_config: AdapterArchConfig = field(default_factory=AdapterArchConfig)
+
+    # sglang-diffusion Adapter-specific parameters
+    prefix: str = ""
+
+    @staticmethod
+    def add_cli_args(parser: Any, prefix: str = "dit-config") -> Any:
+        """Add CLI arguments for AdapterConfig fields"""
+        parser.add_argument(
+            f"--{prefix}.prefix",
+            type=str,
+            dest=f"{prefix.replace('-', '_')}.prefix",
+            default=AdapterConfig.prefix,
+            help="Prefix for the Adapter",
+        )
+
+        return parser

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

@@ -0,0 +1,33 @@
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.adapter.base import (
+    AdapterArchConfig,
+    AdapterConfig,
+)
+
+
+@dataclass
+class LTX2ConnectorArchConfig(AdapterArchConfig):
+    audio_connector_attention_head_dim: int = 128
+    audio_connector_num_attention_heads: int = 30
+    audio_connector_num_layers: int = 2
+    audio_connector_num_learnable_registers: int = 128
+    caption_channels: int = 3840
+    causal_temporal_positioning: bool = False
+    connector_rope_base_seq_len: int = 4096
+    rope_double_precision: bool = True
+    rope_theta: float = 10000.0
+    rope_type: str = "split"
+    text_proj_in_factor: int = 49
+    video_connector_attention_head_dim: int = 128
+    video_connector_num_attention_heads: int = 30
+    video_connector_num_layers: int = 2
+    video_connector_num_learnable_registers: int = 128
+
+
+@dataclass
+class LTX2ConnectorConfig(AdapterConfig):
+
+    arch_config: AdapterArchConfig = field(default_factory=LTX2ConnectorArchConfig)
+
+    prefix: str = "LTX2"

python/sglang/multimodal_gen/configs/models/dits/ltx_2.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from enum import Enum
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+class LTXModelType(Enum):
+    """
+    Model type enum mirroring upstream `LTXModelType`.
+
+    Upstream reference:
+      - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/model.py::LTXModelType`
+    """
+
+    AudioVideo = "ltx av model"
+    VideoOnly = "ltx video only model"
+    AudioOnly = "ltx audio only model"
+
+    def is_video_enabled(self) -> bool:
+        return self in (LTXModelType.AudioVideo, LTXModelType.VideoOnly)
+
+    def is_audio_enabled(self) -> bool:
+        return self in (LTXModelType.AudioVideo, LTXModelType.AudioOnly)
+
+
+class LTX2RopeType(str, Enum):
+    """
+    Minimal RoPE type enum mirroring LTX-2 upstream `LTXRopeType`.
+
+    Upstream reference:
+      - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/rope.py::LTXRopeType`
+    """
+
+    INTERLEAVED = "interleaved"
+    SPLIT = "split"
+
+
+class LTX2AttentionFunction(str, Enum):
+    """
+    Placeholder enum for upstream `AttentionFunction.DEFAULT`.
+
+    Upstream reference:
+      - `LTX-2/packages/ltx-core/src/ltx_core/model/transformer/attention.py`
+    """
+
+    DEFAULT = "default"
+
+
+def is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class LTX2ArchConfig(DiTArchConfig):
+    """Architecture configuration for LTX-2 Video Transformer."""
+
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [is_blocks])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Parameter name mappings from HuggingFace checkpoint keys to SGLang module names.
+            # We use upstream variable names (patchify_proj, adaln_single) but HF uses different keys.
+            #
+            # HF key -> SGLang key (upstream naming)
+            r"^proj_in\.(.*)$": r"patchify_proj.\1",
+            r"^time_embed\.(.*)$": r"adaln_single.\1",
+            r"^audio_proj_in\.(.*)$": r"audio_patchify_proj.\1",
+            r"^audio_time_embed\.(.*)$": r"audio_adaln_single.\1",
+            # FeedForward
+            r"(.*)ff\.net\.0\.proj\.(.*)$": r"\1ff.proj_in.\2",
+            r"(.*)ff\.net\.2\.(.*)$": r"\1ff.proj_out.\2",
+            # Attention Norms
+            r"(.*)\.norm_q\.(.*)$": r"\1.q_norm.\2",
+            r"(.*)\.norm_k\.(.*)$": r"\1.k_norm.\2",
+            # Scale Shift Tables (Global)
+            r"^av_cross_attn_video_scale_shift\.(.*)$": r"av_ca_video_scale_shift_adaln_single.\1",
+            r"^av_cross_attn_audio_scale_shift\.(.*)$": r"av_ca_audio_scale_shift_adaln_single.\1",
+            r"^av_cross_attn_video_a2v_gate\.(.*)$": r"av_ca_a2v_gate_adaln_single.\1",
+            r"^av_cross_attn_audio_v2a_gate\.(.*)$": r"av_ca_v2a_gate_adaln_single.\1",
+            # Scale Shift Tables (Block Level)
+            # HF: scale_shift_table_a2v_ca_video -> SGLang: video_a2v_cross_attn_scale_shift_table
+            r"(.*)scale_shift_table_a2v_ca_video": r"\1video_a2v_cross_attn_scale_shift_table",
+            r"(.*)scale_shift_table_a2v_ca_audio": r"\1audio_a2v_cross_attn_scale_shift_table",
+        }
+    )
+
+    reverse_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # Reverse mapping: SGLang module names -> HF checkpoint keys (for saving).
+            r"^patchify_proj\.(.*)$": r"proj_in.\1",
+            r"^adaln_single\.(.*)$": r"time_embed.\1",
+            r"^audio_patchify_proj\.(.*)$": r"audio_proj_in.\1",
+            r"^audio_adaln_single\.(.*)$": r"audio_time_embed.\1",
+            # FeedForward
+            r"(.*)ff\.proj_in\.(.*)$": r"\1ff.net.0.proj.\2",
+            r"(.*)ff\.proj_out\.(.*)$": r"\1ff.net.2.\2",
+            # Attention Norms
+            r"(.*)\.q_norm\.(.*)$": r"\1.norm_q.\2",
+            r"(.*)\.k_norm\.(.*)$": r"\1.norm_k.\2",
+            # Scale Shift Tables (Global)
+            r"^av_ca_video_scale_shift_adaln_single\.(.*)$": r"av_cross_attn_video_scale_shift.\1",
+            r"^av_ca_audio_scale_shift_adaln_single\.(.*)$": r"av_cross_attn_audio_scale_shift.\1",
+            r"^av_ca_a2v_gate_adaln_single\.(.*)$": r"av_cross_attn_video_a2v_gate.\1",
+            r"^av_ca_v2a_gate_adaln_single\.(.*)$": r"av_cross_attn_audio_v2a_gate.\1",
+            # Scale Shift Tables (Block Level)
+            # SGLang: video_a2v_cross_attn_scale_shift_table -> HF: scale_shift_table_a2v_ca_video
+            r"(.*)video_a2v_cross_attn_scale_shift_table": r"\1scale_shift_table_a2v_ca_video",
+            r"(.*)audio_a2v_cross_attn_scale_shift_table": r"\1scale_shift_table_a2v_ca_audio",
+        }
+    )
+
+    lora_param_names_mapping: dict = field(
+        default_factory=lambda: {
+            # LoRA parameter name mappings from official repo format to HF format.
+            # This is applied before param_names_mapping when loading LoRA adapters.
+            # Will be populated if LoRA adapters use different naming conventions.
+        }
+    )
+
+    # Model type and attention configuration
+    model_type: LTXModelType = LTXModelType.AudioVideo
+    attention_type: LTX2AttentionFunction = LTX2AttentionFunction.DEFAULT
+    rope_type: LTX2RopeType = LTX2RopeType.INTERLEAVED
+    double_precision_rope: bool = False
+
+    # Video parameters
+    num_attention_heads: int = 32
+    attention_head_dim: int = 128
+    in_channels: int = 128
+    out_channels: int = 128
+    num_layers: int = 48
+    cross_attention_dim: int = 4096
+    norm_eps: float = 1e-6
+    caption_channels: int = 3840
+    positional_embedding_theta: float = 10000.0
+    positional_embedding_max_pos: list[int] | None = None
+    timestep_scale_multiplier: int = 1000
+    use_middle_indices_grid: bool = True
+
+    # Audio parameters
+    audio_num_attention_heads: int = 32
+    audio_attention_head_dim: int = 64
+    audio_in_channels: int = 128
+    audio_out_channels: int = 128
+    audio_cross_attention_dim: int = 2048
+    audio_positional_embedding_max_pos: list[int] | None = None
+    av_ca_timestep_scale_multiplier: int = 1
+
+    # SGLang-specific parameters
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    text_len: int = 512
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Video derived values
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+        if self.positional_embedding_max_pos is None:
+            self.positional_embedding_max_pos = [20, 2048, 2048]
+
+        # Audio derived values
+        self.audio_hidden_size = (
+            self.audio_num_attention_heads * self.audio_attention_head_dim
+        )
+        if self.audio_positional_embedding_max_pos is None:
+            self.audio_positional_embedding_max_pos = [2048]
+
+
+@dataclass
+class LTX2Config(DiTConfig):
+    """Configuration for LTX-2 Video Transformer."""
+
+    arch_config: LTX2ArchConfig = field(default_factory=LTX2ArchConfig)
+
+    prefix: str = "ltx2"

python/sglang/multimodal_gen/configs/models/encoders/__init__.py

@@ -10,6 +10,7 @@ from sglang.multimodal_gen.configs.models.encoders.clip import (
     CLIPTextConfig,
     CLIPVisionConfig,
 )
+from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config
 from sglang.multimodal_gen.configs.models.encoders.llama import LlamaConfig
 from sglang.multimodal_gen.configs.models.encoders.qwen3 import Qwen3TextConfig
 from sglang.multimodal_gen.configs.models.encoders.t5 import T5Config
@@ -24,4 +25,5 @@ __all__ = [
     "LlamaConfig",
     "Qwen3TextConfig",
     "T5Config",
+    "Gemma3Config",
 ]

python/sglang/multimodal_gen/configs/models/encoders/gemma_3.py

@@ -0,0 +1,81 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+def _is_transformer_layer(n: str, m) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m) -> bool:
+    return n.endswith("norm")
+
+
+@dataclass
+class Gemma3ArchConfig(TextEncoderArchConfig):
+    """Minimal Gemma text-encoder config for tokenizer kwargs.
+
+    Note: runtime will load the actual `text_encoder/` module from the model repo
+    (e.g. Gemma3Model) via transformers; this config mainly controls tokenization.
+    """
+
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = None
+    hidden_act: str = "gelu_pytorch_tanh"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    pretraining_tp: int = 1
+    tie_word_embeddings: bool = True
+    rope_theta: float = 10000.0
+    rope_scaling: dict | None = None
+    rope_local_base_freq: float = 10000.0
+    sliding_window: int = 4096
+    layer_types: list[str] = field(default_factory=list)
+    query_pre_attn_scalar: int | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int | None = None
+    hidden_state_skip_layer: int = 2
+    text_len: int = 1024
+
+    stacked_params_mapping: list[tuple[str, str, str]] = field(
+        default_factory=lambda: [
+            # (param_name, shard_name, shard_id)
+            (".qkv_proj", ".q_proj", "q"),
+            (".qkv_proj", ".k_proj", "k"),
+            (".qkv_proj", ".v_proj", "v"),
+            (".gate_up_proj", ".gate_proj", "0"),  # type: ignore
+            (".gate_up_proj", ".up_proj", "1"),  # type: ignore
+        ]
+    )
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]
+    )
+
+
+@dataclass
+class Gemma3Config(TextEncoderConfig):
+    arch_config: TextEncoderArchConfig = field(default_factory=Gemma3ArchConfig)
+
+    prefix: str = "gemma_3"

python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import Optional, Tuple
+
+from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class LTXAudioVAEArchConfig(VAEArchConfig):
+    # Architecture params
+    causality_axis: str = "height"
+    attn_resolutions: Optional[Tuple[int, ...]] = None
+    base_channels: int = 128
+    latent_channels: int = 8
+    output_channels: int = 2
+    ch_mult: Tuple[int, ...] = (1, 2, 4)
+    num_res_blocks: int = 2
+    norm_type: str = "pixel"
+    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
+
+
+@dataclass
+class LTXAudioVAEConfig(VAEConfig):
+    arch_config: LTXAudioVAEArchConfig = field(default_factory=LTXAudioVAEArchConfig)

python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import List
+
+from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class LTXVideoVAEArchConfig(VAEArchConfig):
+    # Architecture params
+    in_channels: int = 3
+    latent_channels: int = 128
+    out_channels: int = 3
+    block_out_channels: List[int] = field(
+        default_factory=lambda: [256, 512, 1024, 2048]
+    )
+    down_block_types: List[str] = field(
+        default_factory=lambda: [
+            "LTX2VideoDownBlock3D",
+            "LTX2VideoDownBlock3D",
+            "LTX2VideoDownBlock3D",
+            "LTX2VideoDownBlock3D",
+        ]
+    )
+    spatio_temporal_scaling: List[bool] = field(
+        default_factory=lambda: [True, True, True, True]
+    )
+    layers_per_block: List[int] = field(default_factory=lambda: [4, 6, 6, 2, 2])
+    downsample_type: List[str] = field(
+        default_factory=lambda: [
+            "spatial",
+            "temporal",
+            "spatiotemporal",
+            "spatiotemporal",
+        ]
+    )
+    patch_size: int = 4
+    patch_size_t: int = 1
+    resnet_norm_eps: float = 1e-6
+    encoder_causal: bool = True
+    encoder_spatial_padding_mode: str = "zeros"
+
+    decoder_block_out_channels: List[int] = field(
+        default_factory=lambda: [256, 512, 1024]
+    )
+    decoder_spatio_temporal_scaling: List[bool] = field(
+        default_factory=lambda: [True, True, True]
+    )
+    decoder_layers_per_block: List[int] = field(default_factory=lambda: [5, 5, 5, 5])
+    decoder_causal: bool = False
+    decoder_spatial_padding_mode: str = "reflect"
+
+
+@dataclass
+class LTXVideoVAEConfig(VAEConfig):
+    arch_config: LTXVideoVAEArchConfig = field(default_factory=LTXVideoVAEArchConfig)

python/sglang/multimodal_gen/configs/models/vocoder/__init__.py

@@ -0,0 +1,3 @@
+from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig
+
+__all__ = ["LTXVocoderConfig"]

python/sglang/multimodal_gen/configs/models/vocoder/base.py

@@ -0,0 +1,29 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+import argparse
+import dataclasses
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
+
+
+@dataclass
+class VocoderArchConfig(ArchConfig):
+    in_channels: int = 128
+    hidden_channels: int = 1024
+    out_channels: int = 2
+
+
+@dataclass
+class VocoderConfig(ModelConfig):
+    arch_config: VocoderArchConfig = field(default_factory=VocoderArchConfig)
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "VocoderConfig":
+        kwargs = {}
+        for attr in dataclasses.fields(cls):
+            value = getattr(args, attr.name, None)
+            if value is not None:
+                kwargs[attr.name] = value
+        return cls(**kwargs)

python/sglang/multimodal_gen/configs/models/vocoder/ltx_vocoder.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+from typing import List
+
+from sglang.multimodal_gen.configs.models.vocoder.base import (
+    VocoderArchConfig,
+    VocoderConfig,
+)
+
+
+@dataclass
+class LTXVocoderArchConfig(VocoderArchConfig):
+    # Architecture params
+    in_channels: int = 128
+    hidden_channels: int = 1024
+    out_channels: int = 2
+    upsample_kernel_sizes: List[int] = field(default_factory=lambda: [3, 7, 11])
+    upsample_factors: List[int] = field(default_factory=lambda: [6, 5, 2, 2, 2])
+    resnet_kernel_sizes: List[int] = field(default_factory=lambda: [3, 7, 11])
+    resnet_dilations: List[List[int]] = field(
+        default_factory=lambda: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
+    )
+    leaky_relu_negative_slope: float = 0.1
+    sample_rate: int = 24000
+
+
+@dataclass
+class LTXVocoderConfig(VocoderConfig):
+    arch_config: LTXVocoderArchConfig = field(default_factory=LTXVocoderArchConfig)

python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py

@@ -19,6 +19,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import (
     FastHunyuanConfig,
     HunyuanConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     SelfForcingWanT2V480PConfig,
     WanI2V480PConfig,
@@ -44,4 +45,5 @@ __all__ = [
     "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig",
     "ZImagePipelineConfig",
+    "LTX2PipelineConfig",
 ]

python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py

@@ -0,0 +1,558 @@
+import dataclasses
+from dataclasses import field
+from typing import Callable
+
+import numpy as np
+import torch
+
+from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2Config
+from sglang.multimodal_gen.configs.models.encoders import (
+    BaseEncoderOutput,
+    EncoderConfig,
+)
+from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config
+from sglang.multimodal_gen.configs.models.vaes.ltx_audio import LTXAudioVAEConfig
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ModelTaskType,
+    PipelineConfig,
+    preprocess_text,
+)
+from sglang.multimodal_gen.runtime.distributed import (
+    get_sp_parallel_rank,
+    get_sp_world_size,
+    sequence_model_parallel_all_gather,
+)
+
+
+def pack_text_embeds(
+    text_hidden_states: torch.Tensor,
+    sequence_lengths: torch.Tensor,
+    padding_side: str = "left",
+    scale_factor: int = 8,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """
+    Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and
+    per-layer in a masked fashion (only over non-padded positions).
+
+    Args:
+        text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`):
+            Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`).
+        sequence_lengths (`torch.Tensor of shape `(batch_size,)`):
+            The number of valid (non-padded) tokens for each batch instance.
+        device: (`str` or `torch.device`, *optional*):
+            torch device to place the resulting embeddings on
+        padding_side: (`str`, *optional*, defaults to `"left"`):
+            Whether the text tokenizer performs padding on the `"left"` or `"right"`.
+        scale_factor (`int`, *optional*, defaults to `8`):
+            Scaling factor to multiply the normalized hidden states by.
+        eps (`float`, *optional*, defaults to `1e-6`):
+            A small positive value for numerical stability when performing normalization.
+
+    Returns:
+        `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`:
+            Normed and flattened text encoder hidden states.
+    """
+    batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape
+    original_dtype = text_hidden_states.dtype
+    device = text_hidden_states.device
+
+    # Create padding mask
+    token_indices = torch.arange(seq_len, device=device).unsqueeze(0)
+    if padding_side == "right":
+        mask = token_indices < sequence_lengths[:, None]
+    elif padding_side == "left":
+        start_indices = seq_len - sequence_lengths[:, None]
+        mask = token_indices >= start_indices
+    else:
+        raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}")
+    mask = mask[:, :, None, None]  # [batch_size, seq_len, 1, 1]
+
+    masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0)
+    num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1)
+    masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (
+        num_valid_positions + eps
+    )
+
+    x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(
+        dim=(1, 2), keepdim=True
+    )
+    x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(
+        dim=(1, 2), keepdim=True
+    )
+
+    normalized_hidden_states = (text_hidden_states - masked_mean) / (
+        x_max - x_min + eps
+    )
+    normalized_hidden_states = normalized_hidden_states * scale_factor
+
+    normalized_hidden_states = normalized_hidden_states.flatten(2)
+    mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers)
+    normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0)
+    normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype)
+
+    return normalized_hidden_states
+
+
+def _gemma_postprocess_func(
+    outputs: BaseEncoderOutput, text_inputs: dict
+) -> torch.Tensor:
+    # LTX-2 requires all hidden states concatenated for the connector
+    if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None:
+        # outputs.hidden_states is a tuple of tensors
+        # We need to stack them along the last dimension and pack them
+        hidden_states = torch.stack(outputs.hidden_states, dim=-1)
+        attention_mask = text_inputs["attention_mask"]
+        sequence_lengths = attention_mask.sum(dim=-1)
+        # Assuming left padding for Gemma as per Diffusers
+        return pack_text_embeds(hidden_states, sequence_lengths, padding_side="left")
+    else:
+        raise AttributeError(
+            "Unsupported text encoder output: expected `hidden_states`."
+        )
+
+
+@dataclasses.dataclass
+class LTX2PipelineConfig(PipelineConfig):
+    """Configuration for LTX-Video pipeline."""
+
+    task_type: ModelTaskType = ModelTaskType.T2V
+    dit_config: LTX2Config = field(default_factory=LTX2Config)
+
+    # Model architecture
+    in_channels: int = 128
+    out_channels: int = 128
+    patch_size: int = 1
+    patch_size_t: int = 1
+
+    # Audio VAE configuration
+    audio_vae_config: LTXAudioVAEConfig = field(default_factory=LTXAudioVAEConfig)
+    audio_vae_precision: str = "fp32"
+    audio_vae_temporal_compression_ratio: int = 4
+    audio_vae_mel_compression_ratio: int = 4
+
+    @property
+    def vae_scale_factor(self):
+        return getattr(self.vae_config.arch_config, "spatial_compression_ratio", 32)
+
+    @property
+    def vae_temporal_compression(self):
+        return getattr(self.vae_config.arch_config, "temporal_compression_ratio", 8)
+
+    def prepare_audio_latent_shape(self, batch, batch_size, num_frames):
+        # Adapted from diffusers pipeline prepare_audio_latents
+        duration_s = num_frames / batch.fps
+
+        sample_rate = self.audio_vae_config.arch_config.sample_rate
+        hop_length = self.audio_vae_config.arch_config.mel_hop_length
+        temporal_compression = self.audio_vae_temporal_compression_ratio
+
+        latents_per_second = (
+            float(sample_rate) / float(hop_length) / float(temporal_compression)
+        )
+        latent_length = round(duration_s * latents_per_second)
+
+        num_mel_bins = self.audio_vae_config.arch_config.mel_bins
+        mel_compression_ratio = self.audio_vae_mel_compression_ratio
+        latent_mel_bins = num_mel_bins // mel_compression_ratio
+
+        # Default to 8
+        num_channels_latents = self.audio_vae_config.arch_config.latent_channels
+
+        shape = (batch_size, num_channels_latents, latent_length, latent_mel_bins)
+
+        return shape
+
+    # Text encoding stage (Gemma)
+    # LTX-2 needs separate contexts for video/audio streams. We model this as
+    # two logical encoders sharing the same underlying `text_encoder` module.
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(
+        default_factory=lambda: (Gemma3Config(),)
+    )
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16",))
+    text_encoder_extra_args: list[dict] = field(default_factory=lambda: [{}])
+
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
+        default_factory=lambda: (preprocess_text,)
+    )
+    postprocess_text_funcs: tuple[
+        Callable[[BaseEncoderOutput, dict], torch.Tensor], ...
+    ] = field(default_factory=lambda: (_gemma_postprocess_func,))
+
+    def prepare_sigmas(self, sigmas, num_inference_steps):
+        if sigmas is None:
+            steps = int(num_inference_steps)
+            if steps <= 0:
+                raise ValueError(f"num_inference_steps must be positive, got {steps}")
+            return np.linspace(1.0, 1.0 / float(steps), steps).tolist()
+        return sigmas
+
+    def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict:
+        # Adapted from diffusers_pipeline.py _get_gemma_prompt_embeds
+        # But we only need tokenization here, the embedding happens in TextEncodingStage
+
+        # Gemma expects left padding for chat-style prompts
+        tokenizer.padding_side = "left"
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        max_sequence_length = tok_kwargs.get(
+            "max_length", 1024
+        )  # Default from diffusers pipeline
+
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        return text_inputs
+
+    def maybe_pack_latents(self, latents, batch_size, batch):
+        # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p].
+        # The patch dimensions are then permuted and collapsed into the channel dimension of shape:
+        # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor).
+        # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features
+        batch_size, num_channels, num_frames, height, width = latents.shape
+        post_patch_num_frames = num_frames // self.patch_size_t
+        post_patch_height = height // self.patch_size
+        post_patch_width = width // self.patch_size
+        latents = latents.reshape(
+            batch_size,
+            -1,
+            post_patch_num_frames,
+            self.patch_size_t,
+            post_patch_height,
+            self.patch_size,
+            post_patch_width,
+            self.patch_size,
+        )
+        latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+        return latents
+
+    def _infer_video_latent_frames_and_tokens_per_frame(
+        self, batch, seq_len: int
+    ) -> tuple[int, int]:
+        """Infer latent-frame count and tokens-per-frame for packed token latents [B, S, D].
+
+        Notes:
+        - This assumes `patch_size_t == 1` (no temporal patching).
+        - Tokens are ordered as (frame, height, width) after packing.
+        """
+        if int(self.patch_size_t) != 1:
+            raise ValueError(
+                "LTX-2 SP time-sharding for packed token latents currently requires "
+                f"{self.patch_size_t=}. (Expected 1)"
+            )
+        if int(seq_len) <= 0:
+            raise ValueError(f"Expected {seq_len=} > 0 for packed token latents.")
+        if int(self.vae_scale_factor) <= 0:
+            raise ValueError(f"Invalid {self.vae_scale_factor=}. Must be > 0.")
+        if int(self.patch_size) <= 0:
+            raise ValueError(f"Invalid {self.patch_size=}. Must be > 0.")
+
+        latent_height = int(batch.height) // int(self.vae_scale_factor)
+        latent_width = int(batch.width) // int(self.vae_scale_factor)
+        if latent_height <= 0 or latent_width <= 0:
+            raise ValueError(
+                "Invalid latent H/W computed from batch.height/width: "
+                f"{batch.height=} {batch.width=} {self.vae_scale_factor=}"
+            )
+        if (latent_height % int(self.patch_size)) != 0 or (
+            latent_width % int(self.patch_size)
+        ) != 0:
+            raise ValueError(
+                "Invalid spatial patching for packed token latents. Expected latent H/W "
+                "to be divisible by patch_size, got "
+                f"{latent_height=} {latent_width=} {self.patch_size=}."
+            )
+
+        post_patch_h = latent_height // int(self.patch_size)
+        post_patch_w = latent_width // int(self.patch_size)
+        tokens_per_frame = int(post_patch_h) * int(post_patch_w)
+        if tokens_per_frame <= 0:
+            raise ValueError(
+                f"Invalid tokens_per_frame={tokens_per_frame} from "
+                f"{latent_height=} {latent_width=} {self.patch_size=}"
+            )
+        if int(seq_len) % int(tokens_per_frame) != 0:
+            raise ValueError(
+                f"LTX-2 token latents seq_len={seq_len} is not divisible by "
+                f"tokens_per_frame={tokens_per_frame}. Cannot time-shard for SP."
+            )
+        latent_num_frames = int(seq_len) // int(tokens_per_frame)
+        return int(latent_num_frames), int(tokens_per_frame)
+
+    def shard_latents_for_sp(self, batch, latents):
+        """Shard LTX-2 packed token latents across SP ranks by latent time (frame) dimension."""
+        sp_world_size = get_sp_world_size()
+        if sp_world_size <= 1:
+            return latents, False
+
+        # Default behavior for 5D latents.
+        if isinstance(latents, torch.Tensor) and latents.ndim == 5:
+            return super().shard_latents_for_sp(batch, latents)
+
+        # LTX-2 packed token latents [B, S, D]
+        if not (isinstance(latents, torch.Tensor) and latents.ndim == 3):
+            return latents, False
+
+        sp_rank = get_sp_parallel_rank()
+        seq_len = int(latents.shape[1])
+        latent_frames, tokens_per_frame = (
+            self._infer_video_latent_frames_and_tokens_per_frame(batch, seq_len)
+        )
+
+        # Pad whole frames so `latent_frames` is divisible by `sp_world_size`.
+        pad_frames = (sp_world_size - (latent_frames % sp_world_size)) % sp_world_size
+        if pad_frames:
+            pad_tokens = int(pad_frames) * int(tokens_per_frame)
+            pad = torch.zeros(
+                (latents.shape[0], pad_tokens, latents.shape[2]),
+                device=latents.device,
+                dtype=latents.dtype,
+            )
+            latents = torch.cat([latents, pad], dim=1)
+            latent_frames = int(latent_frames) + int(pad_frames)
+
+        local_frames = int(latent_frames) // int(sp_world_size)
+        start_frame = int(sp_rank) * int(local_frames)
+        start = int(start_frame) * int(tokens_per_frame)
+        end = int(start) + int(local_frames) * int(tokens_per_frame)
+        latents = latents[:, start:end, :]
+
+        # Store SP metadata for denoising (TI2V gating) and model-side RoPE shift.
+        batch.sp_video_latent_num_frames = int(local_frames)
+        batch.sp_video_start_frame = int(start_frame)
+        batch.sp_video_tokens_per_frame = int(tokens_per_frame)
+
+        return latents, True
+
+    def gather_latents_for_sp(self, latents):
+        """Gather latents after SP. For packed token latents [B, S_local, D], gather on dim=1."""
+        if get_sp_world_size() <= 1:
+            return latents
+        if isinstance(latents, torch.Tensor) and latents.ndim == 3:
+            return sequence_model_parallel_all_gather(latents.contiguous(), dim=1)
+        return super().gather_latents_for_sp(latents)
+
+    def maybe_pack_audio_latents(self, latents, batch_size, batch):
+        # Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins
+        # We need to pack them if patch_size/patch_size_t are defined for audio (not standard DiT patch size)
+
+        # So for LTX-2 (unless we change patch sizes), we just do:
+        latents = latents.transpose(1, 2).flatten(
+            2, 3
+        )  # [B, C, L, M] --> [B, L, C * M]
+        return latents
+
+    def get_pos_prompt_embeds(self, batch):
+        # LTX-2 returns multiple prompt embed tensors (video/audio contexts).
+        return (
+            batch.prompt_embeds[0]
+            if isinstance(batch.prompt_embeds, list)
+            else batch.prompt_embeds
+        )
+
+    def get_neg_prompt_embeds(self, batch):
+        return (
+            batch.negative_prompt_embeds[0]
+            if isinstance(batch.negative_prompt_embeds, list)
+            else batch.negative_prompt_embeds
+        )
+
+    def get_decode_scale_and_shift(self, device, dtype, vae):
+        latents_mean = getattr(vae, "latents_mean", None)
+        latents_std = getattr(vae, "latents_std", None)
+
+        scaling_factor = (
+            getattr(getattr(vae, "config", None), "scaling_factor", None)
+            or getattr(vae, "scaling_factor", None)
+            or getattr(self.vae_config.arch_config, "scaling_factor", None)
+            or 1.0
+        )
+        if isinstance(scaling_factor, (int, float)) and float(scaling_factor) == 0.0:
+            scaling_factor = 1.0
+
+        if isinstance(latents_mean, torch.Tensor) and isinstance(
+            latents_std, torch.Tensor
+        ):
+            latents_mean = latents_mean.to(device=device, dtype=dtype).view(
+                1, -1, 1, 1, 1
+            )
+            latents_std = latents_std.to(device=device, dtype=dtype).view(
+                1, -1, 1, 1, 1
+            )
+            sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view(
+                1, 1, 1, 1, 1
+            )
+            return sf / latents_std, latents_mean
+
+        sf = torch.tensor(float(scaling_factor), device=device, dtype=dtype).view(
+            1, 1, 1, 1, 1
+        )
+        return sf, None
+
+    @staticmethod
+    def _unpack_latents(
+        latents: torch.Tensor,
+        num_frames: int,
+        height: int,
+        width: int,
+        patch_size: int = 1,
+        patch_size_t: int = 1,
+    ) -> torch.Tensor:
+        # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions)
+        # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of
+        # what happens in the `_pack_latents` method.
+        batch_size = latents.size(0)
+        latents = latents.reshape(
+            batch_size,
+            num_frames,
+            height,
+            width,
+            -1,
+            patch_size_t,
+            patch_size,
+            patch_size,
+        )
+        latents = (
+            latents.permute(0, 4, 1, 5, 2, 6, 3, 7)
+            .flatten(6, 7)
+            .flatten(4, 5)
+            .flatten(2, 3)
+        )
+        return latents
+
+    @staticmethod
+    def _denormalize_latents(
+        latents: torch.Tensor,
+        latents_mean: torch.Tensor,
+        latents_std: torch.Tensor,
+        scaling_factor: float = 1.0,
+    ) -> torch.Tensor:
+        # Denormalize latents across the channel dimension [B, C, F, H, W]
+        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(
+            latents.device, latents.dtype
+        )
+        latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
+        latents = latents * latents_std / scaling_factor + latents_mean
+        return latents
+
+    @staticmethod
+    def _denormalize_audio_latents(
+        latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor
+    ):
+        latents_mean = latents_mean.to(latents.device, latents.dtype)
+        latents_std = latents_std.to(latents.device, latents.dtype)
+        return (latents * latents_std) + latents_mean
+
+    @staticmethod
+    def _unpack_audio_latents(
+        latents: torch.Tensor,
+        latent_length: int,
+        num_mel_bins: int,
+        patch_size: int | None = None,
+        patch_size_t: int | None = None,
+    ) -> torch.Tensor:
+        # Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M],
+        # where L is the latent audio length and M is the number of mel bins.
+        if patch_size is not None and patch_size_t is not None:
+            batch_size = latents.size(0)
+            latents = latents.reshape(
+                batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size
+            )
+            latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
+        else:
+            # Assume [B, S, D] = [B, L, C * M], which implies that patch_size = M and patch_size_t = 1.
+            latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
+        return latents
+
+    def _unpad_and_unpack_latents(self, latents, audio_latents, batch, vae, audio_vae):
+        # Calculate latent dimensions
+        # Assuming batch has height, width, num_frames
+        height = batch.height
+        width = batch.width
+        num_frames = batch.num_frames
+
+        # Get compression ratios
+        # Default LTX-2 values if not present in config
+        vae_spatial_compression_ratio = getattr(
+            self.vae_config.arch_config, "spatial_compression_ratio", 32
+        )
+        vae_temporal_compression_ratio = getattr(
+            self.vae_config.arch_config, "temporal_compression_ratio", 8
+        )
+
+        latent_height = height // vae_spatial_compression_ratio
+        latent_width = width // vae_spatial_compression_ratio
+        latent_num_frames = (num_frames - 1) // vae_temporal_compression_ratio + 1
+
+        latents = self._unpack_latents(
+            latents,
+            latent_num_frames,
+            latent_height,
+            latent_width,
+            self.patch_size,
+            self.patch_size_t,
+        )
+
+        sample_rate = self.audio_vae_config.arch_config.sample_rate
+        hop_length = self.audio_vae_config.arch_config.mel_hop_length
+        temporal_compression = self.audio_vae_temporal_compression_ratio
+        duration_s = num_frames / batch.fps
+
+        latents_per_second = (
+            float(sample_rate) / float(hop_length) / float(temporal_compression)
+        )
+        audio_num_frames = round(duration_s * latents_per_second)
+
+        num_mel_bins = self.audio_vae_config.arch_config.mel_bins
+        mel_compression_ratio = self.audio_vae_mel_compression_ratio
+        latent_mel_bins = num_mel_bins // mel_compression_ratio
+
+        audio_latents_mean = getattr(audio_vae, "latents_mean", None)
+        audio_latents_std = getattr(audio_vae, "latents_std", None)
+        if (
+            isinstance(audio_latents_mean, torch.Tensor)
+            and isinstance(audio_latents_std, torch.Tensor)
+            and audio_latents_mean.numel() == audio_latents_std.numel()
+        ):
+            audio_latents_mean = audio_latents_mean.to(
+                device=audio_latents.device, dtype=audio_latents.dtype
+            )
+            audio_latents_std = audio_latents_std.to(
+                device=audio_latents.device, dtype=audio_latents.dtype
+            )
+            if audio_latents.ndim == 3:
+                if audio_latents.shape[-1] != audio_latents_mean.numel():
+                    raise ValueError(
+                        f"audio_latents last dim {audio_latents.shape[-1]} "
+                        f"does not match audio_vae stats {audio_latents_mean.numel()}"
+                    )
+                audio_latents = audio_latents * audio_latents_std.view(
+                    1, 1, -1
+                ) + audio_latents_mean.view(1, 1, -1)
+            elif audio_latents.ndim == 2:
+                if audio_latents.shape[-1] != audio_latents_mean.numel():
+                    raise ValueError(
+                        f"audio_latents last dim {audio_latents.shape[-1]} "
+                        f"does not match audio_vae stats {audio_latents_mean.numel()}"
+                    )
+                audio_latents = audio_latents * audio_latents_std.view(
+                    1, -1
+                ) + audio_latents_mean.view(1, -1)
+            else:
+                audio_latents = audio_latents * audio_latents_std + audio_latents_mean
+
+        audio_latents = self._unpack_audio_latents(
+            audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins
+        )
+
+        return latents, audio_latents
+
+
+@dataclasses.dataclass
+class LTX2I2VPipelineConfig(LTX2PipelineConfig):
+    task_type: ModelTaskType = ModelTaskType.TI2V

python/sglang/multimodal_gen/configs/sample/ltx_2.py

@@ -0,0 +1,40 @@
+import dataclasses
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+
+
+@dataclasses.dataclass
+class LTX2SamplingParams(SamplingParams):
+    """Sampling parameters for LTX-2."""
+
+    # Match the reference defaults used by ltx-pipelines (one-stage).
+    # See: LTX-2/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+    seed: int = 10
+
+    # Video parameters
+    height: int = 512
+    width: int = 768
+    num_frames: int = 121
+    fps: int = 24
+
+    # Audio specific
+    generate_audio: bool = True
+
+    # Denoising parameters
+    guidance_scale: float = 4.0
+    num_inference_steps: int = 40
+
+    # Match ltx-pipelines default negative prompt (covers video + audio artifacts).
+    negative_prompt: str = (
+        "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
+        "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
+        "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
+        "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of "
+        "field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent "
+        "lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny "
+        "valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, "
+        "mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, "
+        "off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward "
+        "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, "
+        "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
+    )

python/sglang/multimodal_gen/registry.py

@@ -45,6 +45,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.flux import (
 from sglang.multimodal_gen.configs.pipeline_configs.glm_image import (
     GlmImagePipelineConfig,
 )
+from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImageEditPipelineConfig,
     QwenImageEditPlus_2511_PipelineConfig,
@@ -70,6 +71,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
     FastHunyuanSamplingParam,
     HunyuanSamplingParams,
 )
+from sglang.multimodal_gen.configs.sample.ltx_2 import LTX2SamplingParams
 from sglang.multimodal_gen.configs.sample.qwenimage import (
     QwenImage2512SamplingParams,
     QwenImageEditPlusSamplingParams,
@@ -423,6 +425,16 @@ def get_model_info(
 
 # Registration of model configs
 def _register_configs():
+    # LTX-2
+    register_configs(
+        sampling_param_cls=LTX2SamplingParams,
+        pipeline_config_cls=LTX2PipelineConfig,
+        model_detectors=[
+            lambda path: "ltx" in path.lower() and "video" in path.lower(),
+            lambda path: "ltx-2" in path.lower(),
+        ],
+    )
+
     # Hunyuan
     register_configs(
         sampling_param_cls=HunyuanSamplingParams,

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