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
README.md
+13 -14
View File
@@ -1,7 +1,7 @@
# LTX-2
[![Website](https://img.shields.io/badge/Website-LTX-181717?logo=google-chrome)](https://ltx.io)
[![Model](https://img.shields.io/badge/HuggingFace-Model-orange?logo=huggingface)](https://huggingface.co/Lightricks/LTX-2)
[![Model](https://img.shields.io/badge/HuggingFace-Model-orange?logo=huggingface)](https://huggingface.co/Lightricks/LTX-2.3)
[![Demo](https://img.shields.io/badge/Demo-Try%20Now-brightgreen?logo=vercel)](https://app.ltx.studio/ltx-2-playground/i2v)
[![Paper](https://img.shields.io/badge/Paper-PDF-EC1C24?logo=adobeacrobatreader&logoColor=white)](https://arxiv.org/abs/2601.03233)
[![Discord](https://img.shields.io/badge/Join-Discord-5865F2?logo=discord)](https://discord.gg/ltxplatform)
@@ -26,30 +26,29 @@ source .venv/bin/activate
### Required Models
Download the following models from the [LTX-2 HuggingFace repository](https://huggingface.co/Lightricks/LTX-2):
Download the following models from the [LTX-2.3 HuggingFace repository](https://huggingface.co/Lightricks/LTX-2.3):
**LTX-2 Model Checkpoint** (choose and download one of the following)
* [`ltx-2-19b-dev-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev-fp8.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-dev-fp8.safetensors)
* [`ltx-2-19b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-dev.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-dev.safetensors)
* [`ltx-2-19b-distilled.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-distilled.safetensors)
* [`ltx-2-19b-distilled-fp8.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-fp8.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-distilled-fp8.safetensors)
**LTX-2.3 Model Checkpoint** (choose and download one of the following)
* [`ltx-2.3-22b-dev.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-dev.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-dev.safetensors)
* [`ltx-2.3-22b-distilled.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled.safetensors)
**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository
* [`ltx-2-spatial-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-spatial-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-spatial-upscaler-x2-1.0.safetensors)
* [`ltx-2.3-spatial-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x2-1.0.safetensors)
* [`ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors)
**Temporal Upscaler** - Supported by the model and will be required for future pipeline implementations
* [`ltx-2-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-temporal-upscaler-x2-1.0.safetensors)
* [`ltx-2.3-temporal-upscaler-x2-1.0.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-temporal-upscaler-x2-1.0.safetensors)
**Distilled LoRA** - Required for current two-stage pipeline implementations in this repository (except DistilledPipeline and ICLoraPipeline)
* [`ltx-2-19b-distilled-lora-384.safetensors`](https://huggingface.co/Lightricks/LTX-2/blob/main/ltx-2-19b-distilled-lora-384.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2/resolve/main/ltx-2-19b-distilled-lora-384.safetensors)
* [`ltx-2.3-22b-distilled-lora-384.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-lora-384.safetensors)
**Gemma Text Encoder** (download all assets from the repository)
* [`Gemma 3`](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized/tree/main)
**LoRAs**
* [`LTX-2-19b-IC-LoRA-Canny-Control`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Canny-Control) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Canny-Control/resolve/main/ltx-2-19b-ic-lora-canny-control.safetensors)
* [`LTX-2-19b-IC-LoRA-Depth-Control`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Depth-Control) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Depth-Control/resolve/main/ltx-2-19b-ic-lora-depth-control.safetensors)
* [`LTX-2.3-22b-IC-LoRA-Union-Control`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Union-Control/resolve/main/ltx-2.3-22b-ic-lora-union-control-ref0.5.safetensors)
* [`LTX-2.3-22b-IC-LoRA-Inpainting`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Inpainting) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Inpainting/resolve/main/ltx-2.3-22b-ic-lora-inpainting.safetensors)
* [`LTX-2.3-22b-IC-LoRA-Motion-Track-Control`](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control) - [Download](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Motion-Track-Control/resolve/main/ltx-2.3-22b-ic-lora-motion-track-control-ref0.5.safetensors)
* [`LTX-2-19b-IC-LoRA-Detailer`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Detailer/resolve/main/ltx-2-19b-ic-lora-detailer.safetensors)
* [`LTX-2-19b-IC-LoRA-Pose-Control`](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-IC-LoRA-Pose-Control/resolve/main/ltx-2-19b-ic-lora-pose-control.safetensors)
* [`LTX-2-19b-LoRA-Camera-Control-Dolly-In`](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In) - [Download](https://huggingface.co/Lightricks/LTX-2-19b-LoRA-Camera-Control-Dolly-In/resolve/main/ltx-2-19b-lora-camera-control-dolly-in.safetensors)
@@ -63,7 +62,7 @@ Download the following models from the [LTX-2 HuggingFace repository](https://hu
### Available Pipelines
* **[TI2VidTwoStagesPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py)** - Production-quality text/image-to-video with 2x upsampling (recommended)
* **[TI2VidTwoStagesRes2sPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_res2s.py)** - Same two-stage flow as above but uses the res_2s second-order sampler (fewer steps, different quality/speed trade-off)
* **[TI2VidTwoStagesHQPipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py)** - Same two-stage flow as above but uses the res_2s second-order sampler (fewer steps, better quality)
* **[TI2VidOneStagePipeline](packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py)** - Single-stage generation for quick prototyping
* **[DistilledPipeline](packages/ltx-pipelines/src/ltx_pipelines/distilled.py)** - Fastest inference with 8 predefined sigmas
* **[ICLoraPipeline](packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py)** - Video-to-video and image-to-video transformations (uses distilled model.)
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,
packages/ltx-pipelines/README.md
+5 -5
View File
@@ -56,7 +56,7 @@ python -m ltx_pipelines.ti2vid_two_stages --help
Available pipeline modules:
- `ltx_pipelines.ti2vid_two_stages` - Two-stage text/image-to-video (recommended).
- `ltx_pipelines.ti2vid_two_stages_res2s` - Two-stage text/image-to-video (use 2 times less steps).
- `ltx_pipelines.ti2vid_two_stages_hq` - Two-stage text/image-to-video (different sampler, better quality).
- `ltx_pipelines.ti2vid_one_stage` - Single-stage text/image-to-video.
- `ltx_pipelines.distilled` - Fast text/image-to-video pipeline using only the distilled model.
- `ltx_pipelines.ic_lora` - Video-to-video with IC-LoRA.
@@ -94,14 +94,14 @@ Do you need to condition on existing images/videos?
└─ YES → Use DistilledPipeline (with 8 predefined sigmas)
```
> **Note:** [`TI2VidOneStagePipeline`](src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](src/ltx_pipelines/ti2vid_two_stages.py), [`TI2VidTwoStagesRes2sPipeline`](src/ltx_pipelines/ti2vid_two_stages_res2s.py), [`ICLoraPipeline`](src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](src/ltx_pipelines/keyframe_interpolation.py), [`A2VidPipelineTwoStage`](src/ltx_pipelines/a2vid_two_stage.py), or [`DistilledPipeline`](src/ltx_pipelines/distilled.py)). For editing existing videos, use [`RetakePipeline`](src/ltx_pipelines/retake.py).
> **Note:** [`TI2VidOneStagePipeline`](src/ltx_pipelines/ti2vid_one_stage.py) is primarily for educational purposes. For best quality, use two-stage pipelines ([`TI2VidTwoStagesPipeline`](src/ltx_pipelines/ti2vid_two_stages.py), [`TI2VidTwoStagesHQPipeline`](src/ltx_pipelines/ti2vid_two_stages_hq.py), [`ICLoraPipeline`](src/ltx_pipelines/ic_lora.py), [`KeyframeInterpolationPipeline`](src/ltx_pipelines/keyframe_interpolation.py), [`A2VidPipelineTwoStage`](src/ltx_pipelines/a2vid_two_stage.py), or [`DistilledPipeline`](src/ltx_pipelines/distilled.py)). For editing existing videos, use [`RetakePipeline`](src/ltx_pipelines/retake.py).
### Features Comparison
| Pipeline | Stages | [Multimodal Guidance](#%EF%B8%8F-multimodal-guidance) | Upsampling | Conditioning | Best For |
| -------- | ------ | --- | ---------- | ------------- | -------- |
| **TI2VidTwoStagesPipeline** | 2 | ✅ | ✅ | Image | **Production quality** (recommended) |
| **TI2VidTwoStagesRes2sPipeline** | 2 | ✅ | ✅ | Image | Same as above, res_2s sampler (fewer steps) |
| **TI2VidTwoStagesHQPipeline** | 2 | ✅ | ✅ | Image | Same as above, res_2s sampler (higher quality) |
| **TI2VidOneStagePipeline** | 1 | ✅ | ❌ | Image | Educational, prototyping |
| **DistilledPipeline** | 2 | ❌ | ✅ | Image | Fastest inference (8 sigmas) |
| **ICLoraPipeline** | 2 | ✅ | ✅ | Image + Video | Video-to-video transformations |
@@ -125,11 +125,11 @@ Two-stage generation: Stage 1 generates low-resolution video with [multimodal gu
---
### 2. TI2VidTwoStagesRes2sPipeline
### 2. TI2VidTwoStagesHQPipeline
**Best for:** Same two-stage text/image-to-video as TI2VidTwoStagesPipeline but with a different sampler and step count.
**Source**: [`src/ltx_pipelines/ti2vid_two_stages_res2s.py`](src/ltx_pipelines/ti2vid_two_stages_res2s.py)
**Source**: [`src/ltx_pipelines/ti2vid_two_stages_hq.py`](src/ltx_pipelines/ti2vid_two_stages_hq.py)
Uses the **res_2s** second-order sampler instead of Euler. Same stage structure (stage 1 at target resolution with CFG, stage 2 upsampling with distilled LoRA) and image conditioning support. Typically allows fewer steps for comparable quality; trade-offs differ from the default Euler-based pipeline.
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+44 -44
View File
@@ -14,7 +14,6 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import default_2_stage_arg_parser
@@ -24,10 +23,10 @@ from ltx_pipelines.utils.constants import (
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_video_only,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
simple_denoising_func,
)
@@ -69,7 +68,7 @@ class A2VidPipelineTwoStage:
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
loras=distilled_lora,
)
@@ -103,16 +102,14 @@ class A2VidPipelineTwoStage:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, images[0][0] if len(images) > 0 else None)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, _ = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, _ = ctx_n.video_encoding, ctx_n.audio_encoding
# Encode audio.
decoded_audio = decode_audio_from_file(audio_path, self.device, audio_start_time, audio_max_duration)
@@ -120,8 +117,28 @@ class A2VidPipelineTwoStage:
audio_shape = AudioLatentShape.from_duration(batch=1, duration=num_frames / frame_rate, channels=8, mel_bins=16)
encoded_audio_latent = encoded_audio_latent[:, :, : audio_shape.frames]
# Stage 1: encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del video_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation with audio conditioning.
transformer = self.stage_1_model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
@@ -147,24 +164,6 @@ class A2VidPipelineTwoStage:
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state = denoise_video_only(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -183,12 +182,23 @@ class A2VidPipelineTwoStage:
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
video_encoder = self.stage_1_model_ledger.video_encoder()
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
del video_encoder
torch.cuda.synchronize()
cleanup_memory()
@@ -210,15 +220,6 @@ class A2VidPipelineTwoStage:
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state = denoise_video_only(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
@@ -236,7 +237,6 @@ class A2VidPipelineTwoStage:
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
@@ -278,7 +278,7 @@ def main() -> None:
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+12 -15
View File
@@ -12,7 +12,6 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger, euler_denoising_loop
from ltx_pipelines.utils.args import (
@@ -28,10 +27,10 @@ from ltx_pipelines.utils.constants import (
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.media_io import encode_video
@@ -93,15 +92,13 @@ class DistilledPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, images[0][0] if len(images) > 0 else None)
context_p = encode_text(text_encoder, prompts=[prompt])[0]
video_context, audio_context = context_p
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
(ctx_p,) = encode_prompts(
[prompt],
self.model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
# Stage 1: Initial low resolution video generation.
video_encoder = self.model_ledger.video_encoder()
@@ -130,7 +127,7 @@ class DistilledPipeline:
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
@@ -161,7 +158,7 @@ class DistilledPipeline:
stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
@@ -209,7 +206,7 @@ def main() -> None:
distilled_checkpoint_path=args.distilled_checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+33 -36
View File
@@ -19,17 +19,16 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, VideoEncoder, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoLatentShape, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
simple_denoising_func,
)
from ltx_pipelines.utils.args import (
@@ -169,20 +168,37 @@ class ICLoraPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
video_context, audio_context = encode_text(text_encoder, prompts=[prompt])[0]
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
(ctx_p,) = encode_prompts(
[prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
# Stage 1: Initial low resolution video generation.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
# Encode conditionings before loading transformer to reduce peak VRAM
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = self._create_conditionings(
images=images,
video_conditioning=video_conditioning,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
num_frames=num_frames,
conditioning_attention_strength=conditioning_attention_strength,
conditioning_attention_mask=conditioning_attention_mask,
)
transformer = self.stage_1_model_ledger.transformer()
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
@@ -201,25 +217,6 @@ class ICLoraPipeline:
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = self._create_conditionings(
images=images,
video_conditioning=video_conditioning,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
num_frames=num_frames,
conditioning_attention_strength=conditioning_attention_strength,
conditioning_attention_mask=conditioning_attention_mask,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -278,7 +275,7 @@ class ICLoraPipeline:
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
@@ -340,7 +337,7 @@ class ICLoraPipeline:
Returns:
List of conditioning items. IC-LoRA conditionings are appended last.
"""
conditionings = image_conditionings_by_replacing_latent(
conditionings = combined_image_conditionings(
images=images,
height=height,
width=width,
@@ -510,7 +507,7 @@ def main() -> None:
distilled_checkpoint_path=args.distilled_checkpoint_path,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+12 -16
View File
@@ -18,7 +18,6 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
@@ -27,7 +26,7 @@ from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
denoise_audio_video,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_adding_guiding_latent,
multi_modal_guider_factory_denoising_func,
@@ -71,7 +70,7 @@ class KeyframeInterpolationPipeline:
loras=loras,
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
loras=distilled_lora,
)
self.pipeline_components = PipelineComponents(
@@ -102,18 +101,15 @@ class KeyframeInterpolationPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
@@ -252,7 +248,7 @@ def main() -> None:
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+34 -22
View File
@@ -20,7 +20,6 @@ from ltx_core.model.audio_vae import encode_audio as vae_encode_audio
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.tools import LatentTools
from ltx_core.types import (
Audio,
@@ -30,10 +29,11 @@ from ltx_core.types import (
VideoPixelShape,
)
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.args import QuantizationAction
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.helpers import (
cleanup_memory,
generate_enhanced_prompt,
encode_prompts,
get_device,
multi_modal_guider_denoising_func,
noise_audio_state,
@@ -197,7 +197,6 @@ class RetakePipeline:
# Public entry point #
# --------------------------------------------------------------------- #
@torch.inference_mode()
def __call__( # noqa: PLR0913, PLR0915
self,
video_path: str,
@@ -214,6 +213,7 @@ class RetakePipeline:
regenerate_audio: bool = True,
enhance_prompt: bool = False,
distilled: bool = False,
tiling_config: TilingConfig | None = None,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
"""Regenerate ``[start_time, end_time]`` of the source video (retake).
Parameters
@@ -321,22 +321,17 @@ class RetakePipeline:
del audio_encoder
cleanup_memory()
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(text_encoder, prompt, None, seed=effective_seed)
prompts_to_encode = [prompt] if distilled else [prompt, negative_prompt]
contexts = encode_prompts(
prompts_to_encode,
self.model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_seed=effective_seed,
)
if distilled:
# Distilled mode: single prompt, no negative
context_p = encode_text(text_encoder, prompts=[prompt])[0]
v_context_p, a_context_p = context_p
else:
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
v_context_p, a_context_p = contexts[0].video_encoding, contexts[0].audio_encoding
if not distilled:
v_context_n, a_context_n = contexts[1].video_encoding, contexts[1].audio_encoding
transformer = self.model_ledger.transformer()
@@ -412,7 +407,9 @@ class RetakePipeline:
del transformer
cleanup_memory()
decoded_video = vae_decode_video(video_state.latent, self.model_ledger.video_decoder(), generator=generator)
decoded_video = vae_decode_video(
video_state.latent, self.model_ledger.video_decoder(), tiling_config, generator
)
decoded_audio = vae_decode_audio(
audio_state.latent, self.model_ledger.audio_decoder(), self.model_ledger.vocoder()
)
@@ -420,6 +417,7 @@ class RetakePipeline:
return decoded_video, decoded_audio
@torch.inference_mode()
def main() -> None:
"""CLI entry point for retake (regenerate a time region)."""
logging.getLogger().setLevel(logging.INFO)
@@ -433,6 +431,15 @@ def main() -> None:
parser.add_argument("--gemma-root", type=str, required=True, help="Path to Gemma text encoder weights.")
parser.add_argument("--seed", type=int, default=42, help="Random seed. Use -1 for a random seed.")
parser.add_argument("--loras", nargs="*", default=[], help="LoRA paths (optional).")
parser.add_argument(
"--quantization",
dest="quantization",
action=QuantizationAction,
nargs="+",
metavar=("POLICY", "AMAX_PATH"),
default=None,
help="Quantization policy: fp8-cast or fp8-scaled-mm [AMAX_PATH].",
)
args = parser.parse_args()
if args.start_time >= args.end_time:
@@ -452,16 +459,21 @@ def main() -> None:
pipeline = RetakePipeline(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
loras=args.loras or [],
loras=tuple(args.loras) if args.loras else (),
quantization=args.quantization,
)
params = detect_params(args.checkpoint_path)
tiling_config = TilingConfig.default()
video_iter, audio = pipeline(
video_path=args.video_path,
prompt=args.prompt,
start_time=args.start_time,
end_time=args.end_time,
seed=args.seed,
video_guider_params=params.video_guider_params,
audio_guider_params=params.audio_guider_params,
tiling_config=tiling_config,
)
tiling_config = TilingConfig.default()
video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
encode_video(
video=video_iter,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+25 -25
View File
@@ -16,17 +16,16 @@ from ltx_core.loader import LoraPathStrengthAndSDOps
from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_factory_denoising_func,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_1_stage_arg_parser, detect_checkpoint_path
@@ -91,21 +90,32 @@ class TI2VidOneStagePipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
video_encoder = self.model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del text_encoder
del video_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.model_ledger.video_encoder()
transformer = self.model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
@@ -135,16 +145,6 @@ class TI2VidOneStagePipeline:
),
)
stage_1_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -178,7 +178,7 @@ def main() -> None:
pipeline = TI2VidOneStagePipeline(
checkpoint_path=args.checkpoint_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
video, audio = pipeline(
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+43 -41
View File
@@ -18,17 +18,16 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
euler_denoising_loop,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_factory_denoising_func,
simple_denoising_func,
)
@@ -71,7 +70,7 @@ class TI2VidTwoStagesPipeline:
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
self.stage_2_model_ledger = self.stage_1_model_ledger.with_additional_loras(
loras=distilled_lora,
)
@@ -103,21 +102,38 @@ class TI2VidTwoStagesPipeline:
stepper = EulerDiffusionStep()
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del text_encoder
del video_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
@@ -144,21 +160,6 @@ class TI2VidTwoStagesPipeline:
),
)
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
width=width // 2,
height=height // 2,
fps=frame_rate,
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -176,12 +177,23 @@ class TI2VidTwoStagesPipeline:
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
video_encoder = self.stage_1_model_ledger.video_encoder()
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
del video_encoder
torch.cuda.synchronize()
cleanup_memory()
@@ -203,15 +215,6 @@ class TI2VidTwoStagesPipeline:
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
@@ -229,7 +232,6 @@ class TI2VidTwoStagesPipeline:
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
@@ -253,7 +255,7 @@ def main() -> None:
distilled_lora=args.distilled_lora,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_res2s.py → packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+62 -48
View File
@@ -14,30 +14,29 @@ from ltx_core.model.upsampler import upsample_video
from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
from ltx_core.model.video_vae import decode_video as vae_decode_video
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import encode_text
from ltx_core.tools import VideoLatentShape
from ltx_core.types import Audio, LatentState, VideoPixelShape
from ltx_pipelines.utils import (
ModelLedger,
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
generate_enhanced_prompt,
encode_prompts,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
res2s_audio_video_denoising_loop,
simple_denoising_func,
)
from ltx_pipelines.utils.args import ImageConditioningInput, default_2_stage_arg_parser, detect_checkpoint_path
from ltx_pipelines.utils.constants import STAGE_2_DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.args import ImageConditioningInput, hq_2_stage_arg_parser
from ltx_pipelines.utils.constants import LTX_2_3_HQ_PARAMS, STAGE_2_DISTILLED_SIGMA_VALUES
from ltx_pipelines.utils.media_io import encode_video
from ltx_pipelines.utils.types import PipelineComponents
device = get_device()
class TI2VidTwoStagesRes2sPipeline:
class TI2VidTwoStagesHQPipeline:
"""
Two-stage text/image-to-video generation pipeline using the res_2s sampler.
Same structure as :class:`TI2VidTwoStagesPipeline`: stage 1 generates video at
@@ -53,26 +52,38 @@ class TI2VidTwoStagesRes2sPipeline:
self,
checkpoint_path: str,
distilled_lora: list[LoraPathStrengthAndSDOps],
distilled_lora_strength_stage_1: float,
distilled_lora_strength_stage_2: float,
spatial_upsampler_path: str,
gemma_root: str,
loras: list[LoraPathStrengthAndSDOps],
loras: tuple[LoraPathStrengthAndSDOps, ...],
device: str = device,
quantization: QuantizationPolicy | None = None,
):
self.device = device
self.dtype = torch.bfloat16
distilled_lora_stage_1 = LoraPathStrengthAndSDOps(
path=distilled_lora[0].path,
strength=distilled_lora_strength_stage_1,
sd_ops=distilled_lora[0].sd_ops,
)
distilled_lora_stage_2 = LoraPathStrengthAndSDOps(
path=distilled_lora[0].path,
strength=distilled_lora_strength_stage_2,
sd_ops=distilled_lora[0].sd_ops,
)
self.stage_1_model_ledger = ModelLedger(
dtype=self.dtype,
device=device,
checkpoint_path=checkpoint_path,
gemma_root_path=gemma_root,
spatial_upsampler_path=spatial_upsampler_path,
loras=loras,
loras=(*loras, distilled_lora_stage_1),
quantization=quantization,
)
self.stage_2_model_ledger = self.stage_1_model_ledger.with_loras(
loras=distilled_lora,
loras=(*loras, distilled_lora_stage_2),
)
self.pipeline_components = PipelineComponents(
@@ -103,23 +114,18 @@ class TI2VidTwoStagesRes2sPipeline:
noiser = GaussianNoiser(generator=generator)
dtype = torch.bfloat16
text_encoder = self.stage_1_model_ledger.text_encoder()
if enhance_prompt:
prompt = generate_enhanced_prompt(
text_encoder, prompt, images[0][0] if len(images) > 0 else None, seed=seed
)
context_p, context_n = encode_text(text_encoder, prompts=[prompt, negative_prompt])
v_context_p, a_context_p = context_p
v_context_n, a_context_n = context_n
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
# Stage 1: Initial low resolution video generation.
video_encoder = self.stage_1_model_ledger.video_encoder()
transformer = self.stage_1_model_ledger.transformer()
ctx_p, ctx_n = encode_prompts(
[prompt, negative_prompt],
self.stage_1_model_ledger,
enhance_first_prompt=enhance_prompt,
enhance_prompt_image=images[0][0] if len(images) > 0 else None,
enhance_prompt_seed=seed,
)
v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: encode image conditionings with the VAE encoder, then free it
# before loading the transformer to reduce peak VRAM.
stage_1_output_shape = VideoPixelShape(
batch=1,
frames=num_frames,
@@ -127,6 +133,21 @@ class TI2VidTwoStagesRes2sPipeline:
height=height // 2,
fps=frame_rate,
)
video_encoder = self.stage_1_model_ledger.video_encoder()
stage_1_conditionings = combined_image_conditionings(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del video_encoder
cleanup_memory()
transformer = self.stage_1_model_ledger.transformer()
empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
stepper = Res2sDiffusionStep()
sigmas = (
@@ -158,14 +179,6 @@ class TI2VidTwoStagesRes2sPipeline:
),
)
stage_1_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_1_output_shape.height,
width=stage_1_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_1_output_shape,
conditionings=stage_1_conditionings,
@@ -183,13 +196,24 @@ class TI2VidTwoStagesRes2sPipeline:
cleanup_memory()
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
video_encoder = self.stage_1_model_ledger.video_encoder()
upscaled_video_latent = upsample_video(
latent=video_state.latent[:1],
video_encoder=video_encoder,
upsampler=self.stage_2_model_ledger.spatial_upsampler(),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = combined_image_conditionings(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
torch.cuda.synchronize()
del video_encoder
cleanup_memory()
transformer = self.stage_2_model_ledger.transformer()
@@ -210,15 +234,6 @@ class TI2VidTwoStagesRes2sPipeline:
),
)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = image_conditionings_by_replacing_latent(
images=images,
height=stage_2_output_shape.height,
width=stage_2_output_shape.width,
video_encoder=video_encoder,
dtype=dtype,
device=self.device,
)
video_state, audio_state = denoise_audio_video(
output_shape=stage_2_output_shape,
conditionings=stage_2_conditionings,
@@ -236,7 +251,6 @@ class TI2VidTwoStagesRes2sPipeline:
torch.cuda.synchronize()
del transformer
del video_encoder
cleanup_memory()
decoded_video = vae_decode_video(
@@ -251,16 +265,16 @@ class TI2VidTwoStagesRes2sPipeline:
@torch.inference_mode()
def main() -> None:
logging.getLogger().setLevel(logging.INFO)
checkpoint_path = detect_checkpoint_path()
params = detect_params(checkpoint_path)
parser = default_2_stage_arg_parser(params=params)
parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS)
args = parser.parse_args()
pipeline = TI2VidTwoStagesRes2sPipeline(
pipeline = TI2VidTwoStagesHQPipeline(
checkpoint_path=args.checkpoint_path,
distilled_lora=args.distilled_lora,
distilled_lora_strength_stage_1=args.distilled_lora_strength_stage_1,
distilled_lora_strength_stage_2=args.distilled_lora_strength_stage_2,
spatial_upsampler_path=args.spatial_upsampler_path,
gemma_root=args.gemma_root,
loras=args.lora,
loras=tuple(args.lora) if args.lora else (),
quantization=args.quantization,
)
tiling_config = TilingConfig.default()
packages/ltx-pipelines/src/ltx_pipelines/utils/__init__.py
+4 -2
View File
@@ -1,10 +1,11 @@
from ltx_pipelines.utils.helpers import (
assert_resolution,
cleanup_memory,
combined_image_conditionings,
denoise_audio_video,
encode_prompts,
generate_enhanced_prompt,
get_device,
image_conditionings_by_replacing_latent,
multi_modal_guider_denoising_func,
multi_modal_guider_factory_denoising_func,
simple_denoising_func,
@@ -20,12 +21,13 @@ __all__ = [
"ModelLedger",
"assert_resolution",
"cleanup_memory",
"combined_image_conditionings",
"denoise_audio_video",
"encode_prompts",
"euler_denoising_loop",
"generate_enhanced_prompt",
"get_device",
"gradient_estimating_euler_denoising_loop",
"image_conditionings_by_replacing_latent",
"multi_modal_guider_denoising_func",
"multi_modal_guider_factory_denoising_func",
"res2s_audio_video_denoising_loop",
packages/ltx-pipelines/src/ltx_pipelines/utils/args.py
+18
View File
@@ -8,6 +8,7 @@ from ltx_pipelines.utils.constants import (
DEFAULT_IMAGE_CRF,
DEFAULT_LORA_STRENGTH,
DEFAULT_NEGATIVE_PROMPT,
LTX_2_3_HQ_PARAMS,
LTX_2_3_PARAMS,
PipelineParams,
)
@@ -457,6 +458,23 @@ def default_2_stage_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argpa
return parser
def hq_2_stage_arg_parser(params: PipelineParams = LTX_2_3_HQ_PARAMS) -> argparse.ArgumentParser:
parser = default_2_stage_arg_parser(params=params)
parser.add_argument(
"--distilled-lora-strength-stage-1",
type=float,
default=0.25,
help=(f"Strength of the distilled LoRA used in the first stage (default: {0.25})."),
)
parser.add_argument(
"--distilled-lora-strength-stage-2",
type=float,
default=0.5,
help=(f"Strength of the distilled LoRA used in the second stage (default: {0.5})."),
)
return parser
def default_2_stage_distilled_arg_parser(params: PipelineParams = LTX_2_3_PARAMS) -> argparse.ArgumentParser:
parser = basic_arg_parser(params=params, distilled=True)
parser.set_defaults(height=params.stage_2_height, width=params.stage_2_width)
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+21
View File
@@ -71,6 +71,27 @@ LTX_2_3_PARAMS = replace(
video_guider_params=replace(LTX_2_PARAMS.video_guider_params, stg_blocks=[28]),
audio_guider_params=replace(LTX_2_PARAMS.audio_guider_params, stg_blocks=[28]),
)
LTX_2_3_HQ_PARAMS = PipelineParams(
num_inference_steps=15,
stage_1_height=1088 // 2,
stage_1_width=1920 // 2,
video_guider_params=MultiModalGuiderParams(
cfg_scale=3.0,
stg_scale=0.0,
rescale_scale=0.45,
modality_scale=3.0,
skip_step=0,
stg_blocks=[],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=7.0,
stg_scale=0.0,
rescale_scale=1.0,
modality_scale=3.0,
skip_step=0,
stg_blocks=[],
),
)
DEFAULT_LORA_STRENGTH = 1.0
DEFAULT_IMAGE_CRF = 33
packages/ltx-pipelines/src/ltx_pipelines/utils/helpers.py
+79
View File
@@ -21,6 +21,7 @@ from ltx_core.guidance.perturbations import (
from ltx_core.model.transformer import Modality, X0Model
from ltx_core.model.video_vae import VideoEncoder
from ltx_core.text_encoders.gemma import GemmaTextEncoder
from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessorOutput
from ltx_core.tools import AudioLatentTools, LatentTools, VideoLatentTools
from ltx_core.types import AudioLatentShape, LatentState, VideoLatentShape, VideoPixelShape
from ltx_pipelines.utils.args import ImageConditioningInput
@@ -44,6 +45,84 @@ def cleanup_memory() -> None:
torch.cuda.synchronize()
def encode_prompts(
prompts: list[str],
model_ledger: object,
*,
enhance_prompt_image: str | None = None,
enhance_prompt_seed: int = 42,
enhance_first_prompt: bool = False,
) -> list[EmbeddingsProcessorOutput]:
"""Encode prompts through Gemma → embeddings processor, freeing each after use.
Loads the text encoder from *model_ledger*, optionally enhances the first
prompt, encodes all *prompts*, frees the text encoder, then loads the
embeddings processor to produce the final outputs. Because the text encoder
is loaded and freed entirely within this function, there are no lingering
references that could prevent GPU memory reclamation.
Args:
prompts: Text prompts to encode.
model_ledger: ModelLedger instance (used to load text encoder and embeddings processor).
enhance_prompt_image: Optional image path for prompt enhancement.
enhance_prompt_seed: Seed for prompt enhancement (default 42).
enhance_first_prompt: If True, enhance ``prompts[0]`` before encoding.
Returns:
List of EmbeddingsProcessorOutput, one per prompt.
"""
text_encoder = model_ledger.text_encoder()
if enhance_first_prompt:
prompts = list(prompts)
prompts[0] = generate_enhanced_prompt(text_encoder, prompts[0], enhance_prompt_image, seed=enhance_prompt_seed)
raw_outputs = [text_encoder.encode(p) for p in prompts]
torch.cuda.synchronize()
del text_encoder
cleanup_memory()
embeddings_processor = model_ledger.gemma_embeddings_processor()
results: list[EmbeddingsProcessorOutput] = [
embeddings_processor.process_hidden_states(hs, mask) for hs, mask in raw_outputs
]
del embeddings_processor
cleanup_memory()
return results
def combined_image_conditionings(
images: list[ImageConditioningInput],
height: int,
width: int,
video_encoder: VideoEncoder,
dtype: torch.dtype,
device: torch.device,
) -> list[ConditioningItem]:
"""Create a list of conditionings by replacing the latent at the first frame with the encoded image if present
and using other encoded images as the keyframe conditionings."""
conditionings = []
for img in images:
image = load_image_conditioning(
image_path=img.path,
height=height,
width=width,
dtype=dtype,
device=device,
crf=img.crf,
)
encoded_image = video_encoder(image)
if img.frame_idx == 0:
conditioning = VideoConditionByLatentIndex(
latent=encoded_image,
strength=img.strength,
latent_idx=0,
)
else:
conditioning = VideoConditionByKeyframeIndex(
keyframes=encoded_image,
strength=img.strength,
frame_idx=img.frame_idx,
)
conditionings.append(conditioning)
return conditionings
def image_conditionings_by_replacing_latent(
images: list[ImageConditioningInput],
height: int,
packages/ltx-pipelines/src/ltx_pipelines/utils/model_ledger.py
+40 -11
View File
@@ -33,8 +33,11 @@ from ltx_core.model.video_vae import (
)
from ltx_core.quantization import QuantizationPolicy
from ltx_core.text_encoders.gemma import (
AV_GEMMA_TEXT_ENCODER_KEY_OPS,
EMBEDDINGS_PROCESSOR_KEY_OPS,
GEMMA_LLM_KEY_OPS,
GEMMA_MODEL_OPS,
EmbeddingsProcessor,
EmbeddingsProcessorConfigurator,
GemmaTextEncoder,
GemmaTextEncoderConfigurator,
module_ops_from_gemma_root,
@@ -76,8 +79,8 @@ class ModelLedger:
:meth:`spatial_upsampler` method becomes available; otherwise calling it raises
a :class:`ValueError`.
loras:
Optional collection of LoRA configurations (paths, strengths, and key operations)
that are applied on top of the base transformer weights when building the model.
Tuple of LoRA configurations (path, strength, sd_ops) applied on top of the base
transformer weights. Use ``()`` for none.
registry:
Optional :class:`Registry` instance for weight caching across builders.
Defaults to :class:`DummyRegistry` which performs no cross-builder caching.
@@ -85,8 +88,9 @@ class ModelLedger:
Optional :class:`QuantizationPolicy` controlling how transformer weights
are stored and how matmul is executed. Defaults to None, which means no quantization.
### Creating Variants
Use :meth:`with_loras` to create a new ``ModelLedger`` instance that includes
additional LoRA configurations while sharing the same registry for weight caching.
Use :meth:`with_additional_loras` to create a new ``ModelLedger`` instance that
includes additional LoRA configurations or :meth:`with_loras` to replace existing
lora configurations while sharing the same registry for weight caching.
"""
def __init__(
@@ -96,7 +100,7 @@ class ModelLedger:
checkpoint_path: str | None = None,
gemma_root_path: str | None = None,
spatial_upsampler_path: str | None = None,
loras: LoraPathStrengthAndSDOps | None = None,
loras: tuple[LoraPathStrengthAndSDOps, ...] = (),
registry: Registry | None = None,
quantization: QuantizationPolicy | None = None,
):
@@ -105,7 +109,7 @@ class ModelLedger:
self.checkpoint_path = checkpoint_path
self.gemma_root_path = gemma_root_path
self.spatial_upsampler_path = spatial_upsampler_path
self.loras = loras or ()
self.loras = loras
self.registry = registry or DummyRegistry()
self.quantization = quantization
self.build_model_builders()
@@ -155,15 +159,23 @@ class ModelLedger:
registry=self.registry,
)
# Embeddings processor only needs the LTX checkpoint (no Gemma weights)
self.embeddings_processor_builder = Builder(
model_path=self.checkpoint_path,
model_class_configurator=EmbeddingsProcessorConfigurator,
model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
registry=self.registry,
)
if self.gemma_root_path is not None:
module_ops = module_ops_from_gemma_root(self.gemma_root_path)
model_folder = find_matching_file(self.gemma_root_path, "model*.safetensors").parent
weight_paths = [str(p) for p in model_folder.rglob("*.safetensors")]
self.text_encoder_builder = Builder(
model_path=(str(self.checkpoint_path), *weight_paths),
model_path=tuple(weight_paths),
model_class_configurator=GemmaTextEncoderConfigurator,
model_sd_ops=AV_GEMMA_TEXT_ENCODER_KEY_OPS,
model_sd_ops=GEMMA_LLM_KEY_OPS,
registry=self.registry,
module_ops=(GEMMA_MODEL_OPS, *module_ops),
)
@@ -181,14 +193,19 @@ class ModelLedger:
else:
return torch.device("cpu")
def with_loras(self, loras: LoraPathStrengthAndSDOps) -> "ModelLedger":
def with_additional_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
"""Add new lora configurations to the existing ones."""
return self.with_loras((*self.loras, *loras))
def with_loras(self, loras: tuple[LoraPathStrengthAndSDOps, ...]) -> "ModelLedger":
"""Replace existing lora configurations with new ones."""
return ModelLedger(
dtype=self.dtype,
device=self.device,
checkpoint_path=self.checkpoint_path,
gemma_root_path=self.gemma_root_path,
spatial_upsampler_path=self.spatial_upsampler_path,
loras=(*self.loras, *loras),
loras=loras,
registry=self.registry,
quantization=self.quantization,
)
@@ -244,6 +261,18 @@ class ModelLedger:
return self.text_encoder_builder.build(device=self._target_device(), dtype=self.dtype).to(self.device).eval()
def gemma_embeddings_processor(self) -> EmbeddingsProcessor:
if not hasattr(self, "embeddings_processor_builder"):
raise ValueError(
"Embeddings processor not initialized. Please provide a checkpoint path to the ModelLedger constructor."
)
return (
self.embeddings_processor_builder.build(device=self._target_device(), dtype=self.dtype)
.to(self.device)
.eval()
)
def audio_encoder(self) -> AudioEncoder:
if not hasattr(self, "audio_encoder_builder"):
raise ValueError(
packages/ltx-trainer/AGENTS.md
+8 -6
View File
@@ -14,7 +14,7 @@ This file provides guidance to AI coding assistants (Claude, Cursor, etc.) when
**Supported model versions:**
- **LTX-2** (19B, initial audio-video model)
- **LTX-2.3** (20B, improved text conditioning and audio quality)
- **LTX-2.3** (22B, improved text conditioning and audio quality)
Version detection is fully automatic — ltx-core reads the checkpoint config and selects the correct architecture
components. The trainer does not need version-specific code paths.
@@ -80,9 +80,10 @@ packages/ltx-trainer/
- `ltx_trainer.model_loader` provides component loaders using `ltx-core`
- Individual loaders: `load_transformer()`, `load_video_vae_encoder()`, `load_video_vae_decoder()`,
`load_text_encoder()`, etc.
`load_text_encoder()`, `load_embeddings_processor()`, etc.
- Combined loader: `load_model()` returns `LtxModelComponents` dataclass
- Uses `SingleGPUModelBuilder` from ltx-core internally
- Text encoder and embeddings processor are loaded separately (the text encoder only needs Gemma weights; the embeddings processor only needs the LTX checkpoint)
- 8-bit text encoder loading via `gemma_8bit.py` (bitsandbytes)
**Training Flow:**
@@ -146,7 +147,7 @@ Both model versions share the same latent space interface (see [Latent Space Con
The differences lie in how text conditioning and audio generation work. Version detection is automatic via checkpoint
config — the trainer uses a unified API.
| Component | LTX-2 (19B) | LTX-2.3 (20B) |
| Component | LTX-2 (19B) | LTX-2.3 (22B) |
|-----------------------|---------------------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------|
| Feature extractor | `FeatureExtractorV1`: single `aggregate_embed`, same output for video and audio | `FeatureExtractorV2`: separate `video_aggregate_embed` + `audio_aggregate_embed`, per-token RMSNorm |
| Caption projection | Inside the transformer (`caption_projection`) | Inside the feature extractor (before connector) |
@@ -309,7 +310,7 @@ Key classes:
- Implements distributed training with Accelerate
- Handles mixed precision, gradient accumulation, checkpointing
- `_training_step()` applies embedding connectors then delegates to strategy
- `_load_text_encoder_and_cache_embeddings()` caches validation embeddings and unloads heavy components
- `_load_text_encoder_and_cache_embeddings()` loads the text encoder + embeddings processor, caches validation embeddings, then unloads the Gemma LLM (keeps only the embeddings processor connectors for training)
- Uses training strategies for mode-specific logic
**`src/ltx_trainer/training_strategies/`** - Strategy pattern
@@ -334,7 +335,8 @@ Component loaders:
- `load_video_vae_decoder()` → `VideoDecoder`
- `load_audio_vae_decoder()` → `AudioDecoder`
- `load_vocoder()` → `Vocoder` or `VocoderWithBWE` (auto-detected)
- `load_text_encoder()` → `GemmaTextEncoder` (unified, handles V1/V2 automatically)
- `load_text_encoder(gemma_model_path)` → `GemmaTextEncoder` (pure Gemma LLM, no checkpoint needed)
- `load_embeddings_processor(checkpoint_path)` → `EmbeddingsProcessor` (feature extractor + connectors)
- `load_model()` → `LtxModelComponents` (convenience wrapper)
**`src/ltx_trainer/validation_sampler.py`** - Inference for validation
@@ -506,7 +508,7 @@ packages/ltx-core/src/ltx_core/
│ ├── encoders/
│ │ ├── base_encoder.py # GemmaTextEncoder (unified 3-block pipeline)
│ │ └── encoder_configurator.py # GemmaTextEncoderConfigurator, _create_feature_extractor
│ ├── feature_extractor.py # FeatureExtractorV1 (19B), FeatureExtractorV2 (20B)
│ ├── feature_extractor.py # FeatureExtractorV1 (19B), FeatureExtractorV2 (22B)
│ ├── embeddings_connector.py # Embeddings1DConnector, Embeddings1DConnectorConfigurator,
│ │ # AudioEmbeddings1DConnectorConfigurator
│ ├── embeddings_processor.py # EmbeddingsProcessor (wraps video + audio connectors)
packages/ltx-trainer/scripts/process_captions.py
+13 -7
View File
@@ -34,7 +34,7 @@ from torch.utils.data import DataLoader, Dataset
from transformers.utils.logging import disable_progress_bar
from ltx_trainer import logger
from ltx_trainer.model_loader import load_text_encoder
from ltx_trainer.model_loader import load_embeddings_processor, load_text_encoder
# Disable tokenizers parallelism to avoid warnings
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -264,17 +264,21 @@ def compute_captions_embeddings( # noqa: PLR0913
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Load text encoder
# Load text encoder and embeddings processor
with console.status("[bold]Loading Gemma text encoder...", spinner="dots"):
text_encoder = load_text_encoder(
model_path,
text_encoder_path,
device=device,
dtype=torch.bfloat16,
load_in_8bit=load_in_8bit,
)
embeddings_processor = load_embeddings_processor(
model_path,
device=device,
dtype=torch.bfloat16,
)
logger.info("Text encoder loaded successfully")
logger.info("Text encoder and embeddings processor loaded successfully")
# TODO(batch-tokenization): The current Gemma tokenizer doesn't support batched tokenization.
if batch_size > 1:
@@ -303,14 +307,16 @@ def compute_captions_embeddings( # noqa: PLR0913
) as progress:
task = progress.add_task("Processing captions", total=len(dataloader))
for batch in dataloader:
# Encode prompts using precompute() (returns video/audio features before connector)
# Encode prompts using text_encoder.encode() + feature_extractor
# (returns video/audio features before connector).
# The connector is applied during training via embeddings_processor
with torch.inference_mode():
# TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once.
# For now, process one at a time:
for i in range(len(batch["prompt"])):
video_prompt_embeds, audio_prompt_embeds, prompt_attention_mask = text_encoder.precompute(
batch["prompt"][i], padding_side="left"
hidden_states, prompt_attention_mask = text_encoder.encode(batch["prompt"][i], padding_side="left")
video_prompt_embeds, audio_prompt_embeds = embeddings_processor.feature_extractor(
hidden_states, prompt_attention_mask, "left"
)
output_rel_path = Path(batch["output_path"][i])
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