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

Automated PR - 2026-04-13

Browse Source
This commit is contained in:
github-actions[bot]
2026-04-13 14:29:35 +00:00
parent 59ca828d5a
commit d887bbd1e0
29 changed files with 463 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.clang-format
+3
View File
@@ -0,0 +1,3 @@
BasedOnStyle: LLVM
IndentWidth: 2
ColumnLimit: 120
.gitignore
+3
View File
@@ -39,3 +39,6 @@ tmp
*.png
*.wav
*.webp
# Binary files
*.so
README.md
+3 -3
View File
@@ -30,17 +30,17 @@ Download the following models from the [LTX-2.3 HuggingFace repository](https://
**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)
* [`ltx-2.3-22b-distilled-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-1.1.safetensors)
**Spatial Upscaler** - Required for current two-stage pipeline implementations in this repository
* [`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-x2-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-spatial-upscaler-x2-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x2-1.1.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.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.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)
* [`ltx-2.3-22b-distilled-lora-384-1.1.safetensors`](https://huggingface.co/Lightricks/LTX-2.3/blob/main/ltx-2.3-22b-distilled-lora-384-1.1.safetensors) - [Download](https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-22b-distilled-lora-384-1.1.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)
packages/ltx-core/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-core"
version = "1.1.0"
version = "1.1.1"
description = "Core implementation of Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-core/src/ltx_core/layer_streaming.py
+32 -50
View File
@@ -44,32 +44,26 @@ def _resolve_attr(module: nn.Module, dotted_path: str) -> nn.ModuleList:
class _LayerStore:
"""Manages on-demand pinning of layer parameters for GPU streaming.
Stores references to each layer's source data (which may be file-backed
mmap views or in-memory tensors). When a layer needs to be transferred
to GPU, its source data is pinned on demand and copied; on eviction the
pinned copy is freed and the source data is restored.
"""Manages CPU-pinned copies of layer parameters/buffers.
Tracks which layers currently reside on GPU so the prefetcher and evictor
can make correct decisions.
"""
def __init__(self, layers: nn.ModuleList, target_device: torch.device) -> None:
self.target_device = target_device
self.num_layers = len(layers)
# CPU-pinned copies keyed by (layer_idx, param_name)
self._pinned: list[dict[str, torch.Tensor]] = []
self._on_gpu: set[int] = set()
# Keep a reference to the source data for each layer so we can pin it
# on demand and restore it after eviction.
self._source_data: list[dict[str, torch.Tensor]] = []
for layer in layers:
source: dict[str, torch.Tensor] = {}
pinned: dict[str, torch.Tensor] = {}
for name, tensor in itertools.chain(layer.named_parameters(), layer.named_buffers()):
source[name] = tensor.data
self._source_data.append(source)
# Hold pinned tensors alive until the H2D transfer completes.
# Without this, the CachingHostAllocator can reclaim a pinned tensor
# as soon as its Python reference is dropped, even if an async H2D
# transfer is still reading from it.
self._pinned_in_flight: dict[int, list[torch.Tensor]] = {}
pinned_tensor = tensor.data.pin_memory()
tensor.data = pinned_tensor
pinned[name] = pinned_tensor
self._pinned.append(pinned)
def _check_idx(self, idx: int) -> None:
if idx < 0 or idx >= self.num_layers:
@@ -79,45 +73,34 @@ class _LayerStore:
return idx in self._on_gpu
def move_to_gpu(self, idx: int, layer: nn.Module, *, non_blocking: bool = False) -> None:
"""Pin layer *idx* on demand, then transfer to GPU."""
"""Move layer *idx* parameters from pinned CPU to *target_device*."""
self._check_idx(idx)
if idx in self._on_gpu:
return
source = self._source_data[idx]
pinned_refs: list[torch.Tensor] = []
pinned = self._pinned[idx]
for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
pinned = source[name].pin_memory()
param.data = pinned.to(self.target_device, non_blocking=non_blocking)
pinned_refs.append(pinned)
# Keep pinned tensors alive until eviction — the async H2D transfer
# may still be reading from them.
self._pinned_in_flight[idx] = pinned_refs
param.data = pinned[name].to(self.target_device, non_blocking=non_blocking)
self._on_gpu.add(idx)
def evict_to_cpu(self, idx: int, layer: nn.Module) -> None:
"""Restore source data, freeing the GPU and pinned copies."""
"""Swap layer *idx* parameters back to their pinned CPU copies."""
self._check_idx(idx)
if idx not in self._on_gpu:
return
source = self._source_data[idx]
pinned = self._pinned[idx]
for name, param in itertools.chain(layer.named_parameters(), layer.named_buffers()):
param.data = source[name]
# Release pinned tensors — the H2D transfer is complete by now
# (the compute stream waited on the prefetch event before using
# the layer, and we only evict after compute finishes).
self._pinned_in_flight.pop(idx, None)
param.data = pinned[name]
self._on_gpu.discard(idx)
def cleanup(self) -> None:
"""Release all source data and in-flight pinned references.
After this call, the source tensors can be garbage-collected once
"""Release all pinned memory references.
After this call, the pinned tensors can be garbage-collected once
the layer parameters (which still reference them via ``.data``) are
also released (e.g. via ``.to("meta")``).
"""
for source_dict in self._source_data:
source_dict.clear()
self._source_data.clear()
self._pinned_in_flight.clear()
for pinned_dict in self._pinned:
pinned_dict.clear()
self._pinned.clear()
class _AsyncPrefetcher:
@@ -228,8 +211,6 @@ class LayerStreamingWrapper(nn.Module):
idx_map: dict[int, int] = {id(layer): idx for idx, layer in enumerate(self._layers)}
num_layers = len(self._layers)
compute_stream = torch.cuda.current_stream(self._target_device)
def _pre_hook(
module: nn.Module,
_args: Any, # noqa: ANN401
@@ -246,6 +227,7 @@ class LayerStreamingWrapper(nn.Module):
# caching allocator would allow the prefetch stream to reuse their
# memory immediately after eviction — even if the compute kernel
# that reads them hasn't finished yet.
compute_stream = torch.cuda.current_stream(self._target_device)
for param in itertools.chain(module.parameters(), module.buffers()):
param.data.record_stream(compute_stream)
@@ -270,12 +252,12 @@ class LayerStreamingWrapper(nn.Module):
self._hooks.extend([h1, h2])
def teardown(self) -> None:
"""Remove hooks, release resources, and move parameters back to CPU.
"""Remove hooks, release pinned memory, and move parameters back to CPU.
After this call the wrapper is inert: hooks are removed, the prefetch
stream is drained and destroyed, all parameters reside on CPU, and the
``_LayerStore`` source data references are cleared. Callers should
still follow up with ``.to("meta")`` to release the CPU copies if the
model is no longer needed.
stream is drained and destroyed, all parameters reside on regular
(non-pinned) CPU memory, and the ``_LayerStore`` pinned-tensor cache is
cleared. Callers should still follow up with ``.to("meta")`` to release
the CPU copies if the model is no longer needed.
"""
for h in self._hooks:
h.remove()
@@ -298,10 +280,10 @@ class LayerStreamingWrapper(nn.Module):
for b in self._model.buffers():
b.data = b.data.to("cpu")
# Release source data references. After evict_to_cpu() the layer
# params point to the source data. The caller is expected to follow
# up with .to("meta") to drop the param refs; cleanup() drops the
# store's refs.
# Release pinned memory. After evict_to_cpu() the layer parameters
# still reference the pinned tensors (since .to("cpu") on a pinned
# tensor is a no-op). The caller is expected to follow up with
# .to("meta") to drop the param refs; cleanup() drops the store's refs.
self._store.cleanup()
# ------------------------------------------------------------------
packages/ltx-core/src/ltx_core/modality_tiling.py
+14 -2
View File
@@ -67,11 +67,16 @@ class VideoModalityTilingHelper:
# -- tile modality -----------------------------------------------------
def tile_modality(self, modality: Modality, tile: Tile) -> tuple[Modality, TilingContext]:
def tile_modality(
self, modality: Modality, tile: Tile, *, normalize_positions: bool = True
) -> tuple[Modality, TilingContext]:
"""Slice *modality* to the tokens covered by *tile*.
Selects generated tokens belonging to the tile's spatial region
and conditioning tokens that overlap with the tile (or have
negative time coordinates).
Args:
normalize_positions: When True, shift all positions so the
tile's generated tokens start at zero in every dimension.
Returns:
A ``(tiled_modality, context)`` tuple. Pass *context* to
:meth:`blend` together with the model output.
@@ -83,11 +88,18 @@ class VideoModalityTilingHelper:
keep_indices = keep_mask.nonzero(as_tuple=False).squeeze(1)
tile_attention_mask = modality.attention_mask[:, keep_indices, :][:, :, keep_indices]
positions = modality.positions[:, :, keep_mask, :]
if normalize_positions:
num_tile_gen = self._tile_generated_token_count(tile)
gen_pos = positions[:, :, :num_tile_gen, :] # (B, 3, num_tile_gen, 2)
offset = gen_pos[..., 0].amin(dim=2, keepdim=True).unsqueeze(-1) # (B, 3, 1, 1)
positions = positions - offset
tiled = replace(
modality,
latent=modality.latent[:, keep_mask, :],
timesteps=modality.timesteps[:, keep_mask],
positions=modality.positions[:, :, keep_mask, :],
positions=positions,
attention_mask=tile_attention_mask,
)
packages/ltx-pipelines/CLAUDE.md
+82
View File
@@ -0,0 +1,82 @@
<!--
MAINTENANCE: When modifying any pipeline class in src/ltx_pipelines/,
update this document to reflect changes to:
- __init__ / __call__ signatures
- sigma handling or step counts
- denoiser types or guidance
- new or removed pipelines
Run: ls src/ltx_pipelines/*.py to check for new pipeline files.
-->
# ltx-pipelines
Inference pipelines for LTX-2 audio-video generation. Depends on `ltx-core` for model definitions, diffusion components, and loading. All pipelines live in `packages/ltx-pipelines/src/ltx_pipelines/`.
## Pipeline selection
| Pipeline | File | Stages | Model | Sampler | Use case |
|----------|------|--------|-------|---------|----------|
| `TI2VidOneStagePipeline` | `ti2vid_one_stage.py` | 1 | Full | Euler | Simple text/image-to-video |
| `TI2VidTwoStagesPipeline` | `ti2vid_two_stages.py` | 2 | Full + distilled LoRA | Euler | Production quality |
| `TI2VidTwoStagesHQPipeline` | `ti2vid_two_stages_hq.py` | 2 | Full + distilled LoRA (both stages) | Res2s | Highest quality, fewer steps |
| `A2VidPipelineTwoStage` | `a2vid_two_stage.py` | 2 | Full + distilled LoRA | Euler | Audio-conditioned video |
| `KeyframeInterpolationPipeline` | `keyframe_interpolation.py` | 2 | Full + distilled LoRA | Euler | Keyframe interpolation |
| `DistilledPipeline` | `distilled.py` | 2 | Distilled only | Euler | Fastest inference |
| `ICLoraPipeline` | `ic_lora.py` | 2 | Distilled only | Euler | Video-to-video with IC-LoRA control |
| `RetakePipeline` | `retake.py` | 1 | Full or distilled | Euler | Video region regeneration |
## Guidance
- **CFG**: Blends conditioned/unconditioned predictions. Defaults: `cfg_scale=3.0` (video), `7.0` (audio).
- **STG**: Perturbs self-attention in transformer blocks. Default `stg_scale=1.0`, `stg_blocks=[28]` (LTX-2.3) / `[29]` (LTX-2). HQ disables STG (`stg_scale=0.0`).
- **Modality guidance**: Cross-modal attention scaling (`modality_scale=3.0`).
- All guidance is stage 1 only. Stage 2 always uses `SimpleDenoiser`.
## Sigma schedules and step counts
- **Scheduler-based** (full model): `self._scheduler = LTX2Scheduler()` with `execute(steps=N)` (HQ also passes `latent=` for token-count-dependent shift). Defaults: 30 steps (LTX-2.3), 40 (LTX-2), 15 (HQ).
- **Distilled**: Fixed 8-step `DISTILLED_SIGMA_VALUES` (9 values). Stage 2 uses 3-step `STAGE_2_DISTILLED_SIGMA_VALUES` (4 values). No `num_inference_steps` param.
- **Retake**: `num_inference_steps=40` default; ignored when `distilled=True` (fixed 8-step).
- **Overrides**: All pipelines accept optional sigma tensors in `__call__`: `sigmas` (one-stage), `stage_1_sigmas` + `stage_2_sigmas` (two-stage).
## LoRA conventions
- No default LoRAs. `loras` param defaults to empty list/tuple. `DEFAULT_LORA_STRENGTH = 1.0`.
- Two-stage non-distilled pipelines require `distilled_lora` (applied to stage 2 only in TI2Vid/A2Vid/Keyframe).
- HQ is unique: applies distilled LoRA to **both** stages with separate `distilled_lora_strength_stage_1` / `_stage_2` params.
## Shared building blocks (`utils/blocks.py`)
- `DiffusionStage` -- owns transformer lifecycle; builds model on call, frees on exit via `gpu_model()` context manager (moves params to meta device to release GPU/CPU memory). Accepts optional `stepper` and `loop` overrides.
- `PromptEncoder` -- Gemma text encoder + embeddings processor (video 4096-dim, audio 2048-dim).
- `ImageConditioner` / `AudioConditioner` -- temporary encoder scope; builds encoder, passes to callable, frees.
- `VideoUpsampler` -- 2x spatial upsampling via encoder + upsampler.
- `VideoDecoder` / `AudioDecoder` -- latent-to-pixel decoding (iterator for video, `Audio` for audio).
### Memory management
- **Model lifecycle**: All blocks build their model on call and free it on exit. `gpu_model()` moves params to `"meta"` device on exit, immediately releasing storage. No model persists between calls.
- **Layer streaming**: When `streaming_prefetch_count` is set, `DiffusionStage` wraps the transformer in `LayerStreamingWrapper`. Layers live on pinned CPU memory; only `1 + prefetch_count` layers are on GPU at a time, with async H2D prefetch on a separate CUDA stream.
- **Batch splitting**: `BatchSplitAdapter` wraps the transformer and splits inputs exceeding `max_batch_size` into sequential chunks. If guidance needs B=4 but `max_batch_size=1`, it runs 4 sequential B=1 passes. Higher `max_batch_size` reduces layer-streaming PCIe transfers at the cost of peak memory.
## Denoisers (`utils/denoisers.py`)
- `SimpleDenoiser` -- single forward pass (B=1), no guidance. Used by distilled pipelines and all stage 2.
- `GuidedDenoiser` -- CFG/STG with static `MultiModalGuider` instances (HQ, A2Vid, Retake non-distilled).
- `FactoryGuidedDenoiser` -- per-step guider creation via factory (OneStageTI2Vid, TwoStagesTI2Vid, Keyframe).
Guided denoisers batch all guidance passes into a **single transformer call**: states are repeated along the batch dimension, contexts concatenated, and a `BatchedPerturbationConfig` controls which attention ops are skipped per sample. Pass count is dynamic: B=2 for CFG-only, up to B=4 with CFG+STG+modality isolation. Results are split back and blended by the guider.
## Per-pipeline unique features
- **HQ**: Res2s second-order sampler for **both** stages, latent-dependent sigma schedule, distilled LoRA on both stages with separate strengths.
- **A2Vid**: Audio frozen in both stages (`frozen=True, noise_scale=0.0`). Returns original audio (not VAE-decoded); no `AudioDecoder`.
- **IC-LoRA**: `VideoConditionByReferenceLatent`, `reference_downscale_factor` from LoRA metadata, `skip_stage_2`, attention mask downsampling. Stage 2 is LoRA-free and uses `combined_image_conditionings` (no IC-LoRA conditioning).
- **Keyframe**: Uses `image_conditionings_by_adding_guiding_latent` in both stages (all frames as keyframe guidance, no replacement) -- unlike TI2Vid which uses `combined_image_conditionings` (frame_idx=0 replaces, others guide).
- **Retake**: `TemporalRegionMask` for selective time-window regeneration. `regenerate_video`/`regenerate_audio` flags. Conditional distilled/full behavior.
- **Distilled**: Single `self.stage` reused for both stages (not `stage_1`/`stage_2`).
## Image conditioning helpers (`utils/helpers.py`)
- `combined_image_conditionings()` -- images with `frame_idx==0` replace latent (`VideoConditionByLatentIndex`), others guide (`VideoConditionByKeyframeIndex`).
- `image_conditionings_by_adding_guiding_latent()` -- all images become keyframe guidance regardless of `frame_idx`.
packages/ltx-pipelines/pyproject.toml
+1 -1
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-pipelines"
version = "1.1.0"
version = "1.1.1"
description = "Pipelines implementation for Lightricks' LTX-2 model"
readme = "README.md"
requires-python = ">=3.10"
packages/ltx-pipelines/src/ltx_pipelines/a2vid_two_stage.py
+10 -5
View File
@@ -22,7 +22,7 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
)
from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
@@ -56,6 +56,7 @@ class A2VidPipelineTwoStage:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
@@ -103,6 +104,8 @@ class A2VidPipelineTwoStage:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -148,7 +151,9 @@ class A2VidPipelineTwoStage:
)
)
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (
stage_1_sigmas if stage_1_sigmas is not None else self._scheduler.execute(steps=num_inference_steps)
).to(dtype=torch.float32, device=self.device)
video_state, _ = self.stage_1(
denoiser=GuidedDenoiser(
@@ -185,7 +190,7 @@ class A2VidPipelineTwoStage:
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -200,7 +205,7 @@ class A2VidPipelineTwoStage:
video_state, _ = self.stage_2(
denoiser=SimpleDenoiser(v_context_p, a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -209,7 +214,7 @@ class A2VidPipelineTwoStage:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
packages/ltx-pipelines/src/ltx_pipelines/distilled.py
+7 -5
View File
@@ -23,8 +23,8 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
DISTILLED_SIGMAS,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import SimpleDenoiser
@@ -77,7 +77,7 @@ class DistilledPipeline:
self.video_decoder = VideoDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
self.audio_decoder = AudioDecoder(distilled_checkpoint_path, self.dtype, self.device, registry=registry)
def __call__(
def __call__( # noqa: PLR0913
self,
prompt: str,
seed: int,
@@ -89,6 +89,8 @@ class DistilledPipeline:
tiling_config: TilingConfig | None = None,
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -105,7 +107,7 @@ class DistilledPipeline:
video_context, audio_context = ctx_p.video_encoding, ctx_p.audio_encoding
# Stage 1: Initial low resolution video generation.
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
stage_1_sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
stage_1_w, stage_1_h = width // 2, height // 2
stage_1_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -134,7 +136,7 @@ class DistilledPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
stage_2_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
images=images,
packages/ltx-pipelines/src/ltx_pipelines/ic_lora.py
+9 -7
View File
@@ -32,8 +32,8 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMA_VALUES,
DISTILLED_SIGMAS,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import SimpleDenoiser
@@ -126,6 +126,8 @@ class ICLoraPipeline:
skip_stage_2: bool = False,
conditioning_attention_mask: torch.Tensor | None = None,
streaming_prefetch_count: int | None = None,
stage_1_sigmas: torch.Tensor = DISTILLED_SIGMAS,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
"""
Generate video with IC-LoRA conditioning.
@@ -200,7 +202,7 @@ class ICLoraPipeline:
)
)
stage_1_sigmas = torch.Tensor(DISTILLED_SIGMA_VALUES).to(self.device)
stage_1_sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
video_state, audio_state = self.stage_1(
denoiser=SimpleDenoiser(video_context, audio_context),
@@ -230,7 +232,7 @@ class ICLoraPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -245,7 +247,7 @@ class ICLoraPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(video_context, audio_context),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -254,12 +256,12 @@ class ICLoraPipeline:
video=ModalitySpec(
context=video_context,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=audio_context,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
packages/ltx-pipelines/src/ltx_pipelines/keyframe_interpolation.py
+11 -6
View File
@@ -25,7 +25,7 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
@@ -62,6 +62,7 @@ class KeyframeInterpolationPipeline:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
@@ -107,6 +108,8 @@ class KeyframeInterpolationPipeline:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -125,7 +128,9 @@ class KeyframeInterpolationPipeline:
v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
# Stage 1: Initial low resolution video generation.
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (
stage_1_sigmas if stage_1_sigmas is not None else self._scheduler.execute(steps=num_inference_steps)
).to(dtype=torch.float32, device=self.device)
stage_1_output_shape = VideoPixelShape(
batch=1,
@@ -181,7 +186,7 @@ class KeyframeInterpolationPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LORA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: image_conditionings_by_adding_guiding_latent(
@@ -196,7 +201,7 @@ class KeyframeInterpolationPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(v_context_p, a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -205,12 +210,12 @@ class KeyframeInterpolationPipeline:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=a_context_p,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
packages/ltx-pipelines/src/ltx_pipelines/retake.py
+10 -4
View File
@@ -25,7 +25,7 @@ from ltx_pipelines.utils.blocks import (
PromptEncoder,
VideoDecoder,
)
from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES, detect_params
from ltx_pipelines.utils.constants import DISTILLED_SIGMAS, detect_params
from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
audio_latent_from_file,
@@ -78,6 +78,8 @@ class RetakePipeline:
self.device = device or get_device()
self.dtype = torch.bfloat16
self.distilled = distilled
if not distilled:
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(
checkpoint_path=checkpoint_path,
gemma_root=gemma_root,
@@ -141,6 +143,7 @@ class RetakePipeline:
tiling_config: TilingConfig | None = None,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
sigmas: torch.Tensor | None = None,
) -> tuple[Iterator[torch.Tensor], torch.Tensor]:
"""Regenerate ``[start_time, end_time]`` of the source video (retake).
Parameters
@@ -227,15 +230,18 @@ class RetakePipeline:
initial_latent=initial_audio_latent,
frozen=initial_audio_latent is not None and not regenerate_audio,
)
# Build denoiser
# Build denoiser and resolve sigma schedule.
if sigmas is None:
sigmas = DISTILLED_SIGMAS if self.distilled else self._scheduler.execute(steps=num_inference_steps)
sigmas = sigmas.to(dtype=torch.float32, device=self.device)
if self.distilled:
sigmas = torch.tensor(DISTILLED_SIGMA_VALUES).to(dtype=torch.float32, device=self.device)
denoiser = SimpleDenoiser(
v_context=v_context_p,
a_context=a_context_p,
)
else:
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
v_context_n, a_context_n = contexts[1].video_encoding, contexts[1].audio_encoding
video_guider = MultiModalGuider(
params=video_guider_params,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_one_stage.py
+5 -1
View File
@@ -55,6 +55,7 @@ class TI2VidOneStagePipeline:
):
self.dtype = torch.bfloat16
self.device = device or get_device()
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(
checkpoint_path=checkpoint_path,
gemma_root=gemma_root,
@@ -107,6 +108,7 @@ class TI2VidOneStagePipeline:
streaming_prefetch_count: int | None = None,
tiling_config: TilingConfig | None = None,
max_batch_size: int = 1,
sigmas: torch.Tensor | None = None,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=False)
@@ -135,7 +137,9 @@ class TI2VidOneStagePipeline:
)
)
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (sigmas if sigmas is not None else self._scheduler.execute(steps=num_inference_steps)).to(
dtype=torch.float32, device=self.device
)
video_guider_factory = create_multimodal_guider_factory(
params=video_guider_params,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages.py
+11 -6
View File
@@ -25,7 +25,7 @@ from ltx_pipelines.utils.blocks import (
VideoUpsampler,
)
from ltx_pipelines.utils.constants import (
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
detect_params,
)
from ltx_pipelines.utils.denoisers import FactoryGuidedDenoiser, SimpleDenoiser
@@ -61,6 +61,7 @@ class TI2VidTwoStagesPipeline:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
@@ -106,6 +107,8 @@ class TI2VidTwoStagesPipeline:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -142,7 +145,9 @@ class TI2VidTwoStagesPipeline:
)
)
sigmas = LTX2Scheduler().execute(steps=num_inference_steps).to(dtype=torch.float32, device=self.device)
sigmas = (
stage_1_sigmas if stage_1_sigmas is not None else self._scheduler.execute(steps=num_inference_steps)
).to(dtype=torch.float32, device=self.device)
video_state, audio_state = self.stage_1(
denoiser=FactoryGuidedDenoiser(
@@ -172,7 +177,7 @@ class TI2VidTwoStagesPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.Tensor(STAGE_2_DISTILLED_SIGMA_VALUES).to(self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
images=images,
@@ -186,7 +191,7 @@ class TI2VidTwoStagesPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
width=width,
height=height,
@@ -195,12 +200,12 @@ class TI2VidTwoStagesPipeline:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=a_context_p,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
streaming_prefetch_count=streaming_prefetch_count,
packages/ltx-pipelines/src/ltx_pipelines/ti2vid_two_stages_hq.py
+13 -11
View File
@@ -23,7 +23,7 @@ from ltx_pipelines.utils.blocks import (
)
from ltx_pipelines.utils.constants import (
LTX_2_3_HQ_PARAMS,
STAGE_2_DISTILLED_SIGMA_VALUES,
STAGE_2_DISTILLED_SIGMAS,
)
from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
from ltx_pipelines.utils.helpers import (
@@ -64,6 +64,7 @@ class TI2VidTwoStagesHQPipeline:
):
self.device = device or get_device()
self.dtype = torch.bfloat16
self._scheduler = LTX2Scheduler()
distilled_lora_stage_1 = LoraPathStrengthAndSDOps(
path=distilled_lora[0].path,
@@ -121,6 +122,8 @@ class TI2VidTwoStagesHQPipeline:
enhance_prompt: bool = False,
streaming_prefetch_count: int | None = None,
max_batch_size: int = 1,
stage_1_sigmas: torch.Tensor | None = None,
stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
) -> tuple[Iterator[torch.Tensor], Audio]:
assert_resolution(height=height, width=width, is_two_stage=True)
@@ -157,13 +160,12 @@ class TI2VidTwoStagesHQPipeline:
)
)
empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
stepper = Res2sDiffusionStep()
sigmas = (
LTX2Scheduler()
.execute(latent=empty_latent, steps=num_inference_steps)
.to(dtype=torch.float32, device=self.device)
)
if stage_1_sigmas is None:
empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
stage_1_sigmas = self._scheduler.execute(latent=empty_latent, steps=num_inference_steps)
sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
video_state, audio_state = self.stage_1(
denoiser=GuidedDenoiser(
@@ -195,7 +197,7 @@ class TI2VidTwoStagesHQPipeline:
# Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
upscaled_video_latent = self.upsampler(video_state.latent[:1])
distilled_sigmas = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=self.device)
stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
stage_2_conditionings = self.image_conditioner(
lambda enc: combined_image_conditionings(
@@ -210,7 +212,7 @@ class TI2VidTwoStagesHQPipeline:
video_state, audio_state = self.stage_2(
denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
sigmas=distilled_sigmas,
sigmas=stage_2_sigmas,
noiser=noiser,
stepper=stepper,
width=width,
@@ -220,12 +222,12 @@ class TI2VidTwoStagesHQPipeline:
video=ModalitySpec(
context=v_context_p,
conditionings=stage_2_conditionings,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=upscaled_video_latent,
),
audio=ModalitySpec(
context=a_context_p,
noise_scale=distilled_sigmas[0].item(),
noise_scale=stage_2_sigmas[0].item(),
initial_latent=audio_state.latent,
),
loop=res2s_audio_video_denoising_loop,
packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py
+4
View File
@@ -1,6 +1,7 @@
import logging
from dataclasses import dataclass, field, replace
import torch
from safetensors import safe_open
from ltx_core.components.guiders import MultiModalGuiderParams
@@ -17,6 +18,9 @@ DISTILLED_SIGMA_VALUES = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725,
# Reduced schedule for super-resolution stage 2 (subset of distilled values)
STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
DISTILLED_SIGMAS = torch.tensor(DISTILLED_SIGMA_VALUES)
STAGE_2_DISTILLED_SIGMAS = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES)
# =============================================================================
# Pipeline Parameters
packages/ltx-trainer/docs/custom-training-strategies.md
+4 -5
View File
@@ -307,15 +307,14 @@ class InpaintingStrategy(TrainingStrategy):
audio_pred: Tensor | None,
inputs: ModelInputs,
) -> Tensor:
"""Compute training loss on inpaint regions only."""
"""Compute training loss on inpaint regions only. Returns [B,]."""
# MSE loss
loss = (video_pred - inputs.video_targets).pow(2)
# Apply loss mask
# Apply loss mask and reduce to per-element [B,]
loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
loss = loss.mul(loss_mask).div(loss_mask.mean() + 1e-8)
return loss.mean()
masked = loss.mul(loss_mask)
return masked.mean(dim=[-2, -1]) / loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
```
### Step 5: Register the Strategy
packages/ltx-trainer/docs/utility-scripts.md
+26 -1
View File
@@ -48,6 +48,10 @@ uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --use-
uv run python scripts/caption_videos.py videos_dir/ --output dataset.json \
--captioner-type gemini_flash --api-key YOUR_API_KEY
# Use Gemini Flash with parallel workers for faster throughput
uv run python scripts/caption_videos.py videos_dir/ --output dataset.json \
--captioner-type gemini_flash --num-workers 5
# Caption without audio processing (video-only)
uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --no-audio
@@ -61,9 +65,10 @@ uv run python scripts/caption_videos.py videos_dir/ --output dataset.json --over
- **Multiple backends**:
- `qwen_omni` (default): Local Qwen2.5-Omni model - processes video + audio locally
- `gemini_flash`: Google Gemini Flash API - cloud-based, requires API key
- **Parallel captioning** (Gemini Flash only): Use `--num-workers` to run multiple API calls concurrently for faster throughput on large datasets
- **Structured output**: Captions include visual description, speech transcription, sounds, and on-screen text
- **Memory optimization**: 8-bit quantization option for limited VRAM
- **Incremental processing**: Skips already-captioned files by default
- **Incremental processing**: Skips already-captioned files by default; progress is saved every 5 videos
- **Multiple output formats**: JSON, JSONL, CSV, or TXT
**Caption format:**
@@ -74,6 +79,26 @@ The captioner produces structured captions with four sections:
- `[SOUNDS]`: Description of music, ambient sounds, sound effects
- `[TEXT]`: Any on-screen text visible in the video
**Parallel captioning with Gemini Flash:**
When using `--captioner-type gemini_flash`, you can speed up large dataset captioning by running multiple API calls at the same time using `--num-workers` (accepts 1–10, default is 1):
```bash
export GEMINI_API_KEY="your-key-here"
# Caption a large dataset with 5 workers running concurrently
uv run python scripts/caption_videos.py videos_dir/ \
--output dataset.json \
--captioner-type gemini_flash \
--num-workers 5
```
> [!NOTE]
> `--num-workers` is only supported with `gemini_flash`. Using it with `qwen_omni` or any other local model will raise an error, because local GPU models are not thread-safe.
> [!TIP]
> Keep `--num-workers` between 3–5 for most use cases. Very high values (8–10) may hit Gemini API rate limits depending on your quota tier.
**Environment variables (for Gemini Flash):**
Set one of these to use Gemini Flash without passing `--api-key`:
packages/ltx-trainer/pyproject.toml
+2 -2
View File
@@ -1,6 +1,6 @@
[project]
name = "ltx-trainer"
version = "1.1.0"
version = "1.1.1"
description = "LTX-2 training, democratized."
readme = "README.md"
authors = [
@@ -48,7 +48,7 @@ build-backend = "hatchling.build"
[tool.ruff]
target-version = "1.1.0"
target-version = "1.1.1"
line-length = 120
[tool.ruff.lint]
packages/ltx-trainer/scripts/caption_videos.py
+69 -25
View File
@@ -19,6 +19,8 @@ Basic usage:
Advanced usage:
# Use Gemini Flash API (requires GEMINI_API_KEY or GOOGLE_API_KEY env var)
caption_videos.py videos_dir/ --captioner-type gemini_flash
# Use Gemini Flash with parallel workers (2-10 workers, cloud API only)
caption_videos.py videos_dir/ --captioner-type gemini_flash --num-workers 5
# Disable audio processing (video-only captions)
caption_videos.py videos_dir/ --no-audio
# Process videos with specific extensions and save as JSON
@@ -27,6 +29,7 @@ Advanced usage:
import csv
import json
from concurrent.futures import ThreadPoolExecutor, as_completed
from enum import Enum
from pathlib import Path
@@ -70,7 +73,7 @@ class OutputFormat(str, Enum):
JSONL = "jsonl" # JSON Lines file with one JSON object per line
def caption_media(
def caption_media( # noqa: PLR0913
input_path: Path,
output_path: Path,
captioner: MediaCaptioningModel,
@@ -81,6 +84,7 @@ def caption_media(
clean_caption: bool,
output_format: OutputFormat,
override: bool,
num_workers: int = 1,
) -> None:
"""Caption videos and images using the provided captioning model.
Args:
@@ -94,6 +98,7 @@ def caption_media(
clean_caption: Whether to clean up captions
output_format: Format to save the captions in
override: Whether to override existing captions
num_workers: Number of parallel workers (only for cloud-based captioners like Gemini)
"""
# Get list of media files to process
@@ -121,9 +126,13 @@ def caption_media(
console.print("[bold yellow]All media already have captions. Use --override to recaption.[/]")
return
# Process media files
if num_workers > 1:
console.print(f"Running with [bold cyan]{num_workers}[/] parallel workers.")
captions = existing_captions.copy()
successfully_captioned = 0
completed_since_save = 0
progress = Progress(
SpinnerColumn(),
TextColumn("{task.description}"),
@@ -135,36 +144,47 @@ def caption_media(
console=console,
)
def process_one(media_file: Path) -> tuple[str, str]:
"""Caption a single media file and return (relative_path, caption)."""
caption = captioner.caption(
path=media_file,
fps=fps,
include_audio=include_audio,
clean_caption=clean_caption,
)
rel_path = str(media_file.resolve().relative_to(base_dir))
return rel_path, caption
with progress:
task = progress.add_task("Captioning", total=len(media_to_process))
task = progress.add_task(
f"Captioning (workers: {num_workers})" if num_workers > 1 else "Captioning",
total=len(media_to_process),
)
for i, media_file in enumerate(media_to_process):
progress.update(task, description=f"Captioning [bold blue]{media_file.name}[/]")
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = {executor.submit(process_one, f): f for f in media_to_process}
try:
# Generate caption for the media
caption = captioner.caption(
path=media_file,
fps=fps,
include_audio=include_audio,
clean_caption=clean_caption,
)
for future in as_completed(futures):
media_file = futures[future]
progress.update(task, description=f"Captioning [bold blue]{media_file.name}[/]")
# Convert absolute path to relative path (relative to the output file's directory)
rel_path = str(media_file.resolve().relative_to(base_dir))
# Store the caption with the relative path as key
captions[rel_path] = caption
successfully_captioned += 1
except Exception as e:
console.print(f"[bold red]Error captioning {media_file}: {e}[/]")
try:
rel_path, caption = future.result()
if i % SAVE_INTERVAL == 0:
_save_captions(captions, output_path, output_format)
captions[rel_path] = caption
successfully_captioned += 1
completed_since_save += 1
# Advance progress bar
progress.advance(task)
if completed_since_save >= SAVE_INTERVAL:
_save_captions(captions, output_path, output_format)
completed_since_save = 0
# Save captions to file
except Exception as e:
console.print(f"[bold red]Error captioning {media_file.name}: {e}[/]")
progress.advance(task)
# Final save with everything accumulated
_save_captions(captions, output_path, output_format)
# Print summary
@@ -407,6 +427,18 @@ def main( # noqa: PLR0913
envvar=["GOOGLE_API_KEY", "GEMINI_API_KEY"],
help="API key for Gemini Flash (can also use GOOGLE_API_KEY or GEMINI_API_KEY env var)",
),
num_workers: int = typer.Option(
1,
"--num-workers",
"-w",
min=1,
max=10,
help=(
"Number of parallel workers for captioning (1-10). "
"Values above 1 are only supported for cloud-based captioners (gemini_flash). "
"Using multiple workers with a local model will raise an error."
),
),
) -> None:
"""Auto-caption videos with audio using multimodal models.
This script supports audio-visual captioning using:
@@ -424,6 +456,17 @@ def main( # noqa: PLR0913
caption_videos.py video.mp4 -o captions.json -i "Describe this video in detail"
"""
# Parallel workers are only safe for cloud-based (stateless) captioners.
# Local models like Qwen-Omni hold GPU state and are not thread-safe.
if num_workers > 1 and captioner_type != CaptionerType.GEMINI_FLASH:
console.print(
"[bold red]Error:[/] --num-workers > 1 is only supported with [bold]--captioner-type gemini_flash[/].\n"
"Local models (e.g. qwen_omni) run on GPU and are not thread-safe — "
"parallel calls would cause memory corruption or incorrect results.\n"
"Either set [bold]--num-workers 1[/] (default) or switch to [bold]--captioner-type gemini_flash[/]."
)
raise typer.Exit(code=1)
# Determine device for local models
device_str = device or ("cuda" if torch.cuda.is_available() else "cpu")
@@ -479,6 +522,7 @@ def main( # noqa: PLR0913
clean_caption=clean_caption,
output_format=output_format,
override=override,
num_workers=num_workers,
)
packages/ltx-trainer/src/ltx_trainer/sigma_tracker.py
+59
View File
@@ -0,0 +1,59 @@
"""Sigma-bucketed loss tracking.
Maps each training step's per-element sigmas and losses to buckets.
Smoothing is left to wandb's UI.
"""
import bisect
from collections import defaultdict
class SigmaBucketTracker:
"""Map per-element sigma values to named buckets for per-bucket loss logging.
By default, partitions [0, 1] into four equal-width buckets.
Custom boundaries can be provided for non-uniform bucketing.
Each call to update() receives per-element sigmas and losses (both [B,]),
buckets each element, and computes the mean loss per bucket. This gives
accurate per-sigma loss tracking even for batch_size > 1.
"""
def __init__(
self,
bucket_boundaries: list[float] | None = None,
) -> None:
if bucket_boundaries is None:
bucket_boundaries = [0.0, 0.25, 0.5, 0.75, 1.0]
if len(bucket_boundaries) < 2:
raise ValueError("bucket_boundaries must have at least 2 elements")
if any(bucket_boundaries[i] >= bucket_boundaries[i + 1] for i in range(len(bucket_boundaries) - 1)):
raise ValueError("bucket_boundaries must be strictly increasing")
self._boundaries = list(bucket_boundaries)
self._num_buckets = len(bucket_boundaries) - 1
self._bucket_labels = [
f"{bucket_boundaries[i]:.2f}-{bucket_boundaries[i + 1]:.2f}" for i in range(self._num_buckets)
]
self._last_metrics: dict[str, float] = {}
def _get_bucket_index(self, sigma: float) -> int:
"""Map sigma value to bucket index."""
idx = bisect.bisect_right(self._boundaries, sigma) - 1
return max(0, min(idx, self._num_buckets - 1))
def update(self, sigmas: list[float], losses: list[float]) -> None:
"""Record per-element losses into their sigma buckets.
Args:
sigmas: Per-element sigma values, one per batch element.
losses: Per-element losses, one per batch element.
"""
if not sigmas:
self._last_metrics = {}
return
bucket_losses: dict[int, list[float]] = defaultdict(list)
for sigma, loss in zip(sigmas, losses, strict=True):
bucket_losses[self._get_bucket_index(sigma)].append(loss)
self._last_metrics = {self._bucket_labels[b]: sum(vals) / len(vals) for b, vals in bucket_losses.items()}
def get_metrics(self, prefix: str = "train") -> dict[str, float]:
"""Return the mean loss for each bucket hit on the last update.
Wandb handles smoothing in the UI.
"""
return {f"{prefix}/loss_sigma_{label}": loss for label, loss in self._last_metrics.items()}
packages/ltx-trainer/src/ltx_trainer/trainer.py
+50 -26
View File
@@ -2,8 +2,9 @@ import os
import re
import time
import warnings
from dataclasses import dataclass
from pathlib import Path
from typing import Callable
from typing import Any, Callable
import torch
import wandb
@@ -39,6 +40,7 @@ from ltx_trainer.model_loader import load_embeddings_processor, load_text_encode
from ltx_trainer.model_loader import load_model as load_ltx_model
from ltx_trainer.progress import TrainingProgress
from ltx_trainer.quantization import quantize_model
from ltx_trainer.sigma_tracker import SigmaBucketTracker
from ltx_trainer.timestep_samplers import SAMPLERS
from ltx_trainer.training_state import ConfigFingerprint, RngStates, TrainingState
from ltx_trainer.training_strategies import get_training_strategy
@@ -77,6 +79,14 @@ class TrainingStats(BaseModel):
num_processes: int
@dataclass(frozen=True)
class TrainingStepOutput:
"""Output from a single training step."""
loss: Tensor # [B,] per-element loss (unreduced)
sigma: Tensor # [B,] sampled sigma, detached from computational graph
class LtxvTrainer:
def __init__(self, trainer_config: LtxTrainerConfig) -> None:
self._config = trainer_config
@@ -95,7 +105,8 @@ class LtxvTrainer:
self._checkpoint_paths: list[Path] = []
self._training_state_paths: list[Path] = []
self._training_state_size_warned = False
self._init_wandb()
self._wandb_run = None
self._sigma_tracker = SigmaBucketTracker()
def train( # noqa: PLR0912, PLR0915
self,
@@ -128,6 +139,10 @@ class LtxvTrainer:
initial_step = 0
resuming = False
# Initialize W&B after restore so we only resume the run when state restore succeeds.
resume_run_id = training_state.wandb_run_id if resuming and training_state is not None else None
self._init_wandb(resume_run_id=resume_run_id)
self._init_dataloader()
data_iter = iter(self._dataloader)
self._init_timestep_sampler()
@@ -191,8 +206,8 @@ class LtxvTrainer:
if is_optimization_step:
self._global_step += 1
loss = self._training_step(batch)
self._accelerator.backward(loss)
output = self._training_step(batch)
self._accelerator.backward(output.loss.mean())
if self._accelerator.sync_gradients and cfg.optimization.max_grad_norm > 0:
self._accelerator.clip_grad_norm_(
@@ -244,9 +259,10 @@ class LtxvTrainer:
# Update progress and log metrics
current_lr = self._optimizer.param_groups[0]["lr"]
step_time = (time.time() - step_start_time) * cfg.optimization.gradient_accumulation_steps
step_loss = output.loss.detach().mean().item()
progress.update_training(
loss=loss.item(),
loss=step_loss,
lr=current_lr,
step_time=step_time,
advance=is_optimization_step,
@@ -254,14 +270,16 @@ class LtxvTrainer:
# Log metrics to W&B (only on main process and optimization steps)
if IS_MAIN_PROCESS and is_optimization_step:
self._log_metrics(
{
"train/loss": loss.item(),
"train/learning_rate": current_lr,
"train/step_time": step_time,
"train/global_step": self._global_step,
}
)
# Track per-element loss by sigma bucket
self._sigma_tracker.update(output.sigma.cpu().tolist(), output.loss.detach().cpu().tolist())
metrics = {
"train/loss": step_loss,
"train/learning_rate": current_lr,
"train/step_time": step_time,
"train/global_step": self._global_step,
}
metrics.update(self._sigma_tracker.get_metrics())
self._log_metrics(metrics)
# Fallback logging when progress bars are disabled
if disable_progress_bars and IS_MAIN_PROCESS and self._global_step % 20 == 0:
@@ -274,7 +292,7 @@ class LtxvTrainer:
total_time = "calculating..."
logger.info(
f"Step {self._global_step}/{cfg.optimization.steps} - "
f"Loss: {loss.item():.4f}, LR: {current_lr:.2e}, "
f"Loss: {step_loss:.4f}, LR: {current_lr:.2e}, "
f"Time/Step: {step_time:.2f}s, Total Time: {total_time}",
)
@@ -330,7 +348,7 @@ class LtxvTrainer:
return saved_path, stats
def _training_step(self, batch: dict[str, dict[str, Tensor]]) -> Tensor:
def _training_step(self, batch: dict[str, dict[str, Tensor]]) -> TrainingStepOutput:
"""Perform a single training step using the configured strategy."""
# Apply embedding connectors to transform pre-computed text embeddings
conditions = batch["conditions"]
@@ -366,8 +384,9 @@ class LtxvTrainer:
# Use strategy to compute loss
loss = self._training_strategy.compute_loss(video_pred, audio_pred, model_inputs)
sigma = model_inputs.video.sigma.detach() if model_inputs.video.enabled else model_inputs.audio.sigma.detach()
return loss
return TrainingStepOutput(loss=loss, sigma=sigma)
@free_gpu_memory_context(after=True)
def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
@@ -1064,8 +1083,8 @@ class LtxvTrainer:
def _save_training_state(self, save_dir: Path) -> None:
"""Save training state alongside checkpoint for resume.
Respects checkpoints.save_training_state config:
- "full": optimizer + scheduler + RNG + step
- "minimal": scheduler + RNG + step only
- "full": optimizer + scheduler + RNG + step + wandb_run_id
- "minimal": scheduler + RNG + step + wandb_run_id
- "off": skip entirely
"""
if not IS_MAIN_PROCESS:
@@ -1101,6 +1120,7 @@ class LtxvTrainer:
),
lr_scheduler_state_dict=self._lr_scheduler.state_dict() if self._lr_scheduler is not None else None,
optimizer_state_dict=optimizer_state,
wandb_run_id=self._wandb_run.id if self._wandb_run is not None else None,
)
state_path = save_dir / f"training_state_step_{self._global_step:05d}.pt"
@@ -1166,20 +1186,24 @@ class LtxvTrainer:
logger.info(f"💾 Training configuration saved to: {config_path.relative_to(self._config.output_dir)}")
def _init_wandb(self) -> None:
def _init_wandb(self, resume_run_id: str | None = None) -> None:
"""Initialize Weights & Biases run."""
if not self._config.wandb.enabled or not IS_MAIN_PROCESS:
self._wandb_run = None
return
wandb_config = self._config.wandb
run = wandb.init(
project=wandb_config.project,
entity=wandb_config.entity,
name=Path(self._config.output_dir).name,
tags=wandb_config.tags,
config=self._config.model_dump(),
)
init_kwargs: dict[str, Any] = {
"project": wandb_config.project,
"entity": wandb_config.entity,
"name": Path(self._config.output_dir).name,
"tags": wandb_config.tags,
"config": self._config.model_dump(),
}
if resume_run_id is not None:
init_kwargs["id"] = resume_run_id
init_kwargs["resume"] = "allow"
run = wandb.init(**init_kwargs)
self._wandb_run = run
def _log_metrics(self, metrics: dict[str, float]) -> None:
packages/ltx-trainer/src/ltx_trainer/training_state.py
+2
View File
@@ -28,6 +28,7 @@ class TrainingState(BaseModel):
rng_states: RngStates
lr_scheduler_state_dict: dict[str, Any] | None = None
optimizer_state_dict: dict[str, Any] | None = None
wandb_run_id: str | None = None
def to_save_dict(self) -> dict[str, Any]:
"""Build dict suitable for torch.save -- recurses BaseModel sub-models, passes tensors/dicts through."""
@@ -48,4 +49,5 @@ class TrainingState(BaseModel):
rng_states=RngStates(**data["rng_states"]),
lr_scheduler_state_dict=data.get("lr_scheduler_state_dict"),
optimizer_state_dict=data.get("optimizer_state_dict"),
wandb_run_id=data.get("wandb_run_id"),
)
packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py
+2 -1
View File
@@ -125,7 +125,8 @@ class TrainingStrategy(ABC):
audio_pred: Audio prediction from the transformer model (None for video-only)
inputs: The prepared model inputs containing targets and masks
Returns:
Scalar loss tensor
Per-element loss tensor of shape [B,]. The trainer reduces to a scalar
before backward(). Returning unreduced loss enables per-sigma-bucket tracking.
"""
def get_checkpoint_metadata(self) -> dict[str, Any]:
packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py
+7 -7
View File
@@ -273,19 +273,19 @@ class TextToVideoStrategy(TrainingStrategy):
audio_pred: Tensor | None,
inputs: ModelInputs,
) -> Tensor:
"""Compute masked MSE loss for video and optionally audio."""
# Video loss
"""Compute masked MSE loss for video and optionally audio. Returns [B,]."""
# Video loss: per-element mean over (seq, channels), [B,]
video_loss = (video_pred - inputs.video_targets).pow(2)
video_loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
video_loss = video_loss.mul(video_loss_mask).div(video_loss_mask.mean())
video_loss = video_loss.mean()
masked = video_loss.mul(video_loss_mask)
video_loss = masked.mean(dim=[-2, -1]) / video_loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
# If no audio, return video loss only
if not self.config.with_audio or audio_pred is None or inputs.audio_targets is None:
return video_loss
# Audio loss (no conditioning mask)
audio_loss = (audio_pred - inputs.audio_targets).pow(2).mean()
# Audio loss: per-element mean over (seq, channels), [B,]
audio_loss = (audio_pred - inputs.audio_targets).pow(2).mean(dim=[-2, -1])
# Combined loss
# Combined loss [B,]
return video_loss + audio_loss
packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py
+4 -7
View File
@@ -240,7 +240,7 @@ class VideoToVideoStrategy(TrainingStrategy):
_audio_pred: Tensor | None,
inputs: ModelInputs,
) -> Tensor:
"""Compute masked loss only on target portion."""
"""Compute masked loss only on target portion. Returns [B,]."""
# Extract target portion of prediction
ref_seq_len = inputs.ref_seq_len
target_pred = video_pred[:, ref_seq_len:, :]
@@ -248,14 +248,11 @@ class VideoToVideoStrategy(TrainingStrategy):
# Get target portion of loss mask
target_loss_mask = inputs.video_loss_mask[:, ref_seq_len:]
# Compute loss
# Compute per-element loss [B,]
loss = (target_pred - inputs.video_targets).pow(2)
# Apply loss mask
loss_mask = target_loss_mask.unsqueeze(-1).float()
loss = loss.mul(loss_mask).div(loss_mask.mean())
return loss.mean()
masked = loss.mul(loss_mask)
return masked.mean(dim=[-2, -1]) / loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
def get_checkpoint_metadata(self) -> dict[str, Any]:
"""Get metadata for checkpoint files."""
packages/ltx-trainer/src/ltx_trainer/video_utils.py
+16 -5
View File
@@ -14,6 +14,9 @@ from torch import Tensor
def get_video_frame_count(video_path: str | Path) -> int:
"""Get the number of frames in a video file.
Tries three approaches in order: stream metadata, duration*fps estimate,
full decode. The estimate may be off by a few frames for VFR videos or
containers with edit lists — exact for the min_frames filtering use case.
Args:
video_path: Path to the video file
Returns:
@@ -21,11 +24,19 @@ def get_video_frame_count(video_path: str | Path) -> int:
"""
with av.open(str(video_path)) as container:
video_stream = container.streams.video[0]
frame_count = video_stream.frames
if frame_count == 0:
# Fallback: count frames by decoding
frame_count = sum(1 for _ in container.decode(video=0))
return frame_count
if video_stream.frames > 0:
return video_stream.frames
# Fast estimate from container metadata (avoids full decode).
# Uses Fraction arithmetic to prevent float precision loss.
rate = video_stream.average_rate or video_stream.base_rate
if video_stream.duration and video_stream.time_base and rate:
duration = Fraction(video_stream.duration) * Fraction(video_stream.time_base)
return round(duration * Fraction(rate))
# Last resort: full decode (very slow for 4K)
return sum(1 for _ in container.decode(video=0))
def read_video(video_path: str | Path, max_frames: int | None = None) -> tuple[Tensor, float]:
uv.lock
Generated
+3 -3
View File
@@ -2063,7 +2063,7 @@ wheels = [
[[package]]
name = "ltx-core"
version = "1.1.0"
version = "1.1.1"
source = { editable = "packages/ltx-core" }
dependencies = [
{ name = "accelerate" },
@@ -2121,7 +2121,7 @@ dev = [{ name = "scikit-image", specifier = ">=0.25.2" }]
[[package]]
name = "ltx-pipelines"
version = "1.1.0"
version = "1.1.1"
source = { editable = "packages/ltx-pipelines" }
dependencies = [
{ name = "av" },
@@ -2141,7 +2141,7 @@ requires-dist = [
[[package]]
name = "ltx-trainer"
version = "1.1.0"
version = "1.1.1"
source = { editable = "packages/ltx-trainer" }
dependencies = [
{ name = "accelerate" },