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

Automated PR - 2026-03-04

Browse Source
This commit is contained in:
sync-bot
2026-03-04 19:34:46 +00:00
parent 28c3c73fe5
commit 822ce3c4b1
73 changed files with 4984 additions and 1220 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-trainer/src/ltx_trainer/datasets.py
+7 -3
View File
@@ -67,14 +67,18 @@ class DummyDataset(Dataset):
"fps": self.fps,
},
"text_conditions": {
"prompt_embeds": torch.randn(
"video_prompt_embeds": torch.randn(
self.prompt_sequence_length,
self.prompt_embed_dim,
), # random text embeddings
),
"audio_prompt_embeds": torch.randn(
self.prompt_sequence_length,
self.prompt_embed_dim,
),
"prompt_attention_mask": torch.ones(
self.prompt_sequence_length,
dtype=torch.bool,
), # random attention mask
),
},
}
packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py
+29 -24
View File
@@ -18,28 +18,26 @@ import logging
from collections.abc import Generator
from contextlib import contextmanager
from pathlib import Path
from typing import TYPE_CHECKING
import torch
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
from ltx_core.text_encoders.gemma.embeddings_connector import Embeddings1DConnectorConfigurator
from ltx_core.text_encoders.gemma.encoders.av_encoder import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
AVGemmaTextEncoderModel,
from ltx_core.text_encoders.gemma import AV_GEMMA_TEXT_ENCODER_KEY_OPS
from ltx_core.text_encoders.gemma.embeddings_connector import (
AudioEmbeddings1DConnectorConfigurator,
Embeddings1DConnectorConfigurator,
)
from ltx_core.text_encoders.gemma.feature_extractor import GemmaFeaturesExtractorProjLinear
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
from ltx_core.text_encoders.gemma.encoders.encoder_configurator import _create_feature_extractor
from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
if TYPE_CHECKING:
from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
def load_8bit_gemma(
checkpoint_path: str | Path,
gemma_model_path: str | Path,
dtype: torch.dtype = torch.bfloat16,
) -> "AVGemmaTextEncoderModel":
) -> GemmaTextEncoder:
"""Load the Gemma text encoder in 8-bit precision using bitsandbytes.
This function bypasses ltx-core's standard loading path to enable 8-bit quantization
via the bitsandbytes library. The Gemma model is loaded with load_in_8bit=True and
@@ -50,7 +48,7 @@ def load_8bit_gemma(
gemma_model_path: Path to Gemma model directory
dtype: Data type for non-quantized model weights (feature extractor, connectors)
Returns:
Loaded AVGemmaTextEncoderModel with 8-bit quantized Gemma backbone
Loaded GemmaTextEncoder with 8-bit quantized Gemma backbone
Raises:
ImportError: If bitsandbytes is not installed
FileNotFoundError: If required model files are not found
@@ -88,28 +86,35 @@ def load_8bit_gemma(
def extract_state_dict(prefix: str) -> dict[str, torch.Tensor]:
return {k.replace(prefix, ""): v for k, v in sd.sd.items() if k.startswith(prefix)}
# Create and load feature extractor
feature_extractor = GemmaFeaturesExtractorProjLinear()
feature_extractor.load_state_dict(extract_state_dict("feature_extractor_linear."))
feature_extractor = feature_extractor.to(device=gemma_model.device, dtype=dtype)
# Create and load video embeddings connector
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
embeddings_connector.load_state_dict(extract_state_dict("embeddings_connector."))
embeddings_connector.load_state_dict(extract_state_dict("embeddings_processor.video_connector."))
embeddings_connector = embeddings_connector.to(device=gemma_model.device, dtype=dtype)
# Create and load audio embeddings connector
audio_embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
audio_embeddings_connector.load_state_dict(extract_state_dict("audio_embeddings_connector."))
audio_embeddings_connector = AudioEmbeddings1DConnectorConfigurator.from_config(config)
audio_embeddings_connector.load_state_dict(extract_state_dict("embeddings_processor.audio_connector."))
audio_embeddings_connector = audio_embeddings_connector.to(device=gemma_model.device, dtype=dtype)
# Construct the text encoder
text_encoder = AVGemmaTextEncoderModel(
feature_extractor_linear=feature_extractor,
embeddings_connector=embeddings_connector,
audio_embeddings_connector=audio_embeddings_connector,
# Create embeddings processor
embeddings_processor = EmbeddingsProcessor(
video_connector=embeddings_connector,
audio_connector=audio_embeddings_connector,
)
transformer_config = config.get("transformer", {})
feature_extractor = _create_feature_extractor(transformer_config)
feature_extractor.load_state_dict(
{k.removeprefix("feature_extractor."): v for k, v in sd.sd.items() if k.startswith("feature_extractor.")},
)
feature_extractor = feature_extractor.to(device=gemma_model.device, dtype=dtype)
text_encoder = GemmaTextEncoder(
feature_extractor=feature_extractor,
embeddings_processor=embeddings_processor,
tokenizer=tokenizer,
model=gemma_model,
dtype=dtype,
)
return text_encoder
packages/ltx-trainer/src/ltx_trainer/model_loader.py
+8 -8
View File
@@ -32,7 +32,7 @@ if TYPE_CHECKING:
from ltx_core.model.audio_vae import AudioDecoder, AudioEncoder, Vocoder
from ltx_core.model.transformer import LTXModel
from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
from ltx_core.text_encoders.gemma import GemmaTextEncoder
def _to_torch_device(device: Device) -> torch.device:
@@ -192,7 +192,7 @@ def load_text_encoder(
device: Device = "cpu",
dtype: torch.dtype = torch.bfloat16,
load_in_8bit: bool = False,
) -> "AVGemmaTextEncoderModel":
) -> "GemmaTextEncoder":
"""Load the Gemma text encoder.
Args:
checkpoint_path: Path to the LTX-2 safetensors checkpoint file
@@ -203,7 +203,7 @@ def load_text_encoder(
When True, the model is loaded with device_map="auto" and the device argument
is ignored for the Gemma backbone (feature extractor still uses dtype).
Returns:
Loaded AVGemmaTextEncoderModel
Loaded GemmaTextEncoder (unified encoder handling V1/V2/V3)
"""
if not Path(gemma_model_path).is_dir():
raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
@@ -216,12 +216,12 @@ def load_text_encoder(
# Standard loading path
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
from ltx_core.text_encoders.gemma.encoders.av_encoder import (
from ltx_core.text_encoders.gemma import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
GEMMA_MODEL_OPS,
AVGemmaTextEncoderModelConfigurator,
GemmaTextEncoderConfigurator,
module_ops_from_gemma_root,
)
from ltx_core.text_encoders.gemma.encoders.base_encoder import module_ops_from_gemma_root
from ltx_core.utils import find_matching_file
torch_device = _to_torch_device(device)
@@ -231,7 +231,7 @@ def load_text_encoder(
text_encoder = SingleGPUModelBuilder(
model_path=(str(checkpoint_path), *gemma_weight_paths),
model_class_configurator=AVGemmaTextEncoderModelConfigurator,
model_class_configurator=GemmaTextEncoderConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
).build(device=torch_device, dtype=dtype)
@@ -253,7 +253,7 @@ class LtxModelComponents:
video_vae_decoder: "VideoDecoder | None" = None
audio_vae_decoder: "AudioDecoder | None" = None
vocoder: "Vocoder | None" = None
text_encoder: "AVGemmaTextEncoderModel | None" = None
text_encoder: "GemmaTextEncoder | None" = None
scheduler: "LTX2Scheduler | None" = None
packages/ltx-trainer/src/ltx_trainer/timestep_samplers.py
+42 -10
View File
@@ -45,29 +45,61 @@ class UniformTimestepSampler(TimestepSampler):
return self.sample(batch.shape[0], device=batch.device)
class ShiftedLogitNormalTimestepSampler:
class ShiftedLogitNormalTimestepSampler(TimestepSampler):
"""
Samples timesteps from a shifted logit-normal distribution,
Samples timesteps from a stretched shifted logit-normal distribution,
where the shift is determined by the sequence length.
The stretching normalizes samples between percentile bounds to ensure
the distribution covers [0, 1] more evenly. A uniform fallback prevents
collapse at high token counts.
"""
def __init__(self, std: float = 1.0):
def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1):
self.std = std
self.eps = eps
self.uniform_prob = uniform_prob
# Percentile values for stretching (scaled by std)
# 99.9th percentile of standard normal ≈ 3.0902
# 0.5th percentile of standard normal ≈ -2.5758
self.normal_999_percentile = 3.0902 * std
self.normal_005_percentile = -2.5758 * std
def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:
"""Sample timesteps for a batch from a shifted logit-normal distribution.
"""Sample timesteps for a batch from a stretched shifted logit-normal distribution.
Args:
batch_size: Number of timesteps to sample
seq_length: Length of the sequence being processed, used to determine the shift
device: Device to place the samples on
Returns:
Tensor of shape (batch_size,) containing timesteps sampled from a shifted
logit-normal distribution, where the shift is determined by seq_length
Tensor of shape (batch_size,) containing timesteps sampled from a stretched
shifted logit-normal distribution, where the shift is determined by seq_length
"""
shift = self._get_shift_for_sequence_length(seq_length)
normal_samples = torch.randn((batch_size,), device=device) * self.std + shift
timesteps = torch.sigmoid(normal_samples)
return timesteps
mu = self._get_shift_for_sequence_length(seq_length)
# Sample from shifted logit-normal
normal_samples = torch.randn((batch_size,), device=device) * self.std + mu
logitnormal_samples = torch.sigmoid(normal_samples)
# Compute percentile bounds for stretching
percentile_999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device))
percentile_005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device))
# Stretch to [0, 1] range by normalizing between percentiles
zero_terminal_raw = (logitnormal_samples - percentile_005) / (percentile_999 - percentile_005)
# Reflect small values around eps for numerical stability
stretched_logit = torch.where(
zero_terminal_raw >= self.eps,
zero_terminal_raw,
2 * self.eps - zero_terminal_raw,
)
stretched_logit = torch.clamp(stretched_logit, 0, 1)
# Mix with uniform samples (uniform_prob of the time)
uniform = (1 - self.eps) * torch.rand((batch_size,), device=device) + self.eps
prob = torch.rand((batch_size,), device=device)
return torch.where(prob > self.uniform_prob, stretched_logit, uniform)
def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
"""Sample timesteps for a specific batch tensor.
packages/ltx-trainer/src/ltx_trainer/trainer.py
+17 -4
View File
@@ -309,9 +309,22 @@ class LtxvTrainer:
"""Perform a single training step using the configured strategy."""
# Apply embedding connectors to transform pre-computed text embeddings
conditions = batch["conditions"]
video_embeds, audio_embeds, attention_mask = self._text_encoder._run_connectors(
conditions["prompt_embeds"], conditions["prompt_attention_mask"]
if "video_prompt_embeds" in conditions:
# New format: separate video/audio features from precompute()
video_features = conditions["video_prompt_embeds"]
audio_features = conditions.get("audio_prompt_embeds")
else:
# Legacy format: single prompt_embeds tensor — duplicate for both modalities
video_features = conditions["prompt_embeds"]
audio_features = conditions["prompt_embeds"]
mask = conditions["prompt_attention_mask"]
additive_mask = self._text_encoder._convert_to_additive_mask(mask, video_features.dtype)
video_embeds, audio_embeds, attention_mask = self._text_encoder.embeddings_processor.create_embeddings(
video_features, audio_features, additive_mask
)
conditions["video_prompt_embeds"] = video_embeds
conditions["audio_prompt_embeds"] = audio_embeds
conditions["prompt_attention_mask"] = attention_mask
@@ -375,7 +388,7 @@ class LtxvTrainer:
# Unload heavy components to free VRAM, keeping only the embedding connectors
self._text_encoder.model = None
self._text_encoder.tokenizer = None
self._text_encoder.feature_extractor_linear = None
self._text_encoder.feature_extractor = None
logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
return cached_embeddings
@@ -822,7 +835,7 @@ class LtxvTrainer:
output_path=output_path,
fps=self._config.validation.frame_rate,
audio=audio,
audio_sample_rate=self._vocoder.output_sample_rate if audio is not None else None,
audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
)
video_paths.append(output_path)
packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py
+2
View File
@@ -163,6 +163,7 @@ class TextToVideoStrategy(TrainingStrategy):
# Create video Modality
video_modality = Modality(
enabled=True,
sigma=sigmas,
latent=noisy_video,
timesteps=video_timesteps,
positions=video_positions,
@@ -254,6 +255,7 @@ class TextToVideoStrategy(TrainingStrategy):
audio_modality = Modality(
enabled=True,
latent=noisy_audio,
sigma=sigmas,
timesteps=audio_timesteps,
positions=audio_positions,
context=audio_prompt_embeds,
packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py
+1
View File
@@ -210,6 +210,7 @@ class VideoToVideoStrategy(TrainingStrategy):
video_modality = Modality(
enabled=True,
latent=combined_latents,
sigma=sigmas,
timesteps=timesteps,
positions=positions,
context=prompt_embeds,
packages/ltx-trainer/src/ltx_trainer/validation_sampler.py
+8 -3
View File
@@ -36,7 +36,7 @@ if TYPE_CHECKING:
from ltx_core.model.audio_vae import AudioDecoder, Vocoder
from ltx_core.model.transformer import LTXModel
from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
from ltx_core.text_encoders.gemma import AVGemmaTextEncoderModel
from ltx_core.text_encoders.gemma import GemmaTextEncoder
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
@@ -124,7 +124,7 @@ class ValidationSampler:
transformer: "LTXModel",
vae_decoder: "VideoDecoder",
vae_encoder: "VideoEncoder | None",
text_encoder: "AVGemmaTextEncoderModel | None" = None,
text_encoder: "GemmaTextEncoder | None" = None,
audio_decoder: "AudioDecoder | None" = None,
vocoder: "Vocoder | None" = None,
sampling_context: SamplingContext | None = None,
@@ -497,6 +497,7 @@ class ValidationSampler:
video = Modality(
enabled=True,
latent=video_state.latent,
sigma=sigmas[0].repeat(video_state.latent.shape[0]),
timesteps=video_state.denoise_mask,
positions=video_state.positions,
context=v_ctx_pos,
@@ -509,6 +510,7 @@ class ValidationSampler:
audio = Modality(
enabled=True,
latent=audio_state.latent,
sigma=sigmas[0].repeat(audio_state.latent.shape[0]),
timesteps=audio_state.denoise_mask,
positions=audio_state.positions,
context=a_ctx_pos,
@@ -525,6 +527,7 @@ class ValidationSampler:
video = replace(
video,
latent=video_state.latent,
sigma=sigma.repeat(video_state.latent.shape[0]),
timesteps=sigma * video_state.denoise_mask,
positions=video_state.positions,
)
@@ -533,6 +536,7 @@ class ValidationSampler:
audio = replace(
audio,
latent=audio_state.latent,
sigma=sigma.repeat(audio_state.latent.shape[0]),
timesteps=sigma * audio_state.denoise_mask,
positions=audio_state.positions,
)
@@ -703,7 +707,8 @@ class ValidationSampler:
# Move the base Gemma model to CPU but keep embeddings connectors on GPU
# as this module is also used during training
self._text_encoder.model.to("cpu")
self._text_encoder.feature_extractor_linear.to("cpu")
if self._text_encoder.feature_extractor is not None:
self._text_encoder.feature_extractor.to("cpu")
return v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg