github-actions[bot]
1475 lines
60 KiB
Python
Executable File
1475 lines
60 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
"""
|
|
Compute latent representations for video generation training.
|
|
This module provides functionality for processing video and image files, including:
|
|
- Loading videos/images from various file formats (CSV, JSON, JSONL)
|
|
- Resizing, cropping, and transforming media
|
|
- MediaDataset for video-only preprocessing workflows
|
|
- BucketSampler for grouping videos by resolution
|
|
Can be used as a standalone script:
|
|
python scripts/process_videos.py dataset.csv --resolution-buckets 768x768x25 \
|
|
--output-dir /path/to/output --model-source /path/to/ltx2.safetensors
|
|
"""
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import torch
|
|
import torchaudio
|
|
import typer
|
|
from accelerate import PartialState
|
|
from pillow_heif import register_heif_opener
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
BarColumn,
|
|
MofNCompleteColumn,
|
|
Progress,
|
|
SpinnerColumn,
|
|
TaskProgressColumn,
|
|
TextColumn,
|
|
TimeElapsedColumn,
|
|
TimeRemainingColumn,
|
|
)
|
|
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
|
|
from ltx_core.types import Audio
|
|
from ltx_trainer import logger
|
|
from ltx_trainer.model_loader import load_audio_vae_encoder, load_video_vae_encoder
|
|
from ltx_trainer.utils import open_image_as_srgb
|
|
from ltx_trainer.video_utils import get_video_frame_count, read_video
|
|
|
|
disable_progress_bar()
|
|
|
|
# Register HEIF/HEIC support
|
|
register_heif_opener()
|
|
|
|
# Constants for validation
|
|
VAE_SPATIAL_FACTOR = 32
|
|
VAE_TEMPORAL_FACTOR = 8
|
|
|
|
# Audio constants
|
|
AUDIO_LATENT_CHANNELS = 8
|
|
AUDIO_FREQUENCY_BINS = 16
|
|
|
|
DEFAULT_TILE_SIZE = 512 # Spatial tile size in pixels (must be ≥64 and divisible by 32)
|
|
DEFAULT_TILE_OVERLAP = 128 # Spatial tile overlap in pixels (must be divisible by 32)
|
|
|
|
app = typer.Typer(
|
|
pretty_exceptions_enable=False,
|
|
no_args_is_help=True,
|
|
help="Process videos/images and save latent representations for video generation training.",
|
|
)
|
|
|
|
|
|
def _clamp_01(x: torch.Tensor) -> torch.Tensor:
|
|
return x.clamp_(0, 1)
|
|
|
|
|
|
class MediaDataset(Dataset):
|
|
"""
|
|
Dataset for processing video and image files.
|
|
This dataset is designed for media preprocessing workflows where you need to:
|
|
- Load and preprocess videos/images
|
|
- Apply resizing and cropping transformations
|
|
- Handle different resolution buckets
|
|
- Filter out invalid media files
|
|
- Optionally extract audio from video files
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset_file: str | Path,
|
|
main_media_column: str,
|
|
video_column: str,
|
|
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.
|
|
Args:
|
|
dataset_file: Path to CSV/JSON/JSONL metadata file
|
|
video_column: Column name for video paths in the metadata file
|
|
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__()
|
|
|
|
self.dataset_file = Path(dataset_file)
|
|
self.main_media_column = main_media_column
|
|
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)
|
|
|
|
# Then load reference video paths
|
|
self.video_paths = self._load_video_paths(video_column)
|
|
|
|
# Filter out videos with insufficient frames
|
|
self._filter_valid_videos()
|
|
|
|
self.max_target_frames = max(self.resolution_buckets, key=lambda x: x[0])[0]
|
|
|
|
# Set up video transforms
|
|
self.transforms = transforms.Compose(
|
|
[
|
|
transforms.Lambda(_clamp_01),
|
|
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
|
|
]
|
|
)
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.video_paths)
|
|
|
|
def __getitem__(self, index: int) -> dict[str, Any]:
|
|
"""Get a single video/image with metadata, and optionally audio."""
|
|
if isinstance(index, list):
|
|
# Special case for BucketSampler - return cached data
|
|
return index
|
|
|
|
video_path: Path = self.video_paths[index]
|
|
|
|
# Compute relative path of the video
|
|
data_root = self.dataset_file.parent
|
|
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)
|
|
fps = 1.0
|
|
audio_data = None # Images don't have audio
|
|
else:
|
|
media_tensor, fps = self._preprocess_video(video_path)
|
|
|
|
# Extract audio if enabled
|
|
if self.with_audio:
|
|
# Calculate target duration from the processed video frames
|
|
# This ensures audio is trimmed to match the exact video duration
|
|
# media_tensor is [C, F, H, W] so shape[1] is num_frames
|
|
target_duration = media_tensor.shape[1] / fps
|
|
audio_data = self._extract_audio(video_path, target_duration)
|
|
else:
|
|
audio_data = None
|
|
|
|
# media_tensor is [C, F, H, W] format for VAE compatibility
|
|
_, num_frames, height, width = media_tensor.shape
|
|
|
|
result = {
|
|
"video": media_tensor,
|
|
"relative_path": relative_path,
|
|
"main_media_relative_path": media_relative_path,
|
|
"video_metadata": {
|
|
"num_frames": num_frames,
|
|
"height": height,
|
|
"width": width,
|
|
"fps": fps,
|
|
},
|
|
}
|
|
|
|
# Add audio data if available
|
|
if audio_data is not None:
|
|
result["audio"] = audio_data
|
|
|
|
return result
|
|
|
|
@staticmethod
|
|
def _extract_audio(video_path: Path, target_duration: float) -> dict[str, torch.Tensor | int] | None:
|
|
"""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
|
|
|
|
# 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:
|
|
waveform = audio.waveform
|
|
|
|
return {"waveform": waveform, "sample_rate": audio.sampling_rate}
|
|
|
|
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."""
|
|
original_length = len(self.video_paths)
|
|
valid_video_paths = []
|
|
valid_main_media_paths = []
|
|
min_frames_required = min(self.resolution_buckets, key=lambda x: x[0])[0]
|
|
|
|
for i, video_path in enumerate(self.video_paths):
|
|
if video_path.suffix.lower() in [".png", ".jpg", ".jpeg"]:
|
|
valid_video_paths.append(video_path)
|
|
valid_main_media_paths.append(self.main_media_paths[i])
|
|
continue
|
|
|
|
try:
|
|
frame_count = get_video_frame_count(video_path)
|
|
|
|
if frame_count >= min_frames_required:
|
|
valid_video_paths.append(video_path)
|
|
valid_main_media_paths.append(self.main_media_paths[i])
|
|
else:
|
|
logger.warning(
|
|
f"Skipping video at {video_path} - has {frame_count} frames, "
|
|
f"which is less than the minimum required frames ({min_frames_required})"
|
|
)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to read video at {video_path}: {e!s}")
|
|
|
|
# Update both path lists to maintain synchronization
|
|
self.video_paths = valid_video_paths
|
|
self.main_media_paths = valid_main_media_paths
|
|
|
|
if len(self.video_paths) < original_length:
|
|
logger.warning(
|
|
f"Filtered out {original_length - len(self.video_paths)} videos with insufficient frames. "
|
|
f"Proceeding with {len(self.video_paths)} valid videos."
|
|
)
|
|
|
|
def _preprocess_image(self, path: Path) -> torch.Tensor:
|
|
"""Preprocess a single image by resizing and applying transforms."""
|
|
image = open_image_as_srgb(path)
|
|
image = to_tensor(image)
|
|
image = image.unsqueeze(0) # Add frame dimension [1, C, H, W] for bucket selection
|
|
|
|
# Find nearest resolution bucket and resize
|
|
nearest_bucket = self._get_resolution_bucket_for_item(image)
|
|
_, target_height, target_width = nearest_bucket
|
|
image_resized = self._resize_and_crop(image, target_height, target_width)
|
|
# _resize_and_crop returns [C, H, W] for single-frame input (squeeze removes dim 0)
|
|
|
|
# Apply transforms
|
|
image = self.transforms(image_resized) # [C, H, W] -> [C, H, W]
|
|
|
|
# Add frame dimension in VAE format: [C, H, W] -> [C, 1, H, W]
|
|
image = image.unsqueeze(1)
|
|
return image
|
|
|
|
def _preprocess_video(self, path: Path) -> tuple[torch.Tensor, float]:
|
|
"""Preprocess a video by loading, resizing, and applying transforms.
|
|
Returns:
|
|
Tuple of (video tensor in [C, F, H, W] format, fps)
|
|
"""
|
|
# Load video frames up to max_target_frames
|
|
video, fps = read_video(path, max_frames=self.max_target_frames)
|
|
|
|
nearest_bucket = self._get_resolution_bucket_for_item(video)
|
|
target_num_frames, target_height, target_width = nearest_bucket
|
|
frames_resized = self._resize_and_crop(video, target_height, target_width)
|
|
|
|
# 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)
|
|
|
|
# Permute [F,C,H,W] -> [C,F,H,W] for VAE compatibility
|
|
# After DataLoader batching, this becomes [B,C,F,H,W] which VAE expects
|
|
video = video.permute(1, 0, 2, 3).contiguous()
|
|
|
|
return video, fps
|
|
|
|
def _get_resolution_bucket_for_item(self, media_tensor: torch.Tensor) -> tuple[int, int, int]:
|
|
"""Get the nearest resolution bucket for the given media tensor."""
|
|
num_frames, _, height, width = media_tensor.shape
|
|
|
|
def distance(bucket: tuple[int, int, int]) -> tuple:
|
|
bucket_num_frames, bucket_height, bucket_width = bucket
|
|
# Lexicographic key:
|
|
# 1) minimize aspect-ratio diff (in log-scale, for invariance to shorter/longer ARs)
|
|
# 2) prefer buckets with more frames (by using negative)
|
|
# 3) prefer buckets with larger spatial area (by using negative)
|
|
return (
|
|
abs(math.log(width / height) - math.log(bucket_width / bucket_height)),
|
|
-bucket_num_frames,
|
|
-(bucket_height * bucket_width),
|
|
)
|
|
|
|
# Keep only buckets with <= available frames
|
|
relevant_buckets = [b for b in self.resolution_buckets if b[0] <= num_frames]
|
|
if not relevant_buckets:
|
|
raise ValueError(f"No resolution buckets have <= {num_frames} frames. Available: {self.resolution_buckets}")
|
|
|
|
# Find the bucket with the minimal distance (according to the function above) to the media item's shape.
|
|
nearest_bucket = min(relevant_buckets, key=distance)
|
|
|
|
return nearest_bucket
|
|
|
|
def _resize_and_crop(self, media_tensor: torch.Tensor, target_height: int, target_width: int) -> torch.Tensor:
|
|
"""Resize and crop tensor to target size."""
|
|
# Get current dimensions
|
|
current_height, current_width = media_tensor.shape[2], media_tensor.shape[3]
|
|
|
|
# Calculate aspect ratios to determine which dimension to resize first
|
|
current_aspect = current_width / current_height
|
|
target_aspect = target_width / target_height
|
|
|
|
# Resize while maintaining aspect ratio - scale to make the smaller dimension fit
|
|
if current_aspect > target_aspect:
|
|
# Current is wider than target, so scale by height
|
|
new_width = int(current_width * target_height / current_height)
|
|
media_tensor = resize(
|
|
media_tensor,
|
|
size=[target_height, new_width], # type: ignore
|
|
interpolation=InterpolationMode.BICUBIC,
|
|
)
|
|
else:
|
|
# Current is taller than target, so scale by width
|
|
new_height = int(current_height * target_width / current_width)
|
|
media_tensor = resize(
|
|
media_tensor,
|
|
size=[new_height, target_width],
|
|
interpolation=InterpolationMode.BICUBIC,
|
|
)
|
|
|
|
# Update dimensions after resize
|
|
current_height, current_width = media_tensor.shape[2], media_tensor.shape[3]
|
|
media_tensor = media_tensor.squeeze(0)
|
|
|
|
# Calculate how much we need to crop from each dimension
|
|
delta_h = current_height - target_height
|
|
delta_w = current_width - target_width
|
|
|
|
# Determine crop position based on reshape mode
|
|
if self.reshape_mode == "random":
|
|
# Random crop position
|
|
top = np.random.randint(0, delta_h + 1)
|
|
left = np.random.randint(0, delta_w + 1)
|
|
elif self.reshape_mode == "center":
|
|
# Center crop
|
|
top, left = delta_h // 2, delta_w // 2
|
|
else:
|
|
raise ValueError(f"Unsupported reshape mode: {self.reshape_mode}")
|
|
|
|
# Perform the final crop to exact target dimensions
|
|
media_tensor = crop(media_tensor, top=top, left=left, height=target_height, width=target_width)
|
|
return media_tensor
|
|
|
|
|
|
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]],
|
|
output_dir: str,
|
|
model_path: str,
|
|
main_media_column: str | None = None,
|
|
reshape_mode: str = "center",
|
|
batch_size: int = 1,
|
|
device: str = "cuda",
|
|
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.
|
|
Under ``accelerate launch``, each process handles an interleaved shard of
|
|
the dataset (rank/world read from ``accelerate.PartialState``). Already-
|
|
computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
|
|
atomic so an interrupted run is safe to resume.
|
|
Args:
|
|
dataset_file: Path to metadata file (CSV/JSON/JSONL) containing video paths
|
|
video_column: Column name for video paths in the metadata file
|
|
resolution_buckets: List of (frames, height, width) tuples
|
|
output_dir: Directory to save video latents
|
|
model_path: Path to LTX-2 checkpoint (.safetensors)
|
|
reshape_mode: How to crop videos ("center", "random")
|
|
main_media_column: Column name for main media paths (if different from video_column)
|
|
batch_size: Batch size for processing
|
|
device: Device to use for computation
|
|
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")
|
|
|
|
console = Console()
|
|
torch_device = torch.device(device)
|
|
|
|
dataset = MediaDataset(
|
|
dataset_file=dataset_file,
|
|
main_media_column=main_media_column or video_column,
|
|
video_column=video_column,
|
|
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)
|
|
audio_output_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Audio processing requires batch_size=1; must be applied before the dataloader is built.
|
|
if with_audio and batch_size > 1:
|
|
logger.warning("Audio processing requires batch_size=1. Overriding batch_size to 1.")
|
|
batch_size = 1
|
|
|
|
data_root = Path(dataset_file).parent
|
|
|
|
def _is_done(idx: int) -> bool:
|
|
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()
|
|
|
|
dataloader = _build_sharded_dataloader(
|
|
dataset,
|
|
batch_size=batch_size,
|
|
num_workers=num_dataloader_workers,
|
|
is_done=_is_done,
|
|
overwrite=overwrite,
|
|
)
|
|
if dataloader is None:
|
|
return
|
|
|
|
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)
|
|
|
|
audio_vae_encoder = None
|
|
audio_processor = None
|
|
if with_audio:
|
|
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 VAE needs float32 for quality. TODO: re-test with bfloat16.
|
|
)
|
|
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)
|
|
|
|
# Track audio statistics
|
|
audio_success_count = 0
|
|
audio_skip_count = 0
|
|
|
|
# Process batches
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TaskProgressColumn(),
|
|
MofNCompleteColumn(),
|
|
TimeElapsedColumn(),
|
|
TimeRemainingColumn(),
|
|
console=console,
|
|
) as progress:
|
|
task = progress.add_task("Processing videos", total=len(dataloader))
|
|
|
|
for batch in dataloader:
|
|
# Get video tensor - shape is [B, F, C, H, W] from DataLoader
|
|
video = batch["video"]
|
|
|
|
# Encode video
|
|
with torch.inference_mode():
|
|
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"])):
|
|
output_rel_path = Path(batch["main_media_relative_path"][i]).with_suffix(".pt")
|
|
output_file = output_path / output_rel_path
|
|
|
|
# Create output directory maintaining structure
|
|
output_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 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": effective_fps,
|
|
}
|
|
|
|
_atomic_save(latent_data, output_file)
|
|
|
|
# Process audio if enabled (audio is already extracted by the dataset)
|
|
if with_audio:
|
|
audio_batch = batch.get("audio")
|
|
if audio_batch is not None:
|
|
# Extract the i-th item from batched audio data
|
|
# DataLoader collates [channels, samples] -> [batch, channels, samples]
|
|
audio_data = Audio(
|
|
waveform=audio_batch["waveform"][i],
|
|
sampling_rate=audio_batch["sample_rate"][i].item(),
|
|
)
|
|
|
|
# Encode audio
|
|
with torch.inference_mode():
|
|
audio_latents = _encode_audio(audio_vae_encoder, audio_processor, audio_data)
|
|
|
|
# Save audio latents
|
|
audio_output_file = audio_output_path / output_rel_path
|
|
audio_output_file.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
audio_save_data = {
|
|
"latents": audio_latents["latents"].cpu().contiguous(),
|
|
"num_time_steps": audio_latents["num_time_steps"],
|
|
"frequency_bins": audio_latents["frequency_bins"],
|
|
"duration": audio_latents["duration"],
|
|
}
|
|
|
|
_atomic_save(audio_save_data, audio_output_file)
|
|
audio_success_count += 1
|
|
else:
|
|
# Video has no audio track
|
|
audio_skip_count += 1
|
|
|
|
progress.advance(task)
|
|
|
|
logger.info(f"Processed {len(dataloader.dataset)} videos -> {output_path}") # type: ignore[arg-type]
|
|
if with_audio:
|
|
logger.info(
|
|
f"Audio processing: {audio_success_count} videos with audio, "
|
|
f"{audio_skip_count} videos without audio (skipped)"
|
|
)
|
|
|
|
|
|
def _encode_video(
|
|
vae: torch.nn.Module,
|
|
video: torch.Tensor,
|
|
dtype: torch.dtype | None = None,
|
|
use_tiling: bool = False,
|
|
tile_size: int = DEFAULT_TILE_SIZE,
|
|
tile_overlap: int = DEFAULT_TILE_OVERLAP,
|
|
) -> dict[str, torch.Tensor | int]:
|
|
"""Encode video into non-patchified latent representation.
|
|
Args:
|
|
vae: Video VAE encoder model
|
|
video: Input tensor of shape [B, C, F, H, W] (batch, channels, frames, height, width)
|
|
This is the format expected by the VAE encoder.
|
|
dtype: Target dtype for output latents
|
|
use_tiling: Whether to use spatial tiling for memory efficiency
|
|
tile_size: Tile size in pixels (must be divisible by 32)
|
|
tile_overlap: Overlap between tiles in pixels (must be divisible by 32)
|
|
Returns:
|
|
Dict containing non-patchified latents and shape information:
|
|
{
|
|
"latents": Tensor[B, C, F', H', W'], # Non-patchified format with batch dim
|
|
"num_frames": int, # Latent frame count
|
|
"height": int, # Latent height
|
|
"width": int, # Latent width
|
|
}
|
|
"""
|
|
device = next(vae.parameters()).device
|
|
vae_dtype = next(vae.parameters()).dtype
|
|
|
|
# Add batch dimension if needed
|
|
if video.ndim == 4:
|
|
video = video.unsqueeze(0) # [C, F, H, W] -> [B, C, F, H, W]
|
|
|
|
video = video.to(device=device, dtype=vae_dtype)
|
|
|
|
# Choose encoding method based on tiling flag
|
|
if use_tiling:
|
|
latents = _tiled_encode_video(
|
|
vae=vae,
|
|
video=video,
|
|
tile_size=tile_size,
|
|
tile_overlap=tile_overlap,
|
|
)
|
|
else:
|
|
# Encode video - VAE expects [B, C, F, H, W], returns [B, C, F', H', W']
|
|
latents = vae(video)
|
|
|
|
if dtype is not None:
|
|
latents = latents.to(dtype=dtype)
|
|
|
|
_, _, num_frames, height, width = latents.shape
|
|
|
|
return {
|
|
"latents": latents, # [B, C, F', H', W']
|
|
"num_frames": num_frames,
|
|
"height": height,
|
|
"width": width,
|
|
}
|
|
|
|
|
|
def _tiled_encode_video( # noqa: PLR0912, PLR0915
|
|
vae: torch.nn.Module,
|
|
video: torch.Tensor,
|
|
tile_size: int = DEFAULT_TILE_SIZE,
|
|
tile_overlap: int = DEFAULT_TILE_OVERLAP,
|
|
) -> torch.Tensor:
|
|
"""Encode video using spatial tiling for memory efficiency.
|
|
Splits the video into overlapping spatial tiles, encodes each tile separately,
|
|
and blends the results using linear feathering in the overlap regions.
|
|
Args:
|
|
vae: Video VAE encoder model
|
|
video: Input tensor of shape [B, C, F, H, W]
|
|
tile_size: Tile size in pixels (must be divisible by 32)
|
|
tile_overlap: Overlap between tiles in pixels (must be divisible by 32)
|
|
Returns:
|
|
Encoded latent tensor [B, C_latent, F_latent, H_latent, W_latent]
|
|
"""
|
|
batch, _channels, frames, height, width = video.shape
|
|
device = video.device
|
|
dtype = video.dtype
|
|
|
|
# Validate tile parameters
|
|
if tile_size % VAE_SPATIAL_FACTOR != 0:
|
|
raise ValueError(f"tile_size must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_size}")
|
|
if tile_overlap % VAE_SPATIAL_FACTOR != 0:
|
|
raise ValueError(f"tile_overlap must be divisible by {VAE_SPATIAL_FACTOR}, got {tile_overlap}")
|
|
if tile_overlap >= tile_size:
|
|
raise ValueError(f"tile_overlap ({tile_overlap}) must be less than tile_size ({tile_size})")
|
|
|
|
# If video fits in a single tile, use regular encoding
|
|
if height <= tile_size and width <= tile_size:
|
|
return vae(video)
|
|
|
|
# Calculate output dimensions
|
|
# VAE compresses: H -> H/32, W -> W/32, F -> 1 + (F-1)/8
|
|
output_height = height // VAE_SPATIAL_FACTOR
|
|
output_width = width // VAE_SPATIAL_FACTOR
|
|
output_frames = 1 + (frames - 1) // VAE_TEMPORAL_FACTOR
|
|
|
|
# Latent channels (128 for LTX-2)
|
|
# Get from a small test encode or assume 128
|
|
latent_channels = 128
|
|
|
|
# Initialize output and weight tensors
|
|
output = torch.zeros(
|
|
(batch, latent_channels, output_frames, output_height, output_width),
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
weights = torch.zeros(
|
|
(batch, 1, output_frames, output_height, output_width),
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
# Calculate tile positions with overlap
|
|
# Step size is tile_size - tile_overlap
|
|
step_h = tile_size - tile_overlap
|
|
step_w = tile_size - tile_overlap
|
|
|
|
h_positions = list(range(0, max(1, height - tile_overlap), step_h))
|
|
w_positions = list(range(0, max(1, width - tile_overlap), step_w))
|
|
|
|
# Ensure last tile covers the edge
|
|
if h_positions[-1] + tile_size < height:
|
|
h_positions.append(height - tile_size)
|
|
if w_positions[-1] + tile_size < width:
|
|
w_positions.append(width - tile_size)
|
|
|
|
# Remove duplicates and sort
|
|
h_positions = sorted(set(h_positions))
|
|
w_positions = sorted(set(w_positions))
|
|
|
|
# Overlap in latent space
|
|
overlap_out_h = tile_overlap // VAE_SPATIAL_FACTOR
|
|
overlap_out_w = tile_overlap // VAE_SPATIAL_FACTOR
|
|
|
|
# Process each tile
|
|
for h_pos in h_positions:
|
|
for w_pos in w_positions:
|
|
# Calculate tile boundaries in input space
|
|
h_start = max(0, h_pos)
|
|
w_start = max(0, w_pos)
|
|
h_end = min(h_start + tile_size, height)
|
|
w_end = min(w_start + tile_size, width)
|
|
|
|
# Ensure tile dimensions are divisible by VAE_SPATIAL_FACTOR
|
|
tile_h = ((h_end - h_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
|
|
tile_w = ((w_end - w_start) // VAE_SPATIAL_FACTOR) * VAE_SPATIAL_FACTOR
|
|
|
|
if tile_h < VAE_SPATIAL_FACTOR or tile_w < VAE_SPATIAL_FACTOR:
|
|
continue
|
|
|
|
# Adjust end positions
|
|
h_end = h_start + tile_h
|
|
w_end = w_start + tile_w
|
|
|
|
# Extract tile
|
|
tile = video[:, :, :, h_start:h_end, w_start:w_end]
|
|
|
|
# Encode tile
|
|
encoded_tile = vae(tile)
|
|
|
|
# Get actual encoded dimensions
|
|
_, _, tile_out_frames, tile_out_height, tile_out_width = encoded_tile.shape
|
|
|
|
# Calculate output positions
|
|
out_h_start = h_start // VAE_SPATIAL_FACTOR
|
|
out_w_start = w_start // VAE_SPATIAL_FACTOR
|
|
out_h_end = min(out_h_start + tile_out_height, output_height)
|
|
out_w_end = min(out_w_start + tile_out_width, output_width)
|
|
|
|
# Trim encoded tile if necessary
|
|
actual_tile_h = out_h_end - out_h_start
|
|
actual_tile_w = out_w_end - out_w_start
|
|
encoded_tile = encoded_tile[:, :, :, :actual_tile_h, :actual_tile_w]
|
|
|
|
# Create blending mask with linear feathering at edges
|
|
mask = torch.ones(
|
|
(1, 1, tile_out_frames, actual_tile_h, actual_tile_w),
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
# Apply feathering at edges (linear blend in overlap regions)
|
|
# Left edge
|
|
if h_pos > 0 and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
|
|
fade_in = torch.linspace(0.0, 1.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
|
|
mask[:, :, :, :overlap_out_h, :] *= fade_in.view(1, 1, 1, -1, 1)
|
|
|
|
# Right edge (bottom in height dimension)
|
|
if h_end < height and overlap_out_h > 0 and overlap_out_h < actual_tile_h:
|
|
fade_out = torch.linspace(1.0, 0.0, overlap_out_h + 2, device=device, dtype=dtype)[1:-1]
|
|
mask[:, :, :, -overlap_out_h:, :] *= fade_out.view(1, 1, 1, -1, 1)
|
|
|
|
# Top edge (left in width dimension)
|
|
if w_pos > 0 and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
|
|
fade_in = torch.linspace(0.0, 1.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
|
|
mask[:, :, :, :, :overlap_out_w] *= fade_in.view(1, 1, 1, 1, -1)
|
|
|
|
# Bottom edge (right in width dimension)
|
|
if w_end < width and overlap_out_w > 0 and overlap_out_w < actual_tile_w:
|
|
fade_out = torch.linspace(1.0, 0.0, overlap_out_w + 2, device=device, dtype=dtype)[1:-1]
|
|
mask[:, :, :, :, -overlap_out_w:] *= fade_out.view(1, 1, 1, 1, -1)
|
|
|
|
# Accumulate weighted results
|
|
output[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += encoded_tile * mask
|
|
weights[:, :, :, out_h_start:out_h_end, out_w_start:out_w_end] += mask
|
|
|
|
# Normalize by weights (avoid division by zero)
|
|
output = output / (weights + 1e-8)
|
|
|
|
return output
|
|
|
|
|
|
def _encode_audio(
|
|
audio_vae_encoder: torch.nn.Module,
|
|
audio_processor: torch.nn.Module,
|
|
audio: Audio,
|
|
) -> dict[str, torch.Tensor | int | float]:
|
|
"""Encode audio waveform into latent representation.
|
|
Args:
|
|
audio_vae_encoder: Audio VAE encoder model from ltx-core
|
|
audio_processor: AudioProcessor for waveform-to-spectrogram conversion
|
|
audio: Audio container with waveform tensor and sampling rate.
|
|
Returns:
|
|
Dict containing audio latents and shape information:
|
|
{
|
|
"latents": Tensor[C, T, F], # Non-patchified format
|
|
"num_time_steps": int,
|
|
"frequency_bins": int,
|
|
"duration": float,
|
|
}
|
|
"""
|
|
device = next(audio_vae_encoder.parameters()).device
|
|
dtype = next(audio_vae_encoder.parameters()).dtype
|
|
|
|
waveform = audio.waveform.to(device=device, dtype=dtype)
|
|
|
|
# Add batch dimension if needed: [channels, samples] -> [batch, channels, samples]
|
|
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
|
|
|
|
# Convert waveform to mel spectrogram using AudioProcessor
|
|
mel_spectrogram = audio_processor.waveform_to_mel(Audio(waveform=waveform, sampling_rate=audio.sampling_rate))
|
|
mel_spectrogram = mel_spectrogram.to(dtype=dtype)
|
|
|
|
# Encode mel spectrogram to latents
|
|
latents = audio_vae_encoder(mel_spectrogram)
|
|
|
|
# latents shape: [batch, channels, time, freq] = [1, 8, T, 16]
|
|
_, _channels, time_steps, freq_bins = latents.shape
|
|
|
|
return {
|
|
"latents": latents.squeeze(0), # [C, T, F] - remove batch dim
|
|
"num_time_steps": time_steps,
|
|
"frequency_bins": freq_bins,
|
|
"duration": duration,
|
|
}
|
|
|
|
|
|
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 = []
|
|
for bucket_str in resolution_buckets_str.split(";"):
|
|
w, h, f = map(int, bucket_str.split("x"))
|
|
|
|
if w % VAE_SPATIAL_FACTOR != 0 or h % VAE_SPATIAL_FACTOR != 0:
|
|
raise typer.BadParameter(
|
|
f"Width and height must be multiples of {VAE_SPATIAL_FACTOR}, got {w}x{h}",
|
|
param_hint="resolution-buckets",
|
|
)
|
|
|
|
if f % VAE_TEMPORAL_FACTOR != 1:
|
|
raise typer.BadParameter(
|
|
f"Number of frames must be a multiple of {VAE_TEMPORAL_FACTOR} plus 1, got {f}",
|
|
param_hint="resolution-buckets",
|
|
)
|
|
|
|
resolution_buckets.append((f, h, w))
|
|
return resolution_buckets
|
|
|
|
|
|
def compute_scaled_resolution_buckets(
|
|
resolution_buckets: list[tuple[int, int, int]],
|
|
scale_factor: int,
|
|
) -> list[tuple[int, int, int]]:
|
|
"""Compute scaled resolution buckets and validate the results."""
|
|
if scale_factor == 1:
|
|
return resolution_buckets
|
|
|
|
scaled_buckets = []
|
|
for frames, height, width in resolution_buckets:
|
|
# Validate that scale factor evenly divides the dimensions
|
|
if height % scale_factor != 0:
|
|
raise ValueError(
|
|
f"Height {height} is not evenly divisible by scale factor {scale_factor}. "
|
|
f"Choose a scale factor that divides {height} evenly."
|
|
)
|
|
if width % scale_factor != 0:
|
|
raise ValueError(
|
|
f"Width {width} is not evenly divisible by scale factor {scale_factor}. "
|
|
f"Choose a scale factor that divides {width} evenly."
|
|
)
|
|
|
|
scaled_height = height // scale_factor
|
|
scaled_width = width // scale_factor
|
|
|
|
# Validate scaled dimensions are divisible by VAE spatial factor
|
|
if scaled_height % VAE_SPATIAL_FACTOR != 0:
|
|
raise ValueError(
|
|
f"Scaled height {scaled_height} (from {height} / {scale_factor}) "
|
|
f"is not divisible by {VAE_SPATIAL_FACTOR}. "
|
|
f"Choose a different scale factor or adjust your resolution buckets."
|
|
)
|
|
if scaled_width % VAE_SPATIAL_FACTOR != 0:
|
|
raise ValueError(
|
|
f"Scaled width {scaled_width} (from {width} / {scale_factor}) "
|
|
f"is not divisible by {VAE_SPATIAL_FACTOR}. "
|
|
f"Choose a different scale factor or adjust your resolution buckets."
|
|
)
|
|
|
|
scaled_buckets.append((frames, scaled_height, scaled_width))
|
|
|
|
return scaled_buckets
|
|
|
|
|
|
def _atomic_save(data: Any, out: Path) -> None: # noqa: ANN401
|
|
"""Save to ``out`` atomically via per-PID temp file + replace.
|
|
Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
|
|
ignores. The per-PID suffix makes concurrent writes from multiple ranks
|
|
collision-free.
|
|
"""
|
|
tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
|
|
torch.save(data, tmp)
|
|
tmp.replace(out)
|
|
|
|
|
|
def _build_sharded_dataloader(
|
|
dataset: Dataset,
|
|
*,
|
|
batch_size: int,
|
|
num_workers: int,
|
|
is_done: Callable[[int], bool],
|
|
overwrite: bool,
|
|
) -> DataLoader | None:
|
|
"""Return a DataLoader over this rank's interleaved shard of ``dataset``.
|
|
When ``overwrite`` is False, items whose outputs already exist (per
|
|
``is_done``) are filtered out. Returns ``None`` if this rank has nothing
|
|
to do, so the caller can early-return without loading any models.
|
|
"""
|
|
state = PartialState()
|
|
todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
|
|
if not todo:
|
|
logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
|
|
return None
|
|
logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
|
|
return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
|
|
|
|
|
|
@app.command()
|
|
def main( # noqa: PLR0913
|
|
dataset_file: str = typer.Argument(
|
|
...,
|
|
help="Path to metadata file (CSV/JSON/JSONL) containing video paths",
|
|
),
|
|
resolution_buckets: str = typer.Option(
|
|
...,
|
|
help='Resolution buckets in format "WxHxF;WxHxF;..." (e.g. "768x768x25;512x512x49")',
|
|
),
|
|
output_dir: str = typer.Option(
|
|
...,
|
|
help="Output directory to save video latents",
|
|
),
|
|
model_path: str = typer.Option(
|
|
...,
|
|
help="Path to LTX-2 checkpoint (.safetensors file)",
|
|
),
|
|
video_column: str = typer.Option(
|
|
default="media_path",
|
|
help="Column name in the dataset JSON/JSONL/CSV file containing video paths",
|
|
),
|
|
batch_size: int = typer.Option(
|
|
default=1,
|
|
help="Batch size for processing",
|
|
),
|
|
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",
|
|
),
|
|
reshape_mode: str = typer.Option(
|
|
default="center",
|
|
help="How to crop videos: 'center' or 'random'",
|
|
),
|
|
with_audio: bool = typer.Option(
|
|
default=False,
|
|
help="Extract and encode audio from video files",
|
|
),
|
|
audio_output_dir: str | None = typer.Option(
|
|
default=None,
|
|
help="Output directory for audio latents (required if --with-audio is set)",
|
|
),
|
|
overwrite: bool = typer.Option(
|
|
default=False,
|
|
help="Re-encode every item even if its output exists. Use when rerunning with "
|
|
"changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
|
|
),
|
|
) -> None:
|
|
"""Process videos/images and save latent representations for video generation training.
|
|
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 \\
|
|
--output-dir ./latents --model-path /path/to/ltx2.safetensors
|
|
# Process videos from a JSON file with custom video column
|
|
python scripts/process_videos.py dataset.json --resolution-buckets 768x768x25 \\
|
|
--output-dir ./latents --model-path /path/to/ltx2.safetensors --video-column "video_path"
|
|
# Enable VAE tiling to save GPU VRAM
|
|
python scripts/process_videos.py dataset.csv --resolution-buckets 1024x1024x25 \\
|
|
--output-dir ./latents --model-path /path/to/ltx2.safetensors --vae-tiling
|
|
# Process videos with audio
|
|
python scripts/process_videos.py dataset.csv --resolution-buckets 768x768x25 \\
|
|
--output-dir ./latents --model-path /path/to/ltx2.safetensors \\
|
|
--with-audio --audio-output-dir ./audio_latents
|
|
"""
|
|
|
|
# Validate dataset file exists
|
|
if not Path(dataset_file).is_file():
|
|
raise typer.BadParameter(f"Dataset file not found: {dataset_file}")
|
|
|
|
# Validate audio parameters
|
|
if with_audio and audio_output_dir is None:
|
|
raise typer.BadParameter("--audio-output-dir is required when --with-audio is set")
|
|
|
|
# Parse resolution buckets
|
|
parsed_resolution_buckets = parse_resolution_buckets(resolution_buckets)
|
|
|
|
if len(parsed_resolution_buckets) > 1:
|
|
logger.warning(
|
|
"Using multiple resolution buckets. "
|
|
"When training with multiple resolution buckets, you must use a batch size of 1."
|
|
)
|
|
|
|
# Process latents
|
|
compute_latents(
|
|
dataset_file=dataset_file,
|
|
video_column=video_column,
|
|
resolution_buckets=parsed_resolution_buckets,
|
|
output_dir=output_dir,
|
|
model_path=model_path,
|
|
reshape_mode=reshape_mode,
|
|
batch_size=batch_size,
|
|
device=device,
|
|
vae_tiling=vae_tiling,
|
|
with_audio=with_audio,
|
|
audio_output_dir=audio_output_dir,
|
|
overwrite=overwrite,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app()
|