#!/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.`` 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()