diff --git a/python/sglang/multimodal_gen/configs/models/adapter/base.py b/python/sglang/multimodal_gen/configs/models/adapter/base.py
new file mode 100644
index 000000000..abd22f115
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py
new file mode 100644
index 000000000..d8f6c43ef
--- /dev/null
+++ b/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"
diff --git a/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py b/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py
new file mode 100644
index 000000000..554ab6234
--- /dev/null
+++ b/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"
diff --git a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py b/python/sglang/multimodal_gen/configs/models/encoders/__init__.py
index b66389ed9..f0b25420e 100644
--- a/python/sglang/multimodal_gen/configs/models/encoders/__init__.py
+++ b/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",
]
diff --git a/python/sglang/multimodal_gen/configs/models/encoders/gemma_3.py b/python/sglang/multimodal_gen/configs/models/encoders/gemma_3.py
new file mode 100644
index 000000000..64636985f
--- /dev/null
+++ b/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"
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py b/python/sglang/multimodal_gen/configs/models/vaes/ltx_audio.py
new file mode 100644
index 000000000..dfdb940eb
--- /dev/null
+++ b/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)
diff --git a/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py b/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py
new file mode 100644
index 000000000..3757cce63
--- /dev/null
+++ b/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)
diff --git a/python/sglang/multimodal_gen/configs/models/vocoder/__init__.py b/python/sglang/multimodal_gen/configs/models/vocoder/__init__.py
new file mode 100644
index 000000000..ce94b2036
--- /dev/null
+++ b/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"]
diff --git a/python/sglang/multimodal_gen/configs/models/vocoder/base.py b/python/sglang/multimodal_gen/configs/models/vocoder/base.py
new file mode 100644
index 000000000..0ac8150b6
--- /dev/null
+++ b/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)
diff --git a/python/sglang/multimodal_gen/configs/models/vocoder/ltx_vocoder.py b/python/sglang/multimodal_gen/configs/models/vocoder/ltx_vocoder.py
new file mode 100644
index 000000000..503dc11bd
--- /dev/null
+++ b/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)
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py b/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
index 48e2ffe25..ebb120938 100644
--- a/python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py
+++ b/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",
]
diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py
new file mode 100644
index 000000000..c99a51670
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/configs/sample/ltx_2.py b/python/sglang/multimodal_gen/configs/sample/ltx_2.py
new file mode 100644
index 000000000..5d2e92b58
--- /dev/null
+++ b/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."
+ )
diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py
index 2f09062af..be9032f28 100644
--- a/python/sglang/multimodal_gen/registry.py
+++ b/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,
diff --git a/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py
new file mode 100644
index 000000000..24ea884ae
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
new file mode 100644
index 000000000..e3383865c
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
@@ -0,0 +1,1411 @@
+# Copied and adapted from LTX-2 and WanVideo implementations.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sglang.multimodal_gen.configs.models.dits.ltx_2 import LTX2ArchConfig, LTX2Config
+from sglang.multimodal_gen.runtime.distributed import (
+ get_sp_parallel_rank,
+ get_sp_world_size,
+ get_tp_rank,
+ get_tp_world_size,
+ model_parallel_is_initialized,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+ tensor_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.linear import (
+ ColumnParallelLinear,
+ RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+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)
+ x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
+ return (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+
+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:
+ b = x.shape[0]
+ _, h, t, _ = cos.shape
+ x = x.reshape(b, t, h, -1).swapaxes(1, 2)
+ needs_reshape = True
+
+ 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
+
+ split_x = x.reshape(*x.shape[:-1], 2, r).float()
+ first_x = split_x[..., :1, :]
+ second_x = split_x[..., 1:, :]
+
+ cos_u = cos.unsqueeze(-2)
+ 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)
+ return out.to(dtype=x_dtype)
+
+
+# ==============================================================================
+# Layers and Embeddings
+# ==============================================================================
+
+
+class LTX2AudioVideoRotaryPosEmbed(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ patch_size: int = 1,
+ patch_size_t: int = 1,
+ base_num_frames: int = 20,
+ base_height: int = 2048,
+ base_width: int = 2048,
+ sampling_rate: int = 16000,
+ hop_length: int = 160,
+ scale_factors: Tuple[int, ...] = (8, 32, 32),
+ theta: float = 10000.0,
+ causal_offset: int = 1,
+ modality: str = "video",
+ double_precision: bool = True,
+ rope_type: str = "interleaved",
+ num_attention_heads: int = 32,
+ ) -> None:
+ super().__init__()
+ self.dim = int(dim)
+ self.patch_size = int(patch_size)
+ self.patch_size_t = int(patch_size_t)
+
+ if rope_type not in ["interleaved", "split"]:
+ raise ValueError(
+ f"{rope_type=} not supported. Choose between 'interleaved' and 'split'."
+ )
+ self.rope_type = rope_type
+
+ self.base_num_frames = int(base_num_frames)
+ self.num_attention_heads = int(num_attention_heads)
+
+ self.base_height = int(base_height)
+ self.base_width = int(base_width)
+
+ self.sampling_rate = int(sampling_rate)
+ self.hop_length = int(hop_length)
+ self.audio_latents_per_second = (
+ float(self.sampling_rate) / float(self.hop_length) / float(scale_factors[0])
+ )
+
+ self.scale_factors = tuple(int(x) for x in scale_factors)
+ self.theta = float(theta)
+ self.causal_offset = int(causal_offset)
+
+ self.modality = modality
+ if self.modality not in ["video", "audio"]:
+ raise ValueError(
+ f"Modality {modality} is not supported. Supported modalities are `video` and `audio`."
+ )
+ self.double_precision = bool(double_precision)
+
+ def prepare_video_coords(
+ self,
+ batch_size: int,
+ num_frames: int,
+ height: int,
+ width: int,
+ device: torch.device,
+ fps: float = 24.0,
+ *,
+ start_frame: int = 0,
+ ) -> torch.Tensor:
+ grid_f = torch.arange(
+ start=int(start_frame),
+ end=int(num_frames) + int(start_frame),
+ step=self.patch_size_t,
+ dtype=torch.float32,
+ device=device,
+ )
+ grid_h = torch.arange(
+ start=0,
+ end=height,
+ step=self.patch_size,
+ dtype=torch.float32,
+ device=device,
+ )
+ grid_w = torch.arange(
+ start=0,
+ end=width,
+ step=self.patch_size,
+ dtype=torch.float32,
+ device=device,
+ )
+ grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij")
+ grid = torch.stack(grid, dim=0)
+
+ patch_size = (self.patch_size_t, self.patch_size, self.patch_size)
+ patch_size_delta = torch.tensor(
+ patch_size, dtype=grid.dtype, device=grid.device
+ )
+ patch_ends = grid + patch_size_delta.view(3, 1, 1, 1)
+
+ latent_coords = torch.stack([grid, patch_ends], dim=-1)
+ latent_coords = latent_coords.flatten(1, 3)
+ latent_coords = latent_coords.unsqueeze(0).repeat(batch_size, 1, 1, 1)
+
+ scale_tensor = torch.tensor(self.scale_factors, device=latent_coords.device)
+ broadcast_shape = [1] * latent_coords.ndim
+ broadcast_shape[1] = -1
+ pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape)
+ pixel_coords[:, 0, ...] = (
+ pixel_coords[:, 0, ...] + self.causal_offset - self.scale_factors[0]
+ ).clamp(min=0)
+ pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
+ return pixel_coords
+
+ def prepare_audio_coords(
+ self,
+ batch_size: int,
+ num_frames: int,
+ device: torch.device,
+ *,
+ start_frame: int = 0,
+ ) -> torch.Tensor:
+ grid_f = torch.arange(
+ start=int(start_frame),
+ end=int(num_frames) + int(start_frame),
+ step=self.patch_size_t,
+ dtype=torch.float32,
+ device=device,
+ )
+
+ audio_scale_factor = self.scale_factors[0]
+ grid_start_mel = grid_f * audio_scale_factor
+ grid_start_mel = (
+ grid_start_mel + self.causal_offset - audio_scale_factor
+ ).clip(min=0)
+ grid_start_s = grid_start_mel * self.hop_length / self.sampling_rate
+
+ grid_end_mel = (grid_f + self.patch_size_t) * audio_scale_factor
+ grid_end_mel = (grid_end_mel + self.causal_offset - audio_scale_factor).clip(
+ min=0
+ )
+ grid_end_s = grid_end_mel * self.hop_length / self.sampling_rate
+
+ audio_coords = torch.stack([grid_start_s, grid_end_s], dim=-1)
+ audio_coords = audio_coords.unsqueeze(0).expand(batch_size, -1, -1)
+ audio_coords = audio_coords.unsqueeze(1)
+ return audio_coords
+
+ def prepare_coords(self, *args, **kwargs):
+ if self.modality == "video":
+ return self.prepare_video_coords(*args, **kwargs)
+ return self.prepare_audio_coords(*args, **kwargs)
+
+ def forward(
+ self, coords: torch.Tensor, device: Optional[Union[str, torch.device]] = None
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ device = device or coords.device
+ num_pos_dims = coords.shape[1]
+
+ if coords.ndim == 4:
+ coords_start, coords_end = coords.chunk(2, dim=-1)
+ coords = (coords_start + coords_end) / 2.0
+ coords = coords.squeeze(-1)
+
+ if self.modality == "video":
+ max_positions = (self.base_num_frames, self.base_height, self.base_width)
+ else:
+ max_positions = (self.base_num_frames,)
+
+ grid = torch.stack(
+ [coords[:, i] / max_positions[i] for i in range(num_pos_dims)], dim=-1
+ ).to(device)
+
+ num_rope_elems = num_pos_dims * 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)
+
+ freqs = (grid.unsqueeze(-1) * 2 - 1) * freqs
+ freqs = freqs.transpose(-1, -2).flatten(2)
+
+ 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(
+ cos_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)
+ else:
+ 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.cat([cos_padding, cos_freq], dim=-1)
+ sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
+
+ 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)
+ sin_freqs = torch.swapaxes(sin_freq, 1, 2)
+
+ return cos_freqs, sin_freqs
+
+
+def rms_norm(x: torch.Tensor, eps: float) -> torch.Tensor:
+ return F.rms_norm(x, normalized_shape=(x.shape[-1],), eps=eps)
+
+
+class LTX2TextProjection(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_size: int,
+ out_features: int | None = None,
+ act_fn: str = "gelu_tanh",
+ ) -> None:
+ super().__init__()
+ if out_features is None:
+ out_features = hidden_size
+
+ self.linear_1 = ColumnParallelLinear(
+ in_features, hidden_size, bias=True, gather_output=True
+ )
+ if act_fn == "gelu_tanh":
+ self.act_1 = nn.GELU(approximate="tanh")
+ elif act_fn == "silu":
+ self.act_1 = nn.SiLU()
+ else:
+ raise ValueError(f"Unknown activation function: {act_fn}")
+
+ self.linear_2 = ColumnParallelLinear(
+ hidden_size, out_features, bias=True, gather_output=True
+ )
+
+ def forward(self, caption: torch.Tensor) -> torch.Tensor:
+ hidden_states, _ = self.linear_1(caption)
+ hidden_states = self.act_1(hidden_states)
+ hidden_states, _ = self.linear_2(hidden_states)
+ return hidden_states
+
+
+class LTX2TimestepEmbedder(nn.Module):
+ def __init__(self, embedding_dim: int, in_channels: int = 256) -> None:
+ super().__init__()
+ self.linear_1 = ColumnParallelLinear(
+ in_channels, embedding_dim, bias=True, gather_output=True
+ )
+ self.linear_2 = ColumnParallelLinear(
+ embedding_dim, embedding_dim, bias=True, gather_output=True
+ )
+
+ def forward(self, t_emb: torch.Tensor) -> torch.Tensor:
+ x, _ = self.linear_1(t_emb)
+ x = F.silu(x)
+ x, _ = self.linear_2(x)
+ return x
+
+
+class LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(nn.Module):
+ def __init__(self, embedding_dim: int) -> None:
+ super().__init__()
+ self.timestep_embedder = LTX2TimestepEmbedder(embedding_dim, in_channels=256)
+
+ def forward(
+ self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None
+ ) -> torch.Tensor:
+ t = timestep.reshape(-1).to(dtype=torch.float32)
+ t_emb = timestep_embedding(t, dim=256, max_period=10000, dtype=torch.float32)
+ if hidden_dtype is not None:
+ t_emb = t_emb.to(dtype=hidden_dtype)
+ return self.timestep_embedder(t_emb)
+
+
+class LTX2AdaLayerNormSingle(nn.Module):
+ def __init__(self, embedding_dim: int, embedding_coefficient: int = 6) -> None:
+ super().__init__()
+ self.emb = LTX2PixArtAlphaCombinedTimestepSizeEmbeddings(embedding_dim)
+ self.silu = nn.SiLU()
+ self.linear = ColumnParallelLinear(
+ embedding_dim,
+ embedding_coefficient * embedding_dim,
+ bias=True,
+ gather_output=True,
+ )
+
+ def forward(
+ self, timestep: torch.Tensor, hidden_dtype: torch.dtype | None = None
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype).to(
+ dtype=self.linear.weight.dtype
+ )
+ out, _ = self.linear(self.silu(embedded_timestep))
+ return out, embedded_timestep
+
+
+class LTX2TPRMSNormAcrossHeads(nn.Module):
+ def __init__(
+ self, full_hidden_size: int, local_hidden_size: int, eps: float
+ ) -> None:
+ super().__init__()
+ self.full_hidden_size = full_hidden_size
+ self.local_hidden_size = local_hidden_size
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(local_hidden_size))
+
+ tp_rank = get_tp_rank()
+
+ def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+ shard = loaded_weight.narrow(
+ 0, tp_rank * local_hidden_size, local_hidden_size
+ )
+ param.data.copy_(shard.to(dtype=param.dtype, device=param.device))
+
+ setattr(self.weight, "weight_loader", _weight_loader)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # Keep track of the original dtype. We do the statistics in fp32 for
+ # numerical stability, but cast the output back to the input dtype to
+ orig_dtype = x.dtype
+ if get_tp_world_size() == 1:
+ var = x.float().pow(2).mean(dim=-1, keepdim=True)
+ else:
+ local_sumsq = x.float().pow(2).sum(dim=-1, keepdim=True)
+ global_sumsq = tensor_model_parallel_all_reduce(local_sumsq)
+ var = global_sumsq / float(self.full_hidden_size)
+
+ inv_rms_fp32 = torch.rsqrt(var + self.eps)
+ y = (x.float() * inv_rms_fp32).to(dtype=orig_dtype)
+ return y * self.weight.to(dtype=orig_dtype)
+
+
+class LTX2Attention(nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ context_dim: int | None = None,
+ heads: int = 8,
+ dim_head: int = 64,
+ norm_eps: float = 1e-6,
+ qk_norm: bool = True,
+ supported_attention_backends: set[AttentionBackendEnum] | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+
+ self.query_dim = int(query_dim)
+ self.context_dim = int(query_dim if context_dim is None else context_dim)
+ self.heads = int(heads)
+ self.dim_head = int(dim_head)
+ self.inner_dim = self.heads * self.dim_head
+ self.norm_eps = float(norm_eps)
+ self.qk_norm = bool(qk_norm)
+
+ tp_size = get_tp_world_size()
+ if tp_size <= 0:
+ raise ValueError(f"Invalid {tp_size=}. Expected tp_size >= 1.")
+ if self.heads % tp_size != 0:
+ raise ValueError(
+ f"LTX2Attention requires heads divisible by tp_size, got "
+ f"{self.heads=} {tp_size=}."
+ )
+ if self.inner_dim % tp_size != 0:
+ # This should follow from heads % tp_size, but keep explicit for clarity.
+ raise ValueError(
+ f"LTX2Attention requires inner_dim divisible by tp_size, got "
+ f"{self.inner_dim=} {tp_size=}."
+ )
+ self.local_heads = self.heads // tp_size
+
+ self.to_q = ColumnParallelLinear(
+ self.query_dim, self.inner_dim, bias=True, gather_output=False
+ )
+ self.to_k = ColumnParallelLinear(
+ self.context_dim, self.inner_dim, bias=True, gather_output=False
+ )
+ self.to_v = ColumnParallelLinear(
+ self.context_dim, self.inner_dim, bias=True, gather_output=False
+ )
+
+ self.q_norm: nn.Module | None = None
+ self.k_norm: nn.Module | None = None
+ if self.qk_norm:
+ if tp_size == 1:
+ self.q_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps)
+ self.k_norm = torch.nn.RMSNorm(self.inner_dim, eps=self.norm_eps)
+ else:
+ self.q_norm = LTX2TPRMSNormAcrossHeads(
+ full_hidden_size=self.inner_dim,
+ local_hidden_size=self.inner_dim // tp_size,
+ eps=self.norm_eps,
+ )
+ self.k_norm = LTX2TPRMSNormAcrossHeads(
+ full_hidden_size=self.inner_dim,
+ local_hidden_size=self.inner_dim // tp_size,
+ eps=self.norm_eps,
+ )
+
+ self.to_out = nn.Sequential(
+ RowParallelLinear(
+ self.inner_dim, self.query_dim, bias=True, input_is_parallel=True
+ ),
+ nn.Identity(),
+ )
+
+ self.attn = USPAttention(
+ num_heads=self.local_heads,
+ head_size=self.dim_head,
+ num_kv_heads=self.local_heads,
+ dropout_rate=0,
+ softmax_scale=None,
+ causal=False,
+ supported_attention_backends=supported_attention_backends,
+ prefix=f"{prefix}.attn",
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ context: torch.Tensor | None = None,
+ mask: torch.Tensor | None = None,
+ pe: tuple[torch.Tensor, torch.Tensor] | None = None,
+ k_pe: tuple[torch.Tensor, torch.Tensor] | None = None,
+ ) -> torch.Tensor:
+ q, _ = self.to_q(x)
+ context_ = x if context is None else context
+ k, _ = self.to_k(context_)
+ v, _ = self.to_v(context_)
+
+ if self.qk_norm:
+ assert self.q_norm is not None and self.k_norm is not None
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+
+ if pe is not None:
+ cos, sin = pe
+ k_cos, k_sin = pe if k_pe is None else k_pe
+ tp_size = get_tp_world_size()
+ if tp_size > 1:
+ tp_rank = get_tp_rank()
+ cos, sin = self._slice_rope_for_tp(
+ cos, sin, tp_rank=tp_rank, tp_size=tp_size
+ )
+ k_cos, k_sin = self._slice_rope_for_tp(
+ k_cos, k_sin, tp_rank=tp_rank, tp_size=tp_size
+ )
+ if cos.dim() == 3:
+ q = apply_interleaved_rotary_emb(q, (cos, sin))
+ k = apply_interleaved_rotary_emb(k, (k_cos, k_sin))
+ else:
+ q = apply_split_rotary_emb(q, (cos, sin))
+ k = apply_split_rotary_emb(k, (k_cos, k_sin))
+
+ q = q.view(*q.shape[:-1], self.local_heads, self.dim_head)
+ k = k.view(*k.shape[:-1], self.local_heads, self.dim_head)
+ v = v.view(*v.shape[:-1], self.local_heads, self.dim_head)
+
+ if mask is not None:
+ # Fallback to SDPA for masked attention
+ q_ = q.transpose(1, 2)
+ k_ = k.transpose(1, 2)
+ v_ = v.transpose(1, 2)
+
+ if torch.is_floating_point(mask):
+ m = mask
+ if m.dim() == 2:
+ m = m[:, None, None, :]
+ elif m.dim() == 3:
+ m = m[:, None, :, :]
+ sdpa_mask = m.to(dtype=q_.dtype, device=q_.device)
+ else:
+ m = mask.to(dtype=q_.dtype, device=q_.device)
+ if m.dim() == 2:
+ m = m[:, None, None, :]
+ elif m.dim() == 3:
+ m = m[:, None, :, :]
+ sdpa_mask = (m - 1.0) * torch.finfo(q_.dtype).max
+
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q_, k_, v_, attn_mask=sdpa_mask, dropout_p=0.0, is_causal=False
+ ).transpose(1, 2)
+ else:
+ out = self.attn(q, k, v)
+
+ out = out.flatten(2)
+ out, _ = self.to_out[0](out)
+ return out
+
+ def _slice_rope_for_tp(
+ self,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ *,
+ tp_rank: int,
+ tp_size: int,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """Slice RoPE tensors to the local TP shard.
+
+ - split-rope: cos/sin are shaped [B, H, T, R] (head-major), slice by heads.
+ - interleaved-rope: cos/sin are shaped [B, T, D], where D matches the projected
+ feature dimension and is sharded by TP.
+ """
+ if cos.ndim == 4:
+ # [B, H, T, R]
+ start = tp_rank * self.local_heads
+ end = start + self.local_heads
+ return cos[:, start:end, :, :], sin[:, start:end, :, :]
+ elif cos.ndim == 3:
+ # [B, T, D]
+ d = cos.shape[-1]
+ if d % tp_size != 0:
+ raise ValueError(
+ f"RoPE dim must be divisible by tp_size, got {d=} {tp_size=}."
+ )
+ local_d = d // tp_size
+ start = tp_rank * local_d
+ end = start + local_d
+ return cos[:, :, start:end], sin[:, :, start:end]
+ raise ValueError(f"Unexpected RoPE tensor rank: {cos.ndim}. Expected 3 or 4.")
+
+
+class LTX2FeedForward(nn.Module):
+ def __init__(self, dim: int, dim_out: int | None = None, mult: int = 4) -> None:
+ super().__init__()
+ if dim_out is None:
+ dim_out = dim
+ inner_dim = int(dim * mult)
+
+ self.proj_in = ColumnParallelLinear(
+ dim, inner_dim, bias=True, gather_output=True
+ )
+ self.act = nn.GELU(approximate="tanh")
+ self.proj_out = ColumnParallelLinear(
+ inner_dim, dim_out, bias=True, gather_output=True
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x, _ = self.proj_in(x)
+ x = self.act(x)
+ x, _ = self.proj_out(x)
+ return x
+
+
+class LTX2TransformerBlock(nn.Module):
+ def __init__(
+ self,
+ idx: int,
+ dim: int,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ cross_attention_dim: int,
+ audio_dim: int,
+ audio_num_attention_heads: int,
+ audio_attention_head_dim: int,
+ audio_cross_attention_dim: int,
+ qk_norm: bool = True,
+ norm_eps: float = 1e-6,
+ supported_attention_backends: set[AttentionBackendEnum] | None = None,
+ prefix: str = "",
+ ):
+ super().__init__()
+ self.idx = idx
+ self.norm_eps = norm_eps
+
+ # 1. Self-Attention (video and audio)
+ self.attn1 = LTX2Attention(
+ query_dim=dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ norm_eps=norm_eps,
+ qk_norm=qk_norm,
+ supported_attention_backends=supported_attention_backends,
+ prefix=f"{prefix}.attn1",
+ )
+ self.audio_attn1 = LTX2Attention(
+ query_dim=audio_dim,
+ heads=audio_num_attention_heads,
+ dim_head=audio_attention_head_dim,
+ norm_eps=norm_eps,
+ qk_norm=qk_norm,
+ supported_attention_backends=supported_attention_backends,
+ prefix=f"{prefix}.audio_attn1",
+ )
+
+ # 2. Prompt Cross-Attention
+ self.attn2 = LTX2Attention(
+ query_dim=dim,
+ context_dim=cross_attention_dim,
+ heads=num_attention_heads,
+ dim_head=attention_head_dim,
+ norm_eps=norm_eps,
+ qk_norm=qk_norm,
+ supported_attention_backends=supported_attention_backends,
+ prefix=f"{prefix}.attn2",
+ )
+ self.audio_attn2 = LTX2Attention(
+ query_dim=audio_dim,
+ context_dim=audio_cross_attention_dim,
+ heads=audio_num_attention_heads,
+ dim_head=audio_attention_head_dim,
+ norm_eps=norm_eps,
+ qk_norm=qk_norm,
+ supported_attention_backends=supported_attention_backends,
+ prefix=f"{prefix}.audio_attn2",
+ )
+
+ # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
+ self.audio_to_video_attn = LTX2Attention(
+ query_dim=dim,
+ context_dim=audio_dim,
+ heads=audio_num_attention_heads,
+ dim_head=audio_attention_head_dim,
+ norm_eps=norm_eps,
+ qk_norm=qk_norm,
+ supported_attention_backends=supported_attention_backends,
+ prefix=f"{prefix}.audio_to_video_attn",
+ )
+ self.video_to_audio_attn = LTX2Attention(
+ query_dim=audio_dim,
+ context_dim=dim,
+ heads=audio_num_attention_heads,
+ dim_head=audio_attention_head_dim,
+ norm_eps=norm_eps,
+ qk_norm=qk_norm,
+ supported_attention_backends=supported_attention_backends,
+ prefix=f"{prefix}.video_to_audio_attn",
+ )
+
+ # 4. Feedforward layers
+ self.ff = LTX2FeedForward(dim, dim_out=dim)
+ self.audio_ff = LTX2FeedForward(audio_dim, dim_out=audio_dim)
+
+ # 5. Modulation Parameters
+ self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+ self.audio_scale_shift_table = nn.Parameter(
+ torch.randn(6, audio_dim) / audio_dim**0.5
+ )
+ self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim))
+ self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(
+ torch.randn(5, audio_dim)
+ )
+
+ def get_ada_values(
+ self,
+ scale_shift_table: torch.Tensor,
+ batch_size: int,
+ timestep: torch.Tensor,
+ indices: slice,
+ ) -> tuple[torch.Tensor, ...]:
+ num_ada_params = int(scale_shift_table.shape[0])
+ ada_values = (
+ scale_shift_table[indices]
+ .unsqueeze(0)
+ .unsqueeze(0)
+ .to(device=timestep.device, dtype=timestep.dtype)
+ + timestep.reshape(batch_size, timestep.shape[1], num_ada_params, -1)[
+ :, :, indices, :
+ ]
+ ).unbind(dim=2)
+ return [t.squeeze(2) for t in ada_values]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ audio_hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ audio_encoder_hidden_states: torch.Tensor,
+ temb: torch.Tensor,
+ temb_audio: torch.Tensor,
+ temb_ca_scale_shift: torch.Tensor,
+ temb_ca_audio_scale_shift: torch.Tensor,
+ temb_ca_gate: torch.Tensor,
+ temb_ca_audio_gate: torch.Tensor,
+ video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ca_video_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ ca_audio_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ audio_encoder_attention_mask: Optional[torch.Tensor] = None,
+ a2v_cross_attention_mask: Optional[torch.Tensor] = None,
+ v2a_cross_attention_mask: Optional[torch.Tensor] = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+
+ batch_size = hidden_states.size(0)
+
+ # 1. Video and Audio Self-Attention
+ vshift_msa, vscale_msa, vgate_msa = self.get_ada_values(
+ self.scale_shift_table, batch_size, temb, slice(0, 3)
+ )
+ norm_hidden_states = (
+ rms_norm(hidden_states, self.norm_eps) * (1 + vscale_msa) + vshift_msa
+ )
+ attn_hidden_states = self.attn1(norm_hidden_states, pe=video_rotary_emb)
+ hidden_states = hidden_states + attn_hidden_states * vgate_msa
+
+ ashift_msa, ascale_msa, agate_msa = self.get_ada_values(
+ self.audio_scale_shift_table, batch_size, temb_audio, slice(0, 3)
+ )
+ norm_audio_hidden_states = (
+ rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_msa) + ashift_msa
+ )
+ attn_audio_hidden_states = self.audio_attn1(
+ norm_audio_hidden_states, pe=audio_rotary_emb
+ )
+ audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * agate_msa
+
+ # 2. Prompt Cross-Attention
+ norm_hidden_states = rms_norm(hidden_states, self.norm_eps)
+ attn_hidden_states = self.attn2(
+ norm_hidden_states,
+ context=encoder_hidden_states,
+ mask=encoder_attention_mask,
+ )
+ hidden_states = hidden_states + attn_hidden_states
+
+ norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps)
+ attn_audio_hidden_states = self.audio_attn2(
+ norm_audio_hidden_states,
+ context=audio_encoder_hidden_states,
+ mask=audio_encoder_attention_mask,
+ )
+ audio_hidden_states = audio_hidden_states + attn_audio_hidden_states
+
+ # 3. Audio-to-Video and Video-to-Audio Cross-Attention
+ norm_hidden_states = rms_norm(hidden_states, self.norm_eps)
+ norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps)
+
+ # Compute combined ada params
+ video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[
+ :4, :
+ ]
+ video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :]
+
+ video_ca_scale_shift_table = (
+ video_per_layer_ca_scale_shift[None, None, :, :].to(
+ dtype=temb_ca_scale_shift.dtype, device=temb_ca_scale_shift.device
+ )
+ + temb_ca_scale_shift.reshape(
+ batch_size, temb_ca_scale_shift.shape[1], 4, -1
+ )
+ ).unbind(dim=2)
+ video_ca_gate = (
+ video_per_layer_ca_gate[None, None, :, :].to(
+ dtype=temb_ca_gate.dtype, device=temb_ca_gate.device
+ )
+ + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1)
+ ).unbind(dim=2)
+
+ (
+ video_a2v_ca_scale,
+ video_a2v_ca_shift,
+ video_v2a_ca_scale,
+ video_v2a_ca_shift,
+ ) = [t.squeeze(2) for t in video_ca_scale_shift_table]
+ a2v_gate = video_ca_gate[0].squeeze(2)
+
+ audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[
+ :4, :
+ ]
+ audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :]
+
+ audio_ca_scale_shift_table = (
+ audio_per_layer_ca_scale_shift[None, None, :, :].to(
+ dtype=temb_ca_audio_scale_shift.dtype,
+ device=temb_ca_audio_scale_shift.device,
+ )
+ + temb_ca_audio_scale_shift.reshape(
+ batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1
+ )
+ ).unbind(dim=2)
+ audio_ca_gate = (
+ audio_per_layer_ca_gate[None, None, :, :].to(
+ dtype=temb_ca_audio_gate.dtype, device=temb_ca_audio_gate.device
+ )
+ + temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1)
+ ).unbind(dim=2)
+
+ (
+ audio_a2v_ca_scale,
+ audio_a2v_ca_shift,
+ audio_v2a_ca_scale,
+ audio_v2a_ca_shift,
+ ) = [t.squeeze(2) for t in audio_ca_scale_shift_table]
+ v2a_gate = audio_ca_gate[0].squeeze(2)
+
+ # A2V
+ mod_norm_hidden_states = (
+ norm_hidden_states * (1 + video_a2v_ca_scale) + video_a2v_ca_shift
+ )
+ mod_norm_audio_hidden_states = (
+ norm_audio_hidden_states * (1 + audio_a2v_ca_scale) + audio_a2v_ca_shift
+ )
+
+ a2v_attn_hidden_states = self.audio_to_video_attn(
+ mod_norm_hidden_states,
+ context=mod_norm_audio_hidden_states,
+ pe=ca_video_rotary_emb,
+ k_pe=ca_audio_rotary_emb,
+ mask=a2v_cross_attention_mask,
+ )
+ hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
+
+ # V2A
+ mod_norm_hidden_states = (
+ norm_hidden_states * (1 + video_v2a_ca_scale) + video_v2a_ca_shift
+ )
+ mod_norm_audio_hidden_states = (
+ norm_audio_hidden_states * (1 + audio_v2a_ca_scale) + audio_v2a_ca_shift
+ )
+
+ v2a_attn_hidden_states = self.video_to_audio_attn(
+ mod_norm_audio_hidden_states,
+ context=mod_norm_hidden_states,
+ pe=ca_audio_rotary_emb,
+ k_pe=ca_video_rotary_emb,
+ mask=v2a_cross_attention_mask,
+ )
+ audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
+
+ # 4. Feedforward
+ vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values(
+ self.scale_shift_table, batch_size, temb, slice(3, None)
+ )
+ norm_hidden_states = (
+ rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp
+ )
+ ff_output = self.ff(norm_hidden_states)
+ hidden_states = hidden_states + ff_output * vgate_mlp
+
+ ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values(
+ self.audio_scale_shift_table, batch_size, temb_audio, slice(3, None)
+ )
+ norm_audio_hidden_states = (
+ rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp) + ashift_mlp
+ )
+ audio_ff_output = self.audio_ff(norm_audio_hidden_states)
+ audio_hidden_states = audio_hidden_states + audio_ff_output * agate_mlp
+
+ return hidden_states, audio_hidden_states
+
+
+class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin):
+ _fsdp_shard_conditions = LTX2ArchConfig()._fsdp_shard_conditions
+ _compile_conditions = LTX2ArchConfig()._compile_conditions
+ _supported_attention_backends = LTX2ArchConfig()._supported_attention_backends
+ param_names_mapping = LTX2ArchConfig().param_names_mapping
+ reverse_param_names_mapping = LTX2ArchConfig().reverse_param_names_mapping
+ lora_param_names_mapping = LTX2ArchConfig().lora_param_names_mapping
+
+ def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None:
+ """Validate TP-related dimension constraints (fail-fast)."""
+ if tp_size < 1:
+ raise ValueError(f"Invalid tp_size={tp_size}. Expected tp_size >= 1.")
+
+ if self.hidden_size % self.num_attention_heads != 0:
+ raise ValueError(
+ "video hidden_size must be divisible by num_attention_heads, got "
+ f"{self.hidden_size=} {self.num_attention_heads=}."
+ )
+ if self.audio_hidden_size % self.audio_num_attention_heads != 0:
+ raise ValueError(
+ "audio_hidden_size must be divisible by audio_num_attention_heads, got "
+ f"{self.audio_hidden_size=} {self.audio_num_attention_heads=}."
+ )
+
+ if tp_size == 1:
+ return
+
+ if self.num_attention_heads % tp_size != 0:
+ raise ValueError(
+ "num_attention_heads must be divisible by tp_size, got "
+ f"{self.num_attention_heads=} {tp_size=}."
+ )
+ if self.audio_num_attention_heads % tp_size != 0:
+ raise ValueError(
+ "audio_num_attention_heads must be divisible by tp_size, got "
+ f"{self.audio_num_attention_heads=} {tp_size=}."
+ )
+ if self.hidden_size % tp_size != 0:
+ raise ValueError(
+ "hidden_size must be divisible by tp_size for TP-sharded projections, got "
+ f"{self.hidden_size=} {tp_size=}."
+ )
+ if self.audio_hidden_size % tp_size != 0:
+ raise ValueError(
+ "audio_hidden_size must be divisible by tp_size for TP-sharded projections, got "
+ f"{self.audio_hidden_size=} {tp_size=}."
+ )
+ if int(arch.out_channels) % tp_size != 0:
+ raise ValueError(
+ "out_channels must be divisible by tp_size for TP-sharded output projection, got "
+ f"{arch.out_channels=} {tp_size=}."
+ )
+ if int(arch.audio_out_channels) % tp_size != 0:
+ raise ValueError(
+ "audio_out_channels must be divisible by tp_size for TP-sharded output projection, got "
+ f"{arch.audio_out_channels=} {tp_size=}."
+ )
+
+ def __init__(self, config: LTX2Config, hf_config: dict[str, Any]) -> None:
+ super().__init__(config=config, hf_config=hf_config)
+
+ arch = config.arch_config
+ self.hidden_size = arch.hidden_size
+ self.num_attention_heads = arch.num_attention_heads
+ self.audio_hidden_size = arch.audio_hidden_size
+ self.audio_num_attention_heads = arch.audio_num_attention_heads
+ self.norm_eps = arch.norm_eps
+
+ tp_size = get_tp_world_size()
+ self._validate_tp_config(arch=arch, tp_size=tp_size)
+
+ # 1. Patchification input projections
+ # Matches LTX2Config().param_names_mapping
+ self.patchify_proj = ColumnParallelLinear(
+ arch.in_channels, self.hidden_size, bias=True, gather_output=True
+ )
+ self.audio_patchify_proj = ColumnParallelLinear(
+ arch.audio_in_channels,
+ self.audio_hidden_size,
+ bias=True,
+ gather_output=True,
+ )
+
+ # 2. Prompt embeddings
+ self.caption_projection = LTX2TextProjection(
+ in_features=arch.caption_channels, hidden_size=self.hidden_size
+ )
+ self.audio_caption_projection = LTX2TextProjection(
+ in_features=arch.caption_channels, hidden_size=self.audio_hidden_size
+ )
+
+ # 3. Timestep Modulation Params and Embedding
+ self.adaln_single = LTX2AdaLayerNormSingle(
+ self.hidden_size, embedding_coefficient=6
+ )
+ self.audio_adaln_single = LTX2AdaLayerNormSingle(
+ self.audio_hidden_size, embedding_coefficient=6
+ )
+
+ # Global Cross Attention Modulation Parameters
+ self.av_ca_video_scale_shift_adaln_single = LTX2AdaLayerNormSingle(
+ self.hidden_size, embedding_coefficient=4
+ )
+ self.av_ca_a2v_gate_adaln_single = LTX2AdaLayerNormSingle(
+ self.hidden_size, embedding_coefficient=1
+ )
+ self.av_ca_audio_scale_shift_adaln_single = LTX2AdaLayerNormSingle(
+ self.audio_hidden_size, embedding_coefficient=4
+ )
+ self.av_ca_v2a_gate_adaln_single = LTX2AdaLayerNormSingle(
+ self.audio_hidden_size, embedding_coefficient=1
+ )
+
+ # Output Layer Scale/Shift Modulation parameters
+ self.scale_shift_table = nn.Parameter(
+ torch.randn(2, self.hidden_size) / self.hidden_size**0.5
+ )
+ self.audio_scale_shift_table = nn.Parameter(
+ torch.randn(2, self.audio_hidden_size) / self.audio_hidden_size**0.5
+ )
+
+ hf_patch_size = int(hf_config.get("patch_size", 1))
+ hf_patch_size_t = int(hf_config.get("patch_size_t", 1))
+ self.patch_size = (hf_patch_size_t, hf_patch_size, hf_patch_size)
+
+ hf_audio_patch_size = int(hf_config.get("audio_patch_size", 1))
+ hf_audio_patch_size_t = int(hf_config.get("audio_patch_size_t", 1))
+
+ rope_type = (
+ arch.rope_type.value
+ if hasattr(arch.rope_type, "value")
+ else str(arch.rope_type)
+ )
+ rope_double_precision = bool(getattr(arch, "double_precision_rope", True))
+ causal_offset = int(hf_config.get("causal_offset", 1))
+
+ pos_embed_max_pos = int(arch.positional_embedding_max_pos[0])
+ base_height = int(arch.positional_embedding_max_pos[1])
+ base_width = int(arch.positional_embedding_max_pos[2])
+
+ audio_pos_embed_max_pos = int(arch.audio_positional_embedding_max_pos[0])
+
+ self.video_scale_factors = (8, 32, 32)
+ self.audio_scale_factors = (4,)
+
+ self.rope = LTX2AudioVideoRotaryPosEmbed(
+ dim=self.hidden_size,
+ patch_size=hf_patch_size,
+ patch_size_t=hf_patch_size_t,
+ base_num_frames=pos_embed_max_pos,
+ base_height=base_height,
+ base_width=base_width,
+ scale_factors=self.video_scale_factors,
+ theta=float(arch.positional_embedding_theta),
+ causal_offset=causal_offset,
+ modality="video",
+ double_precision=rope_double_precision,
+ rope_type=rope_type,
+ num_attention_heads=self.num_attention_heads,
+ )
+ self.audio_rope = LTX2AudioVideoRotaryPosEmbed(
+ dim=self.audio_hidden_size,
+ patch_size=hf_audio_patch_size,
+ patch_size_t=hf_audio_patch_size_t,
+ base_num_frames=audio_pos_embed_max_pos,
+ sampling_rate=16000,
+ hop_length=160,
+ scale_factors=self.audio_scale_factors,
+ theta=float(arch.positional_embedding_theta),
+ causal_offset=causal_offset,
+ modality="audio",
+ double_precision=rope_double_precision,
+ rope_type=rope_type,
+ num_attention_heads=self.audio_num_attention_heads,
+ )
+
+ cross_attn_pos_embed_max_pos = max(pos_embed_max_pos, audio_pos_embed_max_pos)
+ self.cross_attn_rope = LTX2AudioVideoRotaryPosEmbed(
+ dim=int(arch.audio_cross_attention_dim),
+ patch_size=hf_patch_size,
+ patch_size_t=hf_patch_size_t,
+ base_num_frames=cross_attn_pos_embed_max_pos,
+ base_height=base_height,
+ base_width=base_width,
+ theta=float(arch.positional_embedding_theta),
+ causal_offset=causal_offset,
+ modality="video",
+ double_precision=rope_double_precision,
+ rope_type=rope_type,
+ num_attention_heads=self.num_attention_heads,
+ )
+ self.cross_attn_audio_rope = LTX2AudioVideoRotaryPosEmbed(
+ dim=int(arch.audio_cross_attention_dim),
+ patch_size=hf_audio_patch_size,
+ patch_size_t=hf_audio_patch_size_t,
+ base_num_frames=cross_attn_pos_embed_max_pos,
+ sampling_rate=16000,
+ hop_length=160,
+ theta=float(arch.positional_embedding_theta),
+ causal_offset=causal_offset,
+ modality="audio",
+ double_precision=rope_double_precision,
+ rope_type=rope_type,
+ num_attention_heads=self.audio_num_attention_heads,
+ )
+
+ self.cross_pe_max_pos = cross_attn_pos_embed_max_pos
+
+ # 5. Transformer Blocks
+ self.transformer_blocks = nn.ModuleList(
+ [
+ LTX2TransformerBlock(
+ idx=idx,
+ dim=self.hidden_size,
+ num_attention_heads=self.num_attention_heads,
+ attention_head_dim=self.hidden_size // self.num_attention_heads,
+ cross_attention_dim=arch.cross_attention_dim,
+ audio_dim=self.audio_hidden_size,
+ audio_num_attention_heads=self.audio_num_attention_heads,
+ audio_attention_head_dim=self.audio_hidden_size
+ // self.audio_num_attention_heads,
+ audio_cross_attention_dim=arch.audio_cross_attention_dim,
+ norm_eps=self.norm_eps,
+ qk_norm=True, # Always True in LTX2
+ supported_attention_backends=self._supported_attention_backends,
+ prefix=config.prefix,
+ )
+ for idx in range(arch.num_layers)
+ ]
+ )
+
+ # 6. Output layers
+ self.norm_out = nn.LayerNorm(
+ self.hidden_size, eps=self.norm_eps, elementwise_affine=False
+ )
+ self.proj_out = ColumnParallelLinear(
+ self.hidden_size, arch.out_channels, bias=True, gather_output=True
+ )
+
+ self.audio_norm_out = nn.LayerNorm(
+ self.audio_hidden_size, eps=self.norm_eps, elementwise_affine=False
+ )
+ self.audio_proj_out = ColumnParallelLinear(
+ self.audio_hidden_size,
+ arch.audio_out_channels,
+ bias=True,
+ gather_output=True,
+ )
+
+ self.out_channels_raw = arch.out_channels // (
+ self.patch_size[0] * self.patch_size[1] * self.patch_size[2]
+ )
+ self.audio_out_channels = arch.audio_out_channels
+ self.timestep_scale_multiplier = arch.timestep_scale_multiplier
+ self.av_ca_timestep_scale_multiplier = arch.av_ca_timestep_scale_multiplier
+
+ self.layer_names = ["transformer_blocks"]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ audio_hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ audio_encoder_hidden_states: torch.Tensor,
+ timestep: torch.LongTensor,
+ audio_timestep: Optional[torch.LongTensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ audio_encoder_attention_mask: Optional[torch.Tensor] = None,
+ num_frames: Optional[int] = None,
+ height: Optional[int] = None,
+ width: Optional[int] = None,
+ fps: float = 24.0,
+ audio_num_frames: Optional[int] = None,
+ video_coords: Optional[torch.Tensor] = None,
+ audio_coords: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+
+ batch_size = hidden_states.size(0)
+ audio_timestep = audio_timestep if audio_timestep is not None else timestep
+
+ if num_frames is None or height is None or width is None:
+ raise ValueError(
+ "num_frames/height/width must be provided for RoPE coordinate generation."
+ )
+ if audio_num_frames is None:
+ raise ValueError(
+ "audio_num_frames must be provided for RoPE coordinate generation."
+ )
+
+ if video_coords is None:
+ # Wan-style SP-RoPE: when SP is enabled, each rank runs on its local
+ # time shard but RoPE positions must be offset to global time.
+ #
+ # We assume equal time sharding across SP ranks.
+ if model_parallel_is_initialized():
+ sp_world_size = get_sp_world_size()
+ sp_rank = get_sp_parallel_rank()
+ else:
+ sp_world_size = 1
+ sp_rank = 0
+
+ video_shift = int(sp_rank) * int(num_frames) if sp_world_size > 1 else 0
+ video_coords = self.rope.prepare_video_coords(
+ batch_size=batch_size,
+ num_frames=num_frames,
+ height=height,
+ width=width,
+ device=hidden_states.device,
+ fps=fps,
+ start_frame=video_shift,
+ )
+ if audio_coords is None:
+ audio_coords = self.audio_rope.prepare_audio_coords(
+ batch_size=batch_size,
+ num_frames=audio_num_frames,
+ device=audio_hidden_states.device,
+ )
+
+ video_rotary_emb = self.rope(video_coords, device=hidden_states.device)
+ audio_rotary_emb = self.audio_rope(
+ audio_coords, device=audio_hidden_states.device
+ )
+ ca_video_rotary_emb = self.cross_attn_rope(
+ video_coords[:, 0:1, :], device=hidden_states.device
+ )
+ ca_audio_rotary_emb = self.cross_attn_audio_rope(
+ audio_coords[:, 0:1, :], device=audio_hidden_states.device
+ )
+
+ # 2. Patchify input projections
+ hidden_states, _ = self.patchify_proj(hidden_states)
+ audio_hidden_states, _ = self.audio_patchify_proj(audio_hidden_states)
+
+ # 3. Prepare timestep embeddings
+ # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters
+ temb, embedded_timestep = self.adaln_single(
+ timestep.flatten(),
+ )
+ temb = temb.view(batch_size, -1, temb.size(-1))
+ embedded_timestep = embedded_timestep.view(
+ batch_size, -1, embedded_timestep.size(-1)
+ )
+
+ temb_audio, audio_embedded_timestep = self.audio_adaln_single(
+ audio_timestep.flatten()
+ )
+ temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1))
+ audio_embedded_timestep = audio_embedded_timestep.view(
+ batch_size, -1, audio_embedded_timestep.size(-1)
+ )
+
+ # 3.2. Prepare global modality cross attention modulation parameters
+ ts_ca_mult = (
+ self.av_ca_timestep_scale_multiplier / self.timestep_scale_multiplier
+ )
+
+ temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single(
+ timestep.flatten()
+ )
+ temb_ca_scale_shift = temb_ca_scale_shift.view(
+ batch_size, -1, temb_ca_scale_shift.shape[-1]
+ )
+
+ temb_ca_gate, _ = self.av_ca_a2v_gate_adaln_single(
+ timestep.flatten() * ts_ca_mult
+ )
+ temb_ca_gate = temb_ca_gate.view(batch_size, -1, temb_ca_gate.shape[-1])
+
+ temb_ca_audio_scale_shift, _ = self.av_ca_audio_scale_shift_adaln_single(
+ audio_timestep.flatten()
+ )
+ temb_ca_audio_scale_shift = temb_ca_audio_scale_shift.view(
+ batch_size, -1, temb_ca_audio_scale_shift.shape[-1]
+ )
+
+ temb_ca_audio_gate, _ = self.av_ca_v2a_gate_adaln_single(
+ audio_timestep.flatten() * ts_ca_mult
+ )
+ temb_ca_audio_gate = temb_ca_audio_gate.view(
+ batch_size, -1, temb_ca_audio_gate.shape[-1]
+ )
+
+ # 4. Prepare prompt embeddings
+ encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+ audio_encoder_hidden_states = self.audio_caption_projection(
+ audio_encoder_hidden_states
+ )
+
+ # 5. Run blocks
+ for block in self.transformer_blocks:
+ hidden_states, audio_hidden_states = block(
+ hidden_states,
+ audio_hidden_states,
+ encoder_hidden_states,
+ audio_encoder_hidden_states,
+ # Keep the first 4 args positional to stay compatible with cache-dit's
+ # LTX2 adapter, which treats `audio_hidden_states` as `encoder_hidden_states`
+ # under ForwardPattern.Pattern_0.
+ temb=temb,
+ temb_audio=temb_audio,
+ temb_ca_scale_shift=temb_ca_scale_shift,
+ temb_ca_audio_scale_shift=temb_ca_audio_scale_shift,
+ temb_ca_gate=temb_ca_gate,
+ temb_ca_audio_gate=temb_ca_audio_gate,
+ video_rotary_emb=video_rotary_emb,
+ audio_rotary_emb=audio_rotary_emb,
+ ca_video_rotary_emb=ca_video_rotary_emb,
+ ca_audio_rotary_emb=ca_audio_rotary_emb,
+ encoder_attention_mask=encoder_attention_mask,
+ audio_encoder_attention_mask=audio_encoder_attention_mask,
+ )
+
+ # 6. Output layers
+ # Video
+ scale_shift_values = self.scale_shift_table[None, None].to(
+ device=hidden_states.device, dtype=hidden_states.dtype
+ ) + embedded_timestep[:, :, None].to(dtype=hidden_states.dtype)
+ shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+ hidden_states = self.norm_out(hidden_states)
+ hidden_states = hidden_states * (1 + scale) + shift
+ hidden_states, _ = self.proj_out(hidden_states)
+
+ # Audio
+ audio_scale_shift_values = self.audio_scale_shift_table[None, None].to(
+ device=audio_hidden_states.device, dtype=audio_hidden_states.dtype
+ ) + audio_embedded_timestep[:, :, None].to(dtype=audio_hidden_states.dtype)
+ audio_shift, audio_scale = (
+ audio_scale_shift_values[:, :, 0],
+ audio_scale_shift_values[:, :, 1],
+ )
+ audio_hidden_states = self.audio_norm_out(audio_hidden_states)
+ audio_hidden_states = audio_hidden_states * (1 + audio_scale) + audio_shift
+ audio_hidden_states, _ = self.audio_proj_out(audio_hidden_states)
+
+ # Unpatchify if requested (default True for pipeline compatibility)
+ return_latents = kwargs.get("return_latents", True)
+
+ if return_latents:
+ # Unpatchify Video
+ # [B, N, C_out_raw*patch_vol] -> [B, C_out_raw, T, H, W]
+ # Requires num_frames, height, width to be known
+ if num_frames is not None and height is not None and width is not None:
+ p_t, p_h, p_w = self.patch_size
+ post_t, post_h, post_w = num_frames // p_t, height // p_h, width // p_w
+ b = batch_size
+ hidden_states = hidden_states.reshape(
+ b, post_t, post_h, post_w, self.out_channels_raw, p_t, p_h, p_w
+ )
+ hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(
+ b, self.out_channels_raw, num_frames, height, width
+ )
+
+ # Unpatchify Audio
+ # [B, N, C_out] -> [B, C_out, T] (or 4D/5D)
+ if audio_num_frames is not None:
+ b = batch_size
+ # simple reshape for 1D patch
+ audio_hidden_states = audio_hidden_states.permute(0, 2, 1) # [B, C, T]
+
+ return hidden_states, audio_hidden_states
+
+
+# Backward-compatible alias (older internal name).
+LTXModel = LTX2VideoTransformer3DModel
+EntryClass = LTX2VideoTransformer3DModel
diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py
new file mode 100644
index 000000000..8dcc27c98
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/encoders/gemma_3.py
@@ -0,0 +1,1149 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from sglang: python/sglang/srt/models/gemma3_causal.py
+
+import logging
+from functools import partial
+from typing import Any, Iterable, Optional, Set, Tuple
+
+import torch
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.encoders.base import BaseEncoderOutput
+from sglang.multimodal_gen.configs.models.encoders.gemma_3 import Gemma3Config
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.activation import GeluAndMul
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention
+from sglang.multimodal_gen.runtime.layers.linear import (
+ ColumnParallelLinear,
+ MergedColumnParallelLinear,
+ QKVParallelLinear,
+ RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import get_rope
+from sglang.multimodal_gen.runtime.loader.weight_utils import default_weight_loader
+from sglang.multimodal_gen.runtime.utils.common import add_prefix
+
+logger = logging.getLogger(__name__)
+
+
+def get_attention_sliding_window_size(config):
+ return config.sliding_window - 1
+
+
+class Gemma3RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.eps = eps
+ self.weight = nn.Parameter(torch.zeros(dim))
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ output = self._norm(x.float())
+ output = output * (1.0 + self.weight.float())
+ return output.type_as(x)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+class Gemma3MLP(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ hidden_act: str,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.gate_up_proj = MergedColumnParallelLinear(
+ input_size=hidden_size,
+ output_sizes=[intermediate_size] * 2,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.gate_up_proj",
+ )
+ self.down_proj = RowParallelLinear(
+ input_size=intermediate_size,
+ output_size=hidden_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.down_proj",
+ )
+ if hidden_act != "gelu_pytorch_tanh":
+ raise ValueError(
+ "Gemma3 uses `gelu_pytorch_tanh` as the hidden activation "
+ "function. Please set `hidden_activation` to "
+ "`gelu_pytorch_tanh`."
+ )
+ self.act_fn = GeluAndMul(approximate="tanh")
+
+ def forward(self, x):
+ x, _ = self.gate_up_proj(x)
+ x = self.act_fn(x)
+ x, _ = self.down_proj(x)
+ return x
+
+
+class Gemma3Attention(nn.Module):
+ def __init__(
+ self,
+ layer_id: int,
+ config: Gemma3Config,
+ hidden_size: int,
+ num_heads: int,
+ num_kv_heads: int,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.layer_id = layer_id
+ self.hidden_size = hidden_size
+ tp_size = get_tp_world_size()
+ self.total_num_heads = num_heads
+ assert self.total_num_heads % tp_size == 0
+ self.num_heads = self.total_num_heads // tp_size
+ self.total_num_kv_heads = num_kv_heads
+ if self.total_num_kv_heads >= tp_size:
+ assert self.total_num_kv_heads % tp_size == 0
+ else:
+ assert tp_size % self.total_num_kv_heads == 0
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+ self.head_dim = getattr(
+ config.text_config, "head_dim", self.hidden_size // self.total_num_heads
+ )
+
+ self.q_size = self.num_heads * self.head_dim
+ self.kv_size = self.num_kv_heads * self.head_dim
+ self.scaling = config.text_config.query_pre_attn_scalar**-0.5
+
+ self.qkv_proj = QKVParallelLinear(
+ hidden_size=hidden_size,
+ head_size=self.head_dim,
+ total_num_heads=self.total_num_heads,
+ total_num_kv_heads=self.total_num_kv_heads,
+ bias=config.text_config.attention_bias,
+ quant_config=quant_config,
+ prefix=f"{prefix}.qkv_proj",
+ )
+
+ self.o_proj = RowParallelLinear(
+ input_size=self.total_num_heads * self.head_dim,
+ output_size=hidden_size,
+ bias=config.text_config.attention_bias,
+ quant_config=quant_config,
+ prefix=f"{prefix}.o_proj",
+ )
+
+ self.is_sliding = (
+ config.text_config.layer_types[layer_id] == "sliding_attention"
+ )
+
+ # Initialize the rotary embedding.
+ if self.is_sliding:
+ # Local attention.
+ self.rope_theta = config.text_config.rope_local_base_freq
+ rope_scaling = None # Default
+ # sliding window
+ self.sliding_window = get_attention_sliding_window_size(config.text_config)
+ # (left, right) = (window, 0) effectively for causal
+ self.window_size = (self.sliding_window, 0)
+ else:
+ # Global attention.
+ self.rope_theta = config.text_config.rope_theta
+ rope_scaling = config.text_config.rope_scaling
+ self.sliding_window = None
+ self.window_size = (-1, -1)
+
+ self.rotary_emb = get_rope(
+ self.head_dim,
+ rotary_dim=self.head_dim,
+ max_position=config.text_config.max_position_embeddings,
+ base=self.rope_theta,
+ rope_scaling=rope_scaling,
+ is_neox_style=True,
+ )
+
+ # Local Attention not support attention mask, we use global attention instead.
+ # self.attn = LocalAttention(
+ # self.num_heads,
+ # self.head_dim,
+ # self.num_kv_heads,
+ # softmax_scale=self.scaling,
+ # causal=True,
+ # supported_attention_backends=config._supported_attention_backends,
+ # window_size=self.window_size,
+ # )
+
+ # Gemma3 adds normalization for q and k
+ self.q_norm = Gemma3RMSNorm(
+ dim=self.head_dim, eps=config.text_config.rms_norm_eps
+ )
+ self.k_norm = Gemma3RMSNorm(
+ dim=self.head_dim, eps=config.text_config.rms_norm_eps
+ )
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ qkv, _ = self.qkv_proj(hidden_states)
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+ batch_size, seq_len, _ = q.shape
+ q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)
+ k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+ v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+ # Apply QK Norm
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+
+ # Apply RoPE
+ q, k = self.rotary_emb(positions, q, k)
+
+ # TODO(FlamingoPg): Support LocalAttention
+ query = q.transpose(1, 2)
+ key = k.transpose(1, 2)
+ value = v.transpose(1, 2)
+
+ attn_mask = torch.zeros(
+ (seq_len, seq_len),
+ device=hidden_states.device,
+ dtype=torch.float32,
+ )
+ causal = torch.triu(
+ torch.ones(
+ (seq_len, seq_len), device=hidden_states.device, dtype=torch.bool
+ ),
+ diagonal=1,
+ )
+ attn_mask = attn_mask.masked_fill(causal, float("-inf"))
+ if self.is_sliding and self.sliding_window is not None:
+ idx = torch.arange(seq_len, device=hidden_states.device)
+ dist = idx[None, :] - idx[:, None]
+ too_far = dist > self.sliding_window
+ attn_mask = attn_mask.masked_fill(too_far, float("-inf"))
+
+ key_pad = ~attention_mask.to(torch.bool)
+ attn_mask = attn_mask[None, None, :, :].expand(batch_size, 1, seq_len, seq_len)
+ attn_mask = attn_mask.masked_fill(
+ key_pad[:, None, None, :].expand(batch_size, 1, seq_len, seq_len),
+ float("-inf"),
+ )
+
+ attn_kwargs = {
+ "attn_mask": attn_mask,
+ "dropout_p": 0.0,
+ "is_causal": False,
+ "scale": self.scaling,
+ }
+ if query.shape[1] != key.shape[1]:
+ attn_kwargs["enable_gqa"] = True
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query, key, value, **attn_kwargs
+ )
+ attn_output = attn_output.transpose(1, 2)
+
+ attn_output = attn_output.reshape(
+ batch_size, seq_len, self.num_heads * self.head_dim
+ )
+
+ output, _ = self.o_proj(attn_output)
+ return output
+
+
+class Gemma3DecoderLayer(nn.Module):
+ def __init__(
+ self,
+ layer_id: int,
+ config: Gemma3Config,
+ quant_config: QuantizationConfig | None = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.hidden_size = config.text_config.hidden_size
+ self.self_attn = Gemma3Attention(
+ layer_id=layer_id,
+ config=config,
+ hidden_size=self.hidden_size,
+ num_heads=config.text_config.num_attention_heads,
+ num_kv_heads=getattr(
+ config.text_config,
+ "num_key_value_heads",
+ config.text_config.num_attention_heads,
+ ),
+ quant_config=quant_config,
+ prefix=f"{prefix}.self_attn",
+ )
+ self.mlp = Gemma3MLP(
+ hidden_size=self.hidden_size,
+ intermediate_size=config.text_config.intermediate_size,
+ hidden_act=config.text_config.hidden_activation,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp",
+ )
+ self.input_layernorm = Gemma3RMSNorm(
+ config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
+ )
+ self.post_attention_layernorm = Gemma3RMSNorm(
+ config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
+ )
+ self.pre_feedforward_layernorm = Gemma3RMSNorm(
+ config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
+ )
+ self.post_feedforward_layernorm = Gemma3RMSNorm(
+ config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
+ )
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ residual: torch.Tensor | None,
+ attention_mask: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # Self Attention
+ # Gemma3 uses "sandwich norm":
+ # x = x + norm(attn(norm(x)))
+ # So we treat input hidden_states as the residual base.
+
+ if residual is not None:
+ hidden_states = hidden_states + residual
+ residual = None
+
+ residual_input = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states = self.self_attn(
+ positions=positions,
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ )
+
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = residual_input + hidden_states
+
+ # MLP
+ residual_mlp = hidden_states
+ hidden_states = self.pre_feedforward_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = self.post_feedforward_layernorm(hidden_states)
+ hidden_states = residual_mlp + hidden_states
+
+ return hidden_states, None
+
+
+class Gemma3TextScaledWordEmbedding(nn.Embedding):
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ padding_idx: int,
+ embed_scale: Optional[float] = 1.0,
+ ):
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
+ self.embed_scale = embed_scale
+
+ def forward(self, input_ids: torch.Tensor):
+ return super().forward(input_ids) * self.embed_scale
+
+
+# --- Siglip Vision Model Implementation ---
+
+
+class QuickGELU(nn.Module):
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return x * torch.sigmoid(1.702 * x)
+
+
+class SiglipVisionEmbeddings(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=config.num_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.patch_size,
+ stride=self.patch_size,
+ padding="valid",
+ )
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches
+ # Use simple Embedding for position embeddings (usually small enough)
+ self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+ self.register_buffer(
+ "position_ids",
+ torch.arange(self.num_positions).expand((1, -1)),
+ persistent=False,
+ )
+
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.patch_embedding.weight.dtype
+ patch_embeds = self.patch_embedding(
+ pixel_values.to(dtype=target_dtype)
+ ) # shape = [*, width, grid, grid]
+ embeddings = patch_embeds.flatten(2).transpose(1, 2)
+ embeddings = embeddings + self.position_embedding(self.position_ids)
+
+ return embeddings
+
+
+class SiglipMLP(nn.Module):
+ def __init__(
+ self,
+ config,
+ act_layer: type[nn.Module] = QuickGELU,
+ quant_config: Optional[QuantizationConfig] = None,
+ prefix: str = "",
+ ):
+ super().__init__()
+ self.fc1 = ColumnParallelLinear(
+ config.hidden_size,
+ config.intermediate_size,
+ quant_config=quant_config,
+ prefix=add_prefix("fc1", prefix),
+ )
+ self.act = act_layer()
+ self.fc2 = RowParallelLinear(
+ config.intermediate_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=add_prefix("fc2", prefix),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x_parallel, _ = self.fc1(x)
+ x_parallel = self.act(x_parallel)
+ x, _ = self.fc2(x_parallel)
+ return x
+
+
+class SiglipAttention(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ num_heads: int,
+ quant_config: Optional[QuantizationConfig] = None,
+ prefix: str = "",
+ ):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.num_heads = num_heads
+ tp_size = get_tp_world_size()
+ self.head_dim = hidden_size // num_heads
+ self.num_heads_per_partition = num_heads // tp_size
+ self.scaling = self.head_dim**-0.5
+
+ self.qkv_proj = QKVParallelLinear(
+ hidden_size=hidden_size,
+ head_size=self.head_dim,
+ total_num_heads=num_heads,
+ total_num_kv_heads=num_heads,
+ bias=True,
+ quant_config=quant_config,
+ prefix=add_prefix("qkv_proj", prefix),
+ )
+
+ self.out_proj = RowParallelLinear(
+ input_size=hidden_size,
+ output_size=hidden_size,
+ bias=True,
+ quant_config=quant_config,
+ prefix=add_prefix("out_proj", prefix),
+ )
+
+ self.attn = LocalAttention(
+ num_heads=self.num_heads_per_partition,
+ head_size=self.head_dim,
+ num_kv_heads=self.num_heads_per_partition,
+ softmax_scale=self.scaling,
+ causal=False, # Bidirectional for Vision
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ qkv, _ = self.qkv_proj(hidden_states)
+ q, k, v = qkv.split([self.hidden_size // get_tp_world_size()] * 3, dim=-1)
+
+ batch_size, seq_len, _ = q.shape
+ q = q.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
+ k = k.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
+ v = v.view(batch_size, seq_len, self.num_heads_per_partition, self.head_dim)
+
+ attn_output = self.attn(q, k, v)
+
+ attn_output = attn_output.reshape(
+ batch_size, seq_len, self.hidden_size // get_tp_world_size()
+ )
+
+ output, _ = self.out_proj(attn_output)
+ return output
+
+
+class SiglipEncoderLayer(nn.Module):
+ def __init__(
+ self,
+ config,
+ act_layer: type[nn.Module] = QuickGELU,
+ norm_layer: type[nn.Module] = None,
+ quant_config: Optional[QuantizationConfig] = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ if norm_layer is None:
+ norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+ self.layer_norm1 = norm_layer(config.hidden_size)
+ self.layer_norm2 = norm_layer(config.hidden_size)
+ self.self_attn = SiglipAttention(
+ hidden_size=config.hidden_size,
+ num_heads=config.num_attention_heads,
+ quant_config=quant_config,
+ prefix=add_prefix("self_attn", prefix),
+ )
+ self.mlp = SiglipMLP(
+ config,
+ act_layer=act_layer,
+ quant_config=quant_config,
+ prefix=add_prefix("mlp", prefix),
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states = self.self_attn(hidden_states)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+class SiglipEncoder(nn.Module):
+ def __init__(
+ self,
+ config,
+ quant_config: Optional[QuantizationConfig] = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.config = config
+ num_hidden_layers = config.num_hidden_layers
+ norm_layer = partial(nn.LayerNorm, eps=config.layer_norm_eps)
+ self.layers = nn.ModuleList(
+ [
+ SiglipEncoderLayer(
+ config=config,
+ norm_layer=norm_layer,
+ quant_config=quant_config,
+ prefix=add_prefix(f"layers.{layer_idx}", prefix),
+ )
+ for layer_idx in range(num_hidden_layers)
+ ]
+ )
+
+ def forward(
+ self,
+ inputs_embeds: torch.Tensor,
+ ) -> torch.Tensor:
+ hidden_states = inputs_embeds
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(hidden_states)
+ return hidden_states
+
+
+class SiglipVisionTransformer(nn.Module):
+ def __init__(
+ self,
+ config,
+ quant_config: Optional[QuantizationConfig] = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.config = config
+ embed_dim = config.hidden_size
+ self.embeddings = SiglipVisionEmbeddings(config)
+ self.encoder = SiglipEncoder(
+ config=config,
+ quant_config=quant_config,
+ prefix=add_prefix("encoder", prefix),
+ )
+ self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ @property
+ def device(self) -> torch.device:
+ return self.encoder.layers[0].layer_norm1.weight.device
+
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.embeddings(pixel_values.to(self.device))
+ last_hidden_state = self.encoder(inputs_embeds=hidden_states)
+ last_hidden_state = self.post_layernorm(last_hidden_state)
+ return last_hidden_state
+
+
+class SiglipVisionModel(nn.Module):
+ def __init__(
+ self,
+ config,
+ quant_config: Optional[QuantizationConfig] = None,
+ prefix: str = "",
+ ):
+ super().__init__()
+ self.vision_model = SiglipVisionTransformer(
+ config, quant_config, prefix=add_prefix("vision_model", prefix)
+ )
+
+ @property
+ def device(self) -> torch.device:
+ return self.vision_model.device
+
+ def forward(self, pixel_values: torch.Tensor):
+ return self.vision_model(pixel_values)
+
+
+class Gemma3MultiModalProjector(nn.Module):
+ """Projector for Gemma3 multimodal."""
+
+ def __init__(self, config: Gemma3Config):
+ super().__init__()
+
+ self.mm_input_projection_weight = nn.Parameter(
+ torch.zeros(
+ config.vision_config.hidden_size, config.text_config.hidden_size
+ )
+ )
+
+ self.mm_soft_emb_norm = Gemma3RMSNorm(
+ config.vision_config.hidden_size, eps=config.vision_config.layer_norm_eps
+ )
+
+ self.patches_per_image = int(
+ config.vision_config.image_size // config.vision_config.patch_size
+ )
+ self.tokens_per_side = int(config.mm_tokens_per_image**0.5)
+ self.kernel_size = self.patches_per_image // self.tokens_per_side
+ self.avg_pool = nn.AvgPool2d(
+ kernel_size=self.kernel_size, stride=self.kernel_size
+ )
+
+ def forward(self, vision_outputs: torch.Tensor) -> torch.Tensor:
+ batch_size, seq_length, hidden_size = vision_outputs.shape
+
+ # Reshape for pooling
+ reshaped_vision_outputs = vision_outputs.transpose(1, 2)
+ reshaped_vision_outputs = reshaped_vision_outputs.reshape(
+ batch_size, hidden_size, self.patches_per_image, self.patches_per_image
+ )
+ reshaped_vision_outputs = reshaped_vision_outputs.contiguous()
+
+ # Apply pooling
+ pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs)
+ pooled_vision_outputs = pooled_vision_outputs.flatten(2)
+ pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2)
+
+ # Apply normalization
+ normed_vision_outputs = self.mm_soft_emb_norm(pooled_vision_outputs)
+
+ # Project to text embedding space
+ projected_vision_outputs = torch.matmul(
+ normed_vision_outputs, self.mm_input_projection_weight
+ )
+
+ return projected_vision_outputs.type_as(vision_outputs)
+
+
+class Gemma3TextModel(nn.Module):
+ def __init__(self, config: Gemma3Config):
+ super().__init__()
+ self.config = config
+ # TODO(yinfan.1024) support text encoding model quant later
+ self.quant_config = None
+
+ # Use VocabParallelEmbedding
+ from sglang.multimodal_gen.runtime.layers.vocab_parallel_embedding import (
+ VocabParallelEmbedding,
+ )
+
+ self.vocab_size = config.text_config.vocab_size
+ self.embed_tokens = VocabParallelEmbedding(
+ self.vocab_size,
+ config.text_config.hidden_size,
+ org_num_embeddings=config.text_config.vocab_size,
+ quant_config=self.quant_config,
+ )
+ self.embed_scale = config.text_config.hidden_size**0.5
+
+ self.layers = nn.ModuleList(
+ [
+ Gemma3DecoderLayer(
+ layer_id=i,
+ config=config,
+ quant_config=self.quant_config,
+ prefix=f"{config.text_config.prefix}.layers.{i}",
+ )
+ for i in range(config.text_config.num_hidden_layers)
+ ]
+ )
+
+ self.norm = Gemma3RMSNorm(
+ config.text_config.hidden_size, eps=config.text_config.rms_norm_eps
+ )
+
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+ return self.embed_tokens(input_ids) * self.embed_scale
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None,
+ position_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ **kwargs,
+ ) -> BaseEncoderOutput:
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+
+ if inputs_embeds is not None:
+ hidden_states = inputs_embeds
+ else:
+ hidden_states = self.get_input_embeddings(input_ids)
+
+ residual = None
+
+ if position_ids is None:
+ position_ids = torch.arange(
+ 0, hidden_states.shape[1], device=hidden_states.device
+ ).unsqueeze(0)
+ position_ids = position_ids + 1
+
+ all_hidden_states: tuple[Any, ...] | None = () if output_hidden_states else None
+
+ for layer in self.layers:
+ if all_hidden_states is not None:
+ all_hidden_states += (hidden_states,)
+
+ hidden_states, residual = layer(
+ position_ids,
+ hidden_states,
+ residual,
+ attention_mask=attention_mask,
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ if all_hidden_states is not None:
+ all_hidden_states += (hidden_states,)
+
+ output = BaseEncoderOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ )
+
+ return output
+
+ def load_weights(self, weights: Any) -> set[str]:
+ # Copied from LlamaModel.load_weights but adapted
+ params_dict = dict(self.named_parameters())
+ loaded_params: set[str] = set()
+
+ def _load_with_shard_id(
+ weight_loader, param, loaded_weight: torch.Tensor, shard_id
+ ) -> None:
+ """Call param.weight_loader with best-effort shard_id normalization.
+
+ Different fused-QKV implementations expect different shard_id types:
+ - Some expect strings: "q"/"k"/"v"
+ - Some expect integer indices: 0/1/2
+ We try the provided shard_id first, then fall back between str/int forms.
+ """
+ try:
+ weight_loader(param, loaded_weight, shard_id)
+ return
+ except (AssertionError, TypeError):
+ pass
+
+ # Fall back between common representations.
+ if isinstance(shard_id, str):
+ mapping = {"q": 0, "k": 1, "v": 2}
+ if shard_id in mapping:
+ weight_loader(param, loaded_weight, mapping[shard_id])
+ return
+ if shard_id.isdigit():
+ weight_loader(param, loaded_weight, int(shard_id))
+ return
+ elif isinstance(shard_id, int):
+ mapping = {0: "q", 1: "k", 2: "v"}
+ if shard_id in mapping:
+ weight_loader(param, loaded_weight, mapping[shard_id])
+ return
+
+ # Re-raise with a clearer message.
+ raise TypeError(
+ f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader} "
+ f"(param={getattr(param, 'name', '')})."
+ )
+
+ stacked_params_mapping = getattr(
+ getattr(self.config, "arch_config", object()),
+ "stacked_params_mapping",
+ None,
+ )
+ if stacked_params_mapping is None:
+ stacked_params_mapping = [
+ # Fused QKV shards; downstream loaders may want "q/k/v" or 0/1/2.
+ (".qkv_proj", ".q_proj", "q"),
+ (".qkv_proj", ".k_proj", "k"),
+ (".qkv_proj", ".v_proj", "v"),
+ (".gate_up_proj", ".gate_proj", 0),
+ (".gate_up_proj", ".up_proj", 1),
+ ]
+
+ for name, loaded_weight in weights:
+ if "rotary_emb.inv_freq" in name:
+ continue
+
+ # The config has stacked_params_mapping
+ for (
+ param_name,
+ weight_name,
+ shard_id,
+ ) in stacked_params_mapping:
+ if weight_name not in name:
+ continue
+ name = name.replace(weight_name, param_name)
+
+ if name not in params_dict:
+ continue
+
+ param = params_dict[name]
+ weight_loader = param.weight_loader
+ _load_with_shard_id(weight_loader, param, loaded_weight, shard_id)
+ break
+ else:
+ if name not in params_dict:
+ continue
+
+ param = params_dict[name]
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
+ weight_loader(param, loaded_weight)
+
+ loaded_params.add(name)
+ return loaded_params
+
+
+class Gemma3ForConditionalGeneration(nn.Module):
+ def __init__(
+ self,
+ config: Gemma3Config,
+ quant_config: Optional[QuantizationConfig] = None,
+ prefix: str = "",
+ ) -> None:
+ super().__init__()
+ self.config = config
+ self.quant_config = quant_config
+ self.text_config = config.text_config
+
+ # Vision Tower
+ self.vision_tower = SiglipVisionModel(
+ config=config.vision_config,
+ quant_config=quant_config,
+ prefix=add_prefix("vision_tower", prefix),
+ )
+
+ # Projector
+ self.multi_modal_projector = Gemma3MultiModalProjector(config)
+
+ # Text Model
+ self.language_model = Gemma3TextModel(config)
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: torch.FloatTensor,
+ ) -> torch.Tensor:
+ image_token_index = int(getattr(self.config, "image_token_index", -1))
+ if image_token_index < 0:
+ image_token_index = int(getattr(self.text_config, "image_token_index", -1))
+ special_image_mask = input_ids == image_token_index
+ n_image_tokens = int(special_image_mask.sum().item())
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds)
+ n_image_features = int(image_features.shape[0] * image_features.shape[1])
+ if inputs_embeds[special_image_mask].numel() != image_features.numel():
+ raise ValueError(
+ f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+ )
+ return special_image_mask
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor,
+ pixel_values: torch.FloatTensor | None = None,
+ **kwargs,
+ ):
+ vocab_size = int(self.language_model.vocab_size)
+ image_token_index = int(getattr(self.config, "image_token_index", -1))
+ if image_token_index < 0:
+ image_token_index = int(getattr(self.text_config, "image_token_index", -1))
+
+ if input_ids is not None and image_token_index >= vocab_size:
+ special_image_mask = input_ids == image_token_index
+ llm_input_ids = input_ids.clone()
+ llm_input_ids[special_image_mask] = 0
+ else:
+ llm_input_ids = input_ids
+
+ inputs_embeds = self.language_model.get_input_embeddings(llm_input_ids)
+
+ if pixel_values is not None:
+ if pixel_values.dim() == 5:
+ pixel_values = pixel_values.reshape(
+ -1,
+ pixel_values.shape[2],
+ pixel_values.shape[3],
+ pixel_values.shape[4],
+ )
+ elif pixel_values.dim() == 3:
+ pixel_values = pixel_values.unsqueeze(0)
+ elif pixel_values.dim() != 4:
+ raise ValueError(f"Unexpected pixel_values shape: {pixel_values.shape}")
+
+ vision_outputs = self.vision_tower(pixel_values)
+ image_features = self.multi_modal_projector(vision_outputs)
+ image_features = image_features.to(
+ device=inputs_embeds.device, dtype=inputs_embeds.dtype
+ )
+ special_image_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(
+ special_image_mask, image_features
+ )
+
+ return self.language_model.forward(
+ llm_input_ids, inputs_embeds=inputs_embeds, **kwargs
+ )
+
+ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
+ loaded_params: Set[str] = set()
+ params_dict = dict(self.named_parameters())
+
+ def _load_with_shard_id(
+ weight_loader, param, loaded_weight: torch.Tensor, shard_id
+ ) -> None:
+ """Call param.weight_loader with best-effort shard_id normalization.
+
+ Different fused-QKV implementations expect different shard_id types:
+ - Some expect strings: "q"/"k"/"v"
+ - Some expect integer indices: 0/1/2
+ We try the provided shard_id first, then fall back between str/int forms.
+ """
+ try:
+ weight_loader(param, loaded_weight, shard_id)
+ return
+ except (AssertionError, TypeError):
+ pass
+
+ # Fall back between common representations.
+ if isinstance(shard_id, str):
+ mapping = {"q": 0, "k": 1, "v": 2}
+ if shard_id in mapping:
+ weight_loader(param, loaded_weight, mapping[shard_id])
+ return
+ if shard_id.isdigit():
+ weight_loader(param, loaded_weight, int(shard_id))
+ return
+ elif isinstance(shard_id, int):
+ mapping = {0: "q", 1: "k", 2: "v"}
+ if shard_id in mapping:
+ weight_loader(param, loaded_weight, mapping[shard_id])
+ return
+
+ raise TypeError(
+ f"Unsupported shard_id={shard_id!r} for weight_loader={weight_loader} "
+ f"(param={getattr(param, 'name', '')})."
+ )
+
+ # Separate weights
+ language_model_weights: list[tuple[str, torch.Tensor]] = []
+ other_weights: list[tuple[str, torch.Tensor]] = []
+
+ for name, loaded_weight in weights:
+ # Handle prefix mapping if needed
+ # HF weights might be "model.vision_tower...", "model.language_model..."
+
+ if "vision_tower" in name or "vision_model" in name:
+ # Load vision tower weights
+ # Map name to local name
+ local_name = name
+ if "model.vision_tower" in name:
+ local_name = name.replace("model.vision_tower", "vision_tower")
+ elif "vision_tower" in name:
+ pass # already correct prefix if matching self.vision_tower
+ elif local_name.startswith("vision_model."):
+ local_name = (
+ "vision_tower.vision_model."
+ + local_name[len("vision_model.") :]
+ )
+
+ # We need to map HF Siglip names to our Siglip implementation
+ # Our Siglip: vision_tower.vision_model.encoder.layers...
+ # HF Siglip: vision_model.encoder.layers...
+
+ # If loading from Gemma3 checkpoint, it usually has "model.vision_tower.vision_model..."
+
+ if local_name in params_dict:
+ param = params_dict[local_name]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
+ loaded_params.add(local_name)
+ else:
+ qkv_shard_id = None
+ fused_name = None
+ if ".self_attn.q_proj." in local_name:
+ fused_name = local_name.replace(
+ ".self_attn.q_proj.", ".self_attn.qkv_proj."
+ )
+ qkv_shard_id = "q"
+ elif ".self_attn.k_proj." in local_name:
+ fused_name = local_name.replace(
+ ".self_attn.k_proj.", ".self_attn.qkv_proj."
+ )
+ qkv_shard_id = "k"
+ elif ".self_attn.v_proj." in local_name:
+ fused_name = local_name.replace(
+ ".self_attn.v_proj.", ".self_attn.qkv_proj."
+ )
+ qkv_shard_id = "v"
+
+ if fused_name is not None and fused_name in params_dict:
+ param = params_dict[fused_name]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ _load_with_shard_id(
+ weight_loader, param, loaded_weight, qkv_shard_id
+ )
+ loaded_params.add(fused_name)
+ continue
+
+ if ".self_attn.proj." in local_name:
+ candidate = local_name.replace(
+ ".self_attn.proj.", ".self_attn.out_proj."
+ )
+ if candidate in params_dict:
+ param = params_dict[candidate]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
+ loaded_params.add(candidate)
+ continue
+ if ".self_attn.out_proj." in local_name:
+ candidate = local_name.replace(
+ ".self_attn.out_proj.", ".self_attn.proj."
+ )
+ if candidate in params_dict:
+ param = params_dict[candidate]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
+ loaded_params.add(candidate)
+ continue
+
+ # Try to find match
+ suffix = local_name.split("vision_tower.")[-1]
+ # Try adding vision_model
+ candidate = f"vision_tower.vision_model.{suffix}"
+ if candidate in params_dict:
+ param = params_dict[candidate]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
+ loaded_params.add(candidate)
+
+ elif "multi_modal_projector" in name:
+ local_name = name
+ if "model.multi_modal_projector" in name:
+ local_name = name.replace(
+ "model.multi_modal_projector", "multi_modal_projector"
+ )
+
+ if local_name in params_dict:
+ param = params_dict[local_name]
+ weight_loader = getattr(
+ param, "weight_loader", default_weight_loader
+ )
+ weight_loader(param, loaded_weight)
+ loaded_params.add(local_name)
+
+ elif "language_model" in name or "model.language_model" in name:
+ # Strip prefix for language model
+ # If name is "model.language_model.model.layers.0...", we want "model.layers.0..." for Gemma3ForCausalLM
+ # Gemma3ForCausalLM has .model (Gemma3TextModel) and .lm_head
+
+ # HF: model.language_model.model.layers...
+ # Ours: language_model.model.layers...
+
+ # We pass (name, weight) to language_model.load_weights
+ # We should strip "model.language_model." or "language_model."
+
+ suffix = name
+ if "model.language_model." in name:
+ suffix = name.replace("model.language_model.", "")
+ elif "language_model." in name:
+ suffix = name.replace("language_model.", "")
+ if suffix.startswith("model."):
+ suffix = suffix[len("model.") :]
+
+ language_model_weights.append((suffix, loaded_weight))
+
+ else:
+ # Fallback for other weights (maybe direct lm_head if not nested?)
+ other_weights.append((name, loaded_weight))
+
+ if language_model_weights:
+ lm_loaded = self.language_model.load_weights(language_model_weights)
+ loaded_params.update({f"language_model.{n}" for n in lm_loaded})
+
+ return loaded_params
+
+ def get_attention_sliding_window_size(self):
+ if self.text_config is not None and hasattr(
+ self.text_config, "get_attention_sliding_window_size"
+ ):
+ return self.text_config.get_attention_sliding_window_size()
+ sliding_window = getattr(self.text_config, "sliding_window", None)
+ if sliding_window is None:
+ sliding_window = getattr(self.config, "sliding_window", None)
+ if sliding_window is None:
+ return None
+ return int(sliding_window) - 1
+
+
+EntryClass = Gemma3ForConditionalGeneration
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_audio.py
new file mode 100644
index 000000000..0341771a0
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py
new file mode 100644
index 000000000..77c4e512b
--- /dev/null
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py
@@ -0,0 +1,1676 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from diffusers.models.activations import get_activation
+from diffusers.models.autoencoders.vae import (
+ DecoderOutput,
+ DiagonalGaussianDistribution,
+)
+from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+
+from sglang.multimodal_gen.configs.models.vaes.ltx_video import LTXVideoVAEConfig
+from sglang.multimodal_gen.runtime.models.vaes.common import ParallelTiledVAE
+
+
+class PerChannelRMSNorm(nn.Module):
+ """
+ Per-pixel (per-location) RMS normalization layer.
+
+ For each element along the chosen dimension, this layer normalizes the tensor by the root-mean-square of its values
+ across that dimension:
+
+ y = x / sqrt(mean(x^2, dim=dim, keepdim=True) + eps)
+ """
+
+ def __init__(self, channel_dim: int = 1, eps: float = 1e-8) -> None:
+ """
+ Args:
+ dim: Dimension along which to compute the RMS (typically channels).
+ eps: Small constant added for numerical stability.
+ """
+ super().__init__()
+ self.channel_dim = channel_dim
+ self.eps = eps
+
+ def forward(
+ self, x: torch.Tensor, channel_dim: Optional[int] = None
+ ) -> torch.Tensor:
+ """
+ Apply RMS normalization along the configured dimension.
+ """
+ channel_dim = channel_dim or self.channel_dim
+ # Compute mean of squared values along `dim`, keep dimensions for broadcasting.
+ mean_sq = torch.mean(x**2, dim=self.channel_dim, keepdim=True)
+ # Normalize by the root-mean-square (RMS).
+ rms = torch.sqrt(mean_sq + self.eps)
+ return x / rms
+
+
+# Like LTXCausalConv3d, but whether causal inference is performed can be specified at runtime
+class LTX2VideoCausalConv3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: Union[int, Tuple[int, int, int]] = 3,
+ stride: Union[int, Tuple[int, int, int]] = 1,
+ dilation: Union[int, Tuple[int, int, int]] = 1,
+ groups: int = 1,
+ spatial_padding_mode: str = "zeros",
+ ):
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.kernel_size = (
+ kernel_size
+ if isinstance(kernel_size, tuple)
+ else (kernel_size, kernel_size, kernel_size)
+ )
+
+ dilation = dilation if isinstance(dilation, tuple) else (dilation, 1, 1)
+ stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
+ height_pad = self.kernel_size[1] // 2
+ width_pad = self.kernel_size[2] // 2
+ padding = (0, height_pad, width_pad)
+
+ self.conv = nn.Conv3d(
+ in_channels,
+ out_channels,
+ self.kernel_size,
+ stride=stride,
+ dilation=dilation,
+ groups=groups,
+ padding=padding,
+ padding_mode=spatial_padding_mode,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor:
+ time_kernel_size = self.kernel_size[0]
+
+ if causal:
+ pad_left = hidden_states[:, :, :1, :, :].repeat(
+ (1, 1, time_kernel_size - 1, 1, 1)
+ )
+ hidden_states = torch.concatenate([pad_left, hidden_states], dim=2)
+ else:
+ pad_left = hidden_states[:, :, :1, :, :].repeat(
+ (1, 1, (time_kernel_size - 1) // 2, 1, 1)
+ )
+ pad_right = hidden_states[:, :, -1:, :, :].repeat(
+ (1, 1, (time_kernel_size - 1) // 2, 1, 1)
+ )
+ hidden_states = torch.concatenate(
+ [pad_left, hidden_states, pad_right], dim=2
+ )
+
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+# Like LTXVideoResnetBlock3d, but uses new causal Conv3d, normal Conv3d for the conv_shortcut, and the spatial padding
+# mode is configurable
+class LTX2VideoResnetBlock3d(nn.Module):
+ r"""
+ A 3D ResNet block used in the LTX 2.0 audiovisual model.
+
+ Args:
+ in_channels (`int`):
+ Number of input channels.
+ out_channels (`int`, *optional*):
+ Number of output channels. If None, defaults to `in_channels`.
+ dropout (`float`, defaults to `0.0`):
+ Dropout rate.
+ eps (`float`, defaults to `1e-6`):
+ Epsilon value for normalization layers.
+ elementwise_affine (`bool`, defaults to `False`):
+ Whether to enable elementwise affinity in the normalization layers.
+ non_linearity (`str`, defaults to `"swish"`):
+ Activation function to use.
+ conv_shortcut (bool, defaults to `False`):
+ Whether or not to use a convolution shortcut.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: Optional[int] = None,
+ dropout: float = 0.0,
+ eps: float = 1e-6,
+ elementwise_affine: bool = False,
+ non_linearity: str = "swish",
+ inject_noise: bool = False,
+ timestep_conditioning: bool = False,
+ spatial_padding_mode: str = "zeros",
+ ) -> None:
+ super().__init__()
+
+ out_channels = out_channels or in_channels
+
+ self.nonlinearity = get_activation(non_linearity)
+
+ self.norm1 = PerChannelRMSNorm()
+ self.conv1 = LTX2VideoCausalConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ self.norm2 = PerChannelRMSNorm()
+ self.dropout = nn.Dropout(dropout)
+ self.conv2 = LTX2VideoCausalConv3d(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ self.norm3 = None
+ self.conv_shortcut = None
+ if in_channels != out_channels:
+ self.norm3 = nn.LayerNorm(
+ in_channels, eps=eps, elementwise_affine=True, bias=True
+ )
+ # LTX 2.0 uses a normal nn.Conv3d here rather than LTXVideoCausalConv3d
+ self.conv_shortcut = nn.Conv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=1,
+ stride=1,
+ )
+
+ self.per_channel_scale1 = None
+ self.per_channel_scale2 = None
+ if inject_noise:
+ self.per_channel_scale1 = nn.Parameter(torch.zeros(in_channels, 1, 1))
+ self.per_channel_scale2 = nn.Parameter(torch.zeros(in_channels, 1, 1))
+
+ self.scale_shift_table = None
+ if timestep_conditioning:
+ self.scale_shift_table = nn.Parameter(
+ torch.randn(4, in_channels) / in_channels**0.5
+ )
+
+ def forward(
+ self,
+ inputs: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ hidden_states = inputs
+
+ hidden_states = self.norm1(hidden_states)
+
+ if self.scale_shift_table is not None:
+ temb = (
+ temb.unflatten(1, (4, -1))
+ + self.scale_shift_table[None, ..., None, None, None]
+ )
+ shift_1, scale_1, shift_2, scale_2 = temb.unbind(dim=1)
+ hidden_states = hidden_states * (1 + scale_1) + shift_1
+
+ hidden_states = self.nonlinearity(hidden_states)
+ hidden_states = self.conv1(hidden_states, causal=causal)
+
+ if self.per_channel_scale1 is not None:
+ spatial_shape = hidden_states.shape[-2:]
+ spatial_noise = torch.randn(
+ spatial_shape,
+ generator=generator,
+ device=hidden_states.device,
+ dtype=hidden_states.dtype,
+ )[None]
+ hidden_states = (
+ hidden_states
+ + (spatial_noise * self.per_channel_scale1)[None, :, None, ...]
+ )
+
+ hidden_states = self.norm2(hidden_states)
+
+ if self.scale_shift_table is not None:
+ hidden_states = hidden_states * (1 + scale_2) + shift_2
+
+ hidden_states = self.nonlinearity(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.conv2(hidden_states, causal=causal)
+
+ if self.per_channel_scale2 is not None:
+ spatial_shape = hidden_states.shape[-2:]
+ spatial_noise = torch.randn(
+ spatial_shape,
+ generator=generator,
+ device=hidden_states.device,
+ dtype=hidden_states.dtype,
+ )[None]
+ hidden_states = (
+ hidden_states
+ + (spatial_noise * self.per_channel_scale2)[None, :, None, ...]
+ )
+
+ if self.norm3 is not None:
+ inputs = self.norm3(inputs.movedim(1, -1)).movedim(-1, 1)
+
+ if self.conv_shortcut is not None:
+ inputs = self.conv_shortcut(inputs)
+
+ hidden_states = hidden_states + inputs
+ return hidden_states
+
+
+# Like LTX 1.0 LTXVideoDownsampler3d, but uses new causal Conv3d
+class LTXVideoDownsampler3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ stride: Union[int, Tuple[int, int, int]] = 1,
+ spatial_padding_mode: str = "zeros",
+ ) -> None:
+ super().__init__()
+
+ self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
+ self.group_size = (
+ in_channels * stride[0] * stride[1] * stride[2]
+ ) // out_channels
+
+ out_channels = out_channels // (
+ self.stride[0] * self.stride[1] * self.stride[2]
+ )
+
+ self.conv = LTX2VideoCausalConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=1,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor:
+ hidden_states = torch.cat(
+ [hidden_states[:, :, : self.stride[0] - 1], hidden_states], dim=2
+ )
+
+ residual = (
+ hidden_states.unflatten(4, (-1, self.stride[2]))
+ .unflatten(3, (-1, self.stride[1]))
+ .unflatten(2, (-1, self.stride[0]))
+ )
+ residual = residual.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4)
+ residual = residual.unflatten(1, (-1, self.group_size))
+ residual = residual.mean(dim=2)
+
+ hidden_states = self.conv(hidden_states, causal=causal)
+ hidden_states = (
+ hidden_states.unflatten(4, (-1, self.stride[2]))
+ .unflatten(3, (-1, self.stride[1]))
+ .unflatten(2, (-1, self.stride[0]))
+ )
+ hidden_states = hidden_states.permute(0, 1, 3, 5, 7, 2, 4, 6).flatten(1, 4)
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+# Like LTX 1.0 LTXVideoUpsampler3d, but uses new causal Conv3d
+class LTXVideoUpsampler3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ stride: Union[int, Tuple[int, int, int]] = 1,
+ residual: bool = False,
+ upscale_factor: int = 1,
+ spatial_padding_mode: str = "zeros",
+ ) -> None:
+ super().__init__()
+
+ self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
+ self.residual = residual
+ self.upscale_factor = upscale_factor
+
+ out_channels = (
+ in_channels * stride[0] * stride[1] * stride[2]
+ ) // upscale_factor
+
+ self.conv = LTX2VideoCausalConv3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=3,
+ stride=1,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, causal: bool = True) -> torch.Tensor:
+ batch_size, num_channels, num_frames, height, width = hidden_states.shape
+
+ if self.residual:
+ residual = hidden_states.reshape(
+ batch_size,
+ -1,
+ self.stride[0],
+ self.stride[1],
+ self.stride[2],
+ num_frames,
+ height,
+ width,
+ )
+ residual = (
+ residual.permute(0, 1, 5, 2, 6, 3, 7, 4)
+ .flatten(6, 7)
+ .flatten(4, 5)
+ .flatten(2, 3)
+ )
+ repeats = (
+ self.stride[0] * self.stride[1] * self.stride[2]
+ ) // self.upscale_factor
+ residual = residual.repeat(1, repeats, 1, 1, 1)
+ residual = residual[:, :, self.stride[0] - 1 :]
+
+ hidden_states = self.conv(hidden_states, causal=causal)
+ hidden_states = hidden_states.reshape(
+ batch_size,
+ -1,
+ self.stride[0],
+ self.stride[1],
+ self.stride[2],
+ num_frames,
+ height,
+ width,
+ )
+ hidden_states = (
+ hidden_states.permute(0, 1, 5, 2, 6, 3, 7, 4)
+ .flatten(6, 7)
+ .flatten(4, 5)
+ .flatten(2, 3)
+ )
+ hidden_states = hidden_states[:, :, self.stride[0] - 1 :]
+
+ if self.residual:
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+# Like LTX 1.0 LTXVideo095DownBlock3D, but with the updated LTX2VideoResnetBlock3d
+class LTX2VideoDownBlock3D(nn.Module):
+ r"""
+ Down block used in the LTXVideo model.
+
+ Args:
+ in_channels (`int`):
+ Number of input channels.
+ out_channels (`int`, *optional*):
+ Number of output channels. If None, defaults to `in_channels`.
+ num_layers (`int`, defaults to `1`):
+ Number of resnet layers.
+ dropout (`float`, defaults to `0.0`):
+ Dropout rate.
+ resnet_eps (`float`, defaults to `1e-6`):
+ Epsilon value for normalization layers.
+ resnet_act_fn (`str`, defaults to `"swish"`):
+ Activation function to use.
+ spatio_temporal_scale (`bool`, defaults to `True`):
+ Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension.
+ Whether or not to downsample across temporal dimension.
+ is_causal (`bool`, defaults to `True`):
+ Whether this layer behaves causally (future frames depend only on past frames) or not.
+ """
+
+ _supports_gradient_checkpointing = True
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: Optional[int] = None,
+ num_layers: int = 1,
+ dropout: float = 0.0,
+ resnet_eps: float = 1e-6,
+ resnet_act_fn: str = "swish",
+ spatio_temporal_scale: bool = True,
+ downsample_type: str = "conv",
+ spatial_padding_mode: str = "zeros",
+ ):
+ super().__init__()
+
+ out_channels = out_channels or in_channels
+
+ resnets = []
+ for _ in range(num_layers):
+ resnets.append(
+ LTX2VideoResnetBlock3d(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ dropout=dropout,
+ eps=resnet_eps,
+ non_linearity=resnet_act_fn,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ self.resnets = nn.ModuleList(resnets)
+
+ self.downsamplers = None
+ if spatio_temporal_scale:
+ self.downsamplers = nn.ModuleList()
+
+ if downsample_type == "conv":
+ self.downsamplers.append(
+ LTX2VideoCausalConv3d(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ kernel_size=3,
+ stride=(2, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ elif downsample_type == "spatial":
+ self.downsamplers.append(
+ LTXVideoDownsampler3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(1, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ elif downsample_type == "temporal":
+ self.downsamplers.append(
+ LTXVideoDownsampler3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(2, 1, 1),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ elif downsample_type == "spatiotemporal":
+ self.downsamplers.append(
+ LTXVideoDownsampler3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ stride=(2, 2, 2),
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ r"""Forward method of the `LTXDownBlock3D` class."""
+
+ for i, resnet in enumerate(self.resnets):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ hidden_states = self._gradient_checkpointing_func(
+ resnet, hidden_states, temb, generator, causal
+ )
+ else:
+ hidden_states = resnet(hidden_states, temb, generator, causal=causal)
+
+ if self.downsamplers is not None:
+ for downsampler in self.downsamplers:
+ hidden_states = downsampler(hidden_states, causal=causal)
+
+ return hidden_states
+
+
+# Adapted from diffusers.models.autoencoders.autoencoder_kl_cogvideox.CogVideoMidBlock3d
+# Like LTX 1.0 LTXVideoMidBlock3d, but with the updated LTX2VideoResnetBlock3d
+class LTX2VideoMidBlock3d(nn.Module):
+ r"""
+ A middle block used in the LTXVideo model.
+
+ Args:
+ in_channels (`int`):
+ Number of input channels.
+ num_layers (`int`, defaults to `1`):
+ Number of resnet layers.
+ dropout (`float`, defaults to `0.0`):
+ Dropout rate.
+ resnet_eps (`float`, defaults to `1e-6`):
+ Epsilon value for normalization layers.
+ resnet_act_fn (`str`, defaults to `"swish"`):
+ Activation function to use.
+ is_causal (`bool`, defaults to `True`):
+ Whether this layer behaves causally (future frames depend only on past frames) or not.
+ """
+
+ _supports_gradient_checkpointing = True
+
+ def __init__(
+ self,
+ in_channels: int,
+ num_layers: int = 1,
+ dropout: float = 0.0,
+ resnet_eps: float = 1e-6,
+ resnet_act_fn: str = "swish",
+ inject_noise: bool = False,
+ timestep_conditioning: bool = False,
+ spatial_padding_mode: str = "zeros",
+ ) -> None:
+ super().__init__()
+
+ self.time_embedder = None
+ if timestep_conditioning:
+ self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
+ in_channels * 4, 0
+ )
+
+ resnets = []
+ for _ in range(num_layers):
+ resnets.append(
+ LTX2VideoResnetBlock3d(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ dropout=dropout,
+ eps=resnet_eps,
+ non_linearity=resnet_act_fn,
+ inject_noise=inject_noise,
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ self.resnets = nn.ModuleList(resnets)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ r"""Forward method of the `LTXMidBlock3D` class."""
+
+ if self.time_embedder is not None:
+ temb = self.time_embedder(
+ timestep=temb.flatten(),
+ resolution=None,
+ aspect_ratio=None,
+ batch_size=hidden_states.size(0),
+ hidden_dtype=hidden_states.dtype,
+ )
+ temb = temb.view(hidden_states.size(0), -1, 1, 1, 1)
+
+ for i, resnet in enumerate(self.resnets):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ hidden_states = self._gradient_checkpointing_func(
+ resnet, hidden_states, temb, generator, causal
+ )
+ else:
+ hidden_states = resnet(hidden_states, temb, generator, causal=causal)
+
+ return hidden_states
+
+
+# Like LTXVideoUpBlock3d but with no conv_in and the updated LTX2VideoResnetBlock3d
+class LTX2VideoUpBlock3d(nn.Module):
+ r"""
+ Up block used in the LTXVideo model.
+
+ Args:
+ in_channels (`int`):
+ Number of input channels.
+ out_channels (`int`, *optional*):
+ Number of output channels. If None, defaults to `in_channels`.
+ num_layers (`int`, defaults to `1`):
+ Number of resnet layers.
+ dropout (`float`, defaults to `0.0`):
+ Dropout rate.
+ resnet_eps (`float`, defaults to `1e-6`):
+ Epsilon value for normalization layers.
+ resnet_act_fn (`str`, defaults to `"swish"`):
+ Activation function to use.
+ spatio_temporal_scale (`bool`, defaults to `True`):
+ Whether or not to use a downsampling layer. If not used, output dimension would be same as input dimension.
+ Whether or not to downsample across temporal dimension.
+ is_causal (`bool`, defaults to `True`):
+ Whether this layer behaves causally (future frames depend only on past frames) or not.
+ """
+
+ _supports_gradient_checkpointing = True
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: Optional[int] = None,
+ num_layers: int = 1,
+ dropout: float = 0.0,
+ resnet_eps: float = 1e-6,
+ resnet_act_fn: str = "swish",
+ spatio_temporal_scale: bool = True,
+ inject_noise: bool = False,
+ timestep_conditioning: bool = False,
+ upsample_residual: bool = False,
+ upscale_factor: int = 1,
+ spatial_padding_mode: str = "zeros",
+ ):
+ super().__init__()
+
+ out_channels = out_channels or in_channels
+
+ self.time_embedder = None
+ if timestep_conditioning:
+ self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
+ in_channels * 4, 0
+ )
+
+ self.conv_in = None
+ if in_channels != out_channels:
+ self.conv_in = LTX2VideoResnetBlock3d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ dropout=dropout,
+ eps=resnet_eps,
+ non_linearity=resnet_act_fn,
+ inject_noise=inject_noise,
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ self.upsamplers = None
+ if spatio_temporal_scale:
+ self.upsamplers = nn.ModuleList(
+ [
+ LTXVideoUpsampler3d(
+ out_channels * upscale_factor,
+ stride=(2, 2, 2),
+ residual=upsample_residual,
+ upscale_factor=upscale_factor,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ ]
+ )
+
+ resnets = []
+ for _ in range(num_layers):
+ resnets.append(
+ LTX2VideoResnetBlock3d(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ dropout=dropout,
+ eps=resnet_eps,
+ non_linearity=resnet_act_fn,
+ inject_noise=inject_noise,
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ )
+ self.resnets = nn.ModuleList(resnets)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ generator: Optional[torch.Generator] = None,
+ causal: bool = True,
+ ) -> torch.Tensor:
+ if self.conv_in is not None:
+ hidden_states = self.conv_in(hidden_states, temb, generator, causal=causal)
+
+ if self.time_embedder is not None:
+ temb = self.time_embedder(
+ timestep=temb.flatten(),
+ resolution=None,
+ aspect_ratio=None,
+ batch_size=hidden_states.size(0),
+ hidden_dtype=hidden_states.dtype,
+ )
+ temb = temb.view(hidden_states.size(0), -1, 1, 1, 1)
+
+ if self.upsamplers is not None:
+ for upsampler in self.upsamplers:
+ hidden_states = upsampler(hidden_states, causal=causal)
+
+ for i, resnet in enumerate(self.resnets):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ hidden_states = self._gradient_checkpointing_func(
+ resnet, hidden_states, temb, generator, causal
+ )
+ else:
+ hidden_states = resnet(hidden_states, temb, generator, causal=causal)
+
+ return hidden_states
+
+
+# Like LTX 1.0 LTXVideoEncoder3d but with different default args - the spatiotemporal downsampling pattern is
+# different, as is the layers_per_block (the 2.0 VAE is bigger)
+class LTX2VideoEncoder3d(nn.Module):
+ r"""
+ The `LTXVideoEncoder3d` layer of a variational autoencoder that encodes input video samples to its latent
+ representation.
+
+ Args:
+ in_channels (`int`, defaults to 3):
+ Number of input channels.
+ out_channels (`int`, defaults to 128):
+ Number of latent channels.
+ block_out_channels (`Tuple[int, ...]`, defaults to `(256, 512, 1024, 2048)`):
+ The number of output channels for each block.
+ spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, True)`:
+ Whether a block should contain spatio-temporal downscaling layers or not.
+ layers_per_block (`Tuple[int, ...]`, defaults to `(4, 6, 6, 2, 2)`):
+ The number of layers per block.
+ downsample_type (`Tuple[str, ...]`, defaults to `("spatial", "temporal", "spatiotemporal", "spatiotemporal")`):
+ The spatiotemporal downsampling pattern per block. Per-layer values can be
+ - `"spatial"` (downsample spatial dims by 2x)
+ - `"temporal"` (downsample temporal dim by 2x)
+ - `"spatiotemporal"` (downsample both spatial and temporal dims by 2x)
+ patch_size (`int`, defaults to `4`):
+ The size of spatial patches.
+ patch_size_t (`int`, defaults to `1`):
+ The size of temporal patches.
+ resnet_norm_eps (`float`, defaults to `1e-6`):
+ Epsilon value for ResNet normalization layers.
+ is_causal (`bool`, defaults to `True`):
+ Whether this layer behaves causally (future frames depend only on past frames) or not.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 128,
+ block_out_channels: Tuple[int, ...] = (256, 512, 1024, 2048),
+ down_block_types: Tuple[str, ...] = (
+ "LTX2VideoDownBlock3D",
+ "LTX2VideoDownBlock3D",
+ "LTX2VideoDownBlock3D",
+ "LTX2VideoDownBlock3D",
+ ),
+ spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True, True),
+ layers_per_block: Tuple[int, ...] = (4, 6, 6, 2, 2),
+ downsample_type: Tuple[str, ...] = (
+ "spatial",
+ "temporal",
+ "spatiotemporal",
+ "spatiotemporal",
+ ),
+ patch_size: int = 4,
+ patch_size_t: int = 1,
+ resnet_norm_eps: float = 1e-6,
+ is_causal: bool = True,
+ spatial_padding_mode: str = "zeros",
+ ):
+ super().__init__()
+
+ self.patch_size = patch_size
+ self.patch_size_t = patch_size_t
+ self.in_channels = in_channels * patch_size**2
+ self.is_causal = is_causal
+
+ output_channel = out_channels
+
+ self.conv_in = LTX2VideoCausalConv3d(
+ in_channels=self.in_channels,
+ out_channels=output_channel,
+ kernel_size=3,
+ stride=1,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ # down blocks
+ num_block_out_channels = len(block_out_channels)
+ self.down_blocks = nn.ModuleList([])
+ for i in range(num_block_out_channels):
+ input_channel = output_channel
+ output_channel = block_out_channels[i]
+
+ if down_block_types[i] == "LTX2VideoDownBlock3D":
+ down_block = LTX2VideoDownBlock3D(
+ in_channels=input_channel,
+ out_channels=output_channel,
+ num_layers=layers_per_block[i],
+ resnet_eps=resnet_norm_eps,
+ spatio_temporal_scale=spatio_temporal_scaling[i],
+ downsample_type=downsample_type[i],
+ spatial_padding_mode=spatial_padding_mode,
+ )
+ else:
+ raise ValueError(f"Unknown down block type: {down_block_types[i]}")
+
+ self.down_blocks.append(down_block)
+
+ # mid block
+ self.mid_block = LTX2VideoMidBlock3d(
+ in_channels=output_channel,
+ num_layers=layers_per_block[-1],
+ resnet_eps=resnet_norm_eps,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ # out
+ self.norm_out = PerChannelRMSNorm()
+ self.conv_act = nn.SiLU()
+ self.conv_out = LTX2VideoCausalConv3d(
+ in_channels=output_channel,
+ out_channels=out_channels + 1,
+ kernel_size=3,
+ stride=1,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self, hidden_states: torch.Tensor, causal: Optional[bool] = None
+ ) -> torch.Tensor:
+ r"""The forward method of the `LTXVideoEncoder3d` class."""
+
+ p = self.patch_size
+ p_t = self.patch_size_t
+
+ batch_size, num_channels, num_frames, height, width = hidden_states.shape
+ post_patch_num_frames = num_frames // p_t
+ post_patch_height = height // p
+ post_patch_width = width // p
+ causal = causal or self.is_causal
+
+ hidden_states = hidden_states.reshape(
+ batch_size,
+ num_channels,
+ post_patch_num_frames,
+ p_t,
+ post_patch_height,
+ p,
+ post_patch_width,
+ p,
+ )
+ # Thanks for driving me insane with the weird patching order :(
+ hidden_states = hidden_states.permute(0, 1, 3, 7, 5, 2, 4, 6).flatten(1, 4)
+ hidden_states = self.conv_in(hidden_states, causal=causal)
+
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ for down_block in self.down_blocks:
+ hidden_states = self._gradient_checkpointing_func(
+ down_block, hidden_states, None, None, causal
+ )
+
+ hidden_states = self._gradient_checkpointing_func(
+ self.mid_block, hidden_states, None, None, causal
+ )
+ else:
+ for down_block in self.down_blocks:
+ hidden_states = down_block(hidden_states, causal=causal)
+
+ hidden_states = self.mid_block(hidden_states, causal=causal)
+
+ hidden_states = self.norm_out(hidden_states)
+ hidden_states = self.conv_act(hidden_states)
+ hidden_states = self.conv_out(hidden_states, causal=causal)
+
+ last_channel = hidden_states[:, -1:]
+ last_channel = last_channel.repeat(1, hidden_states.size(1) - 2, 1, 1, 1)
+ hidden_states = torch.cat([hidden_states, last_channel], dim=1)
+
+ return hidden_states
+
+
+# Like LTX 1.0 LTXVideoDecoder3d, but has only 3 symmetric up blocks which are causal and residual with upsample_factor 2
+class LTX2VideoDecoder3d(nn.Module):
+ r"""
+ The `LTXVideoDecoder3d` layer of a variational autoencoder that decodes its latent representation into an output
+ sample.
+
+ Args:
+ in_channels (`int`, defaults to 128):
+ Number of latent channels.
+ out_channels (`int`, defaults to 3):
+ Number of output channels.
+ block_out_channels (`Tuple[int, ...]`, defaults to `(128, 256, 512, 512)`):
+ The number of output channels for each block.
+ spatio_temporal_scaling (`Tuple[bool, ...], defaults to `(True, True, True, False)`:
+ Whether a block should contain spatio-temporal upscaling layers or not.
+ layers_per_block (`Tuple[int, ...]`, defaults to `(4, 3, 3, 3, 4)`):
+ The number of layers per block.
+ patch_size (`int`, defaults to `4`):
+ The size of spatial patches.
+ patch_size_t (`int`, defaults to `1`):
+ The size of temporal patches.
+ resnet_norm_eps (`float`, defaults to `1e-6`):
+ Epsilon value for ResNet normalization layers.
+ is_causal (`bool`, defaults to `False`):
+ Whether this layer behaves causally (future frames depend only on past frames) or not.
+ timestep_conditioning (`bool`, defaults to `False`):
+ Whether to condition the model on timesteps.
+ """
+
+ def __init__(
+ self,
+ in_channels: int = 128,
+ out_channels: int = 3,
+ block_out_channels: Tuple[int, ...] = (256, 512, 1024),
+ spatio_temporal_scaling: Tuple[bool, ...] = (True, True, True),
+ layers_per_block: Tuple[int, ...] = (5, 5, 5, 5),
+ patch_size: int = 4,
+ patch_size_t: int = 1,
+ resnet_norm_eps: float = 1e-6,
+ is_causal: bool = False,
+ inject_noise: Tuple[bool, ...] = (False, False, False),
+ timestep_conditioning: bool = False,
+ upsample_residual: Tuple[bool, ...] = (True, True, True),
+ upsample_factor: Tuple[bool, ...] = (2, 2, 2),
+ spatial_padding_mode: str = "reflect",
+ ) -> None:
+ super().__init__()
+
+ self.patch_size = patch_size
+ self.patch_size_t = patch_size_t
+ self.out_channels = out_channels * patch_size**2
+ self.is_causal = is_causal
+
+ block_out_channels = tuple(reversed(block_out_channels))
+ spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling))
+ layers_per_block = tuple(reversed(layers_per_block))
+ inject_noise = tuple(reversed(inject_noise))
+ upsample_residual = tuple(reversed(upsample_residual))
+ upsample_factor = tuple(reversed(upsample_factor))
+ output_channel = block_out_channels[0]
+
+ self.conv_in = LTX2VideoCausalConv3d(
+ in_channels=in_channels,
+ out_channels=output_channel,
+ kernel_size=3,
+ stride=1,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ self.mid_block = LTX2VideoMidBlock3d(
+ in_channels=output_channel,
+ num_layers=layers_per_block[0],
+ resnet_eps=resnet_norm_eps,
+ inject_noise=inject_noise[0],
+ timestep_conditioning=timestep_conditioning,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ # up blocks
+ num_block_out_channels = len(block_out_channels)
+ self.up_blocks = nn.ModuleList([])
+ for i in range(num_block_out_channels):
+ input_channel = output_channel // upsample_factor[i]
+ output_channel = block_out_channels[i] // upsample_factor[i]
+
+ up_block = LTX2VideoUpBlock3d(
+ in_channels=input_channel,
+ out_channels=output_channel,
+ num_layers=layers_per_block[i + 1],
+ resnet_eps=resnet_norm_eps,
+ spatio_temporal_scale=spatio_temporal_scaling[i],
+ inject_noise=inject_noise[i + 1],
+ timestep_conditioning=timestep_conditioning,
+ upsample_residual=upsample_residual[i],
+ upscale_factor=upsample_factor[i],
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ self.up_blocks.append(up_block)
+
+ # out
+ self.norm_out = PerChannelRMSNorm()
+ self.conv_act = nn.SiLU()
+ self.conv_out = LTX2VideoCausalConv3d(
+ in_channels=output_channel,
+ out_channels=self.out_channels,
+ kernel_size=3,
+ stride=1,
+ spatial_padding_mode=spatial_padding_mode,
+ )
+
+ # timestep embedding
+ self.time_embedder = None
+ self.scale_shift_table = None
+ self.timestep_scale_multiplier = None
+ if timestep_conditioning:
+ self.timestep_scale_multiplier = nn.Parameter(
+ torch.tensor(1000.0, dtype=torch.float32)
+ )
+ self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
+ output_channel * 2, 0
+ )
+ self.scale_shift_table = nn.Parameter(
+ torch.randn(2, output_channel) / output_channel**0.5
+ )
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ causal: Optional[bool] = None,
+ ) -> torch.Tensor:
+ causal = causal or self.is_causal
+
+ hidden_states = self.conv_in(hidden_states, causal=causal)
+
+ if self.timestep_scale_multiplier is not None:
+ temb = temb * self.timestep_scale_multiplier
+
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ hidden_states = self._gradient_checkpointing_func(
+ self.mid_block, hidden_states, temb, None, causal
+ )
+
+ for up_block in self.up_blocks:
+ hidden_states = self._gradient_checkpointing_func(
+ up_block, hidden_states, temb, None, causal
+ )
+ else:
+ hidden_states = self.mid_block(hidden_states, temb, causal=causal)
+
+ for up_block in self.up_blocks:
+ hidden_states = up_block(hidden_states, temb, causal=causal)
+
+ hidden_states = self.norm_out(hidden_states)
+
+ if self.time_embedder is not None:
+ temb = self.time_embedder(
+ timestep=temb.flatten(),
+ resolution=None,
+ aspect_ratio=None,
+ batch_size=hidden_states.size(0),
+ hidden_dtype=hidden_states.dtype,
+ )
+ temb = temb.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1))
+ temb = temb + self.scale_shift_table[None, ..., None, None, None]
+ shift, scale = temb.unbind(dim=1)
+ hidden_states = hidden_states * (1 + scale) + shift
+
+ hidden_states = self.conv_act(hidden_states)
+ hidden_states = self.conv_out(hidden_states, causal=causal)
+
+ p = self.patch_size
+ p_t = self.patch_size_t
+
+ batch_size, num_channels, num_frames, height, width = hidden_states.shape
+ hidden_states = hidden_states.reshape(
+ batch_size, -1, p_t, p, p, num_frames, height, width
+ )
+ hidden_states = (
+ hidden_states.permute(0, 1, 5, 2, 6, 4, 7, 3)
+ .flatten(6, 7)
+ .flatten(4, 5)
+ .flatten(2, 3)
+ )
+
+ return hidden_states
+
+
+class AutoencoderKLLTX2Video(ParallelTiledVAE):
+ r"""
+ A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos.
+
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
+ for all models (such as downloading or saving).
+ """
+
+ _supports_gradient_checkpointing = False
+
+ def __init__(self, config: LTXVideoVAEConfig):
+ super().__init__(config=config)
+ in_channels = config.arch_config.in_channels
+ latent_channels = config.arch_config.latent_channels
+ out_channels = config.arch_config.out_channels
+ block_out_channels = config.arch_config.block_out_channels
+ down_block_types = config.arch_config.down_block_types
+ spatio_temporal_scaling = config.arch_config.spatio_temporal_scaling
+ layers_per_block = config.arch_config.layers_per_block
+ downsample_type = config.arch_config.downsample_type
+ patch_size = config.arch_config.patch_size
+ patch_size_t = config.arch_config.patch_size_t
+ resnet_norm_eps = config.arch_config.resnet_norm_eps
+ encoder_causal = config.arch_config.encoder_causal
+ encoder_spatial_padding_mode = config.arch_config.encoder_spatial_padding_mode
+
+ decoder_block_out_channels = config.arch_config.decoder_block_out_channels
+ decoder_spatio_temporal_scaling = (
+ config.arch_config.decoder_spatio_temporal_scaling
+ )
+ decoder_layers_per_block = config.arch_config.decoder_layers_per_block
+ decoder_causal = config.arch_config.decoder_causal
+ decoder_spatial_padding_mode = config.arch_config.decoder_spatial_padding_mode
+
+ self.encoder = LTX2VideoEncoder3d(
+ in_channels,
+ latent_channels,
+ block_out_channels,
+ down_block_types,
+ spatio_temporal_scaling,
+ layers_per_block,
+ downsample_type,
+ patch_size,
+ patch_size_t,
+ resnet_norm_eps,
+ encoder_causal,
+ encoder_spatial_padding_mode,
+ )
+
+ self.decoder = LTX2VideoDecoder3d(
+ latent_channels,
+ out_channels,
+ decoder_block_out_channels,
+ decoder_spatio_temporal_scaling,
+ decoder_layers_per_block,
+ patch_size,
+ patch_size_t,
+ resnet_norm_eps,
+ decoder_causal,
+ decoder_spatial_padding_mode,
+ )
+
+ latents_mean = torch.zeros((latent_channels,), requires_grad=False)
+ latents_std = torch.ones((latent_channels,), requires_grad=False)
+ self.register_buffer("latents_mean", latents_mean, persistent=True)
+ self.register_buffer("latents_std", latents_std, persistent=True)
+
+ # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
+ # to perform decoding of a single video latent at a time.
+ self.use_slicing = False
+
+ # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
+ # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
+ # intermediate tiles together, the memory requirement can be lowered.
+ self.use_tiling = False
+
+ # When decoding temporally long video latents, the memory requirement is very high. By decoding latent frames
+ # at a fixed frame batch size (based on `self.num_latent_frames_batch_sizes`), the memory requirement can be lowered.
+ self.use_framewise_encoding = False
+ self.use_framewise_decoding = False
+
+ # This can be configured based on the amount of GPU memory available.
+ # `16` for sample frames and `2` for latent frames are sensible defaults for consumer GPUs.
+ # Setting it to higher values results in higher memory usage.
+ self.num_sample_frames_batch_size = 16
+ self.num_latent_frames_batch_size = 2
+
+ # The minimal tile height and width for spatial tiling to be used
+ self.tile_sample_min_height = 512
+ self.tile_sample_min_width = 512
+ self.tile_sample_min_num_frames = 16
+
+ # The minimal distance between two spatial tiles
+ self.tile_sample_stride_height = 448
+ self.tile_sample_stride_width = 448
+ self.tile_sample_stride_num_frames = 8
+
+ def enable_tiling(
+ self,
+ tile_sample_min_height: Optional[int] = None,
+ tile_sample_min_width: Optional[int] = None,
+ tile_sample_min_num_frames: Optional[int] = None,
+ tile_sample_stride_height: Optional[float] = None,
+ tile_sample_stride_width: Optional[float] = None,
+ tile_sample_stride_num_frames: Optional[float] = None,
+ ) -> None:
+ r"""
+ Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
+ compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
+ processing larger images.
+
+ Args:
+ tile_sample_min_height (`int`, *optional*):
+ The minimum height required for a sample to be separated into tiles across the height dimension.
+ tile_sample_min_width (`int`, *optional*):
+ The minimum width required for a sample to be separated into tiles across the width dimension.
+ tile_sample_stride_height (`int`, *optional*):
+ The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
+ no tiling artifacts produced across the height dimension.
+ tile_sample_stride_width (`int`, *optional*):
+ The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
+ artifacts produced across the width dimension.
+ """
+ self.use_tiling = True
+ self.tile_sample_min_height = (
+ tile_sample_min_height or self.tile_sample_min_height
+ )
+ self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
+ self.tile_sample_min_num_frames = (
+ tile_sample_min_num_frames or self.tile_sample_min_num_frames
+ )
+ self.tile_sample_stride_height = (
+ tile_sample_stride_height or self.tile_sample_stride_height
+ )
+ self.tile_sample_stride_width = (
+ tile_sample_stride_width or self.tile_sample_stride_width
+ )
+ self.tile_sample_stride_num_frames = (
+ tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
+ )
+
+ def _encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> torch.Tensor:
+ batch_size, num_channels, num_frames, height, width = x.shape
+
+ if self.use_framewise_decoding and num_frames > self.tile_sample_min_num_frames:
+ return self._temporal_tiled_encode(x, causal=causal)
+
+ if self.use_tiling and (
+ width > self.tile_sample_min_width or height > self.tile_sample_min_height
+ ):
+ return self.tiled_encode(x, causal=causal)
+
+ enc = self.encoder(x, causal=causal)
+
+ return enc
+
+ def encode(
+ self, x: torch.Tensor, causal: Optional[bool] = None, return_dict: bool = True
+ ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ """
+ Encode a batch of images into latents.
+
+ Args:
+ x (`torch.Tensor`): Input batch of images.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+ Returns:
+ The latent representations of the encoded videos. If `return_dict` is True, a
+ [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+ """
+ if self.use_slicing and x.shape[0] > 1:
+ encoded_slices = [
+ self._encode(x_slice, causal=causal) for x_slice in x.split(1)
+ ]
+ h = torch.cat(encoded_slices)
+ else:
+ h = self._encode(x, causal=causal)
+ posterior = DiagonalGaussianDistribution(h)
+
+ if not return_dict:
+ return (posterior,)
+ return AutoencoderKLOutput(latent_dist=posterior)
+
+ def _decode(
+ self,
+ z: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ causal: Optional[bool] = None,
+ return_dict: bool = True,
+ ) -> Union[DecoderOutput, torch.Tensor]:
+ batch_size, num_channels, num_frames, height, width = z.shape
+ tile_latent_min_height = (
+ self.tile_sample_min_height // self.spatial_compression_ratio
+ )
+ tile_latent_min_width = (
+ self.tile_sample_min_width // self.spatial_compression_ratio
+ )
+ tile_latent_min_num_frames = (
+ self.tile_sample_min_num_frames // self.temporal_compression_ratio
+ )
+
+ if self.use_framewise_decoding and num_frames > tile_latent_min_num_frames:
+ return self._temporal_tiled_decode(
+ z, temb, causal=causal, return_dict=return_dict
+ )
+
+ if self.use_tiling and (
+ width > tile_latent_min_width or height > tile_latent_min_height
+ ):
+ return self.tiled_decode(z, temb, causal=causal, return_dict=return_dict)
+
+ dec = self.decoder(z, temb, causal=causal)
+
+ if not return_dict:
+ return (dec,)
+
+ return DecoderOutput(sample=dec)
+
+ def decode(
+ self,
+ z: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ causal: Optional[bool] = None,
+ return_dict: bool = True,
+ ) -> Union[DecoderOutput, torch.Tensor]:
+ """
+ Decode a batch of images.
+
+ Args:
+ z (`torch.Tensor`): Input batch of latent vectors.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.vae.DecoderOutput`] or `tuple`:
+ If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+ returned.
+ """
+ if self.use_slicing and z.shape[0] > 1:
+ if temb is not None:
+ decoded_slices = [
+ self._decode(z_slice, t_slice, causal=causal).sample
+ for z_slice, t_slice in (z.split(1), temb.split(1))
+ ]
+ else:
+ decoded_slices = [
+ self._decode(z_slice, causal=causal).sample
+ for z_slice in z.split(1)
+ ]
+ decoded = torch.cat(decoded_slices)
+ else:
+ decoded = self._decode(z, temb, causal=causal).sample
+
+ if not return_dict:
+ return (decoded,)
+
+ return DecoderOutput(sample=decoded)
+
+ def blend_v(
+ self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+ ) -> torch.Tensor:
+ blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+ for y in range(blend_extent):
+ b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (
+ 1 - y / blend_extent
+ ) + b[:, :, :, y, :] * (y / blend_extent)
+ return b
+
+ def blend_h(
+ self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+ ) -> torch.Tensor:
+ blend_extent = min(a.shape[4], b.shape[4], blend_extent)
+ for x in range(blend_extent):
+ b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (
+ 1 - x / blend_extent
+ ) + b[:, :, :, :, x] * (x / blend_extent)
+ return b
+
+ def blend_t(
+ self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+ ) -> torch.Tensor:
+ blend_extent = min(a.shape[-3], b.shape[-3], blend_extent)
+ for x in range(blend_extent):
+ b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (
+ 1 - x / blend_extent
+ ) + b[:, :, x, :, :] * (x / blend_extent)
+ return b
+
+ def tiled_encode(
+ self, x: torch.Tensor, causal: Optional[bool] = None
+ ) -> torch.Tensor:
+ r"""Encode a batch of images using a tiled encoder.
+
+ Args:
+ x (`torch.Tensor`): Input batch of videos.
+
+ Returns:
+ `torch.Tensor`:
+ The latent representation of the encoded videos.
+ """
+ batch_size, num_channels, num_frames, height, width = x.shape
+ latent_height = height // self.spatial_compression_ratio
+ latent_width = width // self.spatial_compression_ratio
+
+ tile_latent_min_height = (
+ self.tile_sample_min_height // self.spatial_compression_ratio
+ )
+ tile_latent_min_width = (
+ self.tile_sample_min_width // self.spatial_compression_ratio
+ )
+ tile_latent_stride_height = (
+ self.tile_sample_stride_height // self.spatial_compression_ratio
+ )
+ tile_latent_stride_width = (
+ self.tile_sample_stride_width // self.spatial_compression_ratio
+ )
+
+ blend_height = tile_latent_min_height - tile_latent_stride_height
+ blend_width = tile_latent_min_width - tile_latent_stride_width
+
+ # Split x into overlapping tiles and encode them separately.
+ # The tiles have an overlap to avoid seams between tiles.
+ rows = []
+ for i in range(0, height, self.tile_sample_stride_height):
+ row = []
+ for j in range(0, width, self.tile_sample_stride_width):
+ time = self.encoder(
+ x[
+ :,
+ :,
+ :,
+ i : i + self.tile_sample_min_height,
+ j : j + self.tile_sample_min_width,
+ ],
+ causal=causal,
+ )
+
+ row.append(time)
+ rows.append(row)
+
+ result_rows = []
+ for i, row in enumerate(rows):
+ result_row = []
+ for j, tile in enumerate(row):
+ # blend the above tile and the left tile
+ # to the current tile and add the current tile to the result row
+ if i > 0:
+ tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+ if j > 0:
+ tile = self.blend_h(row[j - 1], tile, blend_width)
+ result_row.append(
+ tile[:, :, :, :tile_latent_stride_height, :tile_latent_stride_width]
+ )
+ result_rows.append(torch.cat(result_row, dim=4))
+
+ enc = torch.cat(result_rows, dim=3)[:, :, :, :latent_height, :latent_width]
+ return enc
+
+ def tiled_decode(
+ self,
+ z: torch.Tensor,
+ temb: Optional[torch.Tensor],
+ causal: Optional[bool] = None,
+ return_dict: bool = True,
+ ) -> Union[DecoderOutput, torch.Tensor]:
+ r"""
+ Decode a batch of images using a tiled decoder.
+
+ Args:
+ z (`torch.Tensor`): Input batch of latent vectors.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+ Returns:
+ [`~models.vae.DecoderOutput`] or `tuple`:
+ If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+ returned.
+ """
+
+ batch_size, num_channels, num_frames, height, width = z.shape
+ sample_height = height * self.spatial_compression_ratio
+ sample_width = width * self.spatial_compression_ratio
+
+ tile_latent_min_height = (
+ self.tile_sample_min_height // self.spatial_compression_ratio
+ )
+ tile_latent_min_width = (
+ self.tile_sample_min_width // self.spatial_compression_ratio
+ )
+ tile_latent_stride_height = (
+ self.tile_sample_stride_height // self.spatial_compression_ratio
+ )
+ tile_latent_stride_width = (
+ self.tile_sample_stride_width // self.spatial_compression_ratio
+ )
+
+ blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
+ blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
+
+ # Split z into overlapping tiles and decode them separately.
+ # The tiles have an overlap to avoid seams between tiles.
+ rows = []
+ for i in range(0, height, tile_latent_stride_height):
+ row = []
+ for j in range(0, width, tile_latent_stride_width):
+ time = self.decoder(
+ z[
+ :,
+ :,
+ :,
+ i : i + tile_latent_min_height,
+ j : j + tile_latent_min_width,
+ ],
+ temb,
+ causal=causal,
+ )
+
+ row.append(time)
+ rows.append(row)
+
+ result_rows = []
+ for i, row in enumerate(rows):
+ result_row = []
+ for j, tile in enumerate(row):
+ # blend the above tile and the left tile
+ # to the current tile and add the current tile to the result row
+ if i > 0:
+ tile = self.blend_v(rows[i - 1][j], tile, blend_height)
+ if j > 0:
+ tile = self.blend_h(row[j - 1], tile, blend_width)
+ result_row.append(
+ tile[
+ :,
+ :,
+ :,
+ : self.tile_sample_stride_height,
+ : self.tile_sample_stride_width,
+ ]
+ )
+ result_rows.append(torch.cat(result_row, dim=4))
+
+ dec = torch.cat(result_rows, dim=3)[:, :, :, :sample_height, :sample_width]
+
+ if not return_dict:
+ return (dec,)
+
+ return DecoderOutput(sample=dec)
+
+ def _temporal_tiled_encode(
+ self, x: torch.Tensor, causal: Optional[bool] = None
+ ) -> AutoencoderKLOutput:
+ batch_size, num_channels, num_frames, height, width = x.shape
+ latent_num_frames = (num_frames - 1) // self.temporal_compression_ratio + 1
+
+ tile_latent_min_num_frames = (
+ self.tile_sample_min_num_frames // self.temporal_compression_ratio
+ )
+ tile_latent_stride_num_frames = (
+ self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+ )
+ blend_num_frames = tile_latent_min_num_frames - tile_latent_stride_num_frames
+
+ row = []
+ for i in range(0, num_frames, self.tile_sample_stride_num_frames):
+ tile = x[:, :, i : i + self.tile_sample_min_num_frames + 1, :, :]
+ if self.use_tiling and (
+ height > self.tile_sample_min_height
+ or width > self.tile_sample_min_width
+ ):
+ tile = self.tiled_encode(tile, causal=causal)
+ else:
+ tile = self.encoder(tile, causal=causal)
+ if i > 0:
+ tile = tile[:, :, 1:, :, :]
+ row.append(tile)
+
+ result_row = []
+ for i, tile in enumerate(row):
+ if i > 0:
+ tile = self.blend_t(row[i - 1], tile, blend_num_frames)
+ result_row.append(tile[:, :, :tile_latent_stride_num_frames, :, :])
+ else:
+ result_row.append(tile[:, :, : tile_latent_stride_num_frames + 1, :, :])
+
+ enc = torch.cat(result_row, dim=2)[:, :, :latent_num_frames]
+ return enc
+
+ def _temporal_tiled_decode(
+ self,
+ z: torch.Tensor,
+ temb: Optional[torch.Tensor],
+ causal: Optional[bool] = None,
+ return_dict: bool = True,
+ ) -> Union[DecoderOutput, torch.Tensor]:
+ batch_size, num_channels, num_frames, height, width = z.shape
+ num_sample_frames = (num_frames - 1) * self.temporal_compression_ratio + 1
+
+ tile_latent_min_height = (
+ self.tile_sample_min_height // self.spatial_compression_ratio
+ )
+ tile_latent_min_width = (
+ self.tile_sample_min_width // self.spatial_compression_ratio
+ )
+ tile_latent_min_num_frames = (
+ self.tile_sample_min_num_frames // self.temporal_compression_ratio
+ )
+ tile_latent_stride_num_frames = (
+ self.tile_sample_stride_num_frames // self.temporal_compression_ratio
+ )
+ blend_num_frames = (
+ self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames
+ )
+
+ row = []
+ for i in range(0, num_frames, tile_latent_stride_num_frames):
+ tile = z[:, :, i : i + tile_latent_min_num_frames + 1, :, :]
+ if self.use_tiling and (
+ tile.shape[-1] > tile_latent_min_width
+ or tile.shape[-2] > tile_latent_min_height
+ ):
+ decoded = self.tiled_decode(
+ tile, temb, causal=causal, return_dict=True
+ ).sample
+ else:
+ decoded = self.decoder(tile, temb, causal=causal)
+ if i > 0:
+ decoded = decoded[:, :, :-1, :, :]
+ row.append(decoded)
+
+ result_row = []
+ for i, tile in enumerate(row):
+ if i > 0:
+ tile = self.blend_t(row[i - 1], tile, blend_num_frames)
+ tile = tile[:, :, : self.tile_sample_stride_num_frames, :, :]
+ result_row.append(tile)
+ else:
+ result_row.append(
+ tile[:, :, : self.tile_sample_stride_num_frames + 1, :, :]
+ )
+
+ dec = torch.cat(result_row, dim=2)[:, :, :num_sample_frames]
+
+ if not return_dict:
+ return (dec,)
+ return DecoderOutput(sample=dec)
+
+ def forward(
+ self,
+ sample: torch.Tensor,
+ temb: Optional[torch.Tensor] = None,
+ sample_posterior: bool = False,
+ encoder_causal: Optional[bool] = None,
+ decoder_causal: Optional[bool] = None,
+ return_dict: bool = True,
+ generator: Optional[torch.Generator] = None,
+ ) -> Union[torch.Tensor, torch.Tensor]:
+ x = sample
+ posterior = self.encode(x, causal=encoder_causal).latent_dist
+ if sample_posterior:
+ z = posterior.sample(generator=generator)
+ else:
+ z = posterior.mode()
+ dec = self.decode(z, temb, causal=decoder_causal)
+ if not return_dict:
+ return (dec.sample,)
+ return dec
+
+
+EntryClass = AutoencoderKLLTX2Video
diff --git a/python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py b/python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py
new file mode 100644
index 000000000..82ad20d2a
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
new file mode 100644
index 000000000..29f3195aa
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py
index 7cfb17854..34ff8658e 100644
--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py
+++ b/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",
]
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py
new file mode 100644
index 000000000..744688ca5
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
new file mode 100644
index 000000000..c18e6a4d4
--- /dev/null
+++ b/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)
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py
new file mode 100644
index 000000000..4440a3060
--- /dev/null
+++ b/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
diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py
new file mode 100644
index 000000000..93baeba6e
--- /dev/null
+++ b/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