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

Automated PR - 2026-05-11

Browse Source
This commit is contained in:
github-actions[bot]
2026-05-11 13:14:05 +00:00
parent 41d9243716
commit 7df34dfa83
72 changed files with 3299 additions and 911 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-trainer/scripts/process_captions.py
+73 -20
View File
@@ -13,12 +13,14 @@ Can be used as a standalone script:
import json
import os
from collections.abc import Callable
from pathlib import Path
from typing import Any
import pandas as pd
import torch
import typer
from accelerate import PartialState
from rich.console import Console
from rich.progress import (
BarColumn,
@@ -30,7 +32,7 @@ from rich.progress import (
TimeElapsedColumn,
TimeRemainingColumn,
)
from torch.utils.data import DataLoader, Dataset
from torch.utils.data import DataLoader, Dataset, Subset
from transformers.utils.logging import disable_progress_bar
from ltx_trainer import logger
@@ -232,9 +234,14 @@ def compute_captions_embeddings( # noqa: PLR0913
batch_size: int = 8,
device: str = "cuda",
load_in_8bit: bool = False,
overwrite: bool = False,
) -> None:
"""
Process captions and save text embeddings.
Under ``accelerate launch``, each process handles an interleaved shard of
the dataset (rank/world read from ``accelerate.PartialState``). Already-
computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
atomic so an interrupted run is safe to resume.
Args:
dataset_file: Path to metadata file (CSV/JSON/JSONL) containing captions and media paths
output_dir: Directory to save embeddings
@@ -247,11 +254,12 @@ def compute_captions_embeddings( # noqa: PLR0913
batch_size: Batch size for processing
device: Device to use for computation
load_in_8bit: Whether to load the Gemma text encoder in 8-bit precision
overwrite: Re-encode every item even if its output exists. Use when rerunning with
changed parameters (different text encoder, lora_trigger, etc.) so stale
outputs are replaced.
"""
console = Console()
# Create dataset
dataset = CaptionsDataset(
dataset_file=dataset_file,
caption_column=caption_column,
@@ -264,6 +272,24 @@ def compute_captions_embeddings( # noqa: PLR0913
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# TODO(batch-tokenization): The current Gemma tokenizer doesn't support batched tokenization.
if batch_size > 1:
logger.warning(
"Batch size greater than 1 is not currently supported with the Gemma tokenizer. "
"Overriding batch_size to 1. This will be fixed in a future update."
)
batch_size = 1
dataloader = _build_sharded_dataloader(
dataset,
batch_size=batch_size,
num_workers=2,
is_done=lambda idx: (output_path / dataset.output_paths[idx]).is_file(),
overwrite=overwrite,
)
if dataloader is None:
return
# Load text encoder and embeddings processor
with console.status("[bold]Loading Gemma text encoder...", spinner="dots"):
text_encoder = load_text_encoder(
@@ -279,21 +305,7 @@ def compute_captions_embeddings( # noqa: PLR0913
)
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:
logger.warning(
"Batch size greater than 1 is not currently supported with the Gemma tokenizer. "
"Overriding batch_size to 1. This will be fixed in a future update."
)
batch_size = 1
# Create dataloader
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=2)
# Process batches
total_batches = len(dataloader)
logger.info(f"Processing captions in {total_batches:,} batches...")
logger.info(f"Processing captions in {len(dataloader):,} batches...")
with Progress(
SpinnerColumn(),
@@ -333,11 +345,44 @@ def compute_captions_embeddings( # noqa: PLR0913
embedding_data["audio_prompt_embeds"] = audio_prompt_embeds[0].cpu().contiguous()
output_file = output_path / output_rel_path
torch.save(embedding_data, output_file)
_atomic_save(embedding_data, output_file)
progress.advance(task)
logger.info(f"Processed {len(dataset):,} captions. Embeddings saved to {output_path}")
logger.info(f"Processed {len(dataloader.dataset):,} captions -> {output_path}") # type: ignore[arg-type]
def _atomic_save(data: Any, out: Path) -> None: # noqa: ANN401
"""Save to ``out`` atomically via per-PID temp file + replace.
Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
ignores. The per-PID suffix makes concurrent writes from multiple ranks
collision-free.
"""
tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
torch.save(data, tmp)
tmp.replace(out)
def _build_sharded_dataloader(
dataset: Dataset,
*,
batch_size: int,
num_workers: int,
is_done: Callable[[int], bool],
overwrite: bool,
) -> DataLoader | None:
"""Return a DataLoader over this rank's interleaved shard of ``dataset``.
When ``overwrite`` is False, items whose outputs already exist (per
``is_done``) are filtered out. Returns ``None`` if this rank has nothing
to do, so the caller can early-return without loading any models.
"""
state = PartialState()
todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
if not todo:
logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
return None
logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
@app.command()
@@ -387,8 +432,15 @@ def main( # noqa: PLR0913
default=False,
help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
),
overwrite: bool = typer.Option(
default=False,
help="Re-encode every caption even if its output exists. Use when rerunning with "
"changed parameters (different text encoder, lora_trigger, etc.) so stale outputs are replaced.",
),
) -> None:
"""Process text captions and save embeddings for video generation training.
For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
will handle an interleaved shard of the dataset.
This script processes captions from metadata files and saves text embeddings
that can be used for training video generation models. The output embeddings
will maintain the same folder structure and naming as the corresponding media files.
@@ -428,6 +480,7 @@ def main( # noqa: PLR0913
batch_size=batch_size,
device=device,
load_in_8bit=load_text_encoder_in_8bit,
overwrite=overwrite,
)
packages/ltx-trainer/scripts/process_dataset.py
+12
View File
@@ -50,6 +50,7 @@ def preprocess_dataset( # noqa: PLR0913
reference_downscale_factor: int = 1,
with_audio: bool = False,
load_text_encoder_in_8bit: bool = False,
overwrite: bool = False,
) -> None:
"""Run the preprocessing pipeline with the given arguments."""
# Validate dataset file
@@ -77,6 +78,7 @@ def preprocess_dataset( # noqa: PLR0913
batch_size=batch_size,
device=device,
load_in_8bit=load_text_encoder_in_8bit,
overwrite=overwrite,
)
# Process videos using the dedicated function
@@ -97,6 +99,7 @@ def preprocess_dataset( # noqa: PLR0913
vae_tiling=vae_tiling,
with_audio=with_audio,
audio_output_dir=str(audio_latents_dir) if audio_latents_dir else None,
overwrite=overwrite,
)
# Process reference videos if reference_column is provided
@@ -133,6 +136,7 @@ def preprocess_dataset( # noqa: PLR0913
batch_size=batch_size,
device=device,
vae_tiling=vae_tiling,
overwrite=overwrite,
)
# Handle decoding if requested (for verification)
@@ -252,8 +256,15 @@ def main( # noqa: PLR0913
help="Downscale factor for reference video resolution. When > 1, reference videos are processed at "
"1/n resolution (e.g., 2 means half resolution). Used for efficient IC-LoRA training.",
),
overwrite: bool = typer.Option(
default=False,
help="Re-compute every item even if its output exists. Use when rerunning with "
"changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
),
) -> None:
"""Preprocess a video dataset by computing and saving latents and text embeddings.
For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
will handle an interleaved shard of the dataset.
The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
This script is designed for LTX-2 models which use the Gemma text encoder.
Examples:
@@ -310,6 +321,7 @@ def main( # noqa: PLR0913
reference_downscale_factor=reference_downscale_factor,
with_audio=with_audio,
load_text_encoder_in_8bit=load_text_encoder_in_8bit,
overwrite=overwrite,
)
packages/ltx-trainer/scripts/process_videos.py
+79 -18
View File
@@ -14,6 +14,8 @@ Can be used as a standalone script:
import json
import math
import os
from collections.abc import Callable
from pathlib import Path
from typing import Any
@@ -22,6 +24,7 @@ import pandas as pd
import torch
import torchaudio
import typer
from accelerate import PartialState
from pillow_heif import register_heif_opener
from rich.console import Console
from rich.progress import (
@@ -34,7 +37,7 @@ from rich.progress import (
TimeElapsedColumn,
TimeRemainingColumn,
)
from torch.utils.data import DataLoader, Dataset
from torch.utils.data import DataLoader, Dataset, Subset
from torchvision import transforms
from torchvision.transforms import InterpolationMode
from torchvision.transforms.functional import crop, resize, to_tensor
@@ -444,9 +447,14 @@ def compute_latents( # noqa: PLR0913, PLR0915
vae_tiling: bool = False,
with_audio: bool = False,
audio_output_dir: str | None = None,
overwrite: bool = False,
) -> None:
"""
Process videos and save latent representations.
Under ``accelerate launch``, each process handles an interleaved shard of
the dataset (rank/world read from ``accelerate.PartialState``). Already-
computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
atomic so an interrupted run is safe to resume.
Args:
dataset_file: Path to metadata file (CSV/JSON/JSONL) containing video paths
video_column: Column name for video paths in the metadata file
@@ -460,15 +468,15 @@ def compute_latents( # noqa: PLR0913, PLR0915
vae_tiling: Whether to enable VAE tiling
with_audio: Whether to extract and encode audio from videos
audio_output_dir: Directory to save audio latents (required if with_audio=True)
overwrite: Re-process every item even if its output exists. Use when rerunning with
changed parameters (different model, resolution, etc.) so stale outputs are replaced.
"""
# Validate audio parameters
if with_audio and audio_output_dir is None:
raise ValueError("audio_output_dir must be provided when with_audio=True")
console = Console()
torch_device = torch.device(device)
# Create dataset
dataset = MediaDataset(
dataset_file=dataset_file,
main_media_column=main_media_column or video_column,
@@ -481,13 +489,34 @@ def compute_latents( # noqa: PLR0913, PLR0915
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Set up audio output directory if needed
audio_output_path = None
audio_output_path: Path | None = None
if with_audio:
audio_output_path = Path(audio_output_dir)
audio_output_path.mkdir(parents=True, exist_ok=True)
# Audio processing requires batch_size=1; must be applied before the dataloader is built.
if with_audio and batch_size > 1:
logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
batch_size = 1
data_root = dataset.dataset_file.parent
def _is_done(idx: int) -> bool:
rel = dataset.main_media_paths[idx].relative_to(data_root).with_suffix(".pt")
if not (output_path / rel).is_file():
return False
return audio_output_path is None or (audio_output_path / rel).is_file()
dataloader = _build_sharded_dataloader(
dataset,
batch_size=batch_size,
num_workers=4,
is_done=_is_done,
overwrite=overwrite,
)
if dataloader is None:
return
# Load video VAE encoder
with console.status(f"[bold]Loading video VAE encoder from [cyan]{model_path}[/]...", spinner="dots"):
vae = load_video_vae_encoder(model_path, device=torch_device, dtype=torch.bfloat16)
@@ -510,14 +539,6 @@ def compute_latents( # noqa: PLR0913, PLR0915
n_fft=audio_vae_encoder.n_fft,
).to(torch_device)
# Create dataloader
# Note: batch_size=1 required when with_audio because audio extraction can fail for some videos,
# and the default collate function can't handle mixed None/dict values across a batch.
if with_audio and batch_size > 1:
logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
batch_size = 1
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)
# Track audio statistics
audio_success_count = 0
audio_skip_count = 0
@@ -560,7 +581,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
"fps": batch["video_metadata"]["fps"][i].item(),
}
torch.save(latent_data, output_file)
_atomic_save(latent_data, output_file)
# Process audio if enabled (audio is already extracted by the dataset)
if with_audio:
@@ -588,7 +609,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
"duration": audio_latents["duration"],
}
torch.save(audio_save_data, audio_output_file)
_atomic_save(audio_save_data, audio_output_file)
audio_success_count += 1
else:
# Video has no audio track
@@ -596,8 +617,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
progress.advance(task)
# Log summary
logger.info(f"Processed {len(dataset)} videos. Latents saved to {output_path}")
logger.info(f"Processed {len(dataloader.dataset)} videos -> {output_path}") # type: ignore[arg-type]
if with_audio:
logger.info(
f"Audio processing: {audio_success_count} videos with audio, "
@@ -935,6 +955,39 @@ def compute_scaled_resolution_buckets(
return scaled_buckets
def _atomic_save(data: Any, out: Path) -> None: # noqa: ANN401
"""Save to ``out`` atomically via per-PID temp file + replace.
Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
ignores. The per-PID suffix makes concurrent writes from multiple ranks
collision-free.
"""
tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
torch.save(data, tmp)
tmp.replace(out)
def _build_sharded_dataloader(
dataset: Dataset,
*,
batch_size: int,
num_workers: int,
is_done: Callable[[int], bool],
overwrite: bool,
) -> DataLoader | None:
"""Return a DataLoader over this rank's interleaved shard of ``dataset``.
When ``overwrite`` is False, items whose outputs already exist (per
``is_done``) are filtered out. Returns ``None`` if this rank has nothing
to do, so the caller can early-return without loading any models.
"""
state = PartialState()
todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
if not todo:
logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
return None
logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
@app.command()
def main( # noqa: PLR0913
dataset_file: str = typer.Argument(
@@ -981,8 +1034,15 @@ def main( # noqa: PLR0913
default=None,
help="Output directory for audio latents (required if --with-audio is set)",
),
overwrite: bool = typer.Option(
default=False,
help="Re-encode every item even if its output exists. Use when rerunning with "
"changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
),
) -> None:
"""Process videos/images and save latent representations for video generation training.
For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
will handle an interleaved shard of the dataset.
This script processes videos and images from metadata files and saves latent representations
that can be used for training video generation models. The output latents will maintain
the same folder structure and naming as the corresponding media files.
@@ -1032,6 +1092,7 @@ def main( # noqa: PLR0913
vae_tiling=vae_tiling,
with_audio=with_audio,
audio_output_dir=audio_output_dir,
overwrite=overwrite,
)