Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 56 additions & 1 deletion fastvideo/configs/models/dits/ltx2.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,15 @@ class LTX2VideoArchConfig(DiTArchConfig):
_fsdp_shard_conditions: list = field(default_factory=lambda: [is_ltx2_blocks])
_compile_conditions: list = field(default_factory=lambda: [is_ltx2_blocks])

# Parameter name mapping for weight conversion (hf/comfy -> FastVideo)
# Parameter name mapping for weight conversion (hf/comfy -> FastVideo).
# The ``to_gate_compress`` -> ``to_gate_logits`` rules for the LTX-2.3
# gated-attention path are inserted at the front of this dict in
# ``__post_init__`` only when ``apply_gated_attention=True``. Without
# that flag the target model has no ``to_gate_logits`` slot, *and* the
# same-named ``to_gate_compress`` already lives on the LTX-2.0 VSA-QAT
# gate path (plus it is in the default ``lora_target_modules`` list).
# Applying the rename unconditionally silently breaks LTX-2.0 VSA
# checkpoints and default-target LoRAs.
param_names_mapping: dict = field(
default_factory=lambda: {
r"^model\.diffusion_model\.(.*)$": r"model.\1",
Expand All @@ -44,6 +52,9 @@ class LTX2VideoArchConfig(DiTArchConfig):
attention_type: str = "default"
rope_type: str = "split"
double_precision_rope: bool = True
# LTX-2.3 gated extensions. All default OFF == LTX-2.0 behavior.
cross_attention_adaln: bool = False
caption_proj_before_connector: bool = False

positional_embedding_theta: float = 10000.0
positional_embedding_max_pos: list[int] = field(default_factory=lambda: [20, 2048, 2048])
Expand All @@ -64,6 +75,27 @@ class LTX2VideoArchConfig(DiTArchConfig):
audio_cross_attention_dim: int = 2048
audio_positional_embedding_max_pos: list[int] = field(default_factory=lambda: [20])
av_ca_timestep_scale_multiplier: int = 1
# LTX-2.3 gated self-attention (distinct from the VSA-QAT to_gate_compress
# gate). Default OFF == LTX-2.0 behavior.
apply_gated_attention: bool = False

# Text connector/feature extractor compatibility fields carried in some
# transformer configs (used by the LTX-2.3 text stack). Defaults match
# the LTX-2.0 connector layout.
caption_projection_first_linear: bool = True
caption_proj_input_norm: bool = True
caption_projection_second_linear: bool = True
connector_num_attention_heads: int = 30
connector_attention_head_dim: int = 128
connector_num_layers: int = 2
audio_connector_num_attention_heads: int = 30
audio_connector_attention_head_dim: int = 128
audio_connector_num_layers: int = 2

# STG perturbation block index differs across model versions.
# LTX-2.0 defaults to block 29; LTX-2.3 (caption_proj_before_connector)
# uses block 28. ``None`` resolves in __post_init__.
stg_block_idx: int | None = None

def __post_init__(self):
super().__post_init__()
Expand All @@ -72,6 +104,29 @@ def __post_init__(self):
self.in_channels = self.num_channels_latents * patch_volume
if self.out_channels is None:
self.out_channels = self.in_channels
if self.stg_block_idx is None:
self.stg_block_idx = 28 if self.caption_proj_before_connector else 29

# LTX-2.3 stores the gated-attention weight under ``to_gate_compress``
# upstream; FastVideo's internal name is ``to_gate_logits``. Only
# enable the rename when the gated path is actually configured: the
# LTX-2.0 attention module's own ``to_gate_compress`` parameter
# (created when the backend is ``VIDEO_SPARSE_ATTN``) and the default
# ``to_gate_compress`` LoRA target both share the upstream name, so
# an unconditional rename would silently retarget them. Inserted at
# the front so first-match-wins matching fires the rename before the
# generic prefix-strip rules.
if self.apply_gated_attention:
gate_rules = {
r"^model\.diffusion_model\.(.*)\.to_gate_compress\.(.*)$": r"model.\1.to_gate_logits.\2",
r"^diffusion_model\.(.*)\.to_gate_compress\.(.*)$": r"model.\1.to_gate_logits.\2",
r"^model\.(.*)\.to_gate_compress\.(.*)$": r"model.\1.to_gate_logits.\2",
r"^(.*)\.to_gate_compress\.(.*)$": r"model.\1.to_gate_logits.\2",
}
self.param_names_mapping = {
**gate_rules,
**self.param_names_mapping,
}


@dataclass
Expand Down
20 changes: 19 additions & 1 deletion fastvideo/configs/models/encoders/gemma.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,13 @@


def _is_feature_extractor_linear(n: str, m) -> bool:
return n.endswith("feature_extractor_linear")
# LTX-2.3 (caption_proj_before_connector) introduces separate
# video/audio feature extractor linears; keep the LTX-2.0 name too.
return (
n.endswith("feature_extractor_linear")
or n.endswith("video_feature_extractor_linear")
or n.endswith("audio_feature_extractor_linear")
)


def _is_embeddings(n: str, m) -> bool:
Expand All @@ -35,14 +41,26 @@ class LTX2GemmaArchConfig(TextEncoderArchConfig):

feature_extractor_in_features: int = 3840 * 49
feature_extractor_out_features: int = 3840
# LTX-2.3 text-stack connector fields (default OFF == LTX-2.0 behavior).
video_feature_extractor_out_features: int | None = None
audio_feature_extractor_out_features: int | None = None
caption_proj_before_connector: bool = False
caption_projection_first_linear: bool = True
caption_proj_input_norm: bool = True
caption_projection_second_linear: bool = True

connector_num_attention_heads: int = 30
connector_attention_head_dim: int = 128
connector_num_layers: int = 2
# Separate audio connector geometry (None falls back to the video values).
audio_connector_num_attention_heads: int | None = None
audio_connector_attention_head_dim: int | None = None
audio_connector_num_layers: int | None = None
connector_positional_embedding_theta: float = 10000.0
connector_positional_embedding_max_pos: list[int] = field(default_factory=lambda: [4096])
connector_rope_type: str = "split"
connector_double_precision_rope: bool = False
connector_apply_gated_attention: bool = False
connector_num_learnable_registers: int | None = 128

_fsdp_shard_conditions: list = field(
Expand Down
Loading
Loading