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-trainer/src/ltx_trainer/gemma_8bit.py
+8 -63
View File
@@ -6,10 +6,7 @@ This module provides functionality for loading the Gemma text encoder in 8-bit p
using bitsandbytes, which significantly reduces GPU memory usage.
Example usage:
from ltx_trainer.gemma_8bit import load_8bit_gemma
text_encoder = load_8bit_gemma(
checkpoint_path="/path/to/ltx2.safetensors",
gemma_model_path="/path/to/gemma",
)
text_encoder = load_8bit_gemma(gemma_model_path="/path/to/gemma")
"""
from __future__ import annotations
@@ -21,34 +18,20 @@ from pathlib import Path
import torch
from ltx_core.loader.sft_loader import SafetensorsModelStateDictLoader
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.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
def load_8bit_gemma(
checkpoint_path: str | Path,
gemma_model_path: str | Path,
dtype: torch.dtype = torch.bfloat16,
) -> GemmaTextEncoder:
def load_8bit_gemma(gemma_model_path: str | Path, dtype: torch.dtype = torch.bfloat16) -> 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
torch_dtype=bfloat16, while the feature extractor and connector weights are loaded
from the LTX-2 checkpoint.
Only the Gemma LLM backbone is loaded here. The embeddings processor
(feature extractor + connectors) should be loaded separately via
:func:`ltx_trainer.model_loader.load_embeddings_processor`.
Args:
checkpoint_path: Path to the LTX-2 safetensors checkpoint file
gemma_model_path: Path to Gemma model directory
dtype: Data type for non-quantized model weights (feature extractor, connectors)
dtype: Data type for non-quantized model weights
Returns:
Loaded GemmaTextEncoder with 8-bit quantized Gemma backbone
GemmaTextEncoder with 8-bit quantized Gemma backbone
Raises:
ImportError: If bitsandbytes is not installed
FileNotFoundError: If required model files are not found
@@ -60,7 +43,6 @@ def load_8bit_gemma(
"8-bit text encoder loading requires bitsandbytes. Install it with: uv pip install bitsandbytes"
) from e
# Find paths within gemma_model_path
gemma_path = _find_gemma_subpath(gemma_model_path, "model*.safetensors")
tokenizer_path = _find_gemma_subpath(gemma_model_path, "tokenizer.model")
@@ -74,51 +56,14 @@ def load_8bit_gemma(
local_files_only=True,
)
# Load tokenizer
tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024)
# Load config and weights from the LTX-2 checkpoint
loader = SafetensorsModelStateDictLoader()
config = loader.metadata(str(checkpoint_path))
sd = loader.load(str(checkpoint_path), sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS)
# Helper to extract state dict for a given prefix
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 video embeddings connector
embeddings_connector = Embeddings1DConnectorConfigurator.from_config(config)
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 = 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)
# 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,
return GemmaTextEncoder(
tokenizer=tokenizer,
model=gemma_model,
dtype=dtype,
)
return text_encoder
def _find_gemma_subpath(root_path: str | Path, pattern: str) -> str:
"""Find a file matching a glob pattern and return its parent directory."""
packages/ltx-trainer/src/ltx_trainer/model_loader.py
+37 -10
View File
@@ -8,7 +8,7 @@ Example usage:
# Load individual components
vae_encoder = load_video_vae_encoder("/path/to/checkpoint.safetensors", device="cuda")
vae_decoder = load_video_vae_decoder("/path/to/checkpoint.safetensors", device="cuda")
text_encoder = load_text_encoder("/path/to/checkpoint.safetensors", "/path/to/gemma", device="cuda")
text_encoder = load_text_encoder("/path/to/gemma", device="cuda")
# Load all components at once
components = load_model("/path/to/checkpoint.safetensors", text_encoder_path="/path/to/gemma")
"""
@@ -33,6 +33,7 @@ if TYPE_CHECKING:
from ltx_core.model.transformer import LTXModel
from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
from ltx_core.text_encoders.gemma import GemmaTextEncoder
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
def _to_torch_device(device: Device) -> torch.device:
@@ -187,7 +188,6 @@ def load_vocoder(
def load_text_encoder(
checkpoint_path: str | Path,
gemma_model_path: str | Path,
device: Device = "cpu",
dtype: torch.dtype = torch.bfloat16,
@@ -195,15 +195,14 @@ def load_text_encoder(
) -> "GemmaTextEncoder":
"""Load the Gemma text encoder.
Args:
checkpoint_path: Path to the LTX-2 safetensors checkpoint file
gemma_model_path: Path to Gemma model directory
device: Device to load model on
dtype: Data type for model weights
load_in_8bit: Whether to load the Gemma model in 8-bit precision using bitsandbytes.
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).
is ignored for the Gemma backbone.
Returns:
Loaded GemmaTextEncoder (unified encoder handling V1/V2/V3)
Loaded GemmaTextEncoder
"""
if not Path(gemma_model_path).is_dir():
raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
@@ -212,12 +211,12 @@ def load_text_encoder(
if load_in_8bit:
from ltx_trainer.gemma_8bit import load_8bit_gemma
return load_8bit_gemma(checkpoint_path, gemma_model_path, dtype)
return load_8bit_gemma(gemma_model_path, dtype)
# Standard loading path
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
from ltx_core.text_encoders.gemma import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
GEMMA_LLM_KEY_OPS,
GEMMA_MODEL_OPS,
GemmaTextEncoderConfigurator,
module_ops_from_gemma_root,
@@ -230,15 +229,43 @@ def load_text_encoder(
gemma_weight_paths = [str(p) for p in gemma_model_folder.rglob("*.safetensors")]
text_encoder = SingleGPUModelBuilder(
model_path=(str(checkpoint_path), *gemma_weight_paths),
model_path=tuple(gemma_weight_paths),
model_class_configurator=GemmaTextEncoderConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
model_sd_ops=GEMMA_LLM_KEY_OPS,
module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
).build(device=torch_device, dtype=dtype)
return text_encoder
def load_embeddings_processor(
checkpoint_path: str | Path,
device: Device = "cpu",
dtype: torch.dtype = torch.bfloat16,
) -> "EmbeddingsProcessor":
"""Load the embeddings processor (feature extractor + video/audio connectors).
Args:
checkpoint_path: Path to the LTX-2 safetensors checkpoint file
device: Device to load model on
dtype: Data type for model weights
Returns:
Loaded EmbeddingsProcessor with feature extractor and connectors
"""
from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
from ltx_core.text_encoders.gemma import (
EMBEDDINGS_PROCESSOR_KEY_OPS,
EmbeddingsProcessorConfigurator,
)
torch_device = _to_torch_device(device)
return SingleGPUModelBuilder(
model_path=str(checkpoint_path),
model_class_configurator=EmbeddingsProcessorConfigurator,
model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
).build(device=torch_device, dtype=dtype)
# =============================================================================
# Combined Component Loader
# =============================================================================
@@ -337,7 +364,7 @@ def load_model(
if text_encoder_path is None:
raise ValueError("text_encoder_path must be provided when with_text_encoder=True")
logger.debug("Loading Gemma text encoder...")
text_encoder = load_text_encoder(checkpoint_path, text_encoder_path, torch_device, dtype)
text_encoder = load_text_encoder(text_encoder_path, torch_device, dtype)
# Create scheduler (stateless, no loading needed)
scheduler = LTX2Scheduler()
packages/ltx-trainer/src/ltx_trainer/trainer.py
+34 -24
View File
@@ -27,14 +27,15 @@ from torch.optim.lr_scheduler import (
from torch.utils.data import DataLoader
from torchvision.transforms import functional as F # noqa: N812
from ltx_core.text_encoders.gemma import convert_to_additive_mask
from ltx_trainer import logger
from ltx_trainer.config import LtxTrainerConfig
from ltx_trainer.config_display import print_config
from ltx_trainer.datasets import PrecomputedDataset
from ltx_trainer.gpu_utils import free_gpu_memory, free_gpu_memory_context, get_gpu_memory_gb
from ltx_trainer.hf_hub_utils import push_to_hub
from ltx_trainer.model_loader import load_embeddings_processor, load_text_encoder
from ltx_trainer.model_loader import load_model as load_ltx_model
from ltx_trainer.model_loader import load_text_encoder
from ltx_trainer.progress import TrainingProgress
from ltx_trainer.quantization import quantize_model
from ltx_trainer.timestep_samplers import SAMPLERS
@@ -320,8 +321,8 @@ class LtxvTrainer:
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(
additive_mask = convert_to_additive_mask(mask, video_features.dtype)
video_embeds, audio_embeds, attention_mask = self._embeddings_processor.create_embeddings(
video_features, audio_features, additive_mask
)
@@ -346,26 +347,31 @@ class LtxvTrainer:
@free_gpu_memory_context(after=True)
def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
"""Load text encoder, computes and returns validation embeddings."""
"""Load text encoder + embeddings processor, compute and cache validation embeddings."""
# This method:
# 1. Loads the text encoder on GPU
# 2. If validation prompts are configured, computes and caches their embeddings
# 3. Unloads the heavy Gemma model while keeping the lightweight embedding connectors
# The text encoder is kept (as self._text_encoder) but with model/tokenizer/feature_extractor
# set to None. Only the embedding connectors remain for use during training.
# 1. Loads the pure Gemma text encoder on GPU
# 2. Loads the embeddings processor (feature extractor + connectors)
# 3. If validation prompts are configured, computes and caches their embeddings
# 4. Unloads the Gemma model entirely, keeps the embeddings processor for training
# Load text encoder on GPU
# Load text encoder (pure Gemma LLM) on GPU
logger.debug("Loading text encoder...")
self._text_encoder = load_text_encoder(
checkpoint_path=self._config.model.model_path,
text_encoder = load_text_encoder(
gemma_model_path=self._config.model.text_encoder_path,
device="cuda",
dtype=torch.bfloat16,
load_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
)
# Load embeddings processor (feature extractor + connectors)
logger.debug("Loading embeddings processor...")
self._embeddings_processor = load_embeddings_processor(
checkpoint_path=self._config.model.model_path,
device="cuda",
dtype=torch.bfloat16,
)
# Cache validation embeddings if prompts are configured
cached_embeddings = None
if self._config.validation.prompts:
@@ -373,22 +379,26 @@ class LtxvTrainer:
cached_embeddings = []
with torch.inference_mode():
for prompt in self._config.validation.prompts:
v_ctx_pos, a_ctx_pos, _ = self._text_encoder(prompt)
v_ctx_neg, a_ctx_neg, _ = self._text_encoder(self._config.validation.negative_prompt)
pos_hs, pos_mask = text_encoder.encode(prompt)
pos_out = self._embeddings_processor.process_hidden_states(pos_hs, pos_mask)
neg_hs, neg_mask = text_encoder.encode(self._config.validation.negative_prompt)
neg_out = self._embeddings_processor.process_hidden_states(neg_hs, neg_mask)
cached_embeddings.append(
CachedPromptEmbeddings(
video_context_positive=v_ctx_pos.cpu(),
audio_context_positive=a_ctx_pos.cpu(),
video_context_negative=v_ctx_neg.cpu() if v_ctx_neg is not None else None,
audio_context_negative=a_ctx_neg.cpu() if a_ctx_neg is not None else None,
video_context_positive=pos_out.video_encoding.cpu(),
audio_context_positive=pos_out.audio_encoding.cpu(),
video_context_negative=neg_out.video_encoding.cpu(),
audio_context_negative=(
neg_out.audio_encoding.cpu() if neg_out.audio_encoding is not None else None
),
)
)
# 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 = None
# Unload Gemma model and feature extractor, keep only connectors for training
del text_encoder
self._embeddings_processor.feature_extractor = None
logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
return cached_embeddings
@@ -426,7 +436,7 @@ class LtxvTrainer:
self._scheduler = components.scheduler
self._audio_vae = components.audio_vae_decoder
self._vocoder = components.vocoder
# Note: self._text_encoder was set in _load_text_encoder_and_cache_embeddings
# Note: self._embeddings_processor was set in _load_text_encoder_and_cache_embeddings
# Determine initial dtype based on training mode.
# Note: For FSDP + LoRA, we'll cast to FP32 later in _prepare_models_for_training()
packages/ltx-trainer/src/ltx_trainer/validation_sampler.py
+18 -8
View File
@@ -37,6 +37,7 @@ if TYPE_CHECKING:
from ltx_core.model.transformer import LTXModel
from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
from ltx_core.text_encoders.gemma import GemmaTextEncoder
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
@@ -128,21 +129,24 @@ class ValidationSampler:
audio_decoder: "AudioDecoder | None" = None,
vocoder: "Vocoder | None" = None,
sampling_context: SamplingContext | None = None,
embeddings_processor: "EmbeddingsProcessor | None" = None,
):
"""Initialize the validation sampler.
Args:
transformer: LTX-2 transformer model
vae_decoder: Video VAE decoder
vae_encoder: Video VAE encoder (for image/video conditioning), can be None if not needed
text_encoder: Gemma text encoder with embeddings connector (optional if cached_embeddings in config)
text_encoder: Gemma text encoder (optional if cached_embeddings in config)
audio_decoder: Optional audio VAE decoder (for audio generation)
vocoder: Optional vocoder (for audio generation)
sampling_context: Optional SamplingContext for progress display during denoising
embeddings_processor: Optional embeddings processor (required if text_encoder provided)
"""
self._transformer = transformer
self._vae_decoder = vae_decoder
self._vae_encoder = vae_encoder
self._text_encoder = text_encoder
self._embeddings_processor = embeddings_processor
self._audio_decoder = audio_decoder
self._vocoder = vocoder
self._sampling_context = sampling_context
@@ -677,6 +681,8 @@ class ValidationSampler:
# Validate prompt embedding source
if config.cached_embeddings is None and self._text_encoder is None:
raise ValueError("Either text_encoder or config.cached_embeddings must be provided")
if config.cached_embeddings is None and self._embeddings_processor is None:
raise ValueError("embeddings_processor is required when encoding prompts on-the-fly")
def _get_prompt_embeddings(
self, config: GenerationConfig, device: torch.device
@@ -697,18 +703,22 @@ class ValidationSampler:
def _encode_prompts(
self, config: GenerationConfig, device: torch.device
) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
"""Encode positive and negative prompts using the text encoder."""
"""Encode positive and negative prompts using the text encoder + embeddings processor."""
self._text_encoder.to(device)
v_ctx_pos, a_ctx_pos, _ = self._text_encoder(config.prompt)
self._embeddings_processor.to(device)
pos_hs, pos_mask = self._text_encoder.encode(config.prompt)
pos_out = self._embeddings_processor.process_hidden_states(pos_hs, pos_mask)
v_ctx_pos, a_ctx_pos = pos_out.video_encoding, pos_out.audio_encoding
v_ctx_neg, a_ctx_neg = None, None
if config.guidance_scale != 1.0:
v_ctx_neg, a_ctx_neg, _ = self._text_encoder(config.negative_prompt)
neg_hs, neg_mask = self._text_encoder.encode(config.negative_prompt)
neg_out = self._embeddings_processor.process_hidden_states(neg_hs, neg_mask)
v_ctx_neg, a_ctx_neg = neg_out.video_encoding, neg_out.audio_encoding
# Move the base Gemma model to CPU but keep embeddings connectors on GPU
# as this module is also used during training
# Move the base Gemma model to CPU
self._text_encoder.model.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