indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-03-05

Browse Source
This commit is contained in:
sync-bot
2026-03-05 15:47:20 +00:00
parent 3b6d09d7b6
commit d230aec5cd
29 changed files with 739 additions and 540 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-core/src/ltx_core/loader/sft_loader.py
+4 -1
View File
@@ -57,7 +57,10 @@ class SafetensorsModelStateDictLoader(StateDictLoader):
def metadata(self, path: str) -> dict:
with safetensors.safe_open(path, framework="pt") as f:
return json.loads(f.metadata()["config"])
meta = f.metadata()
if meta is None or "config" not in meta:
return {}
return json.loads(meta["config"])
def load(self, path: str | list[str], sd_ops: SDOps | None = None, device: torch.device | None = None) -> StateDict:
return self.weight_loader.load(path, sd_ops, device)
packages/ltx-core/src/ltx_core/model/transformer/model_configurator.py
+2 -2
View File
@@ -128,12 +128,12 @@ def _build_caption_projections(
) -> tuple[torch.nn.Module | None, torch.nn.Module | None]:
"""Build caption projections for the transformer when projection is NOT in the text encoder.
19B models: projection is in the transformer (caption_proj_before_connector=False).
20B models: projection is in the text encoder, so no projections are created here.
22B models: projection is in the text encoder, so no projections are created here.
Args:
config: Full model config dict (must contain "transformer" key).
is_av: Whether this is an audio-video model. When False, audio projection is skipped.
Returns:
Tuple of (video_caption_projection, audio_caption_projection), both None for 20B models.
Tuple of (video_caption_projection, audio_caption_projection), both None for 22B models.
"""
transformer_config = config.get("transformer", {})
if transformer_config.get("caption_proj_before_connector", False):
packages/ltx-core/src/ltx_core/model/transformer/transformer.py
+45 -39
View File
@@ -140,19 +140,20 @@ class BasicAVTransformerBlock(torch.nn.Module):
batch_size: int,
scale_shift_timestep: torch.Tensor,
gate_timestep: torch.Tensor,
scale_shift_indices: slice,
num_scale_shift_values: int = 4,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
scale_shift_ada_values = self.get_ada_values(
scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, slice(None, None)
scale_shift_table[:num_scale_shift_values, :], batch_size, scale_shift_timestep, scale_shift_indices
)
gate_ada_values = self.get_ada_values(
scale_shift_table[num_scale_shift_values:, :], batch_size, gate_timestep, slice(None, None)
)
scale_shift_chunks = [t.squeeze(2) for t in scale_shift_ada_values]
gate_ada_values = [t.squeeze(2) for t in gate_ada_values]
scale, shift = (t.squeeze(2) for t in scale_shift_ada_values)
(gate,) = (t.squeeze(2) for t in gate_ada_values)
return (*scale_shift_chunks, *gate_ada_values)
return scale, shift, gate
def _apply_text_cross_attention(
self,
@@ -287,37 +288,26 @@ class BasicAVTransformerBlock(torch.nn.Module):
vx_norm3 = rms_norm(vx, eps=self.norm_eps)
ax_norm3 = rms_norm(ax, eps=self.norm_eps)
(
scale_ca_audio_hidden_states_a2v,
shift_ca_audio_hidden_states_a2v,
scale_ca_audio_hidden_states_v2a,
shift_ca_audio_hidden_states_v2a,
gate_out_v2a,
) = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_audio,
ax.shape[0],
audio.cross_scale_shift_timestep,
audio.cross_gate_timestep,
)
(
scale_ca_video_hidden_states_a2v,
shift_ca_video_hidden_states_a2v,
scale_ca_video_hidden_states_v2a,
shift_ca_video_hidden_states_v2a,
gate_out_a2v,
) = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_video,
vx.shape[0],
video.cross_scale_shift_timestep,
video.cross_gate_timestep,
)
if run_a2v and not perturbations.all_in_batch(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx):
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v) + shift_ca_video_hidden_states_a2v
del scale_ca_video_hidden_states_a2v, shift_ca_video_hidden_states_a2v
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
del scale_ca_audio_hidden_states_a2v, shift_ca_audio_hidden_states_a2v
scale_ca_video_a2v, shift_ca_video_a2v, gate_out_a2v = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_video,
vx.shape[0],
video.cross_scale_shift_timestep,
video.cross_gate_timestep,
slice(0, 2),
)
vx_scaled = vx_norm3 * (1 + scale_ca_video_a2v) + shift_ca_video_a2v
del scale_ca_video_a2v, shift_ca_video_a2v
scale_ca_audio_a2v, shift_ca_audio_a2v, _ = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_audio,
ax.shape[0],
audio.cross_scale_shift_timestep,
audio.cross_gate_timestep,
slice(0, 2),
)
ax_scaled = ax_norm3 * (1 + scale_ca_audio_a2v) + shift_ca_audio_a2v
del scale_ca_audio_a2v, shift_ca_audio_a2v
a2v_mask = perturbations.mask_like(PerturbationType.SKIP_A2V_CROSS_ATTN, self.idx, vx)
vx = vx + (
self.audio_to_video_attn(
@@ -330,11 +320,26 @@ class BasicAVTransformerBlock(torch.nn.Module):
* a2v_mask
)
del gate_out_a2v, a2v_mask, vx_scaled, ax_scaled
if run_v2a and not perturbations.all_in_batch(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx):
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
del scale_ca_audio_hidden_states_v2a, shift_ca_audio_hidden_states_v2a
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
del scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a
scale_ca_audio_v2a, shift_ca_audio_v2a, gate_out_v2a = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_audio,
ax.shape[0],
audio.cross_scale_shift_timestep,
audio.cross_gate_timestep,
slice(2, 4),
)
ax_scaled = ax_norm3 * (1 + scale_ca_audio_v2a) + shift_ca_audio_v2a
del scale_ca_audio_v2a, shift_ca_audio_v2a
scale_ca_video_v2a, shift_ca_video_v2a, _ = self.get_av_ca_ada_values(
self.scale_shift_table_a2v_ca_video,
vx.shape[0],
video.cross_scale_shift_timestep,
video.cross_gate_timestep,
slice(2, 4),
)
vx_scaled = vx_norm3 * (1 + scale_ca_video_v2a) + shift_ca_video_v2a
del scale_ca_video_v2a, shift_ca_video_v2a
v2a_mask = perturbations.mask_like(PerturbationType.SKIP_V2A_CROSS_ATTN, self.idx, ax)
ax = ax + (
self.video_to_audio_attn(
@@ -347,6 +352,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
* v2a_mask
)
del gate_out_v2a, v2a_mask, ax_scaled, vx_scaled
del vx_norm3, ax_norm3
if run_vx:
packages/ltx-core/src/ltx_core/text_encoders/gemma/__init__.py
+16 -8
View File
@@ -1,25 +1,33 @@
"""Gemma text encoder components."""
from ltx_core.text_encoders.gemma.embeddings_processor import (
EmbeddingsProcessor,
EmbeddingsProcessorOutput,
convert_to_additive_mask,
)
from ltx_core.text_encoders.gemma.encoders.base_encoder import (
GemmaEncoderOutput,
GemmaTextEncoder,
encode_text,
module_ops_from_gemma_root,
)
from ltx_core.text_encoders.gemma.encoders.encoder_configurator import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
EMBEDDINGS_PROCESSOR_KEY_OPS,
GEMMA_LLM_KEY_OPS,
GEMMA_MODEL_OPS,
VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS,
VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS,
EmbeddingsProcessorConfigurator,
GemmaTextEncoderConfigurator,
)
__all__ = [
"AV_GEMMA_TEXT_ENCODER_KEY_OPS",
"EMBEDDINGS_PROCESSOR_KEY_OPS",
"GEMMA_LLM_KEY_OPS",
"GEMMA_MODEL_OPS",
"VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS",
"GemmaEncoderOutput",
"VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS",
"EmbeddingsProcessor",
"EmbeddingsProcessorConfigurator",
"EmbeddingsProcessorOutput",
"GemmaTextEncoder",
"GemmaTextEncoderConfigurator",
"encode_text",
"convert_to_additive_mask",
"module_ops_from_gemma_root",
]
packages/ltx-core/src/ltx_core/text_encoders/gemma/embeddings_processor.py
+49 -3
View File
@@ -1,9 +1,24 @@
from typing import NamedTuple
import torch
from torch import nn
from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnector
class EmbeddingsProcessorOutput(NamedTuple):
video_encoding: torch.Tensor
audio_encoding: torch.Tensor | None
attention_mask: torch.Tensor
def convert_to_additive_mask(attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
"""Convert binary attention mask to additive form for transformer masking."""
return (attention_mask.to(torch.int64) - 1).to(dtype).reshape(
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
) * torch.finfo(dtype).max
def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Convert connector output mask to binary mask and apply to encoded tensor."""
binary_mask = (encoded_mask < 0.000001).to(torch.int64)
@@ -13,12 +28,21 @@ def _to_binary_mask(encoded: torch.Tensor, encoded_mask: torch.Tensor) -> tuple[
class EmbeddingsProcessor(nn.Module):
"""Wraps video connector + optional audio connector.
Returns (video_encoded, audio_encoded | None, binary_mask).
"""Wraps feature extractor + video connector + optional audio connector.
Can operate in two modes:
1. create_embeddings(): Takes pre-computed features + additive mask (backward compat, used by trainer)
2. process_hidden_states(): Takes raw Gemma hidden states, runs feature extraction + connectors
"""
def __init__(self, video_connector: Embeddings1DConnector, audio_connector: Embeddings1DConnector | None = None):
def __init__(
self,
*,
feature_extractor: nn.Module | None = None,
video_connector: Embeddings1DConnector,
audio_connector: Embeddings1DConnector | None = None,
):
super().__init__()
self.feature_extractor = feature_extractor
self.video_connector = video_connector
self.audio_connector = audio_connector
@@ -41,3 +65,25 @@ class EmbeddingsProcessor(nn.Module):
audio_encoded, _ = self.audio_connector(audio_features, additive_attention_mask)
return video_encoded, audio_encoded, binary_mask.squeeze(-1)
def process_hidden_states(
self,
hidden_states: tuple[torch.Tensor, ...],
attention_mask: torch.Tensor,
padding_side: str = "left",
) -> EmbeddingsProcessorOutput:
"""Full pipeline: feature extraction -> connectors -> final embeddings.
Args:
hidden_states: Raw Gemma hidden states (tuple of tensors per layer).
attention_mask: Binary attention mask [B, seq_len].
padding_side: Padding side used during tokenization.
Returns:
EmbeddingsProcessorOutput with video_encoding, audio_encoding, and attention_mask.
"""
if self.feature_extractor is None:
raise ValueError("feature_extractor is required for process_hidden_states()")
video_feats, audio_feats = self.feature_extractor(hidden_states, attention_mask, padding_side)
additive_mask = convert_to_additive_mask(attention_mask, video_feats.dtype)
video_enc, audio_enc, binary_mask = self.create_embeddings(video_feats, audio_feats, additive_mask)
return EmbeddingsProcessorOutput(video_enc, audio_enc, binary_mask)
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/base_encoder.py
+17 -54
View File
@@ -1,33 +1,22 @@
import functools
from pathlib import Path
from typing import NamedTuple
import torch
from transformers import AutoImageProcessor, Gemma3ForConditionalGeneration, Gemma3Processor
from ltx_core.loader.module_ops import ModuleOps
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
from ltx_core.utils import find_matching_file
class GemmaEncoderOutput(NamedTuple):
video_encoding: torch.Tensor
audio_encoding: torch.Tensor | None
attention_mask: torch.Tensor
class GemmaTextEncoder(torch.nn.Module):
"""Unified Gemma text encoder with 3-block pipeline.
Block 1: Gemma model (runs LLM, gets hidden states)
Block 2: Feature extractor
Block 3: Embeddings processor (connector with optional audio)
"""Pure Gemma text encoder — runs the LLM and returns raw hidden states.
Prompt enhancement (generate) is also supported since the full
Gemma3ForConditionalGeneration model (including lm_head) is loaded.
"""
def __init__(
self,
feature_extractor: torch.nn.Module,
embeddings_processor: EmbeddingsProcessor,
model: Gemma3ForConditionalGeneration | None = None,
tokenizer: LTXVGemmaTokenizer | None = None,
processor: Gemma3Processor | None = None,
@@ -37,39 +26,25 @@ class GemmaTextEncoder(torch.nn.Module):
self.model = model
self.tokenizer = tokenizer
self.processor = processor
self.feature_extractor = feature_extractor.to(dtype=dtype)
self.embeddings_processor = embeddings_processor.to(dtype=dtype)
self._dtype = dtype
def _convert_to_additive_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
return (attention_mask.to(torch.int64) - 1).to(dtype).reshape(
(attention_mask.shape[0], 1, -1, attention_mask.shape[-1])
) * torch.finfo(dtype).max
def precompute(
self, text: str, padding_side: str = "left"
) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
"""Blocks 1+2: Gemma model -> feature extraction.
Used by process_captions.py for offline precomputation.
Returns (video_features, audio_features | None, attention_mask).
def encode(
self,
text: str,
padding_side: str = "left", # noqa: ARG002
) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
"""Run Gemma LLM and return raw hidden states + attention mask.
Calls the inner model (self.model.model) to skip lm_head logits computation (~500 MiB saving).
Returns:
(hidden_states, attention_mask) where hidden_states is a tuple of per-layer tensors.
"""
# Block 1: Run Gemma
token_pairs = self.tokenizer.tokenize_with_weights(text)["gemma"]
input_ids = torch.tensor([[t[0] for t in token_pairs]], device=self.model.device)
attention_mask = torch.tensor([[w[1] for w in token_pairs]], device=self.model.device)
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
# Block 2: Feature extraction
video_feats, audio_feats = self.feature_extractor(outputs.hidden_states, attention_mask, padding_side)
return video_feats, audio_feats, attention_mask
def forward(self, text: str, padding_side: str = "left") -> GemmaEncoderOutput:
"""Full pipeline: precompute -> embeddings processor."""
video_feats, audio_feats, attention_mask = self.precompute(text, padding_side)
additive_mask = self._convert_to_additive_mask(attention_mask, video_feats.dtype)
video_enc, audio_enc, binary_mask = self.embeddings_processor.create_embeddings(
video_feats, audio_feats, additive_mask
)
return GemmaEncoderOutput(video_enc, audio_enc, binary_mask)
outputs = self.model.model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
hidden_states = outputs.hidden_states
del outputs
return hidden_states, attention_mask
# --- Prompt enhancement methods ---
@@ -225,15 +200,3 @@ def module_ops_from_gemma_root(gemma_root: str) -> tuple[ModuleOps, ...]:
mutator=load_processor,
)
return (tokenizer_load_ops, processor_load_ops)
def encode_text(text_encoder: GemmaTextEncoder, prompts: list[str]) -> list[tuple[torch.Tensor, torch.Tensor]]:
"""Encode a list of prompts using the provided Gemma text encoder.
Returns:
List of tuples, each containing (v_context, a_context) tensors for each prompt.
"""
result = []
for prompt in prompts:
v_context, a_context, _ = text_encoder(prompt)
result.append((v_context, a_context))
return result
packages/ltx-core/src/ltx_core/text_encoders/gemma/encoders/encoder_configurator.py
+42 -35
View File
@@ -22,31 +22,32 @@ from ltx_core.text_encoders.gemma.feature_extractor import (
class GemmaTextEncoderConfigurator(ModelConfigurator[GemmaTextEncoder]):
@classmethod
def from_config(cls, config: dict) -> GemmaTextEncoder:
transformer_config = config.get("transformer", {})
def from_config(cls, config: dict) -> GemmaTextEncoder: # noqa: ARG003
gemma_config = Gemma3Config.from_dict(GEMMA3_CONFIG_FOR_LTX.to_dict())
with torch.device("meta"):
model = Gemma3ForConditionalGeneration(gemma_config)
return GemmaTextEncoder(model=model)
class EmbeddingsProcessorConfigurator(ModelConfigurator[EmbeddingsProcessor]):
@classmethod
def from_config(cls, config: dict) -> EmbeddingsProcessor:
transformer_config = config.get("transformer", {})
# Create video embeddings connector (always needed)
video_connector = Embeddings1DConnectorConfigurator.from_config(config)
# Create audio embeddings connector
audio_connector = AudioEmbeddings1DConnectorConfigurator.from_config(config)
# Create embeddings processor with both connectors
embeddings_processor = EmbeddingsProcessor(
video_connector=video_connector,
audio_connector=audio_connector,
)
# Create feature extractor
feature_extractor = _create_feature_extractor(transformer_config)
return GemmaTextEncoder(
return EmbeddingsProcessor(
video_connector=video_connector,
audio_connector=audio_connector,
feature_extractor=feature_extractor,
embeddings_processor=embeddings_processor,
model=model,
)
@@ -97,8 +98,31 @@ def _create_feature_extractor(transformer_config: dict) -> torch.nn.Module:
)
AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
SDOps("AV_GEMMA_TEXT_ENCODER_KEY_OPS")
# --- Split SDOps: Gemma LLM keys vs Embeddings Processor keys ---
GEMMA_LLM_KEY_OPS = (
SDOps("GEMMA_LLM_KEY_OPS")
# 1. Map language model layers (note the double .model prefix)
.with_matching(prefix="language_model.model.")
.with_replacement("language_model.model.", "model.model.language_model.")
# 2. Map the Vision Tower
.with_matching(prefix="vision_tower.")
.with_replacement("vision_tower.", "model.model.vision_tower.")
# 3. Map the Multi-Modal Projector
.with_matching(prefix="multi_modal_projector.")
.with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
# 4. Duplicate embed_tokens to lm_head (needed for prompt enhancement via generate())
.with_kv_operation(
operation=lambda key, value: [
KeyValueOperationResult(key, value),
KeyValueOperationResult("model.lm_head.weight", value),
],
key_prefix="model.model.language_model.embed_tokens.weight",
)
)
EMBEDDINGS_PROCESSOR_KEY_OPS = (
SDOps("EMBEDDINGS_PROCESSOR_KEY_OPS")
# 1. Map the feature extractor (V1: aggregate_embed inside feature_extractor)
.with_matching(prefix="text_embedding_projection.aggregate_embed.")
.with_replacement("text_embedding_projection.aggregate_embed.", "feature_extractor.aggregate_embed.")
@@ -109,30 +133,13 @@ AV_GEMMA_TEXT_ENCODER_KEY_OPS = (
.with_replacement("text_embedding_projection.audio_aggregate_embed.", "feature_extractor.audio_aggregate_embed.")
# 2. Map the connectors
.with_matching(prefix="model.diffusion_model.video_embeddings_connector.")
.with_replacement("model.diffusion_model.video_embeddings_connector.", "embeddings_processor.video_connector.")
.with_replacement("model.diffusion_model.video_embeddings_connector.", "video_connector.")
.with_matching(prefix="model.diffusion_model.audio_embeddings_connector.")
.with_replacement("model.diffusion_model.audio_embeddings_connector.", "embeddings_processor.audio_connector.")
# 3. Map language model layers (note the double .model prefix)
.with_matching(prefix="language_model.model.")
.with_replacement("language_model.model.", "model.model.language_model.")
# 4. Map the Vision Tower
.with_matching(prefix="vision_tower.")
.with_replacement("vision_tower.", "model.model.vision_tower.")
# 5. Map the Multi-Modal Projector
.with_matching(prefix="multi_modal_projector.")
.with_replacement("multi_modal_projector.", "model.model.multi_modal_projector.")
.with_kv_operation(
operation=lambda key, value: [
KeyValueOperationResult(key, value),
KeyValueOperationResult("model.lm_head.weight", value),
],
key_prefix="model.model.language_model.embed_tokens.weight",
)
.with_replacement("model.diffusion_model.audio_embeddings_connector.", "audio_connector.")
)
VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS = (
SDOps("VIDEO_ONLY_GEMMA_TEXT_ENCODER_KEY_OPS")
VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS = (
SDOps("VIDEO_ONLY_EMBEDDINGS_PROCESSOR_KEY_OPS")
# 1. Map the feature extractor (V1: aggregate_embed inside feature_extractor)
.with_matching(prefix="text_embedding_projection.aggregate_embed.")
.with_replacement("text_embedding_projection.aggregate_embed.", "feature_extractor.aggregate_embed.")
packages/ltx-core/src/ltx_core/text_encoders/gemma/feature_extractor.py
+1 -1
View File
@@ -110,7 +110,7 @@ class FeatureExtractorV1(nn.Module):
class FeatureExtractorV2(nn.Module):
"""20B: per-token RMS norm → rescale → dual aggregate embeds"""
"""22B: per-token RMS norm → rescale → dual aggregate embeds"""
def __init__(
self,