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

Automated PR - 2026-06-17

Browse Source
This commit is contained in:
github-actions[bot]
2026-06-17 14:06:32 +00:00
parent d6053703e0
commit 0d3d3a3855
90 changed files with 8265 additions and 4511 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-trainer/scripts/caption_videos.py
+81 -75
View File
@@ -2,29 +2,28 @@
"""
Auto-caption videos with audio using multimodal models.
This script provides a command-line interface for generating captions for videos
(including audio) using multimodal models. It supports:
- Qwen2.5-Omni: Local model for audio-visual captioning (default)
- Gemini Flash: Cloud-based API for audio-visual captioning
The paths to videos in the generated dataset/captions file will be RELATIVE to the
directory where the output file is stored. This makes the dataset more portable and
easier to use in different environments.
Backends:
- Qwen3-Omni-30B-A3B-Thinking via a local vLLM HTTP server (default,
``qwen_omni``). Launch the server once with ``scripts/serve_captioner.py``.
- Gemini Flash 3.5 via Google's API (``gemini_flash``).
The paths in the output file are RELATIVE to the output file's directory,
making the dataset portable.
Basic usage:
# Caption a single video (includes audio by default)
caption_videos.py video.mp4 --output captions.json
# Caption all videos in a directory
caption_videos.py videos_dir/ --output captions.csv
# Caption with custom instruction
caption_videos.py video.mp4 --instruction "Describe what happens in this video in detail."
# Launch the captioner server once (separate terminal)
uv run python scripts/serve_captioner.py
# Caption a directory
caption_videos.py videos_dir/ --output captions.json
# Caption a single video with a custom prompt
caption_videos.py video.mp4 --output cap.json --instruction "Describe in detail."
Advanced usage:
# Use Gemini Flash API (requires GEMINI_API_KEY or GOOGLE_API_KEY env var)
# Use Gemini Flash 3.5 (cloud, requires GEMINI_API_KEY)
caption_videos.py videos_dir/ --captioner-type gemini_flash
# Use Gemini Flash with parallel workers (2-10 workers, cloud API only)
# Gemini with parallel workers
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
caption_videos.py videos_dir/ --extensions mp4,mov,avi --output captions.json
# Talk to a remote vLLM server
caption_videos.py videos_dir/ --vllm-url http://192.168.1.10:8001/v1
# Enable Qwen3 chain-of-thought (slower, more detail)
caption_videos.py videos_dir/ --enable-thinking
"""
import csv
@@ -33,7 +32,6 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
from enum import Enum
from pathlib import Path
import torch
import typer
from rich.console import Console
from rich.progress import (
@@ -45,9 +43,14 @@ from rich.progress import (
TimeElapsedColumn,
TimeRemainingColumn,
)
from transformers.utils.logging import disable_progress_bar
from ltx_trainer.captioning import CaptionerType, MediaCaptioningModel, create_captioner
from ltx_trainer.captioning import (
DEFAULT_QWEN_MODEL,
DEFAULT_VLLM_BASE_URL,
CaptionerType,
MediaCaptioningModel,
create_captioner,
)
VIDEO_EXTENSIONS = ["mp4", "avi", "mov", "mkv", "webm"]
IMAGE_EXTENSIONS = ["jpg", "jpeg", "png"]
@@ -61,8 +64,6 @@ app = typer.Typer(
help="Auto-caption videos with audio using multimodal models.",
)
disable_progress_bar()
class OutputFormat(str, Enum):
"""Available output formats for captions."""
@@ -73,15 +74,13 @@ class OutputFormat(str, Enum):
JSONL = "jsonl" # JSON Lines file with one JSON object per line
def caption_media( # noqa: PLR0913
def caption_media(
input_path: Path,
output_path: Path,
captioner: MediaCaptioningModel,
extensions: list[str],
recursive: bool,
fps: int,
include_audio: bool,
clean_caption: bool,
output_format: OutputFormat,
override: bool,
num_workers: int = 1,
@@ -94,8 +93,6 @@ def caption_media( # noqa: PLR0913
extensions: List of media file extensions to include
recursive: Whether to search subdirectories recursively
fps: Frames per second to sample from videos (ignored for images)
include_audio: Whether to include audio in captioning
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)
@@ -149,10 +146,10 @@ def caption_media( # noqa: PLR0913
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))
# Don't resolve the file itself, so a symlinked clip keeps its logical path under the
# dataset dir instead of jumping to its (possibly external) link target.
rel_path = str((media_file.parent.resolve() / media_file.name).relative_to(base_dir))
return rel_path, caption
with progress:
@@ -371,16 +368,31 @@ def main( # noqa: PLR0913
help="Type of captioner to use. Valid values: 'qwen_omni' (local), 'gemini_flash' (API)",
case_sensitive=False,
),
device: str | None = typer.Option(
None,
"--device",
"-d",
help="Device to use for inference (e.g., 'cuda', 'cuda:0', 'cpu'). Only for local models.",
vllm_url: str = typer.Option(
DEFAULT_VLLM_BASE_URL,
"--vllm-url",
help=(
"Base URL of the vLLM OpenAI-compatible server (qwen_omni only). "
"Launch the server with `uv run python scripts/serve_captioner.py`."
),
),
use_8bit: bool = typer.Option(
vllm_model: str = typer.Option(
DEFAULT_QWEN_MODEL,
"--vllm-model",
help="Served model identifier on the vLLM server (qwen_omni only).",
),
enable_thinking: bool = typer.Option(
False,
"--use-8bit",
help="Whether to use 8-bit precision for the captioning model (reduces memory usage)",
"--enable-thinking/--no-thinking",
help=(
"Let Qwen3-Omni produce a <think>...</think> chain-of-thought before the caption. "
"Off by default: ~5x slower with marginal quality benefit and occasional hallucinations."
),
),
max_tokens: int = typer.Option(
4096,
"--max-tokens",
help="Maximum new tokens to generate per caption (qwen_omni only).",
),
instruction: str | None = typer.Option(
None,
@@ -401,20 +413,14 @@ def main( # noqa: PLR0913
help="Search for media files in subdirectories recursively",
),
fps: int = typer.Option(
3,
2,
"--fps",
"-f",
help="Frames per second to sample from videos (ignored for images)",
),
include_audio: bool = typer.Option(
True,
"--audio/--no-audio",
help="Whether to include audio in captioning (for videos with audio tracks)",
),
clean_caption: bool = typer.Option(
True,
"--clean-caption/--raw-caption",
help="Whether to clean up captions by removing common VLM patterns",
help=(
"Frames per second to sample from videos. 2 is a typical default; "
"lower values use less compute per video. Ignored for images and for the "
"Gemini backend (which decides its own sampling rate)."
),
),
override: bool = typer.Option(
False,
@@ -441,35 +447,36 @@ def main( # noqa: PLR0913
),
) -> None:
"""Auto-caption videos with audio using multimodal models.
This script supports audio-visual captioning using:
- Qwen2.5-Omni: Local model (default) - processes both video and audio
- Gemini Flash: Cloud API - requires GOOGLE_API_KEY environment variable
Backends:
- ``qwen_omni`` (default): Qwen3-Omni-30B-A3B-Thinking via a local vLLM
HTTP server. Launch the server once in a separate terminal with
``uv run python scripts/serve_captioner.py``. The server stays loaded
across script invocations.
- ``gemini_flash``: Google Gemini (``gemini-3.5-flash``) via the google-genai SDK.
Auth is automatic -- ``GEMINI_API_KEY``/``GOOGLE_API_KEY`` for the Developer API,
or Google Cloud credentials (gcloud / service account) for Vertex AI with no env vars.
The paths in the output file will be relative to the output file's directory.
Examples:
# Caption videos with audio using Qwen2.5-Omni (default)
# Caption videos using the local vLLM server (default)
caption_videos.py videos_dir/ -o captions.json
# Caption using Gemini Flash API
# Point at a remote vLLM server
caption_videos.py videos_dir/ -o captions.json --vllm-url http://other-host:8001/v1
# Caption using Gemini Flash 3.5
caption_videos.py videos_dir/ -o captions.json -c gemini_flash
# Caption without audio (video-only)
caption_videos.py videos_dir/ -o captions.json --no-audio
# Caption with custom instruction
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.
# Parallel workers are only supported for the cloud Gemini backend; qwen_omni
# drives a single shared vLLM server and is captioned serially from here.
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[/]."
"[bold red]Error:[/] --num-workers > 1 is only supported with "
"[bold]--captioner-type gemini_flash[/]. Use --num-workers 1 (default) "
"for the qwen_omni backend."
)
raise typer.Exit(code=1)
# Determine device for local models
device_str = device or ("cuda" if torch.cuda.is_available() else "cpu")
# Parse extensions
ext_list = [ext.strip() for ext in extensions.split(",")]
@@ -490,14 +497,15 @@ def main( # noqa: PLR0913
output = Path(output).resolve()
console.print(f"Output will be saved to [bold blue]{output}[/]")
# Initialize captioning model
with console.status("Loading captioning model...", spinner="dots"):
with console.status("Initializing captioner...", spinner="dots"):
if captioner_type == CaptionerType.QWEN_OMNI:
captioner = create_captioner(
captioner_type=captioner_type,
device=device_str,
use_8bit=use_8bit,
base_url=vllm_url,
model=vllm_model,
instruction=instruction,
max_tokens=max_tokens,
enable_thinking=enable_thinking,
)
elif captioner_type == CaptionerType.GEMINI_FLASH:
captioner = create_captioner(
@@ -508,7 +516,7 @@ def main( # noqa: PLR0913
else:
raise ValueError(f"Unsupported captioner type: {captioner_type}")
console.print(f"[bold green]✓[/] {captioner_type.value} captioning model loaded successfully")
console.print(f"[bold green]✓[/] {captioner_type.value} captioner ready")
# Caption media files
caption_media(
@@ -518,8 +526,6 @@ def main( # noqa: PLR0913
extensions=ext_list,
recursive=recursive,
fps=fps,
include_audio=include_audio,
clean_caption=clean_caption,
output_format=output_format,
override=override,
num_workers=num_workers,
packages/ltx-trainer/scripts/compute_reference.py
+33 -17
View File
@@ -2,7 +2,7 @@
Compute reference videos for IC-LoRA training.
This script provides a command-line interface for generating reference videos to be used for IC-LoRA training.
Note that it reads and writes to the same file (the output of caption_videos.py),
where it adds the "reference_path" field to the JSON.
where it adds the "reference_video" field to the JSON.
Basic usage:
# Compute reference videos for all videos in a directory
compute_reference.py videos_dir/ --output videos_dir/captions.json
@@ -11,7 +11,7 @@ Basic usage:
# Standard library imports
import json
from pathlib import Path
from typing import Dict
from typing import Any
# Third-party imports
import cv2
@@ -37,6 +37,10 @@ from ltx_trainer.video_utils import read_video, save_video
console = Console()
disable_progress_bar()
VIDEO_COLUMNS = ("video", "media_path")
REFERENCE_VIDEO_COLUMN = "reference_video"
LEGACY_REFERENCE_COLUMN = "reference_path"
def compute_reference(
images: torch.Tensor,
@@ -79,15 +83,15 @@ def compute_reference(
def _get_meta_data(
output_path: Path,
) -> Dict[str, str]:
) -> list[dict[str, Any]]:
"""Get set of existing reference video paths without loading the actual files.
Args:
output_path: Path to the reference video paths file
Returns:
Dictionary mapping media paths to reference video paths
Dataset rows with media paths and captions
"""
if not output_path.exists():
return {}
return []
console.print(f"[bold blue]Reading meta data from [cyan]{output_path}[/]...[/]")
@@ -98,11 +102,18 @@ def _get_meta_data(
except Exception as e:
console.print(f"[bold yellow]Warning: Could not check meta data: {e}[/]")
return {}
return []
def _get_media_path(item: dict[str, Any]) -> str:
for column in VIDEO_COLUMNS:
if column in item:
return item[column]
raise KeyError(f"Dataset row must contain one of {VIDEO_COLUMNS}")
def _save_dataset_json(
reference_paths: Dict[str, str],
reference_paths: dict[str, str],
output_path: Path,
) -> None:
"""Save dataset json with reference video paths.
@@ -115,17 +126,17 @@ def _save_dataset_json(
json_data = json.load(f)
new_json_data = json_data.copy()
for i, item in enumerate(json_data):
media_path = item["media_path"]
media_path = _get_media_path(item)
reference_path = reference_paths[media_path]
new_json_data[i]["reference_path"] = reference_path
new_json_data[i].pop(LEGACY_REFERENCE_COLUMN, None)
new_json_data[i][REFERENCE_VIDEO_COLUMN] = reference_path
with output_path.open("w", encoding="utf-8") as f:
json.dump(new_json_data, f, indent=2, ensure_ascii=False)
console.print(f"[bold green]✓[/] Reference video paths saved to [cyan]{output_path}[/]")
console.print("[bold yellow]Note:[/] Use these files with ImageOrVideoDataset by setting:")
console.print(" reference_column='[cyan]reference_path[/]'")
console.print(" video_column='[cyan]media_path[/]'")
console.print("[bold yellow]Note:[/] Reference videos were written to the '[cyan]reference_video[/]' column.")
console.print(" [cyan]process_dataset.py[/] detects this column automatically for IC-LoRA preprocessing.")
def process_media(
@@ -158,7 +169,7 @@ def process_media(
def media_path_to_reference_path(media_file: Path) -> Path:
return media_file.parent / (media_file.stem + "_reference" + media_file.suffix)
media_files = [base_dir / Path(sample["media_path"]) for sample in meta_data]
media_files = [base_dir / Path(_get_media_path(sample)) for sample in meta_data]
for media_file in media_files:
reference_path = media_path_to_reference_path(media_file)
media_to_process.append(media_file)
@@ -178,18 +189,23 @@ def process_media(
)
# Process media files
media_paths = [item["media_path"] for item in meta_data]
media_paths = [_get_media_path(item) for item in meta_data]
reference_paths = {rel_path: str(media_path_to_reference_path(Path(rel_path))) for rel_path in media_paths}
with progress:
task = progress.add_task("Computing condition on videos", total=len(media_to_process))
for media_file in media_to_process:
for media_file, rel_path in zip(media_to_process, media_paths, strict=True):
progress.update(task, description=f"Processing [bold blue]{media_file.name}[/]")
rel_path = str(media_file.resolve().relative_to(base_dir))
# Key by the original media-path string (matches the dict seeded above). Avoid
# resolve()/relative_to here — they crash on symlinked or absolute media paths.
reference_path = media_path_to_reference_path(media_file)
reference_paths[rel_path] = str(reference_path.relative_to(base_dir))
try:
ref_stored = str(reference_path.relative_to(base_dir))
except ValueError:
ref_stored = str(reference_path) # absolute/out-of-tree: keep it next to the source
reference_paths[rel_path] = ref_stored
if not reference_path.resolve().exists() or override:
try:
packages/ltx-trainer/scripts/decode_latents.py
+1 -1
View File
@@ -310,7 +310,7 @@ def main(
help="Device to use for computation",
),
vae_tiling: bool = typer.Option(
default=False,
default=True,
help="Enable VAE tiling for larger video resolutions",
),
seed: int | None = typer.Option(
packages/ltx-trainer/scripts/inference.py
-443
View File
@@ -1,443 +0,0 @@
#!/usr/bin/env python3
# ruff: noqa: T201
"""
CLI script for running LTX video/audio generation inference.
Usage:
# Text-to-Video + Audio (default behavior)
python scripts/inference.py --checkpoint path/to/model.safetensors \
--text-encoder-path path/to/gemma \
--prompt "A cat playing with a ball" --output output.mp4
# Video only (skip audio)
python scripts/inference.py --checkpoint path/to/model.safetensors \
--text-encoder-path path/to/gemma \
--prompt "A cat playing with a ball" --skip-audio --output output.mp4
# Image-to-Video
python scripts/inference.py --checkpoint path/to/model.safetensors \
--text-encoder-path path/to/gemma \
--prompt "A cat walking" --condition-image first_frame.png --output output.mp4
# Video-to-Video (IC-LoRA style)
python scripts/inference.py --checkpoint path/to/model.safetensors \
--text-encoder-path path/to/gemma \
--prompt "A cat turning into a dog" --reference-video input.mp4 --output output.mp4
# With LoRA weights
python scripts/inference.py --checkpoint path/to/model.safetensors \
--text-encoder-path path/to/gemma \
--lora-path path/to/lora.safetensors \
--prompt "A cat in my custom style" --output output.mp4
"""
import argparse
import re
from pathlib import Path
import torch
import torchaudio
from peft import LoraConfig, get_peft_model, set_peft_model_state_dict
from safetensors.torch import load_file
from torchvision import transforms
from ltx_trainer.model_loader import load_model
from ltx_trainer.progress import StandaloneSamplingProgress
from ltx_trainer.utils import open_image_as_srgb
from ltx_trainer.validation_sampler import GenerationConfig, ValidationSampler
from ltx_trainer.video_utils import read_video, save_video
def load_image(image_path: str) -> torch.Tensor:
"""Load an image and convert to tensor [C, H, W] in [0, 1]."""
image = open_image_as_srgb(image_path)
transform = transforms.ToTensor()
return transform(image)
def extract_lora_target_modules(state_dict: dict[str, torch.Tensor]) -> list[str]:
"""Extract target module names from LoRA checkpoint keys.
LoRA keys follow the pattern (after removing "diffusion_model." prefix):
- transformer_blocks.0.attn1.to_k.lora_A.weight
- transformer_blocks.0.ff.net.0.proj.lora_B.weight
This extracts the full module path like "transformer_blocks.0.attn1.to_k".
Using full paths is more robust than partial patterns.
"""
target_modules = set()
# Pattern to extract everything before .lora_A or .lora_B
pattern = re.compile(r"(.+)\.lora_[AB]\.")
for key in state_dict:
match = pattern.match(key)
if match:
module_path = match.group(1)
target_modules.add(module_path)
return sorted(target_modules)
def load_lora_weights(transformer: torch.nn.Module, lora_path: str | Path) -> torch.nn.Module:
"""Load LoRA weights into the transformer model.
The LoRA rank and target modules are automatically detected from the checkpoint.
Alpha is set equal to rank (standard practice for inference).
Args:
transformer: The base transformer model
lora_path: Path to the LoRA weights (.safetensors)
Returns:
The transformer model with LoRA weights applied
"""
print(f"Loading LoRA weights from {lora_path}...")
# Load the LoRA state dict
state_dict = load_file(str(lora_path))
# Remove "diffusion_model." prefix (ComfyUI-compatible format)
state_dict = {k.replace("diffusion_model.", "", 1): v for k, v in state_dict.items()}
# Extract target modules from the checkpoint
target_modules = extract_lora_target_modules(state_dict)
if not target_modules:
raise ValueError(f"Could not extract target modules from LoRA checkpoint: {lora_path}")
print(f" Detected {len(target_modules)} target modules")
# Auto-detect rank from the first lora_A weight shape
lora_rank = None
for key, value in state_dict.items():
if "lora_A" in key and value.ndim == 2:
lora_rank = value.shape[0]
break
if lora_rank is None:
raise ValueError("Could not auto-detect LoRA rank from weights")
print(f" LoRA rank: {lora_rank}")
# Create LoRA config and wrap the model
# Alpha = rank is standard for inference (maintains the trained scale)
lora_config = LoraConfig(
r=lora_rank,
lora_alpha=lora_rank,
target_modules=target_modules,
lora_dropout=0.0,
init_lora_weights=True,
)
# Wrap the transformer with PEFT to add LoRA layers
transformer = get_peft_model(transformer, lora_config)
# Load the LoRA weights
base_model = transformer.get_base_model()
set_peft_model_state_dict(base_model, state_dict)
print("✓ LoRA weights loaded successfully")
return transformer
def main() -> None: # noqa: PLR0912, PLR0915
parser = argparse.ArgumentParser(
description="LTX Video/Audio Generation",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
# Model arguments
parser.add_argument(
"--checkpoint",
type=str,
required=True,
help="Path to model checkpoint (.safetensors)",
)
parser.add_argument(
"--text-encoder-path",
type=str,
required=True,
help="Path to Gemma text encoder directory",
)
# LoRA arguments
parser.add_argument(
"--lora-path",
type=str,
default=None,
help="Path to LoRA weights (.safetensors)",
)
# Generation arguments
parser.add_argument(
"--prompt",
type=str,
required=True,
help="Text prompt for generation",
)
parser.add_argument(
"--negative-prompt",
type=str,
default="",
help="Negative prompt",
)
parser.add_argument(
"--height",
type=int,
default=544,
help="Video height (must be divisible by 32)",
)
parser.add_argument(
"--width",
type=int,
default=960,
help="Video width (must be divisible by 32)",
)
parser.add_argument(
"--num-frames",
type=int,
default=97,
help="Number of video frames (must be k*8 + 1)",
)
parser.add_argument(
"--frame-rate",
type=float,
default=25.0,
help="Video frame rate",
)
parser.add_argument(
"--num-inference-steps",
type=int,
default=30,
help="Number of denoising steps",
)
parser.add_argument(
"--guidance-scale",
type=float,
default=4.0,
help="Classifier-free guidance scale (CFG)",
)
parser.add_argument(
"--stg-scale",
type=float,
default=1.0,
help="STG (Spatio-Temporal Guidance) scale. 0.0 disables STG. Default: 1.0",
)
parser.add_argument(
"--stg-blocks",
type=int,
nargs="*",
default=[29],
help="Which transformer blocks to perturb for STG. Default: 29 (single block).",
)
parser.add_argument(
"--stg-mode",
type=str,
default="stg_av",
choices=["stg_av", "stg_v"],
help="STG mode: 'stg_av' perturbs both audio and video, 'stg_v' perturbs video only",
)
parser.add_argument(
"--seed",
type=int,
default=42,
help="Random seed for reproducibility",
)
# Conditioning arguments
parser.add_argument(
"--condition-image",
type=str,
default=None,
help="Path to conditioning image for image-to-video generation",
)
parser.add_argument(
"--reference-video",
type=str,
default=None,
help="Path to reference video for video-to-video generation (IC-LoRA style)",
)
parser.add_argument(
"--include-reference-in-output",
action="store_true",
help="Include reference video side-by-side with generated output (only for V2V)",
)
# Audio arguments
parser.add_argument(
"--skip-audio",
action="store_true",
help="Skip audio generation (by default, audio is generated alongside video)",
)
# Output arguments
parser.add_argument(
"--output",
type=str,
required=True,
help="Output video path (.mp4)",
)
parser.add_argument(
"--audio-output",
type=str,
default=None,
help="Output audio path (.wav, optional - if not provided, audio will be embedded in video)",
)
# Device arguments
parser.add_argument(
"--device",
type=str,
default="cuda",
help="Device to run on (cuda/cpu)",
)
args = parser.parse_args()
# Validate conditioning arguments
if args.include_reference_in_output and args.reference_video is None:
parser.error("--include-reference-in-output requires --reference-video")
# Validate arguments
generate_audio = not args.skip_audio
print("=" * 80)
print("LTX Video/Audio Generation")
print("=" * 80)
# Determine if we need VAE encoder (for image or video conditioning)
need_vae_encoder = args.condition_image is not None or args.reference_video is not None
components = load_model(
checkpoint_path=args.checkpoint,
device="cpu", # Load to CPU first, sampler will move to device as needed
dtype=torch.bfloat16,
with_video_vae_encoder=need_vae_encoder,
with_video_vae_decoder=True,
with_audio_vae_decoder=generate_audio,
with_vocoder=generate_audio,
with_text_encoder=True,
text_encoder_path=args.text_encoder_path,
)
# Apply LoRA weights if provided
transformer = components.transformer
if args.lora_path is not None:
transformer = load_lora_weights(transformer, args.lora_path)
# Load conditioning image if provided
condition_image = None
if args.condition_image:
print(f"Loading conditioning image from {args.condition_image}...")
condition_image = load_image(args.condition_image)
# Load reference video if provided
reference_video = None
if args.reference_video:
print(f"Loading reference video from {args.reference_video}...")
reference_video, ref_fps = read_video(args.reference_video, max_frames=args.num_frames)
print(f" Loaded {reference_video.shape[0]} frames @ {ref_fps:.1f} fps")
# Determine generation mode
if args.reference_video is not None and args.condition_image is not None:
mode = "Video-to-Video + Image Conditioning (V2V+I2V)"
elif args.reference_video is not None:
mode = "Video-to-Video (V2V)"
elif args.condition_image is not None:
mode = "Image-to-Video (I2V)"
else:
mode = "Text-to-Video (T2V)"
print("\n" + "=" * 80)
print("Generation Parameters")
print("=" * 80)
print(f"Mode: {mode}")
print(f"Prompt: {args.prompt}")
if args.negative_prompt:
print(f"Negative prompt: {args.negative_prompt}")
print(f"Resolution: {args.width}x{args.height}")
print(f"Frames: {args.num_frames} @ {args.frame_rate} fps")
print(f"Inference steps: {args.num_inference_steps}")
print(f"CFG scale: {args.guidance_scale}")
if args.stg_scale > 0:
blocks_str = args.stg_blocks if args.stg_blocks else "all"
print(f"STG scale: {args.stg_scale} (mode: {args.stg_mode}, blocks: {blocks_str})")
else:
print("STG: disabled")
print(f"Seed: {args.seed}")
if args.lora_path:
print(f"LoRA: {args.lora_path}")
if condition_image is not None:
print(f"Conditioning: Image ({args.condition_image})")
if reference_video is not None:
print(f"Reference: Video ({args.reference_video})")
if args.include_reference_in_output:
print(" → Will include reference side-by-side in output")
if generate_audio:
video_duration = args.num_frames / args.frame_rate
print(f"Audio: Enabled (duration will match video: {video_duration:.2f}s)")
print("=" * 80)
print(f"\nGenerating {'video + audio' if generate_audio else 'video'}...")
# Create generation config
gen_config = GenerationConfig(
prompt=args.prompt,
negative_prompt=args.negative_prompt,
height=args.height,
width=args.width,
num_frames=args.num_frames,
frame_rate=args.frame_rate,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
condition_image=condition_image,
reference_video=reference_video,
generate_audio=generate_audio,
include_reference_in_output=args.include_reference_in_output,
stg_scale=args.stg_scale,
stg_blocks=args.stg_blocks,
stg_mode=args.stg_mode,
)
# Generate with progress bar
with StandaloneSamplingProgress(num_steps=args.num_inference_steps) as progress:
# Create sampler with progress context
sampler = ValidationSampler(
transformer=transformer,
vae_decoder=components.video_vae_decoder,
vae_encoder=components.video_vae_encoder,
text_encoder=components.text_encoder,
audio_decoder=components.audio_vae_decoder if generate_audio else None,
vocoder=components.vocoder if generate_audio else None,
sampling_context=progress,
)
video, audio = sampler.generate(
config=gen_config,
device=args.device,
)
# Save video
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
# Get audio sample rate from vocoder if audio was generated
audio_sample_rate = None
if audio is not None and components.vocoder is not None:
audio_sample_rate = components.vocoder.output_sampling_rate
save_video(
video_tensor=video,
output_path=output_path,
fps=args.frame_rate,
audio=audio,
audio_sample_rate=audio_sample_rate,
)
print(f"✓ Video saved to {args.output}")
# Save separate audio file if requested
if audio is not None and args.audio_output is not None:
audio_output_path = Path(args.audio_output)
audio_output_path.parent.mkdir(parents=True, exist_ok=True)
torchaudio.save(
str(audio_output_path),
audio.cpu(),
sample_rate=audio_sample_rate,
)
duration = audio.shape[1] / audio_sample_rate
print(f"✓ Audio saved: {duration:.2f}s at {audio_sample_rate}Hz")
print("\n" + "=" * 80)
print("Generation complete!")
print("=" * 80)
if __name__ == "__main__":
main()
packages/ltx-trainer/scripts/process_captions.py
+16 -7
View File
@@ -154,6 +154,17 @@ class CaptionsDataset(Dataset):
else:
raise ValueError("Expected `dataset_file` to be a path to a CSV, JSON, or JSONL file.")
def _embedding_output_path(self, media_path: Path) -> str:
"""Output `.pt` path relative to the dataset dir; mirrors `process_videos._output_relative`
so caption keys match video/audio latent keys (and absolute paths don't escape output_dir)."""
data_root = self.dataset_file.parent
resolved = data_root / media_path # pathlib: an absolute media_path overrides data_root
try:
rel = resolved.relative_to(data_root)
except ValueError:
rel = Path(*resolved.parts[1:]) if resolved.is_absolute() else resolved
return str(rel.with_suffix(".pt"))
def _load_caption_data_from_csv(self) -> dict[str, str]:
"""Load captions from a CSV file and compute output embedding paths."""
df = pd.read_csv(self.dataset_file)
@@ -166,8 +177,7 @@ class CaptionsDataset(Dataset):
caption_data = {}
for _, row in df.iterrows():
media_path = Path(row[self.media_column].strip())
# Convert media path to embedding output path (same structure, .pt extension)
output_path = str(media_path.with_suffix(".pt"))
output_path = self._embedding_output_path(media_path)
caption_data[output_path] = row[self.caption_column]
return caption_data
@@ -188,8 +198,7 @@ class CaptionsDataset(Dataset):
raise ValueError(f"Key '{self.media_column}' not found in JSON entry: {entry}")
media_path = Path(entry[self.media_column].strip())
# Convert media path to embedding output path (same structure, .pt extension)
output_path = str(media_path.with_suffix(".pt"))
output_path = self._embedding_output_path(media_path)
caption_data[output_path] = entry[self.caption_column]
return caption_data
@@ -206,8 +215,7 @@ class CaptionsDataset(Dataset):
raise ValueError(f"Key '{self.media_column}' not found in JSONL entry: {entry}")
media_path = Path(entry[self.media_column].strip())
# Convert media path to embedding output path (same structure, .pt extension)
output_path = str(media_path.with_suffix(".pt"))
output_path = self._embedding_output_path(media_path)
caption_data[output_path] = entry[self.caption_column]
return caption_data
@@ -326,7 +334,8 @@ def compute_captions_embeddings( # noqa: PLR0913
# 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"])):
hidden_states, prompt_attention_mask = text_encoder.encode(batch["prompt"][i], padding_side="left")
encoded = text_encoder.encode([batch["prompt"][i]], padding_side="left")
hidden_states, prompt_attention_mask = encoded[0]
video_prompt_embeds, audio_prompt_embeds = embeddings_processor.feature_extractor(
hidden_states, prompt_attention_mask, "left"
)
packages/ltx-trainer/scripts/process_dataset.py
+268 -165
View File
@@ -1,14 +1,21 @@
#!/usr/bin/env python3
"""
Preprocess a video dataset by computing video clips latents and text captions embeddings.
This script provides a command-line interface for preprocessing video datasets by computing
latent representations of video clips and text embeddings of their captions. The preprocessed
data can be used to accelerate training of video generation models and to save GPU memory.
Preprocess a media dataset for LTX-2 training.
Automatically detects dataset columns and processes each according to a convention table.
Column names determine what gets encoded and where outputs go — no per-role CLI flags needed.
Convention table:
video → Video VAE → latents/
audio → Audio VAE → audio_latents/
reference_video → Video VAE → reference_latents/
reference_audio → Audio VAE → reference_audio_latents/
video_mask → (downsample) → video_masks/
audio_mask → (downsample) → audio_masks/
caption → Text encoder → conditions/
Legacy aliases: media_path → video, ref_media_path → reference_video
Basic usage:
python scripts/process_dataset.py /path/to/dataset.json --resolution-buckets 768x768x49 \
python scripts/process_dataset.py /path/to/dataset.json --resolution-buckets 768x768x49 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma
The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.
"""
from pathlib import Path
@@ -16,7 +23,15 @@ from pathlib import Path
import typer
from decode_latents import LatentsDecoder
from process_captions import compute_captions_embeddings
from process_videos import compute_latents, compute_scaled_resolution_buckets, parse_resolution_buckets
from process_videos import (
compute_audio_latents,
compute_audio_masks,
compute_latents,
compute_scaled_resolution_buckets,
compute_video_masks,
detect_dataset_columns,
parse_resolution_buckets,
)
from rich.console import Console
from ltx_trainer import logger
@@ -27,52 +42,82 @@ console = Console()
app = typer.Typer(
pretty_exceptions_enable=False,
no_args_is_help=True,
help="Preprocess a video dataset by computing video clips latents and text captions embeddings. "
"The dataset must be a CSV, JSON, or JSONL file with columns for captions and video paths.",
help="Preprocess a media dataset for LTX-2 training. "
"Automatically detects columns (video, audio, reference_video, reference_audio, caption) "
"and processes each with the appropriate encoder.",
)
_KNOWN_ROLES = {"video", "audio", "reference_video", "reference_audio", "video_mask", "audio_mask", "caption"}
_LEGACY_ALIASES = {"media_path": "video", "ref_media_path": "reference_video"}
def preprocess_dataset( # noqa: PLR0913
def preprocess_dataset( # noqa: PLR0912, PLR0913, PLR0915
dataset_file: str,
caption_column: str,
video_column: str,
resolution_buckets: list[tuple[int, int, int]],
batch_size: int,
output_dir: str | None,
lora_trigger: str | None,
vae_tiling: bool,
decode: bool,
resolution_buckets: list[tuple[int, int, int]] | None,
model_path: str,
text_encoder_path: str,
device: str,
output_dir: str | None = None,
video_column: str | None = None,
caption_column: str | None = None,
batch_size: int = 1,
lora_trigger: str | None = None,
vae_tiling: bool = False,
decode: bool = False,
remove_llm_prefixes: bool = False,
reference_column: str | None = None,
reference_downscale_factor: int = 1,
with_audio: bool = False,
reference_temporal_scale_factor: int = 1,
skip_audio: bool = False,
audio_durations: list[float] | None = None,
load_text_encoder_in_8bit: bool = False,
overwrite: bool = False,
) -> None:
"""Run the preprocessing pipeline with the given arguments."""
# Validate dataset file
"""Run the preprocessing pipeline with convention-based column detection."""
_validate_dataset_file(dataset_file)
# Set up output directories
# Detect columns and resolve roles
dataset_columns = detect_dataset_columns(dataset_file)
roles = _resolve_columns(dataset_columns, video_column, caption_column)
# Log detected roles
for role, col in sorted(roles.items()):
alias_note = f" (alias for '{role}')" if col != role else ""
logger.info(f"Detected column '{col}'{alias_note} → {role}")
# Validate: need at least caption
if "caption" not in roles:
raise ValueError(
f"No caption column found. Dataset has columns: {dataset_columns}. "
f"Expected 'caption' or use --caption-column to specify."
)
# Validate: need video or audio
has_video = "video" in roles
has_audio = "audio" in roles
if not has_video and not has_audio:
raise ValueError(
f"No media column found. Dataset has columns: {dataset_columns}. "
f"Expected 'video', 'audio', or 'media_path' (legacy)."
)
# Validate: video modes need resolution buckets
if has_video and not resolution_buckets:
raise ValueError("--resolution-buckets is required when the dataset has a video column.")
output_base = Path(output_dir) if output_dir else Path(dataset_file).parent / ".precomputed"
conditions_dir = output_base / "conditions"
latents_dir = output_base / "latents"
if lora_trigger:
logger.info(f'LoRA trigger word "{lora_trigger}" will be prepended to all captions')
# --- Phase 1: Text encoder ---
with free_gpu_memory_context():
# Process captions using the dedicated function
compute_captions_embeddings(
dataset_file=dataset_file,
output_dir=str(conditions_dir),
output_dir=str(output_base / "conditions"),
model_path=model_path,
text_encoder_path=text_encoder_path,
caption_column=caption_column,
media_column=video_column,
caption_column=roles["caption"],
media_column=roles.get("video") or roles.get("audio") or roles["caption"],
lora_trigger=lora_trigger,
remove_llm_prefixes=remove_llm_prefixes,
batch_size=batch_size,
@@ -81,119 +126,177 @@ def preprocess_dataset( # noqa: PLR0913
overwrite=overwrite,
)
# Process videos using the dedicated function
audio_latents_dir = None
if with_audio:
logger.info("Audio preprocessing enabled - will extract and encode audio from videos")
audio_latents_dir = output_base / "audio_latents"
# --- Phase 2: Video VAE (video, reference_video) ---
if has_video and resolution_buckets:
# Determine if audio should be auto-extracted from video files
auto_audio = not skip_audio and "audio" not in roles
with free_gpu_memory_context():
compute_latents(
dataset_file=dataset_file,
video_column=video_column,
resolution_buckets=resolution_buckets,
output_dir=str(latents_dir),
model_path=model_path,
batch_size=batch_size,
device=device,
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
if reference_column:
# Validate: scaled references with multiple buckets can cause ambiguous bucket matching
if reference_downscale_factor > 1 and len(resolution_buckets) > 1:
raise ValueError(
"When using --reference-downscale-factor > 1, only a single resolution bucket is supported. "
"Using multiple buckets with scaled references can cause ambiguous bucket matching "
"(e.g., a 512x256 reference could match either the scaled-down 1024x512 bucket or the 512x256 "
"bucket). Please use a single resolution bucket or set --reference-downscale-factor to 1."
)
# Calculate and validate scaled resolution buckets for reference videos
reference_buckets = compute_scaled_resolution_buckets(resolution_buckets, reference_downscale_factor)
if reference_downscale_factor > 1:
logger.info(
f"Processing reference videos for IC-LoRA training at 1/{reference_downscale_factor} resolution..."
)
logger.info(f"Reference resolution buckets: {reference_buckets}")
else:
logger.info("Processing reference videos for IC-LoRA training...")
reference_latents_dir = output_base / "reference_latents"
audio_latents_dir = str(output_base / "audio_latents") if auto_audio else None
if auto_audio:
logger.info("Audio will be auto-extracted from video files (use --skip-audio to disable)")
with free_gpu_memory_context():
compute_latents(
dataset_file=dataset_file,
main_media_column=video_column,
video_column=reference_column,
resolution_buckets=reference_buckets,
output_dir=str(reference_latents_dir),
video_column=roles["video"],
resolution_buckets=resolution_buckets,
output_dir=str(output_base / "latents"),
model_path=model_path,
batch_size=batch_size,
device=device,
vae_tiling=vae_tiling,
with_audio=auto_audio,
audio_output_dir=audio_latents_dir,
overwrite=overwrite,
)
# Handle decoding if requested (for verification)
# Process reference video if present
if "reference_video" in roles:
if reference_downscale_factor > 1 and len(resolution_buckets) > 1:
raise ValueError(
"When using --reference-downscale-factor > 1, only a single resolution bucket is supported."
)
if reference_temporal_scale_factor > 1 and len(resolution_buckets) > 1:
raise ValueError(
"When using --reference-temporal-scale-factor > 1, only a single resolution bucket is supported."
)
reference_buckets = compute_scaled_resolution_buckets(resolution_buckets, reference_downscale_factor)
if reference_downscale_factor > 1:
logger.info(f"Processing reference videos at 1/{reference_downscale_factor} resolution...")
if reference_temporal_scale_factor > 1:
logger.info(
f"Temporally subsampling reference videos by {reference_temporal_scale_factor}x "
f"(VAE-aligned pattern)..."
)
with free_gpu_memory_context():
compute_latents(
dataset_file=dataset_file,
main_media_column=roles["video"],
video_column=roles["reference_video"],
resolution_buckets=reference_buckets,
output_dir=str(output_base / "reference_latents"),
model_path=model_path,
batch_size=batch_size,
device=device,
vae_tiling=vae_tiling,
overwrite=overwrite,
temporal_subsample_factor=reference_temporal_scale_factor,
)
# --- Phase 2b: Masks (video_mask, audio_mask) — processed after video latents for alignment ---
if "video_mask" in roles and has_video:
compute_video_masks(
dataset_file=dataset_file,
mask_column=roles["video_mask"],
latents_dir=str(output_base / "latents"),
output_dir=str(output_base / "video_masks"),
main_media_column=roles["video"],
)
# --- Phase 3: Audio VAE (audio, reference_audio) ---
audio_roles_to_process = [
("audio", "audio_latents"),
("reference_audio", "reference_audio_latents"),
]
active_audio_roles = [(role, subdir) for role, subdir in audio_roles_to_process if role in roles]
if active_audio_roles:
# Determine audio duration constraint: video bucket → max_duration, or explicit buckets
max_audio_duration = None
audio_duration_buckets = None
if has_video and resolution_buckets:
max_audio_duration = max(f for f, _h, _w in resolution_buckets) / 25.0
elif audio_durations:
audio_duration_buckets = audio_durations
for role, output_subdir in active_audio_roles:
with free_gpu_memory_context():
compute_audio_latents(
dataset_file=dataset_file,
audio_column=roles[role],
output_dir=str(output_base / output_subdir),
model_path=model_path,
main_media_column=roles.get("video"),
max_duration=max_audio_duration,
duration_buckets=audio_duration_buckets,
device=device,
overwrite=overwrite,
)
# --- Phase 4: Audio masks (after audio latents exist for temporal alignment) ---
if "audio_mask" in roles:
audio_latents_source = output_base / "audio_latents"
if audio_latents_source.exists():
compute_audio_masks(
dataset_file=dataset_file,
mask_column=roles["audio_mask"],
audio_latents_dir=str(audio_latents_source),
output_dir=str(output_base / "audio_masks"),
main_media_column=roles.get("video") or roles.get("audio"),
)
else:
logger.warning("audio_mask column found but no audio_latents/ — run with audio first")
# --- Decode for verification ---
if decode:
logger.info("Decoding latents for verification...")
decoder = LatentsDecoder(model_path=model_path, device=device, vae_tiling=vae_tiling, with_audio=has_audio)
if has_video:
decoder.decode(output_base / "latents", output_base / "decoded_videos")
if "reference_video" in roles and (output_base / "reference_latents").exists():
decoder.decode(output_base / "reference_latents", output_base / "decoded_reference_videos")
decoder = LatentsDecoder(
model_path=model_path,
device=device,
vae_tiling=vae_tiling,
with_audio=with_audio,
)
decoder.decode(latents_dir, output_base / "decoded_videos")
# Also decode reference videos if they exist
if reference_column:
reference_latents_dir = output_base / "reference_latents"
if reference_latents_dir.exists():
logger.info("Decoding reference videos...")
decoder.decode(reference_latents_dir, output_base / "decoded_reference_videos")
# Decode audio latents if they exist
if with_audio and audio_latents_dir and audio_latents_dir.exists():
logger.info("Decoding audio latents...")
decoder.decode_audio(audio_latents_dir, output_base / "decoded_audio")
# Print summary
# --- Summary ---
logger.info(f"Dataset preprocessing complete! Results saved to {output_base}")
if reference_column:
logger.info("Reference videos processed and saved to reference_latents/ directory for IC-LoRA training")
if with_audio:
logger.info("Audio latents saved to audio_latents/ directory for audio-video training")
produced = [d.name for d in output_base.iterdir() if d.is_dir() and not d.name.startswith("decoded")]
logger.info(f"Output directories: {', '.join(sorted(produced))}")
def _validate_dataset_file(dataset_path: str) -> None:
"""Validate that the dataset file exists and has the correct format."""
dataset_file = Path(dataset_path)
if not dataset_file.exists():
raise FileNotFoundError(f"Dataset file does not exist: {dataset_file}")
if not dataset_file.is_file():
raise ValueError(f"Dataset path must be a file, not a directory: {dataset_file}")
if dataset_file.suffix.lower() not in [".csv", ".json", ".jsonl"]:
raise ValueError(f"Dataset file must be CSV, JSON, or JSONL format: {dataset_file}")
def _resolve_columns(
dataset_columns: set[str],
video_column_override: str | None = None,
caption_column_override: str | None = None,
) -> dict[str, str]:
"""Map canonical role names to actual dataset column names.
Returns a dict of role → column_name for recognized roles found in the dataset.
"""
roles: dict[str, str] = {}
for col in dataset_columns:
role = _LEGACY_ALIASES.get(col, col)
if role in _KNOWN_ROLES:
roles[role] = col
if video_column_override and video_column_override in dataset_columns:
roles["video"] = video_column_override
if caption_column_override and caption_column_override in dataset_columns:
roles["caption"] = caption_column_override
return roles
@app.command()
def main( # noqa: PLR0913
dataset_path: str = typer.Argument(
...,
help="Path to metadata file (CSV/JSON/JSONL) containing captions and video paths",
help="Path to metadata file (CSV/JSON/JSONL) with columns matching the convention table",
),
resolution_buckets: str = typer.Option(
...,
help='Resolution buckets in format "WxHxF;WxHxF;..." (e.g. "768x768x25;512x512x49")',
resolution_buckets: str | None = typer.Option(
default=None,
help='Resolution buckets in format "WxHxF;WxHxF;..." (e.g. "768x768x25"). '
"Required when dataset has a video column.",
),
model_path: str = typer.Option(
...,
@@ -203,13 +306,13 @@ def main( # noqa: PLR0913
...,
help="Path to Gemma text encoder directory",
),
caption_column: str = typer.Option(
default="caption",
help="Column name containing captions in the dataset JSON/JSONL/CSV file",
caption_column: str | None = typer.Option(
default=None,
help="Override: treat this column as 'caption' (default: auto-detect 'caption')",
),
video_column: str = typer.Option(
default="media_path",
help="Column name containing video paths in the dataset JSON/JSONL/CSV file",
video_column: str | None = typer.Option(
default=None,
help="Override: treat this column as 'video' (default: auto-detect 'video' or 'media_path')",
),
batch_size: int = typer.Option(
default=1,
@@ -229,32 +332,43 @@ def main( # noqa: PLR0913
),
lora_trigger: str | None = typer.Option(
default=None,
help="Optional trigger word to prepend to each caption (activates the LoRA during inference)",
help="Optional trigger word to prepend to each caption",
),
decode: bool = typer.Option(
default=False,
help="Decode and save latents after encoding (videos and audio) for verification",
help="Decode and save latents after encoding for verification",
),
remove_llm_prefixes: bool = typer.Option(
default=False,
help="Remove LLM prefixes from captions",
),
reference_column: str | None = typer.Option(
skip_audio: bool = typer.Option(
default=False,
help="Don't extract audio from video files (audio extraction is on by default)",
),
audio_durations: str | None = typer.Option(
default=None,
help="Column name containing reference video paths (for video-to-video training)",
help='Audio duration buckets in seconds for audio-only datasets (e.g. "2.0;4.0;8.0"). '
"When set, audio files are trimmed to the best matching duration. "
"Not needed when a video column is present (audio duration derived from video bucket).",
),
with_audio: bool = typer.Option(
default=False,
help="Extract and encode audio from video files",
hidden=True,
help="[DEPRECATED: audio is now on by default, use --skip-audio to disable]",
),
load_text_encoder_in_8bit: bool = typer.Option(
default=False,
help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
help="Load the Gemma text encoder in 8-bit precision to save GPU memory",
),
reference_downscale_factor: int = typer.Option(
default=1,
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.",
help="Downscale factor for reference video resolution (e.g., 2 = half resolution for IC-LoRA)",
),
reference_temporal_scale_factor: int = typer.Option(
default=1,
help="Temporal subsampling factor for reference videos (e.g., 2 = half frame rate, "
"VAE-aligned: keeps frame 0, then every Nth frame from frame 1 onwards)",
),
overwrite: bool = typer.Option(
default=False,
@@ -262,64 +376,53 @@ def main( # noqa: PLR0913
"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
"""Preprocess a media dataset for LTX-2 training.
See module docstring for the convention table. Audio is auto-extracted from
video files by default — use --skip-audio to disable.
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:
# Process a dataset with LTX-2 model
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma
# Process dataset with custom column names
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
--caption-column "text" --video-column "video_path"
# Process dataset with reference videos for IC-LoRA training
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
--reference-column "reference_path"
# Process dataset with scaled reference videos (half resolution) for efficient IC-LoRA
python scripts/process_dataset.py dataset.json --resolution-buckets 768x768x25 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
--reference-column "reference_path" --reference-downscale-factor 2
# Process dataset with audio for audio-video training
python scripts/process_dataset.py dataset.json --resolution-buckets 768x512x97 \\
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
--with-audio
"""
parsed_resolution_buckets = parse_resolution_buckets(resolution_buckets)
if len(parsed_resolution_buckets) > 1:
# Handle deprecated --with-audio flag
if with_audio:
logger.warning(
"Using multiple resolution buckets. "
"When training with multiple resolution buckets, you must use a batch size of 1."
"--with-audio is deprecated. Audio extraction is now on by default. Use --skip-audio to disable."
)
# Validate reference_downscale_factor
parsed_buckets = parse_resolution_buckets(resolution_buckets) if resolution_buckets else None
if parsed_buckets and len(parsed_buckets) > 1:
logger.warning("Using multiple resolution buckets. Training batch size must be 1.")
if reference_downscale_factor < 1:
raise typer.BadParameter("--reference-downscale-factor must be >= 1")
if reference_downscale_factor > 1 and not reference_column:
logger.warning("--reference-downscale-factor specified but no --reference-column provided. Ignoring.")
if reference_temporal_scale_factor < 1:
raise typer.BadParameter("--reference-temporal-scale-factor must be >= 1")
parsed_audio_durations = None
if audio_durations:
parsed_audio_durations = [float(d) for d in audio_durations.split(";")]
if any(d <= 0 for d in parsed_audio_durations):
raise typer.BadParameter("All audio durations must be positive")
preprocess_dataset(
dataset_file=dataset_path,
caption_column=caption_column,
video_column=video_column,
resolution_buckets=parsed_resolution_buckets,
batch_size=batch_size,
output_dir=output_dir,
lora_trigger=lora_trigger,
vae_tiling=vae_tiling,
decode=decode,
resolution_buckets=parsed_buckets,
model_path=model_path,
text_encoder_path=text_encoder_path,
device=device,
output_dir=output_dir,
video_column=video_column,
caption_column=caption_column,
batch_size=batch_size,
lora_trigger=lora_trigger,
vae_tiling=vae_tiling,
decode=decode,
remove_llm_prefixes=remove_llm_prefixes,
reference_column=reference_column,
reference_downscale_factor=reference_downscale_factor,
with_audio=with_audio,
reference_temporal_scale_factor=reference_temporal_scale_factor,
skip_audio=skip_audio,
audio_durations=parsed_audio_durations,
load_text_encoder_in_8bit=load_text_encoder_in_8bit,
overwrite=overwrite,
)
packages/ltx-trainer/scripts/process_videos.py
+480 -106
View File
@@ -41,6 +41,7 @@ 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
from torchvision.transforms.functional import resize as tv_resize
from transformers.utils.logging import disable_progress_bar
from ltx_core.model.audio_vae import AudioProcessor
@@ -73,6 +74,10 @@ app = typer.Typer(
)
def _clamp_01(x: torch.Tensor) -> torch.Tensor:
return x.clamp_(0, 1)
class MediaDataset(Dataset):
"""
Dataset for processing video and image files.
@@ -92,6 +97,7 @@ class MediaDataset(Dataset):
resolution_buckets: list[tuple[int, int, int]],
reshape_mode: str = "center",
with_audio: bool = False,
temporal_subsample_factor: int = 1,
) -> None:
"""
Initialize the media dataset.
@@ -101,6 +107,8 @@ class MediaDataset(Dataset):
resolution_buckets: List of (frames, height, width) tuples
reshape_mode: How to crop videos ("center", "random")
with_audio: Whether to extract audio from video files
temporal_subsample_factor: Factor for VAE-aligned temporal subsampling.
When > 1, keeps frame 0 then takes every Nth frame from frame 1 onwards.
"""
super().__init__()
@@ -109,6 +117,7 @@ class MediaDataset(Dataset):
self.resolution_buckets = resolution_buckets
self.reshape_mode = reshape_mode
self.with_audio = with_audio
self.temporal_subsample_factor = temporal_subsample_factor
# First load main media paths
self.main_media_paths = self._load_video_paths(main_media_column)
@@ -124,7 +133,7 @@ class MediaDataset(Dataset):
# Set up video transforms
self.transforms = transforms.Compose(
[
transforms.Lambda(lambda x: x.clamp_(0, 1)),
transforms.Lambda(_clamp_01),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
]
)
@@ -142,8 +151,8 @@ class MediaDataset(Dataset):
# Compute relative path of the video
data_root = self.dataset_file.parent
relative_path = str(video_path.relative_to(data_root))
media_relative_path = str(self.main_media_paths[index].relative_to(data_root))
relative_path = str(_output_relative(video_path, data_root))
media_relative_path = str(_output_relative(self.main_media_paths[index], data_root))
if video_path.suffix.lower() in [".png", ".jpg", ".jpeg"]:
media_tensor = self._preprocess_image(video_path)
@@ -185,97 +194,29 @@ class MediaDataset(Dataset):
@staticmethod
def _extract_audio(video_path: Path, target_duration: float) -> dict[str, torch.Tensor | int] | None:
"""Extract audio track from a video file, trimmed to match video duration."""
try:
# torchaudio can extract audio from video files directly
# waveform shape: [channels, samples]
waveform, sample_rate = torchaudio.load(str(video_path))
# Trim or pad to target duration
target_samples = int(target_duration * sample_rate)
current_samples = waveform.shape[-1]
if current_samples > target_samples:
# Trim to target duration
waveform = waveform[..., :target_samples]
elif current_samples < target_samples:
# Pad with zeros to target duration
padding = target_samples - current_samples
waveform = torch.nn.functional.pad(waveform, (0, padding))
logger.warning(f"Padded audio to {target_duration:.2f} seconds for {video_path}")
return {"waveform": waveform, "sample_rate": sample_rate}
except Exception as e:
logger.debug(f"Could not extract audio from {video_path}: {e}")
"""Extract audio track from a video file, trimmed/padded to match video duration."""
audio = _load_audio_from_file(video_path, max_duration=target_duration)
if audio is None:
return None
def _load_video_paths(self, column: str) -> list[Path]:
"""Load video paths from the specified data source."""
if self.dataset_file.suffix == ".csv":
return self._load_video_paths_from_csv(column)
elif self.dataset_file.suffix == ".json":
return self._load_video_paths_from_json(column)
elif self.dataset_file.suffix == ".jsonl":
return self._load_video_paths_from_jsonl(column)
# Pad if shorter than target (_load_audio_from_file only trims, doesn't pad)
target_samples = int(target_duration * audio.sampling_rate)
if audio.waveform.shape[-1] < target_samples:
padding = target_samples - audio.waveform.shape[-1]
waveform = torch.nn.functional.pad(audio.waveform, (0, padding))
logger.warning(f"Padded audio to {target_duration:.2f} seconds for {video_path}")
else:
raise ValueError("Expected `dataset_file` to be a path to a CSV, JSON, or JSONL file.")
waveform = audio.waveform
def _load_video_paths_from_csv(self, column: str) -> list[Path]:
"""Load video paths from a CSV file."""
df = pd.read_csv(self.dataset_file)
if column not in df.columns:
raise ValueError(f"Column '{column}' not found in CSV file")
return {"waveform": waveform, "sample_rate": audio.sampling_rate}
data_root = self.dataset_file.parent
video_paths = [data_root / Path(line.strip()) for line in df[column].tolist()]
# Validate that all paths exist
invalid_paths = [path for path in video_paths if not path.is_file()]
if invalid_paths:
raise ValueError(f"Found {len(invalid_paths)} invalid video paths. First few: {invalid_paths[:5]}")
return video_paths
def _load_video_paths_from_json(self, column: str) -> list[Path]:
"""Load video paths from a JSON file."""
with open(self.dataset_file, "r", encoding="utf-8") as file:
data = json.load(file)
if not isinstance(data, list):
raise ValueError("JSON file must contain a list of objects")
data_root = self.dataset_file.parent
video_paths = []
for entry in data:
if column not in entry:
raise ValueError(f"Key '{column}' not found in JSON entry")
video_paths.append(data_root / Path(entry[column].strip()))
# Validate that all paths exist
invalid_paths = [path for path in video_paths if not path.is_file()]
if invalid_paths:
raise ValueError(f"Found {len(invalid_paths)} invalid video paths. First few: {invalid_paths[:5]}")
return video_paths
def _load_video_paths_from_jsonl(self, column: str) -> list[Path]:
"""Load video paths from a JSONL file."""
data_root = self.dataset_file.parent
video_paths = []
with open(self.dataset_file, "r", encoding="utf-8") as file:
for line in file:
entry = json.loads(line)
if column not in entry:
raise ValueError(f"Key '{column}' not found in JSONL entry")
video_paths.append(data_root / Path(entry[column].strip()))
# Validate that all paths exist
invalid_paths = [path for path in video_paths if not path.is_file()]
if invalid_paths:
raise ValueError(f"Found {len(invalid_paths)} invalid video paths. First few: {invalid_paths[:5]}")
return video_paths
def _load_video_paths(self, column: str) -> list[Path]:
"""Load video paths from the specified data source, validating existence."""
paths = _load_paths_from_dataset(self.dataset_file, column)
invalid = [p for p in paths if not p.is_file()]
if invalid:
raise ValueError(f"Found {len(invalid)} invalid paths in '{column}'. First few: {invalid[:5]}")
return paths
def _filter_valid_videos(self) -> None:
"""Filter out videos with insufficient frames."""
@@ -348,6 +289,11 @@ class MediaDataset(Dataset):
# Trim video to target number of frames
frames_resized = frames_resized[:target_num_frames]
# VAE-aligned temporal subsampling: keep frame 0, then every Nth frame
if self.temporal_subsample_factor > 1:
indices = _compute_temporal_subsample_indices(target_num_frames, self.temporal_subsample_factor)
frames_resized = frames_resized[indices]
# Apply transforms to each frame and stack
video = torch.stack([self.transforms(frame) for frame in frames_resized], dim=0)
@@ -434,7 +380,18 @@ class MediaDataset(Dataset):
return media_tensor
def compute_latents( # noqa: PLR0913, PLR0915
def _compute_temporal_subsample_indices(num_frames: int, factor: int) -> list[int]:
"""Compute VAE-aligned temporal subsample indices.
Keeps frame 0 (the VAE's standalone first-frame latent), then takes every
``factor``-th frame from frame 1 onwards. This ensures each resulting
8-frame VAE group spans ``factor`` groups of the original video.
"""
if factor == 1:
return list(range(num_frames))
return [0, *list(range(1, num_frames, factor))]
def compute_latents( # noqa: PLR0912, PLR0913, PLR0915
dataset_file: str | Path,
video_column: str,
resolution_buckets: list[tuple[int, int, int]],
@@ -447,7 +404,9 @@ def compute_latents( # noqa: PLR0913, PLR0915
vae_tiling: bool = False,
with_audio: bool = False,
audio_output_dir: str | None = None,
num_dataloader_workers: int = 4,
overwrite: bool = False,
temporal_subsample_factor: int = 1,
) -> None:
"""
Process videos and save latent representations.
@@ -468,9 +427,28 @@ 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)
num_dataloader_workers: Number of DataLoader worker processes (0 for in-process loading)
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.
temporal_subsample_factor: Factor for VAE-aligned temporal subsampling of reference videos
"""
# Validate temporal subsampling compatibility with resolution buckets
if temporal_subsample_factor > 1:
for frames, _h, _w in resolution_buckets:
pixel_frames_minus_one = frames - 1
if pixel_frames_minus_one % temporal_subsample_factor != 0:
raise ValueError(
f"Frame count {frames} is not compatible with "
f"temporal_subsample_factor={temporal_subsample_factor}. "
f"(frames - 1) must be divisible by the factor."
)
subsampled = 1 + pixel_frames_minus_one // temporal_subsample_factor
if (subsampled - 1) % VAE_TEMPORAL_FACTOR != 0:
raise ValueError(
f"After temporal subsampling {frames} → {subsampled} frames, "
f"result does not satisfy (frames - 1) % {VAE_TEMPORAL_FACTOR} == 0."
)
if with_audio and audio_output_dir is None:
raise ValueError("audio_output_dir must be provided when with_audio=True")
@@ -484,11 +462,13 @@ def compute_latents( # noqa: PLR0913, PLR0915
resolution_buckets=resolution_buckets,
reshape_mode=reshape_mode,
with_audio=with_audio,
temporal_subsample_factor=temporal_subsample_factor,
)
logger.info(f"Loaded {len(dataset)} valid media files")
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
audio_output_path: Path | None = None
if with_audio:
audio_output_path = Path(audio_output_dir)
@@ -499,10 +479,10 @@ def compute_latents( # noqa: PLR0913, PLR0915
logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
batch_size = 1
data_root = dataset.dataset_file.parent
data_root = Path(dataset_file).parent
def _is_done(idx: int) -> bool:
rel = dataset.main_media_paths[idx].relative_to(data_root).with_suffix(".pt")
rel = _output_relative(dataset.main_media_paths[idx], 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()
@@ -510,18 +490,16 @@ def compute_latents( # noqa: PLR0913, PLR0915
dataloader = _build_sharded_dataloader(
dataset,
batch_size=batch_size,
num_workers=4,
num_workers=num_dataloader_workers,
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)
# Load audio VAE encoder and audio processor if needed
audio_vae_encoder = None
audio_processor = None
if with_audio:
@@ -531,7 +509,6 @@ def compute_latents( # noqa: PLR0913, PLR0915
device=torch_device,
dtype=torch.float32, # Audio VAE needs float32 for quality. TODO: re-test with bfloat16.
)
# Create audio processor for waveform-to-spectrogram conversion
audio_processor = AudioProcessor(
target_sample_rate=audio_vae_encoder.sample_rate,
mel_bins=audio_vae_encoder.mel_bins,
@@ -562,7 +539,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
# Encode video
with torch.inference_mode():
video_latent_data = encode_video(vae=vae, video=video, use_tiling=vae_tiling)
video_latent_data = _encode_video(vae=vae, video=video, use_tiling=vae_tiling)
# Save latents for each item in batch
for i in range(len(batch["relative_path"])):
@@ -572,13 +549,15 @@ def compute_latents( # noqa: PLR0913, PLR0915
# Create output directory maintaining structure
output_file.parent.mkdir(parents=True, exist_ok=True)
# Index into batch to get this item's latents
# Store the latent's effective fps (= source_fps / subsample factor).
# Downstream position math expects the rate the saved latents actually have.
effective_fps = batch["video_metadata"]["fps"][i].item() / temporal_subsample_factor
latent_data = {
"latents": video_latent_data["latents"][i].cpu().contiguous(), # [C, F', H', W']
"num_frames": video_latent_data["num_frames"],
"height": video_latent_data["height"],
"width": video_latent_data["width"],
"fps": batch["video_metadata"]["fps"][i].item(),
"fps": effective_fps,
}
_atomic_save(latent_data, output_file)
@@ -596,7 +575,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
# Encode audio
with torch.inference_mode():
audio_latents = encode_audio(audio_vae_encoder, audio_processor, audio_data)
audio_latents = _encode_audio(audio_vae_encoder, audio_processor, audio_data)
# Save audio latents
audio_output_file = audio_output_path / output_rel_path
@@ -625,7 +604,7 @@ def compute_latents( # noqa: PLR0913, PLR0915
)
def encode_video(
def _encode_video(
vae: torch.nn.Module,
video: torch.Tensor,
dtype: torch.dtype | None = None,
@@ -662,7 +641,7 @@ def encode_video(
# Choose encoding method based on tiling flag
if use_tiling:
latents = tiled_encode_video(
latents = _tiled_encode_video(
vae=vae,
video=video,
tile_size=tile_size,
@@ -685,7 +664,7 @@ def encode_video(
}
def tiled_encode_video( # noqa: PLR0912, PLR0915
def _tiled_encode_video( # noqa: PLR0912, PLR0915
vae: torch.nn.Module,
video: torch.Tensor,
tile_size: int = DEFAULT_TILE_SIZE,
@@ -840,7 +819,7 @@ def tiled_encode_video( # noqa: PLR0912, PLR0915
return output
def encode_audio(
def _encode_audio(
audio_vae_encoder: torch.nn.Module,
audio_processor: torch.nn.Module,
audio: Audio,
@@ -868,6 +847,35 @@ def encode_audio(
if waveform.dim() == 2:
waveform = waveform.unsqueeze(0)
# Convert to stereo if needed (audio VAE expects 2 channels)
# Channel order for surround: 5.1=[L,R,C,LFE,Ls,Rs], 7.1=[L,R,C,LFE,Ls,Rs,Lb,Rb]
num_channels = waveform.shape[1]
if num_channels == 1:
# Mono to stereo: duplicate the channel
waveform = waveform.repeat(1, 2, 1)
elif num_channels == 6:
# 5.1 downmix with normalized weights (sum to 1.0)
# Original: L = L + 0.707*C + 0.707*Ls, weights sum = 2.414
w_main = 1.0 / 2.414 # ~0.414
w_other = 0.707 / 2.414 # ~0.293
left = w_main * waveform[:, 0, :] + w_other * waveform[:, 2, :] + w_other * waveform[:, 4, :]
right = w_main * waveform[:, 1, :] + w_other * waveform[:, 2, :] + w_other * waveform[:, 5, :]
waveform = torch.stack([left, right], dim=1)
elif num_channels == 8:
# 7.1 downmix with normalized weights (sum to 1.0)
# Original: L = L + 0.707*C + 0.707*Ls + 0.707*Lb, weights sum = 3.121
w_main = 1.0 / 3.121 # ~0.320
w_other = 0.707 / 3.121 # ~0.227
center = waveform[:, 2, :]
left = w_main * waveform[:, 0, :] + w_other * (center + waveform[:, 4, :] + waveform[:, 6, :])
right = w_main * waveform[:, 1, :] + w_other * (center + waveform[:, 5, :] + waveform[:, 7, :])
waveform = torch.stack([left, right], dim=1)
elif num_channels > 2:
# Unknown layout: average all channels to mono, then duplicate to stereo
logger.warning(f"Unknown audio channel layout ({num_channels} channels), using mean downmix")
mono = waveform.mean(dim=1, keepdim=True)
waveform = mono.repeat(1, 2, 1)
# Calculate duration
duration = waveform.shape[-1] / audio.sampling_rate
@@ -889,6 +897,372 @@ def encode_audio(
}
AUDIO_FILE_EXTENSIONS = {".wav", ".mp3", ".flac", ".ogg", ".aac", ".m4a"}
VIDEO_FILE_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
IMAGE_FILE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".heic", ".heif", ".bmp", ".tiff", ".webp"}
def compute_video_masks(
dataset_file: str | Path,
mask_column: str,
latents_dir: str,
output_dir: str,
main_media_column: str | None = None,
overwrite: bool = False,
) -> None:
"""Preprocess video mask files to latent-space binary masks.
For each sample, loads the mask video/image, applies the same spatial
resize/crop as the target video (read from saved latent metadata), downsamples
to latent dimensions, binarizes, and saves as a .pt tensor.
Args:
dataset_file: Path to metadata file (CSV/JSON/JSONL).
mask_column: Column name containing mask video/image paths.
latents_dir: Directory containing the target video latents (for reading
spatial/temporal metadata to ensure mask alignment).
output_dir: Directory to save mask .pt files.
main_media_column: Column for output file naming (defaults to mask_column).
"""
dataset_path = Path(dataset_file)
data_root = dataset_path.parent
latents_path = Path(latents_dir)
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
naming_column = main_media_column or mask_column
mask_paths = _load_paths_from_dataset(dataset_path, mask_column)
naming_paths = _load_paths_from_dataset(dataset_path, naming_column) if naming_column != mask_column else mask_paths
success = 0
for mask_file, naming_file in zip(mask_paths, naming_paths, strict=True):
rel_path = _output_relative(naming_file, data_root)
latent_file = latents_path / rel_path.with_suffix(".pt")
out_file = output_path / rel_path.with_suffix(".pt")
if not latent_file.exists():
logger.warning(f"No target latent found at {latent_file}, skipping mask {mask_file}")
continue
if not overwrite and out_file.is_file():
continue
target_meta = torch.load(latent_file, map_location="cpu", weights_only=True)
latent_f = target_meta["num_frames"]
latent_h = target_meta["height"]
latent_w = target_meta["width"]
pixel_h = latent_h * VAE_SPATIAL_FACTOR
pixel_w = latent_w * VAE_SPATIAL_FACTOR
pixel_f = (latent_f - 1) * VAE_TEMPORAL_FACTOR + 1
# Load mask as video or image
if mask_file.suffix.lower() in IMAGE_FILE_EXTENSIONS:
img = to_tensor(open_image_as_srgb(mask_file)).mean(dim=0, keepdim=True) # [1, H, W]
img = tv_resize(img.unsqueeze(0), [pixel_h, pixel_w]).squeeze(0) # [1, H, W]
mask_pixels = img.expand(pixel_f, -1, -1) # tile across frames → [F, H, W]
else:
frames, _ = read_video(str(mask_file), max_frames=pixel_f) # [F, C, H, W]
frames = frames[:pixel_f].mean(dim=1) # grayscale → [F, H, W]
frames = torch.nn.functional.interpolate(
frames.unsqueeze(1), size=(pixel_h, pixel_w), mode="nearest"
).squeeze(1) # [F, H, W]
mask_pixels = frames
# Downsample to latent dims: [F, H, W] → [F', H', W']
mask_latent = torch.nn.functional.avg_pool2d(mask_pixels.unsqueeze(1), kernel_size=VAE_SPATIAL_FACTOR).squeeze(
1
) # [F, H', W'] → spatial done
# Temporal: max-pool over groups of VAE_TEMPORAL_FACTOR frames (any masked frame masks the group)
f_spatial = mask_latent.shape[0]
pad_f = (VAE_TEMPORAL_FACTOR - f_spatial % VAE_TEMPORAL_FACTOR) % VAE_TEMPORAL_FACTOR
if pad_f > 0:
mask_latent = torch.nn.functional.pad(mask_latent, (0, 0, 0, 0, 0, pad_f))
h_prime, w_prime = mask_latent.shape[1], mask_latent.shape[2]
mask_latent = mask_latent.reshape(-1, VAE_TEMPORAL_FACTOR, h_prime, w_prime).amax(dim=1)[:latent_f]
# Binarize
mask_latent = (mask_latent > 0.5).float()
out_file.parent.mkdir(parents=True, exist_ok=True)
_atomic_save({"mask": mask_latent}, out_file)
success += 1
logger.info(f"Mask preprocessing complete: {success} masks saved to {output_path}")
def compute_audio_masks(
dataset_file: str | Path,
mask_column: str,
audio_latents_dir: str,
output_dir: str,
main_media_column: str | None = None,
overwrite: bool = False,
) -> None:
"""Preprocess audio mask files to latent-space binary masks.
For each sample, loads the mask (a 1D waveform-like signal or a simple tensor),
resamples it to match the target audio latent temporal length, binarizes, and saves.
Args:
dataset_file: Path to metadata file (CSV/JSON/JSONL).
mask_column: Column name containing mask file paths (.wav or .pt).
audio_latents_dir: Directory containing the target audio latents (for reading
temporal metadata to ensure mask alignment).
output_dir: Directory to save mask .pt files.
main_media_column: Column for output file naming (defaults to mask_column).
"""
dataset_path = Path(dataset_file)
data_root = dataset_path.parent
audio_latents_path = Path(audio_latents_dir)
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
naming_column = main_media_column or mask_column
mask_paths = _load_paths_from_dataset(dataset_path, mask_column)
naming_paths = _load_paths_from_dataset(dataset_path, naming_column) if naming_column != mask_column else mask_paths
success = 0
for mask_file, naming_file in zip(mask_paths, naming_paths, strict=True):
rel_path = _output_relative(naming_file, data_root)
latent_file = audio_latents_path / rel_path.with_suffix(".pt")
out_file = output_path / rel_path.with_suffix(".pt")
if not latent_file.exists():
logger.warning(f"No target audio latent found at {latent_file}, skipping mask {mask_file}")
continue
if not overwrite and out_file.is_file():
continue
target_meta = torch.load(latent_file, map_location="cpu", weights_only=True)
latent_t = target_meta["num_time_steps"]
# Load mask: .pt file (raw tensor) or .wav (use amplitude envelope)
if mask_file.suffix == ".pt":
raw_mask = torch.load(mask_file, map_location="cpu", weights_only=True)
if isinstance(raw_mask, dict):
raw_mask = raw_mask.get("mask", next(iter(raw_mask.values())))
raw_mask = raw_mask.float().flatten()
else:
audio = _load_audio_from_file(mask_file)
if audio is None:
logger.warning(f"Could not load audio mask from {mask_file}")
continue
raw_mask = audio.waveform.abs().mean(dim=0) # mono amplitude envelope
# Resample to target audio latent length
mask_resampled = torch.nn.functional.interpolate(
raw_mask.unsqueeze(0).unsqueeze(0), size=latent_t, mode="nearest"
).squeeze() # [latent_t]
mask_binary = (mask_resampled > 0.5).float()
out_file.parent.mkdir(parents=True, exist_ok=True)
_atomic_save({"mask": mask_binary}, out_file)
success += 1
logger.info(f"Audio mask preprocessing complete: {success} masks saved to {output_path}")
def compute_audio_latents( # noqa: PLR0915
dataset_file: str | Path,
audio_column: str,
output_dir: str,
model_path: str,
main_media_column: str | None = None,
max_duration: float | None = None,
duration_buckets: list[float] | None = None,
device: str = "cuda",
overwrite: bool = False,
) -> None:
"""Encode audio files into latent representations.
Supports standalone audio files (.wav, .mp3, etc.) and audio tracks
extracted from video files (.mp4, etc.).
Args:
dataset_file: Path to metadata file (CSV/JSON/JSONL).
audio_column: Column name containing audio file paths.
output_dir: Directory to save audio latents.
model_path: Path to LTX-2 checkpoint (.safetensors).
main_media_column: Column for output file naming (defaults to audio_column).
Ensures alignment with other latent directories.
max_duration: Maximum audio duration in seconds. Audio is trimmed if longer.
Mutually exclusive with duration_buckets.
duration_buckets: List of allowed durations in seconds (e.g. [2.0, 4.0, 8.0]).
Each audio file is matched to the largest bucket that fits its duration,
then trimmed to exactly that length. Files shorter than the smallest
bucket are skipped. Ensures uniform lengths for batched training.
device: Device to use for computation.
"""
console = Console()
torch_device = torch.device(device)
dataset_path = Path(dataset_file)
data_root = dataset_path.parent
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
naming_column = main_media_column or audio_column
audio_paths = _load_paths_from_dataset(dataset_path, audio_column)
naming_paths = (
_load_paths_from_dataset(dataset_path, naming_column) if naming_column != audio_column else audio_paths
)
with console.status(f"[bold]Loading audio VAE encoder from [cyan]{model_path}[/]...", spinner="dots"):
audio_vae_encoder = load_audio_vae_encoder(
checkpoint_path=model_path,
device=torch_device,
dtype=torch.float32,
)
audio_processor = AudioProcessor(
target_sample_rate=audio_vae_encoder.sample_rate,
mel_bins=audio_vae_encoder.mel_bins,
mel_hop_length=audio_vae_encoder.mel_hop_length,
n_fft=audio_vae_encoder.n_fft,
).to(torch_device)
sorted_buckets = sorted(duration_buckets, reverse=True) if duration_buckets else None
success_count = 0
skip_count = 0
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
MofNCompleteColumn(),
TimeElapsedColumn(),
TimeRemainingColumn(),
console=console,
) as progress:
task = progress.add_task("Encoding audio", total=len(audio_paths))
for audio_path, naming_path in zip(audio_paths, naming_paths, strict=True):
rel_path = _output_relative(naming_path, data_root)
output_file = output_path / rel_path.with_suffix(".pt")
output_file.parent.mkdir(parents=True, exist_ok=True)
if not overwrite and output_file.is_file():
success_count += 1
progress.advance(task)
continue
# Load audio (no trimming yet — need full duration for bucket matching)
audio = _load_audio_from_file(audio_path)
if audio is None:
skip_count += 1
progress.advance(task)
continue
file_duration = audio.waveform.shape[-1] / audio.sampling_rate
# Determine target duration: bucket matching, max_duration cap, or full file
target_duration = file_duration
if sorted_buckets:
bucket = next((b for b in sorted_buckets if b <= file_duration), None)
if bucket is None:
logger.warning(
f"Skipping {audio_path.name} ({file_duration:.1f}s) — shorter than "
f"smallest bucket ({sorted_buckets[-1]:.1f}s)"
)
skip_count += 1
progress.advance(task)
continue
target_duration = bucket
elif max_duration is not None:
target_duration = min(file_duration, max_duration)
# Trim to target duration
target_samples = int(target_duration * audio.sampling_rate)
trimmed_waveform = audio.waveform[:, :target_samples]
audio = Audio(waveform=trimmed_waveform, sampling_rate=audio.sampling_rate)
with torch.inference_mode():
audio_latents = _encode_audio(audio_vae_encoder, audio_processor, audio)
_atomic_save(
{
"latents": audio_latents["latents"].cpu().contiguous(),
"num_time_steps": audio_latents["num_time_steps"],
"frequency_bins": audio_latents["frequency_bins"],
"duration": audio_latents["duration"],
},
output_file,
)
success_count += 1
progress.advance(task)
logger.info(f"Audio encoding complete: {success_count} encoded, {skip_count} skipped. Saved to {output_path}")
def _output_relative(path: Path, data_root: Path) -> Path:
"""Relative path used to name a sample's cached output, mirroring the input layout.
Normally media lives under the dataset directory and this is just the path relative to it.
If a media path is absolute or otherwise outside the dataset directory (e.g. a one-off
metadata file that references media elsewhere), mirror its absolute structure under the
output directory instead of raising, so out-of-tree media stays collision-free.
"""
try:
return path.relative_to(data_root)
except ValueError:
return Path(*path.parts[1:]) if path.is_absolute() else path
def _load_paths_from_dataset(dataset_file: Path, column: str) -> list[Path]:
"""Load file paths from a dataset column, resolving relative to the dataset file's directory."""
data_root = dataset_file.parent
if dataset_file.suffix == ".csv":
df = pd.read_csv(dataset_file)
if column not in df.columns:
raise ValueError(f"Column '{column}' not found in CSV file")
return [data_root / Path(str(v).strip()) for v in df[column].tolist()]
if dataset_file.suffix == ".json":
with open(dataset_file, encoding="utf-8") as f:
data = json.load(f)
if not isinstance(data, list):
raise ValueError("JSON file must contain a list of objects")
return [data_root / Path(entry[column].strip()) for entry in data]
if dataset_file.suffix == ".jsonl":
paths = []
with open(dataset_file, encoding="utf-8") as f:
for line in f:
entry = json.loads(line)
paths.append(data_root / Path(entry[column].strip()))
return paths
raise ValueError(f"Unsupported dataset format: {dataset_file.suffix}")
def _load_audio_from_file(audio_path: Path, max_duration: float | None = None) -> Audio | None:
"""Load audio from an audio or video file, optionally trimming to max_duration."""
try:
waveform, sample_rate = torchaudio.load(str(audio_path))
except Exception:
logger.debug(f"Could not load audio from {audio_path}")
return None
if max_duration is not None:
max_samples = int(max_duration * sample_rate)
if waveform.shape[-1] > max_samples:
waveform = waveform[:, :max_samples]
return Audio(waveform=waveform, sampling_rate=sample_rate)
def detect_dataset_columns(dataset_file: str | Path) -> set[str]:
"""Read column names from a dataset file without loading all data."""
path = Path(dataset_file)
if path.suffix == ".csv":
df = pd.read_csv(path, nrows=0)
return set(df.columns)
if path.suffix == ".json":
with open(path, encoding="utf-8") as f:
data = json.load(f)
return set(data[0].keys()) if isinstance(data, list) and data else set()
if path.suffix == ".jsonl":
with open(path, encoding="utf-8") as f:
return set(json.loads(f.readline()).keys())
return set()
def parse_resolution_buckets(resolution_buckets_str: str) -> list[tuple[int, int, int]]:
"""Parse resolution buckets from string format to list of tuples (frames, height, width)"""
resolution_buckets = []
@@ -1041,11 +1415,11 @@ def main( # noqa: PLR0913
),
) -> 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.
For multi-GPU preprocessing, invoke under ``accelerate launch`` -- each process
will handle an interleaved shard of the dataset.
Examples:
# Process videos from a CSV file
python scripts/process_videos.py dataset.csv --resolution-buckets 768x768x25 \\
packages/ltx-trainer/scripts/serve_captioner.py
Executable
+203
View File
@@ -0,0 +1,203 @@
#!/usr/bin/env python3
"""Launch a vLLM server for Qwen3-Omni captioning.
Runs the actual server via ``uvx`` so that vLLM and its CUDA-tied
dependencies live in their own isolated environment (no impact on this
package's dependency tree).
The captioning script (``caption_videos.py``) talks to the server over its
OpenAI-compatible HTTP API. Once the server is up it stays loaded across
captioning runs; no per-script model warmup cost.
Typical usage::
# Default: dynamic FP8 quantization, listen on 127.0.0.1:8001
uv run python scripts/serve_captioner.py
# Just print the chosen `uvx vllm serve ...` command without running it
uv run python scripts/serve_captioner.py --print-cmd
# Use full bf16 on a GPU with >= 66 GiB free VRAM (slightly more reliable
# numerics but 2x the weight memory)
uv run python scripts/serve_captioner.py --quantization bf16
# Use a different port or expose on all interfaces
uv run python scripts/serve_captioner.py --port 9000 --host 0.0.0.0
"""
import os
import shutil
import subprocess
import sys
from pathlib import Path
import typer
from rich.console import Console
console = Console()
# Model identifier we serve. The captioner client must use the same string.
DEFAULT_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
# Pinned vLLM version known to support Qwen3-Omni on CUDA 12.x.
# vLLM 0.20+ requires CUDA 13. Update both as the environment evolves.
# The ``[audio]`` extra is required for Qwen3-Omni to decode audio at all.
DEFAULT_VLLM_SPEC = "vllm[audio]==0.11.2"
# Approximate disk needed for the model download (HF cache structure).
MODEL_DISK_GIB = 65.0
app = typer.Typer(
pretty_exceptions_enable=False,
no_args_is_help=False,
help="Launch a local vLLM server for Qwen3-Omni captioning.",
)
def _query_disk_free_gib(path: Path) -> float:
return shutil.disk_usage(str(path)).free / 1024**3
def _build_vllm_args(
*,
model: str,
host: str,
port: int,
quantization: str,
max_model_len: int,
gpu_memory_utilization: float,
extra_args: list[str],
) -> list[str]:
"""Construct the `vllm serve ...` argv."""
args = [
"vllm",
"serve",
model,
"--host",
host,
"--port",
str(port),
"--dtype",
"bfloat16",
"--max-model-len",
str(max_model_len),
"--gpu-memory-utilization",
str(gpu_memory_utilization),
# Let the server accept ``file://`` URLs pointing at local videos.
"--allowed-local-media-path",
"/",
# The model is a multimodal MoE; cap each input to one of each
# modality to match what our captioner sends.
"--limit-mm-per-prompt",
'{"image": 1, "video": 1, "audio": 1}',
# Small concurrent-sequence cap so KV cache headroom isn't fragmented.
"--max-num-seqs",
"4",
]
if quantization == "fp8":
args += ["--quantization", "fp8"]
args += extra_args
return args
@app.command()
def main(
model: str = typer.Option(DEFAULT_MODEL, "--model", help="Model identifier to serve."),
host: str = typer.Option("127.0.0.1", "--host", help="Listen address. Use 0.0.0.0 for remote access."),
port: int = typer.Option(8001, "--port", help="HTTP port."),
quantization: str = typer.Option(
"fp8",
"--quantization",
"-q",
help=(
"Weight precision. 'fp8' (default, dynamic FP8 -- ~31 GiB weights) is "
"the recommended choice; it fits on 40 GiB GPUs and runs at the same "
"speed as bf16 on H100. 'bf16' uses ~60 GiB of weights -- pick it if "
"you have abundant VRAM and want minimal numerical drift."
),
),
max_model_len: int = typer.Option(
32768,
"--max-model-len",
help="Maximum context length the server accepts (must fit input video tokens + max_tokens).",
),
gpu_memory_utilization: float = typer.Option(
0.9,
"--gpu-memory-utilization",
help="Fraction of GPU memory vLLM may reserve (model + KV cache).",
),
hf_home: Path | None = typer.Option( # noqa: B008
None,
"--hf-home",
help=(
"Override HF_HOME (where the model is downloaded). The model is ~65 GB; "
"by default this follows your environment's HF_HOME or HuggingFace's default."
),
),
vllm_spec: str = typer.Option(
DEFAULT_VLLM_SPEC,
"--vllm-spec",
help="pip-style spec passed to `uvx --from`. Pin a version that matches your CUDA.",
),
print_cmd: bool = typer.Option(
False,
"--print-cmd",
help="Print the chosen command without running it.",
),
extra_args: list[str] | None = typer.Argument( # noqa: B008
None,
help="Additional args passed through to `vllm serve` after `--`.",
),
) -> None:
"""Launch the vLLM server for Qwen3-Omni."""
extra = extra_args or []
if quantization not in ("bf16", "fp8"):
console.print(f"[red]--quantization must be 'bf16' or 'fp8'; got {quantization!r}.[/]")
raise typer.Exit(code=1)
# Disk check (only meaningful before first download).
cache_root = hf_home or Path(os.environ.get("HF_HOME", str(Path.home() / ".cache" / "huggingface")))
cache_root.mkdir(parents=True, exist_ok=True)
free_disk = _query_disk_free_gib(cache_root)
if free_disk < MODEL_DISK_GIB:
console.print(
f"[yellow]\u26a0 Only {free_disk:.1f} GiB free on disk under {cache_root} but the "
f"model needs ~{MODEL_DISK_GIB:.0f} GiB. Either free up space, set --hf-home "
f"to a larger volume, or expect the download to fail mid-way.[/]"
)
vllm_args = _build_vllm_args(
model=model,
host=host,
port=port,
quantization=quantization,
max_model_len=max_model_len,
gpu_memory_utilization=gpu_memory_utilization,
extra_args=extra,
)
# Use ``uvx --from vllm==...`` so vLLM lives in its own throwaway venv
# (or a cached tool venv). The `--` separates uvx args from the command's.
uvx_cmd = ["uvx", "--from", vllm_spec, *vllm_args]
env = os.environ.copy()
# vLLM 0.11.x requires the V0 engine for Qwen3-Omni's multimodal pipeline.
env.setdefault("VLLM_USE_V1", "0")
if hf_home is not None:
env["HF_HOME"] = str(hf_home)
console.print("\n[bold]Command:[/]")
console.print(" " + " ".join(uvx_cmd))
if hf_home is not None:
console.print(f" [dim](with HF_HOME={hf_home})[/]")
if print_cmd:
return
console.print("\n[dim]Launching... (first run downloads the model -- ~5 min on a fast link)[/]\n")
try:
completed = subprocess.run(uvx_cmd, env=env, check=False)
except KeyboardInterrupt:
console.print("\n[yellow]Interrupted.[/]")
return
sys.exit(completed.returncode)
if __name__ == "__main__":
app()