177 lines
6.8 KiB
Python
177 lines
6.8 KiB
Python
# 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 = [20]
|
|
|
|
|
|
@dataclass
|
|
class LTX2Config(DiTConfig):
|
|
"""Configuration for LTX-2 Video Transformer."""
|
|
|
|
arch_config: LTX2ArchConfig = field(default_factory=LTX2ArchConfig)
|
|
|
|
prefix: str = "ltx2"
|