github-actions[bot]
433 lines
17 KiB
Python
Executable File
433 lines
17 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
"""
|
|
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 \\
|
|
--model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma
|
|
"""
|
|
|
|
from pathlib import Path
|
|
|
|
import typer
|
|
from decode_latents import LatentsDecoder
|
|
from process_captions import compute_captions_embeddings
|
|
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
|
|
from ltx_trainer.gpu_utils import free_gpu_memory_context
|
|
|
|
console = Console()
|
|
|
|
app = typer.Typer(
|
|
pretty_exceptions_enable=False,
|
|
no_args_is_help=True,
|
|
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: PLR0912, PLR0913, PLR0915
|
|
dataset_file: str,
|
|
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_downscale_factor: int = 1,
|
|
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 convention-based column detection."""
|
|
_validate_dataset_file(dataset_file)
|
|
|
|
# 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"
|
|
|
|
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():
|
|
compute_captions_embeddings(
|
|
dataset_file=dataset_file,
|
|
output_dir=str(output_base / "conditions"),
|
|
model_path=model_path,
|
|
text_encoder_path=text_encoder_path,
|
|
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,
|
|
device=device,
|
|
load_in_8bit=load_text_encoder_in_8bit,
|
|
overwrite=overwrite,
|
|
)
|
|
|
|
# --- 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
|
|
|
|
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,
|
|
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,
|
|
)
|
|
|
|
# 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")
|
|
|
|
# --- Summary ---
|
|
logger.info(f"Dataset preprocessing complete! Results saved to {output_base}")
|
|
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) with columns matching the convention table",
|
|
),
|
|
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(
|
|
...,
|
|
help="Path to LTX-2 checkpoint (.safetensors file)",
|
|
),
|
|
text_encoder_path: str = typer.Option(
|
|
...,
|
|
help="Path to Gemma text encoder directory",
|
|
),
|
|
caption_column: str | None = typer.Option(
|
|
default=None,
|
|
help="Override: treat this column as 'caption' (default: auto-detect 'caption')",
|
|
),
|
|
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,
|
|
help="Batch size for preprocessing",
|
|
),
|
|
device: str = typer.Option(
|
|
default="cuda",
|
|
help="Device to use for computation",
|
|
),
|
|
vae_tiling: bool = typer.Option(
|
|
default=False,
|
|
help="Enable VAE tiling for larger video resolutions",
|
|
),
|
|
output_dir: str | None = typer.Option(
|
|
default=None,
|
|
help="Output directory (defaults to .precomputed in dataset directory)",
|
|
),
|
|
lora_trigger: str | None = typer.Option(
|
|
default=None,
|
|
help="Optional trigger word to prepend to each caption",
|
|
),
|
|
decode: bool = typer.Option(
|
|
default=False,
|
|
help="Decode and save latents after encoding for verification",
|
|
),
|
|
remove_llm_prefixes: bool = typer.Option(
|
|
default=False,
|
|
help="Remove LLM prefixes from captions",
|
|
),
|
|
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='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,
|
|
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",
|
|
),
|
|
reference_downscale_factor: int = typer.Option(
|
|
default=1,
|
|
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,
|
|
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 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.
|
|
"""
|
|
# Handle deprecated --with-audio flag
|
|
if with_audio:
|
|
logger.warning(
|
|
"--with-audio is deprecated. Audio extraction is now on by default. Use --skip-audio to disable."
|
|
)
|
|
|
|
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_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,
|
|
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_downscale_factor=reference_downscale_factor,
|
|
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,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|