Upload folder using huggingface_hub (part 2)

packages/ltx-trainer/configs/a2v_lora.yaml

@@ -0,0 +1,344 @@
+# =============================================================================
+# LTX-2 Audio-to-Video LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# audio-to-video generation. The model learns to generate videos conditioned
+# on a frozen audio signal via the transformer's built-in cross-modal attention.
+#
+# In this mode, audio is provided as a frozen (clean, no noise, no loss)
+# conditioning signal. The video modality is the only generated output.
+# Audio influences video generation through the transformer's audio-to-video
+# cross-attention mechanism.
+#
+# Use this configuration when you want to:
+#   - Generate videos driven by audio content (e.g., music visualizations)
+#   - Train audio-reactive video generation
+#   - Create models that synchronize video with given audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (frozen conditioning input)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the audio-to-video training approach using the unified flexible strategy.
+# Audio is frozen (no noise, sigma=0, excluded from loss) and conditions video
+# generation via the transformer's built-in cross-modal attention.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  # Video is the generated (denoised) output
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+  # Audio modality configuration
+  # Audio is frozen — it acts as conditioning for video generation
+  # Frozen modalities get sigma=0, timestep=0, no noise, and no loss
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    # When false, audio is passed through the transformer clean and influences
+    # video via cross-modal attention
+    is_generated: false
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, audio_latents/, and conditions/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A musician plays an acoustic guitar in a dimly lit recording studio, fingers moving
+        across the fretboard with practiced ease. The warm amber light from a desk lamp
+        illuminates the wooden guitar body. Sound-absorbing panels line the walls, and a
+        microphone stands nearby on a boom arm.
+      conditions:
+        - type: audio_to_video
+          audio: "/path/to/conditioning_audio_1.wav"
+    - prompt: >-
+        Rain falls steadily on a cobblestone street in a European town at dusk. Puddles form
+        between the stones, creating ripples as new drops land. Old brick buildings line both
+        sides of the narrow street, their facades glistening with moisture under warm
+        streetlights.
+      conditions:
+        - type: audio_to_video
+          audio: "/path/to/conditioning_audio_2.wav"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_v"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  # For this audio-to-video config, false matches frozen audio conditioning (video-only synthesis).
+  generate_audio: false
+
+  # Generate video from the frozen audio condition
+  generate_video: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "a2v" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/a2v_lora"

packages/ltx-trainer/configs/accelerate/ddp.yaml

@@ -0,0 +1,16 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

packages/ltx-trainer/configs/accelerate/ddp_compile.yaml

@@ -0,0 +1,21 @@
+compute_environment: LOCAL_MACHINE
+dynamo_config:
+  dynamo_backend: INDUCTOR
+  dynamo_mode: default
+  dynamo_use_fullgraph: false
+  dynamo_use_dynamic: true
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: [ ]
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

packages/ltx-trainer/configs/accelerate/fsdp.yaml

@@ -0,0 +1,29 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_activation_checkpointing: false
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch: BACKWARD_PRE
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_transformer_layer_cls_to_wrap: BasicAVTransformerBlock
+  fsdp_use_orig_params: true
+  fsdp_version: 1
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

packages/ltx-trainer/configs/accelerate/fsdp_compile.yaml

@@ -0,0 +1,34 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+dynamo_config:
+  dynamo_backend: INDUCTOR
+  dynamo_mode: default
+  dynamo_use_fullgraph: false
+  dynamo_use_dynamic: true
+enable_cpu_affinity: false
+fsdp_config:
+  fsdp_activation_checkpointing: false
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch: BACKWARD_PRE
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_forward_prefetch: false
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: FULL_SHARD
+  fsdp_state_dict_type: SHARDED_STATE_DICT
+  fsdp_sync_module_states: true
+  fsdp_transformer_layer_cls_to_wrap: BasicAVTransformerBlock
+  fsdp_use_orig_params: true
+  fsdp_version: 1
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

packages/ltx-trainer/configs/audio_extend_lora.yaml

@@ -0,0 +1,343 @@
+# =============================================================================
+# LTX-2 Audio Extension (Forward) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# forward audio extension. The model learns to continue audio forward in time
+# by conditioning on a prefix of existing audio latent timesteps.
+#
+# Prefix conditioning works by providing the first N audio latent timesteps as
+# clean conditioning signals during training — they receive no noise,
+# timestep=0, and are excluded from the loss. The model learns to generate
+# audio that seamlessly continues from the given prefix.
+#
+# This is an audio-only training mode — no video modality is configured.
+#
+# Use this configuration when you want to:
+#   - Train the model to extend/continue existing audio
+#   - Fine-tune audio temporal continuation capabilities
+#   - Create seamless audio extension models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   └── audio_latents/   # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for audio extension LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES (not used for audio-only modes):
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (not used for audio-only modes):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For audio-only extension, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the audio extension training approach using the unified flexible
+# strategy. Prefix conditioning provides the first N audio latent timesteps as
+# clean conditioning, teaching the model to generate temporal continuations.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      - type: prefix
+        # Number of audio latent timesteps to use as conditioning prefix
+        # Each audio latent timestep = 1 patchified token
+        temporal_boundary: 8
+        # Probability of applying prefix conditioning per training sample
+        # At 1.0, all training samples use audio extension mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/ and audio_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Each sample includes an audio prefix condition for forward audio extension.
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+      conditions:
+        - type: prefix
+          audio: "/path/to/prefix_audio_1.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+      conditions:
+        - type: prefix
+          audio: "/path/to/prefix_audio_2.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "audio-extension", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/audio_extend_lora"

packages/ltx-trainer/configs/audio_inpainting_lora.yaml

@@ -0,0 +1,324 @@
+# =============================================================================
+# LTX-2 Audio Inpainting LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# audio inpainting. The model learns to fill in masked regions of audio
+# using per-sample binary masks that define which regions are conditioning
+# (provided clean) and which regions the model must generate.
+#
+# Mask conditioning works by loading per-sample binary masks from disk.
+# Masked regions receive clean latents (no noise, timestep=0) and are excluded
+# from the loss. Unmasked regions are noised and trained normally.
+#
+# This is an audio-only training mode — no video modality is configured.
+#
+# Use this configuration when you want to:
+#   - Train the model to fill in or replace regions of existing audio
+#   - Fine-tune audio inpainting capabilities on custom datasets
+#   - Create audio restoration or editing models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   ├── audio_latents/   # Audio latents (VAE-encoded audio)
+#   └── audio_masks/     # Per-sample binary masks defining conditioning regions
+#
+# Dataset metadata columns: audio, audio_mask, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for audio inpainting LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # For audio-only inpainting, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the audio inpainting training approach using the unified flexible
+# strategy. Per-sample binary masks define which audio regions are provided
+# as clean conditioning and which regions the model must learn to generate.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      - type: mask
+        # Directory name (within preprocessed_data_root) containing binary masks
+        # Each mask file corresponds to a training sample and defines the
+        # conditioning region (mask=1 means conditioning, mask=0 means generate)
+        mask_dir: "audio_masks"
+        # Probability of applying mask conditioning per training sample
+        # At 1.0, all training samples use inpainting mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/, audio_latents/, and audio_masks/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # For audio inpainting, each sample includes mask conditioning (audio + mask paths).
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+      conditions:
+        - type: mask
+          audio: "/path/to/inpainting_audio_1.wav"
+          mask: "/path/to/inpainting_mask_1.pt"
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+      conditions:
+        - type: mask
+          audio: "/path/to/inpainting_audio_2.wav"
+          mask: "/path/to/inpainting_mask_2.pt"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "audio-inpainting", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/audio_inpainting_lora"

packages/ltx-trainer/configs/audio_suffix_lora.yaml

@@ -0,0 +1,343 @@
+# =============================================================================
+# LTX-2 Audio Extension (Backward) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# backward audio extension. The model learns to generate audio content leading into existing audio
+# by conditioning on a suffix of existing audio latent timesteps.
+#
+# Suffix conditioning works by providing the last N audio latent timesteps as
+# clean conditioning signals during training — they receive no noise,
+# timestep=0, and are excluded from the loss. The model learns to generate
+# audio that seamlessly leads into the given suffix.
+#
+# This is an audio-only training mode — no video modality is configured.
+#
+# Use this configuration when you want to:
+#   - Train the model to generate preceding content for existing audio
+#   - Fine-tune audio temporal continuation capabilities
+#   - Create seamless audio extension models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   └── audio_latents/   # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for audio extension LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES (not used for audio-only modes):
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (not used for audio-only modes):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For audio-only extension, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the backward audio extension training approach using the unified flexible
+# strategy. Suffix conditioning provides the last N audio latent timesteps as
+# clean conditioning, teaching the model to generate content leading into existing audio.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      - type: suffix
+        # Number of audio latent timesteps to use as conditioning suffix
+        # Each audio latent timestep = 1 patchified token
+        temporal_boundary: 8
+        # Probability of applying suffix conditioning per training sample
+        # At 1.0, all training samples use audio extension mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/ and audio_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Each sample includes an audio suffix condition for backward audio extension.
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+      conditions:
+        - type: suffix
+          audio: "/path/to/suffix_audio_1.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+      conditions:
+        - type: suffix
+          audio: "/path/to/suffix_audio_2.wav"
+          # Duration is in seconds; choose a value that covers a comparable audio context
+          # to the training temporal_boundary measured in audio latent timesteps.
+          duration: 1.0
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "audio-suffix", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/audio_suffix_lora"

packages/ltx-trainer/configs/av2av_ic_lora.yaml

@@ -0,0 +1,342 @@
+# =============================================================================
+# LTX-2 AV2AV IC-LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training In-Context LoRA (IC-LoRA) adapters that
+# enable joint audio-video-to-audio-video transformations. IC-LoRA learns to
+# apply transformations to both the video and audio modalities simultaneously
+# by conditioning on paired reference video and audio.
+#
+# Both modalities use reference conditioning: pre-encoded reference latents
+# are concatenated to each modality's target sequence. Reference tokens
+# participate in bidirectional self-attention but receive no noise and are
+# excluded from the loss.
+#
+# Key differences from video-only IC-LoRA (v2v_ic_lora.yaml):
+#   - Both video AND audio have reference conditions
+#   - Requires preprocessed reference latents for BOTH modalities
+#   - LoRA targets all modules (video, audio, and cross-modal attention)
+#   - Validation uses both video and audio reference conditions
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/                 # Target video latents
+#   ├── audio_latents/           # Target audio latents
+#   ├── conditions/              # Text embeddings
+#   ├── reference_latents/       # Reference video latents (conditioning input)
+#   └── reference_audio_latents/ # Reference audio latents (conditioning input)
+#
+# Dataset metadata columns: video, audio, reference_video, reference_audio, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # IC-LoRA reference conditioning is intended for LoRA adapter training.
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 16-32 for IC-LoRA.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # For AV2AV IC-LoRA, we target ALL modules — video, audio, and cross-modal attention.
+  # Using short patterns matches all branches simultaneously.
+  target_modules:
+    # Attention layers (matches video, audio, and cross-modal branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the AV2AV IC-LoRA training approach using the unified flexible
+# strategy. Both video and audio modalities have reference conditioning,
+# enabling joint audiovisual transformations.
+training_strategy:
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing target video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    conditions:
+      # Reference conditioning (IC-LoRA): concatenates pre-encoded reference video
+      # latents to the target sequence
+      - type: reference
+        latents_dir: "reference_latents"
+        probability: 1.0
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing target audio latents
+    latents_dir: "audio_latents"
+
+    # Conditions applied to the audio modality during training
+    conditions:
+      # Reference conditioning (IC-LoRA): concatenates pre-encoded reference audio
+      # latents to the target sequence
+      - type: reference
+        latents_dir: "reference_audio_latents"
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, audio_latents/, conditions/, reference_latents/, and reference_audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # For AV2AV IC-LoRA, each sample includes reference video and reference audio conditions.
+  samples:
+    - prompt: >-
+        A man in a casual blue jacket walks along a winding path through a lush green park on a
+        bright sunny afternoon. Tall oak trees line the pathway, their leaves rustling gently in
+        the breeze. Dappled sunlight creates shifting patterns on the ground as he strolls at a
+        relaxed pace, occasionally looking up at the scenery around him. The audio captures
+        footsteps on gravel, birds singing in the trees, distant children playing, and the soft
+        whisper of wind through the foliage.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_1.mp4"
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
+        - type: reference
+          audio: "/path/to/reference_audio_1.wav"
+    - prompt: >-
+        A fluffy orange tabby cat sits perfectly still on a wooden windowsill, its green eyes
+        intently tracking small birds hopping on a branch just outside the glass. The cat's ears
+        twitch and rotate, following every movement. Warm afternoon light illuminates its fur,
+        creating a soft golden glow. Behind the cat, a cozy living room with a bookshelf and
+        houseplants is visible. The audio features gentle purring, occasional soft meows, muffled
+        bird chirps through the window, and quiet ambient room sounds.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_2.mp4"
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
+        - type: reference
+          audio: "/path/to/reference_audio_2.wav"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # Both video and audio modalities are trained
+
+  # Whether to generate audio in validation samples
+  # Can be enabled even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "ic-lora", "av2av" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/av2av_ic_lora"

packages/ltx-trainer/configs/i2v_lora.yaml

@@ -0,0 +1,348 @@
+# =============================================================================
+# LTX-2 Image-to-Video LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# image-to-video generation. The model learns to generate videos conditioned
+# on a starting image (first frame), with optional audio generation.
+#
+# First-frame conditioning works by providing the first frame as a clean
+# conditioning signal during training — it receives no noise, timestep=0,
+# and is excluded from the loss. This teaches the model to animate from
+# a given image.
+#
+# Use this configuration when you want to:
+#   - Train the model to generate videos starting from a given image
+#   - Fine-tune image-to-video capabilities on custom datasets
+#   - Create image animation models with optional audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the image-to-video training approach using the unified flexible strategy.
+# First-frame conditioning provides the first frame as clean conditioning signal
+# during training, teaching the model to animate from a given image.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # First-frame conditioning: the first frame of each video is provided as a clean
+    # conditioning signal (no noise, timestep=0, excluded from loss)
+    conditions:
+      - type: first_frame
+        # Probability of applying first-frame conditioning per training sample
+        # At 0.5, half the training samples use I2V mode, half use pure T2V
+        # Higher values improve I2V quality but may reduce T2V diversity
+        probability: 0.5
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+      conditions:
+        - type: first_frame
+          image_or_video: "/path/to/conditioning_image_1.png"
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+      conditions:
+        - type: first_frame
+          image_or_video: "/path/to/conditioning_image_2.png"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "i2v" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/i2v_lora"

packages/ltx-trainer/configs/t2a_lora.yaml

@@ -0,0 +1,317 @@
+# =============================================================================
+# LTX-2 Text-to-Audio LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# text-to-audio generation. The model learns to generate audio from text
+# prompts without any additional conditioning.
+#
+# This is the simplest audio-only training mode — no reference audio, no
+# video modality. Only the audio branch of the transformer is trained.
+#
+# Use this configuration when you want to:
+#   - Fine-tune audio generation for specific sound styles or domains
+#   - Train custom audio generation capabilities
+#   - Create audio LoRAs that can be combined with video LoRAs
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── conditions/      # Text embeddings for each sample
+#   └── audio_latents/   # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general audio LoRA training.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES (not used for audio-only modes):
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (not used for audio-only modes):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For audio-only training, we explicitly target audio modules.
+  # Including audio FFN layers can increase the LoRA's capacity.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward (often improves transformation quality)
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the text-to-audio training approach using the unified flexible
+# strategy. Audio-only training with no additional conditions — the model
+# learns to generate audio purely from text prompts.
+training_strategy:
+  name: "flexible"
+
+  # Audio modality configuration (audio-only, no video)
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: conditions/ and audio_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Text-to-audio validation samples do not need additional conditions.
+  samples:
+    - prompt: >-
+        A warm and soothing piano melody with soft ambient textures, reminiscent of a quiet
+        evening by the fireplace. Gentle reverb creates a sense of intimate space.
+    - prompt: >-
+        Energetic electronic beats with pulsing synthesizer leads and crisp hi-hat patterns.
+        Deep bass tones provide a driving rhythm underneath bright melodic arpeggios.
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Generation length control [width, height, frames]
+  # With generate_video=false, width/height are unused; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" skips both audio and video self-attention (suited to audio validation)
+
+  # Whether to generate audio in validation samples
+  # Enabled because this audio-only config generates audio
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled because no video modality is configured
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "t2a", "audio-only" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/t2a_lora"

packages/ltx-trainer/configs/t2v_lora.yaml

@@ -0,0 +1,327 @@
+# =============================================================================
+# LTX-2 Text-to-Video LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# text-to-video generation with joint audio-video support.
+#
+# Use this configuration when you want to:
+#   - Fine-tune LTX-2 on your own video dataset
+#   - Train joint audio-video generation from text prompts
+#   - Create custom video generation styles or audiovisual concepts
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the training approach using the unified flexible strategy.
+# This configuration trains both video and audio generation from text prompts.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  # When is_generated is true, the model learns to generate (denoise) video
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+  # Audio modality configuration
+  # When is_generated is true, the model learns to generate (denoise) audio
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "t2v" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/t2v_lora"

packages/ltx-trainer/configs/t2v_lora_low_vram.yaml

@@ -0,0 +1,339 @@
+# =============================================================================
+# LTX-2 Text-to-Video LoRA Training Configuration (Low VRAM)
+# =============================================================================
+#
+# This is a memory-optimized variant of the standard text-to-video LoRA config.
+# It uses 8-bit optimizer, int8 quantization, and reduced LoRA rank to minimize
+# GPU memory usage while maintaining good training quality.
+#
+# Memory optimizations applied:
+#   - 8-bit AdamW optimizer (reduces optimizer state memory by ~75%)
+#   - INT8 model quantization (reduces model memory by ~50%)
+#   - Lower LoRA rank (16 vs 32, reduces trainable parameters)
+#   - Gradient checkpointing enabled
+#
+# Recommended for GPUs with 32GB VRAM (e.g., RTX 5090).
+#
+# Use this configuration when you want to:
+#   - Fine-tune LTX-2 on your own video dataset with limited GPU memory
+#   - Train joint audio-video generation from text prompts
+#   - Create custom video generation styles or audiovisual concepts
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+# Using a lower rank (16) to reduce trainable parameters and memory usage.
+# This still provides good capacity for many fine-tuning tasks.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Using 16 for low VRAM configuration.
+  rank: 16
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 16
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the training approach using the unified flexible strategy.
+# This configuration trains both video and audio generation from text prompts.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  # When is_generated is true, the model learns to generate (denoise) video
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+  # Audio modality configuration
+  # When is_generated is true, the model learns to generate (denoise) audio
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  # Using 8-bit AdamW to reduce optimizer state memory by ~75%
+  optimizer_type: "adamw8bit"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  # Using INT8 quantization to reduce base model memory consumption by ~50%
+  quantization: "int8-quanto"
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: true
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: true
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 49 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [ 29 ]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy audio modality settings — you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "t2v", "low-vram" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/t2v_lora_low_vram"

packages/ltx-trainer/configs/v2a_lora.yaml

@@ -0,0 +1,349 @@
+# =============================================================================
+# LTX-2 Video-to-Audio (Foley) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video-to-audio (Foley) generation. The model learns to generate audio
+# conditioned on a frozen video signal via the transformer's built-in
+# cross-modal attention.
+#
+# In this mode, video is provided as a frozen (clean, no noise, no loss)
+# conditioning signal. The audio modality is the only generated output.
+# Video influences audio generation through the transformer's video-to-audio
+# cross-attention mechanism.
+#
+# Use this configuration when you want to:
+#   - Generate sound effects (Foley) for existing videos
+#   - Train audio generation conditioned on visual content
+#   - Create models that produce audio matching video content
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (frozen conditioning input)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio, generated output)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For video-to-audio training, we target audio modules and the cross-modal
+  # attention that allows audio to attend to video features.
+  target_modules:
+    # Audio self-attention
+    - "audio_attn1.to_k"
+    - "audio_attn1.to_q"
+    - "audio_attn1.to_v"
+    - "audio_attn1.to_out.0"
+    # Audio cross-attention to text
+    - "audio_attn2.to_k"
+    - "audio_attn2.to_q"
+    - "audio_attn2.to_v"
+    - "audio_attn2.to_out.0"
+    # Audio feed-forward
+    - "audio_ff.net.0.proj"
+    - "audio_ff.net.2"
+    # Cross-modal attention: allows audio to attend to video features
+    - "video_to_audio_attn.to_k"
+    - "video_to_audio_attn.to_q"
+    - "video_to_audio_attn.to_v"
+    - "video_to_audio_attn.to_out.0"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video-to-audio (Foley) training approach using the unified
+# flexible strategy. Video is frozen (no noise, sigma=0, excluded from loss)
+# and conditions audio generation via the transformer's cross-modal attention.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  # Video is frozen — it acts as conditioning for audio generation
+  # Frozen modalities get sigma=0, timestep=0, no noise, and no loss
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    # When false, video is passed through the transformer clean and influences
+    # audio via cross-modal attention
+    is_generated: false
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+  # Audio modality configuration
+  # Audio is the generated (denoised) output
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, audio_latents/, and conditions/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        The sound of ocean waves crashing against rocky cliffs, with seagulls calling in the
+        distance and wind whistling through coastal grass.
+      conditions:
+        - type: video_to_audio
+          video: "/path/to/conditioning_video_1.mp4"
+    - prompt: >-
+        Footsteps echo in a marble hallway as a person walks steadily, with the distant hum of
+        air conditioning and occasional door closing sounds.
+      conditions:
+        - type: video_to_audio
+          video: "/path/to/conditioning_video_2.mp4"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Validation dimensions [width, height, frames]
+  # Width/height resize the frozen video condition; frames and frame_rate set audio duration.
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Whether to generate video in validation samples
+  # Disabled for V2A since video is frozen conditioning, not generated
+  generate_video: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "v2a", "foley" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/v2a_lora"

packages/ltx-trainer/configs/v2v_ic_lora.yaml

@@ -0,0 +1,359 @@
+# =============================================================================
+# LTX-2 Video-to-Video (IC-LoRA) Training Configuration
+# =============================================================================
+#
+# This configuration is for training In-Context LoRA (IC-LoRA) adapters that
+# enable video-to-video transformations. IC-LoRA learns to apply visual
+# transformations (e.g., depth-to-video, pose control, style transfer, etc.)
+# by conditioning on reference videos.
+#
+# Key differences from text-to-video LoRA:
+#   - Uses reference videos as conditioning input alongside text prompts
+#   - Requires preprocessed reference latents in addition to target latents
+#   - Validation requires reference videos to demonstrate the transformation
+#
+# Dataset structure for IC-LoRA training:
+#   preprocessed_data_root/
+#   ├── latents/           # Target video latents (what the model learns to generate)
+#   ├── conditions/        # Text embeddings for each video
+#   └── reference_latents/ # Reference video latents (conditioning input)
+#
+# Dataset metadata columns: video, reference_video, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # IC-LoRA reference conditioning is intended for LoRA adapter training.
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 16-32 for IC-LoRA.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES (not used for video-only IC-LoRA):
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction, not used for video-only IC-LoRA):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # For IC-LoRA (video-only), we explicitly target video modules.
+  # Including FFN layers often improves transformation quality.
+  target_modules:
+    # Video self-attention
+    - "attn1.to_k"
+    - "attn1.to_q"
+    - "attn1.to_v"
+    - "attn1.to_out.0"
+    # Video cross-attention
+    - "attn2.to_k"
+    - "attn2.to_q"
+    - "attn2.to_v"
+    - "attn2.to_out.0"
+    # Video feed-forward (often improves transformation quality)
+    - "ff.net.0.proj"
+    - "ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video-to-video (IC-LoRA) training approach using the unified
+# flexible strategy. Reference conditioning concatenates pre-encoded reference
+# video latents to the target sequence. Reference tokens participate in
+# bidirectional self-attention but receive no noise and are excluded from loss.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing target video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    conditions:
+      # Reference conditioning (IC-LoRA): concatenates pre-encoded reference video
+      # latents to the target sequence. The model learns to transform the reference
+      # into the target video based on the text prompt.
+      - type: reference
+        # Directory name (within preprocessed_data_root) containing reference video latents
+        # These are the conditioning inputs that guide the transformation
+        latents_dir: "reference_latents"
+        # Probability of applying reference conditioning per training sample
+        probability: 1.0
+
+      # Optional first-frame conditioning to improve I2V capabilities
+      # At low probability, this teaches the model to also accept first-frame input
+      - type: first_frame
+        probability: 0.2
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 2e-4
+
+  # Total number of training steps
+  steps: 3000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and reference_latents/ subdirectories
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # For IC-LoRA, each sample includes a reference condition pointing to the conditioning video.
+  samples:
+    - prompt: >-
+        A man in a casual blue jacket walks along a winding path through a lush green park on a
+        bright sunny afternoon. Tall oak trees line the pathway, their leaves rustling gently in
+        the breeze. Dappled sunlight creates shifting patterns on the ground as he strolls at a
+        relaxed pace, occasionally looking up at the scenery around him. The audio captures
+        footsteps on gravel, birds singing in the trees, distant children playing, and the soft
+        whisper of wind through the foliage.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_1.mp4"
+          # Set these to match --reference-downscale-factor / --reference-temporal-scale-factor
+          # if reference latents were preprocessed at reduced spatial or temporal resolution.
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
+    - prompt: >-
+        A fluffy orange tabby cat sits perfectly still on a wooden windowsill, its green eyes
+        intently tracking small birds hopping on a branch just outside the glass. The cat's ears
+        twitch and rotate, following every movement. Warm afternoon light illuminates its fur,
+        creating a soft golden glow. Behind the cat, a cozy living room with a bookshelf and
+        houseplants is visible. The audio features gentle purring, occasional soft meows, muffled
+        bird chirps through the window, and quiet ambient room sounds.
+      conditions:
+        - type: reference
+          video: "/path/to/reference_video_2.mp4"
+          # Set these to match --reference-downscale-factor / --reference-temporal-scale-factor
+          # if reference latents were preprocessed at reduced spatial or temporal resolution.
+          downscale_factor: 1
+          temporal_scale_factor: 1
+          include_in_output: true
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 512, 512, 81 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_v"  # "stg_v" for video-only (no audio training)
+
+  # Whether to generate audio in validation samples
+  # Can be enabled even when not training the audio branch
+  generate_audio: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: 3
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-ic-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "ic-lora", "v2v" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/v2v_ic_lora"

packages/ltx-trainer/configs/video_extend_lora.yaml

@@ -0,0 +1,355 @@
+# =============================================================================
+# LTX-2 Video Extension LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video extension (temporal continuation). The model learns to extend a video
+# forward in time by conditioning on a prefix of existing frames.
+#
+# Prefix conditioning works by providing the first N temporal units as clean
+# conditioning signals during training — they receive no noise, timestep=0,
+# and are excluded from the loss. The model learns to generate continuation
+# frames that seamlessly follow the given prefix.
+#
+# Use this configuration when you want to:
+#   - Train the model to extend/continue existing videos
+#   - Fine-tune temporal continuation capabilities on custom datasets
+#   - Create seamless video extension models with optional audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video extension training approach using the unified flexible strategy.
+# Prefix conditioning provides the first N latent frames as clean conditioning,
+# teaching the model to generate temporally coherent continuations.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Prefix conditioning: the first N temporal units are provided as clean
+    # conditioning signals (no noise, timestep=0, excluded from loss)
+    conditions:
+      - type: prefix
+        # Number of latent frames to use as conditioning prefix
+        # For prefix conditioning, N latent frames correspond to (N - 1) * 8 + 1 pixel frames.
+        # temporal_boundary=8 means 57 pixel frames are used as prefix.
+        temporal_boundary: 8
+        # Probability of applying prefix conditioning per training sample
+        # At 1.0, all training samples use video extension mode
+        probability: 1.0
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+      conditions:
+        - type: prefix
+          video: "/path/to/prefix_video_1.mp4"
+          # Matches temporal_boundary=8 during training: (8 - 1) * 8 + 1 = 57 frames
+          num_frames: 57
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+      conditions:
+        - type: prefix
+          video: "/path/to/prefix_video_2.mp4"
+          # Matches temporal_boundary=8 during training: (8 - 1) * 8 + 1 = 57 frames
+          num_frames: 57
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "video-extension" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_extend_lora"

packages/ltx-trainer/configs/video_inpainting_lora.yaml

@@ -0,0 +1,343 @@
+# =============================================================================
+# LTX-2 Video Inpainting LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video inpainting. The model learns to fill in masked regions of a video
+# using per-sample binary masks that define which regions are conditioning
+# (provided clean) and which regions the model must generate.
+#
+# Mask conditioning works by loading per-sample binary masks from disk.
+# Masked regions receive clean latents (no noise, timestep=0) and are excluded
+# from the loss. Unmasked regions are noised and trained normally.
+#
+# Use this configuration when you want to:
+#   - Train the model to fill in or replace regions of existing videos
+#   - Fine-tune video inpainting capabilities on custom datasets
+#   - Create object removal or region editing models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── video_masks/    # Per-sample binary masks defining conditioning regions
+#
+# Dataset metadata columns: video, video_mask, caption
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # This is a video-only config, so explicitly target video attention modules.
+  target_modules:
+    - "attn1.to_k"
+    - "attn1.to_q"
+    - "attn1.to_v"
+    - "attn1.to_out.0"
+    - "attn2.to_k"
+    - "attn2.to_q"
+    - "attn2.to_v"
+    - "attn2.to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video inpainting training approach using the unified flexible
+# strategy. Per-sample binary masks define which regions are provided as clean
+# conditioning and which regions the model must learn to generate.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Mask conditioning: per-sample binary masks loaded from disk define which
+    # tokens are conditioning (clean, timestep=0, no loss) vs generated
+    conditions:
+      - type: mask
+        # Directory name (within preprocessed_data_root) containing binary masks
+        # Each mask file corresponds to a training sample and defines the
+        # conditioning region (mask=1 means conditioning, mask=0 means generate)
+        mask_dir: "video_masks"
+        # Probability of applying mask conditioning per training sample
+        # At 1.0, all training samples use inpainting mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and video_masks/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration.
+      conditions:
+        - type: mask
+          video: "/path/to/inpainting_video_1.mp4"
+          mask: "/path/to/inpainting_mask_1.mp4"
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him.
+      conditions:
+        - type: mask
+          video: "/path/to/inpainting_video_2.mp4"
+          mask: "/path/to/inpainting_mask_2.mp4"
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_v"  # "stg_v" for video-only validation
+
+  # Video inpainting trains video only; do not generate validation audio.
+  generate_audio: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "inpainting" ]
+
+  # Log validation outputs to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_inpainting_lora"

packages/ltx-trainer/configs/video_outpainting_lora.yaml

@@ -0,0 +1,341 @@
+# =============================================================================
+# LTX-2 Video Outpainting LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# video outpainting (spatial extension). The model learns to generate content
+# surrounding a known rectangular region of the video.
+#
+# Spatial crop conditioning works by providing a rectangular pixel region as
+# clean conditioning during training — those tokens receive no noise,
+# timestep=0, and are excluded from the loss. The model learns to generate
+# the content outside the given region, seamlessly extending the scene.
+#
+# Use this configuration when you want to:
+#   - Extend video content beyond its original boundaries
+#   - Train spatial outpainting capabilities on custom datasets
+#   - Create video aspect ratio conversion models
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   └── conditions/     # Text embeddings for each video
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # This is a video-only config, so explicitly target video attention modules.
+  target_modules:
+    - "attn1.to_k"
+    - "attn1.to_q"
+    - "attn1.to_v"
+    - "attn1.to_out.0"
+    - "attn2.to_k"
+    - "attn2.to_q"
+    - "attn2.to_v"
+    - "attn2.to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the video outpainting training approach using the unified flexible
+# strategy. Spatial crop conditioning provides a rectangular pixel region as
+# clean conditioning, teaching the model to generate surrounding content.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Spatial crop conditioning provides a rectangular region as clean context;
+    # the model learns to generate the surrounding video tokens.
+    conditions:
+      - type: spatial_crop
+        # Rectangular pixel region provided as clean conditioning (y1, x1, y2, x2)
+        # Pixels within this region are conditioning (clean, timestep=0, no loss)
+        # Pixels outside this region are generated by the model
+        # Coordinates are in pixel space — automatically converted to latent space
+        spatial_region: [0, 0, 288, 576]
+        # Probability of applying spatial crop conditioning per training sample
+        # At 1.0, all training samples use outpainting mode
+        probability: 1.0
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/ and conditions/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration.
+      conditions:
+        - type: spatial_crop
+          video: "/path/to/outpainting_video_1.mp4"
+          spatial_region: [0, 0, 288, 576]
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him.
+      conditions:
+        - type: spatial_crop
+          video: "/path/to/outpainting_video_2.mp4"
+          spatial_region: [0, 0, 288, 576]
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_v"  # "stg_v" for video-only validation
+
+  # Video outpainting trains video only; do not generate validation audio.
+  generate_audio: false
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "outpainting", "spatial-crop" ]
+
+  # Log validation outputs to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_outpainting_lora"

packages/ltx-trainer/configs/video_suffix_lora.yaml

@@ -0,0 +1,355 @@
+# =============================================================================
+# LTX-2 Video Backward Extension (Suffix) LoRA Training Configuration
+# =============================================================================
+#
+# This configuration is for training LoRA adapters on the LTX-2 model for
+# backward video extension. The model learns to generate video content that
+# leads into a given suffix of existing frames.
+#
+# Suffix conditioning works by providing the last N temporal units as clean
+# conditioning signals during training — they receive no noise, timestep=0,
+# and are excluded from the loss. The model learns to generate preceding
+# frames that seamlessly lead into the given suffix.
+#
+# Use this configuration when you want to:
+#   - Train the model to generate "prequel" content for existing videos
+#   - Fine-tune backward temporal continuation capabilities
+#   - Create video backward extension models with optional audio
+#
+# Dataset structure:
+#   preprocessed_data_root/
+#   ├── latents/        # Video latents (VAE-encoded videos)
+#   ├── conditions/     # Text embeddings for each video
+#   └── audio_latents/  # Audio latents (VAE-encoded audio)
+#
+# =============================================================================
+
+# -----------------------------------------------------------------------------
+# Model Configuration
+# -----------------------------------------------------------------------------
+# Specifies the base model to fine-tune and the training mode.
+model:
+  # Path to the LTX-2 model checkpoint (.safetensors file)
+  # This should be a local path to your downloaded model
+  model_path: "path/to/ltx-2-model.safetensors"
+
+  # Path to the text encoder model directory
+  # For LTX-2, this is typically the Gemma-based text encoder
+  text_encoder_path: "path/to/gemma-text-encoder"
+
+  # Training mode: "lora" for efficient adapter training, "full" for full fine-tuning
+  # LoRA is recommended for most use cases (faster, less memory, prevents overfitting)
+  training_mode: "lora"
+
+  # Optional: Path to resume training from a checkpoint
+  # Can be a checkpoint file (.safetensors) or directory (uses latest checkpoint)
+  load_checkpoint: null
+
+# -----------------------------------------------------------------------------
+# LoRA Configuration
+# -----------------------------------------------------------------------------
+# Controls the Low-Rank Adaptation parameters for efficient fine-tuning.
+lora:
+  # Rank of the LoRA matrices (higher = more capacity but more parameters)
+  # Typical values: 8, 16, 32, 64. Start with 32 for general fine-tuning.
+  rank: 32
+
+  # Alpha scaling factor (usually set equal to rank)
+  # The effective scaling is alpha/rank, so alpha=rank means scaling of 1.0
+  alpha: 32
+
+  # Dropout probability for LoRA layers (0.0 = no dropout)
+  # Can help with regularization if overfitting occurs
+  dropout: 0.0
+
+  # Which transformer modules to apply LoRA to
+  # The LTX-2 transformer has separate attention and FFN blocks for video and audio:
+  #
+  # VIDEO MODULES:
+  #   - attn1.to_k, attn1.to_q, attn1.to_v, attn1.to_out.0  (video self-attention)
+  #   - attn2.to_k, attn2.to_q, attn2.to_v, attn2.to_out.0  (video cross-attention to text)
+  #   - ff.net.0.proj, ff.net.2                             (video feed-forward)
+  #
+  # AUDIO MODULES:
+  #   - audio_attn1.to_k, audio_attn1.to_q, audio_attn1.to_v, audio_attn1.to_out.0  (audio self-attention)
+  #   - audio_attn2.to_k, audio_attn2.to_q, audio_attn2.to_v, audio_attn2.to_out.0  (audio cross-attention to text)
+  #   - audio_ff.net.0.proj, audio_ff.net.2                                         (audio feed-forward)
+  #
+  # AUDIO-VIDEO CROSS-ATTENTION MODULES (for cross-modal interaction):
+  #   - audio_to_video_attn.to_k, audio_to_video_attn.to_q, audio_to_video_attn.to_v, audio_to_video_attn.to_out.0
+  #       (Q from video, K/V from audio - allows video to attend to audio features)
+  #   - video_to_audio_attn.to_k, video_to_audio_attn.to_q, video_to_audio_attn.to_v, video_to_audio_attn.to_out.0
+  #       (Q from audio, K/V from video - allows audio to attend to video features)
+  #
+  # Using short patterns like "to_k" matches ALL attention modules (video, audio, and cross-modal).
+  # For audio-video training, this is the recommended approach.
+  target_modules:
+    # Attention layers (matches both video and audio branches)
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+    # Uncomment below to also train feed-forward layers (can increase the LoRA's capacity):
+    # - "ff.net.0.proj"
+    # - "ff.net.2"
+    # - "audio_ff.net.0.proj"
+    # - "audio_ff.net.2"
+
+# -----------------------------------------------------------------------------
+# Training Strategy Configuration
+# -----------------------------------------------------------------------------
+# Defines the backward video extension training approach using the unified
+# flexible strategy. Suffix conditioning provides the last N latent frames as
+# clean conditioning, teaching the model to generate content leading into them.
+training_strategy:
+  # Strategy name: "flexible" for the unified conditioning framework
+  # Supports all training modes (T2V, I2V, V2V, A2V, V2A, etc.) through
+  # modality-specific configuration blocks.
+  name: "flexible"
+
+  # Video modality configuration
+  video:
+    # Whether the model generates video (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing video latents
+    latents_dir: "latents"
+
+    # Conditions applied to the video modality during training
+    # Suffix conditioning: the last N temporal units are provided as clean
+    # conditioning signals (no noise, timestep=0, excluded from loss)
+    conditions:
+      - type: suffix
+        # Number of latent frames to use as conditioning suffix
+        # For suffix conditioning, N latent frames correspond to N * 8 pixel frames.
+        # temporal_boundary=8 means the last 64 pixel frames are used as suffix.
+        temporal_boundary: 8
+        # Probability of applying suffix conditioning per training sample
+        # At 1.0, all training samples use backward extension mode
+        probability: 1.0
+
+  # Audio modality configuration
+  audio:
+    # Whether the model generates audio (true) or uses it as frozen conditioning (false)
+    is_generated: true
+    # Directory name (within preprocessed_data_root) containing audio latents
+    latents_dir: "audio_latents"
+
+# -----------------------------------------------------------------------------
+# Optimization Configuration
+# -----------------------------------------------------------------------------
+# Controls the training optimization parameters.
+optimization:
+  # Learning rate for the optimizer
+  # Typical range for LoRA: 1e-5 to 1e-4
+  learning_rate: 1e-4
+
+  # Total number of training steps
+  steps: 2000
+
+  # Batch size per GPU
+  # Reduce if running out of memory
+  batch_size: 1
+
+  # Number of gradient accumulation steps
+  # Effective batch size = batch_size * gradient_accumulation_steps * num_gpus
+  gradient_accumulation_steps: 1
+
+  # Maximum gradient norm for clipping (helps training stability)
+  max_grad_norm: 1.0
+
+  # Optimizer type: "adamw" (standard) or "adamw8bit" (memory-efficient)
+  optimizer_type: "adamw"
+
+  # Learning rate scheduler type
+  # Options: "constant", "linear", "cosine", "cosine_with_restarts", "polynomial"
+  scheduler_type: "linear"
+
+  # Additional scheduler parameters (depends on scheduler_type)
+  scheduler_params: { }
+
+  # Enable gradient checkpointing to reduce memory usage
+  # Recommended for training with limited GPU memory
+  enable_gradient_checkpointing: true
+
+# -----------------------------------------------------------------------------
+# Acceleration Configuration
+# -----------------------------------------------------------------------------
+# Hardware acceleration and memory optimization settings.
+acceleration:
+  # Mixed precision training mode
+  # Options: "no" (fp32), "fp16" (half precision), "bf16" (bfloat16, recommended)
+  mixed_precision_mode: "bf16"
+
+  # Model quantization for reduced memory usage
+  # Options: null (none), "int8-quanto", "int4-quanto", "int2-quanto", "fp8-quanto", "fp8uz-quanto"
+  quantization: null
+
+  # Load text encoder in 8-bit precision to save memory
+  # Useful when GPU memory is limited
+  load_text_encoder_in_8bit: false
+
+  # Offload optimizer state to CPU during validation video sampling and restore it after.
+  # Frees VRAM for the VAE decoder when optimizer state is large (full fine-tune, high-rank
+  # LoRA). No effect under FSDP (sharded state).
+  offload_optimizer_during_validation: false
+
+
+# -----------------------------------------------------------------------------
+# Data Configuration
+# -----------------------------------------------------------------------------
+# Specifies the training data location and loading parameters.
+data:
+  # Root directory containing preprocessed training data
+  # Should contain: latents/, conditions/, and audio_latents/
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+  # Number of worker processes for data loading
+  # Used for parallel data loading to speed up data loading
+  num_dataloader_workers: 2
+
+# -----------------------------------------------------------------------------
+# Validation Configuration
+# -----------------------------------------------------------------------------
+# Controls validation sampling during training.
+# NOTE: Validation sampling use simplified inference pipelines and prioritizes speed over
+# maximum quality. For production-quality inference, use `packages/ltx-pipelines`.
+validation:
+  # Validation samples — each sample describes a self-contained generation request.
+  # Use 'conditions' to add conditioning (first_frame, prefix, suffix, reference, video_to_audio, etc.)
+  # See docs/configuration-reference.md#validation-condition-types for the full list of condition types.
+  samples:
+    - prompt: >-
+        A woman with long brown hair sits at a wooden desk in a cozy home office, typing on a
+        laptop while occasionally glancing at notes beside her. Soft natural light streams through
+        a large window, casting warm shadows across the room. She pauses to take a sip from a
+        ceramic mug, then continues working with focused concentration. The audio captures the
+        gentle clicking of keyboard keys, the soft rustle of papers, and ambient room tone with
+        occasional distant bird chirps from outside.
+      conditions:
+        - type: suffix
+          video: "/path/to/suffix_video_1.mp4"
+          # Matches temporal_boundary=8 during training: 8 * 8 = 64 frames
+          num_frames: 64
+    - prompt: >-
+        A chef in a white uniform stands in a professional kitchen, carefully plating a gourmet
+        dish with precise movements. Steam rises from freshly cooked vegetables as he arranges
+        them with tweezers. The stainless steel surfaces gleam under bright overhead lights, and
+        various pots simmer on the stove behind him. The audio features the sizzling of pans,
+        the clinking of utensils against plates, and the ambient hum of kitchen ventilation.
+      conditions:
+        - type: suffix
+          video: "/path/to/suffix_video_2.mp4"
+          # Matches temporal_boundary=8 during training: 8 * 8 = 64 frames
+          num_frames: 64
+
+  # Negative prompt to avoid unwanted artifacts
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"
+
+  # Output video dimensions [width, height, frames]
+  # Width and height must be divisible by 32
+  # Frames must satisfy: frames % 8 == 1 (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, ...)
+  video_dims: [ 576, 576, 89 ]
+
+  # Frame rate for generated videos
+  frame_rate: 25.0
+
+  # Random seed for reproducible validation outputs
+  seed: 42
+
+  # Number of denoising steps for validation inference
+  # Higher values = better quality but slower generation
+  inference_steps: 30
+
+  # Generate validation videos every N training steps
+  # Set to null to disable validation during training
+  interval: 100
+
+  # Classifier-free guidance scale
+  # Higher values = stronger adherence to prompt but may introduce artifacts
+  guidance_scale: 4.0
+
+  # STG (Spatio-Temporal Guidance) parameters for improved video quality
+  # STG is combined with CFG for better temporal coherence
+  stg_scale: 1.0  # Recommended: 1.0 (0.0 disables STG)
+  stg_blocks: [29]  # Recommended: single block 29
+  stg_mode: "stg_av"  # "stg_av" perturbs both audio and video, "stg_v" video only
+
+  # Whether to generate audio in validation samples
+  # Independent of training_strategy.audio.is_generated - you can generate audio
+  # in validation even when not training the audio branch
+  generate_audio: true
+
+  # Skip validation at the beginning of training (step 0)
+  skip_initial_validation: false
+
+# -----------------------------------------------------------------------------
+# Checkpoint Configuration
+# -----------------------------------------------------------------------------
+# Controls model checkpoint saving during training.
+checkpoints:
+  # Save a checkpoint every N steps
+  # Set to null to disable intermediate checkpoints
+  interval: 250
+
+  # Number of most recent checkpoints to keep
+  # Set to -1 to keep all checkpoints
+  keep_last_n: -1
+
+  # Precision to use when saving checkpoint weights
+  # Options: "bfloat16" (default, smaller files) or "float32" (full precision)
+  precision: "bfloat16"
+
+# -----------------------------------------------------------------------------
+# Flow Matching Configuration
+# -----------------------------------------------------------------------------
+# Parameters for the flow matching training objective.
+flow_matching:
+  # Timestep sampling mode
+  # "shifted_logit_normal" is recommended for LTX-2 models
+  timestep_sampling_mode: "shifted_logit_normal"
+
+  # Additional parameters for timestep sampling
+  timestep_sampling_params: { }
+
+# -----------------------------------------------------------------------------
+# Hugging Face Hub Configuration
+# -----------------------------------------------------------------------------
+# Settings for uploading trained models to the Hugging Face Hub.
+hub:
+  # Whether to push the trained model to the Hub
+  push_to_hub: false
+
+  # Repository ID on Hugging Face Hub (e.g., "username/my-lora-model")
+  # Required if push_to_hub is true
+  hub_model_id: null
+
+# -----------------------------------------------------------------------------
+# Weights & Biases Configuration
+# -----------------------------------------------------------------------------
+# Settings for experiment tracking with W&B.
+wandb:
+  # Enable W&B logging
+  enabled: false
+
+  # W&B project name
+  project: "ltx-2-trainer"
+
+  # W&B username or team (null uses default account)
+  entity: null
+
+  # Tags to help organize runs
+  tags: [ "ltx2", "lora", "video-suffix", "backward-extension" ]
+
+  # Log validation media (video/audio) to W&B
+  log_validation_videos: true
+
+# -----------------------------------------------------------------------------
+# General Configuration
+# -----------------------------------------------------------------------------
+# Global settings for the training run.
+
+# Random seed for reproducibility
+seed: 42
+
+# Directory to save outputs (checkpoints, validation videos, logs)
+output_dir: "outputs/video_suffix_lora"

packages/ltx-trainer/docs/configuration-reference.md

@@ -0,0 +1,464 @@
+# Configuration Reference
+
+The trainer uses structured Pydantic models for configuration, making it easy to customize training parameters.
+This guide covers all available configuration options and their usage.
+
+## 📋 Overview
+
+The main configuration class is [`LtxTrainerConfig`](../src/ltx_trainer/config.py), which includes the following
+sub-configurations:
+
+- **ModelConfig**: Base model and training mode settings
+- **LoraConfig**: LoRA training parameters
+- **TrainingStrategyConfig**: Training strategy settings (flexible conditioning framework)
+- **OptimizationConfig**: Learning rate, batch sizes, and scheduler settings
+- **AccelerationConfig**: Mixed precision and quantization settings
+- **DataConfig**: Data loading parameters
+- **ValidationConfig**: Validation and inference settings
+- **CheckpointsConfig**: Checkpoint saving frequency and retention settings
+- **HubConfig**: Hugging Face Hub integration settings
+- **WandbConfig**: Weights & Biases logging settings
+- **FlowMatchingConfig**: Timestep sampling parameters
+
+## 📄 Example Configuration Files
+
+Check out our example configurations in the `configs` directory:
+
+- 📄 [Text-to-Video LoRA](../configs/t2v_lora.yaml) - Text-to-video LoRA training
+- 📄 [Image-to-Video LoRA](../configs/i2v_lora.yaml) - Image-to-video LoRA training
+- 📄 [IC-LoRA Video-to-Video](../configs/v2v_ic_lora.yaml) - IC-LoRA video-to-video training
+- 📄 [Audio-to-Video LoRA](../configs/a2v_lora.yaml) - Audio-to-video LoRA training
+- 📄 [Video-to-Audio LoRA](../configs/v2a_lora.yaml) - Video-to-audio (Foley) LoRA training
+- 📄 [Video Extension LoRA](../configs/video_extend_lora.yaml) - Video extension (forward) LoRA training
+- 📄 [Video Suffix LoRA](../configs/video_suffix_lora.yaml) - Video extension (backward) LoRA training
+- 📄 [Video Inpainting LoRA](../configs/video_inpainting_lora.yaml) - Video inpainting LoRA training
+- 📄 [Video Outpainting LoRA](../configs/video_outpainting_lora.yaml) - Video outpainting (spatial crop) LoRA training
+- 📄 [Text-to-Audio LoRA](../configs/t2a_lora.yaml) - Text-to-audio LoRA training
+- 📄 [Audio Extension LoRA](../configs/audio_extend_lora.yaml) - Audio extension (forward) LoRA training
+- 📄 [Audio Suffix LoRA](../configs/audio_suffix_lora.yaml) - Audio extension (backward) LoRA training
+- 📄 [Audio Inpainting LoRA](../configs/audio_inpainting_lora.yaml) - Audio inpainting LoRA training
+- 📄 [Audio-to-Audio IC-LoRA](../configs/a2a_ic_lora.yaml) - Audio IC-LoRA transformation training
+- 📄 [AV2AV IC-LoRA](../configs/av2av_ic_lora.yaml) - Audio+video IC-LoRA transformation training
+- 📄 [T2V LoRA (Low VRAM)](../configs/t2v_lora_low_vram.yaml) - Memory-optimized config for 32GB GPUs
+
+## ⚙️ Configuration Sections
+
+> [!NOTE]
+> The YAML snippets below show **recommended starting values**, not necessarily the code defaults.
+> Fields you omit from your config file will use the code defaults from [`config.py`](../src/ltx_trainer/config.py).
+
+### ModelConfig
+
+Controls the base model and training mode settings.
+
+```yaml
+model:
+  model_path: "/path/to/ltx-2-model.safetensors"  # Local path to model checkpoint
+  text_encoder_path: "/path/to/gemma-model"       # Path to Gemma text encoder directory
+  training_mode: "lora"                           # "lora" or "full"
+  load_checkpoint: null                           # Path to checkpoint to resume from
+```
+
+**Key parameters:**
+
+| Parameter           | Description                                                                                                                                                    |
+|---------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `model_path`        | **Required.** Local path to the LTX-2 model checkpoint (`.safetensors` file). URLs are not supported.                                                          |
+| `text_encoder_path` | **Required.** Path to the Gemma text encoder model directory. Download from [HuggingFace](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized/). |
+| `training_mode`     | Training approach - `"lora"` for LoRA training or `"full"` for full-rank fine-tuning.                                                                          |
+| `load_checkpoint`   | Optional path to resume training from a checkpoint file or directory.                                                                                          |
+
+> [!NOTE]
+> LTX-2 requires both a model checkpoint and a Gemma text encoder. Both must be local paths.
+
+### LoraConfig
+
+LoRA-specific fine-tuning parameters (only used when `training_mode: "lora"`).
+
+```yaml
+lora:
+  rank: 32         # LoRA rank (higher = more parameters)
+  alpha: 32        # LoRA alpha scaling factor
+  dropout: 0.0     # Dropout probability (0.0-1.0)
+  target_modules: # Modules to apply LoRA to
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+```
+
+**Key parameters:**
+
+| Parameter        | Description                                                                     |
+|------------------|---------------------------------------------------------------------------------|
+| `rank`           | LoRA rank - higher values mean more trainable parameters (typical range: 8-128) |
+| `alpha`          | Alpha scaling factor - typically set equal to rank                              |
+| `dropout`        | Dropout probability for regularization                                          |
+| `target_modules` | List of transformer modules to apply LoRA adapters to (see below)               |
+
+#### Understanding Target Modules
+
+The LTX-2 transformer has separate attention and feed-forward blocks for video and audio, as well as cross-attention
+modules that enable the two modalities to exchange information. Choosing the right `target_modules` is critical for
+achieving good results, especially when training with audio.
+
+**Video-only modules:**
+
+| Module Pattern                                             | Description                     |
+|------------------------------------------------------------|---------------------------------|
+| `attn1.to_k`, `attn1.to_q`, `attn1.to_v`, `attn1.to_out.0` | Video self-attention            |
+| `attn2.to_k`, `attn2.to_q`, `attn2.to_v`, `attn2.to_out.0` | Video cross-attention (to text) |
+| `ff.net.0.proj`, `ff.net.2`                                | Video feed-forward network      |
+
+**Audio-only modules:**
+
+| Module Pattern                                                                     | Description                     |
+|------------------------------------------------------------------------------------|---------------------------------|
+| `audio_attn1.to_k`, `audio_attn1.to_q`, `audio_attn1.to_v`, `audio_attn1.to_out.0` | Audio self-attention            |
+| `audio_attn2.to_k`, `audio_attn2.to_q`, `audio_attn2.to_v`, `audio_attn2.to_out.0` | Audio cross-attention (to text) |
+| `audio_ff.net.0.proj`, `audio_ff.net.2`                                            | Audio feed-forward network      |
+
+**Audio-video cross-attention modules:**
+
+These modules enable bidirectional information flow between the audio and video modalities:
+
+| Module Pattern                                                                                                     | Description                                           |
+|--------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------|
+| `audio_to_video_attn.to_k`, `audio_to_video_attn.to_q`, `audio_to_video_attn.to_v`, `audio_to_video_attn.to_out.0` | Video attends to audio (Q from video, K/V from audio) |
+| `video_to_audio_attn.to_k`, `video_to_audio_attn.to_q`, `video_to_audio_attn.to_v`, `video_to_audio_attn.to_out.0` | Audio attends to video (Q from audio, K/V from video) |
+
+**Recommended configurations:**
+
+For **video-only training**, target the video attention layers:
+
+```yaml
+target_modules:
+  - "attn1.to_k"
+  - "attn1.to_q"
+  - "attn1.to_v"
+  - "attn1.to_out.0"
+  - "attn2.to_k"
+  - "attn2.to_q"
+  - "attn2.to_v"
+  - "attn2.to_out.0"
+```
+
+For **audio-video training**, use patterns that match both branches:
+
+```yaml
+target_modules:
+  - "to_k"
+  - "to_q"
+  - "to_v"
+  - "to_out.0"
+```
+
+> [!NOTE]
+> Using shorter patterns like `"to_k"` will match all attention modules including `attn1.to_k`, `audio_attn1.to_k`,
+> `audio_to_video_attn.to_k`, and `video_to_audio_attn.to_k`, effectively training video, audio, and cross-modal
+> attention branches together.
+
+> [!TIP]
+> You can also target the feed-forward (FFN) modules (`ff.net.0.proj`, `ff.net.2` for video,
+> `audio_ff.net.0.proj`, `audio_ff.net.2` for audio) to increase the LoRA's capacity and potentially
+> help it capture the target distribution better.
+
+### TrainingStrategyConfig
+
+Configures the training strategy. The recommended strategy is `"flexible"`, which supports all conditioning scenarios through configuration.
+
+#### Flexible Strategy
+
+The flexible strategy provides a unified conditioning framework. Each modality (video, audio) is configured
+independently with its own latents directory, generation flag, and list of conditions.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true                 # Video is denoised during training
+    latents_dir: "latents"             # Directory containing precomputed video latents
+    conditions:
+      - type: first_frame              # Use first frame as conditioning
+        probability: 0.5               # Apply this condition 50% of the time
+  audio:
+    is_generated: true                 # Audio is denoised during training
+    latents_dir: "audio_latents"       # Directory containing precomputed audio latents
+    conditions: []                     # No additional audio conditions (text-only)
+```
+
+**ModalityConfig parameters:**
+
+| Parameter      | Description                                                                                                      |
+|----------------|------------------------------------------------------------------------------------------------------------------|
+| `is_generated` | `true` = modality is denoised (contributes to loss). `false` = frozen conditioning (sigma=0, no loss).           |
+| `latents_dir`  | Directory name within `preprocessed_data_root` containing precomputed latents for this modality.                 |
+| `conditions`   | List of conditioning configs applied during training (see condition types below). Text conditioning is implicit. |
+
+**Condition types:**
+
+| Type           | Parameters                                          | Description                                                                           |
+|----------------|-----------------------------------------------------|---------------------------------------------------------------------------------------|
+| `first_frame`  | `probability`                                       | First latent frame is clean, excluded from loss. **Video only.**                      |
+| `prefix`       | `temporal_boundary`, `probability`                  | First N latent temporal units are clean. For extension forward.                       |
+| `suffix`       | `temporal_boundary`, `probability`                  | Last N latent temporal units are clean. For extension backward.                       |
+| `spatial_crop` | `spatial_region` (y1, x1, y2, x2 in px), `probability` | Rectangular region is clean, excluded from loss. For outpainting. **Video only.** |
+| `mask`         | `mask_dir`, `probability`                           | Per-sample mask directory. Masks are thresholded at `0.5`; `1` means conditioning, `0` means generate. |
+| `reference`    | `latents_dir`, `probability`                        | IC-LoRA style concatenation. Reference tokens are prepended, clean (timestep=0), no loss. |
+
+> [!NOTE]
+> The `prefix`, `suffix`, `mask`, and `reference` condition types work on both video and audio modalities —
+> place them in the `video.conditions` or `audio.conditions` list as appropriate.
+> `first_frame` and `spatial_crop` are video-only conditions.
+
+> [!NOTE]
+> Training conditions reference **directories** of precomputed data (within `preprocessed_data_root`),
+> while validation conditions reference **individual files** (images, videos, masks) that are encoded
+> on-the-fly during validation. The condition `type` names are the same, but the fields differ.
+
+> [!NOTE]
+> The legacy `text_to_video` and `video_to_video` strategies are deprecated but remain forward-compatible.
+> New configs should use `name: "flexible"`.
+
+### OptimizationConfig
+
+Training optimization parameters including learning rates, batch sizes, and schedulers.
+
+```yaml
+optimization:
+  learning_rate: 1e-4                  # Learning rate
+  steps: 2000                          # Total training steps
+  batch_size: 1                        # Batch size per GPU
+  gradient_accumulation_steps: 1       # Steps to accumulate gradients
+  max_grad_norm: 1.0                   # Gradient clipping threshold
+  optimizer_type: "adamw"              # "adamw" or "adamw8bit"
+  scheduler_type: "linear"             # Scheduler type
+  scheduler_params: { }                # Additional scheduler parameters
+  enable_gradient_checkpointing: true  # Memory optimization
+```
+
+**Key parameters:**
+
+| Parameter                       | Description                                                                                  |
+|---------------------------------|----------------------------------------------------------------------------------------------|
+| `learning_rate`                 | Learning rate for optimization (typical range: 1e-5 to 1e-3)                                 |
+| `steps`                         | Total number of training steps                                                               |
+| `batch_size`                    | Batch size per GPU (reduce if running out of memory)                                         |
+| `gradient_accumulation_steps`   | Accumulate gradients over multiple steps                                                     |
+| `scheduler_type`                | LR scheduler: `"constant"`, `"linear"`, `"cosine"`, `"cosine_with_restarts"`, `"polynomial"`, `"step"` |
+| `enable_gradient_checkpointing` | Trade training speed for GPU memory savings (recommended for large models)                   |
+
+### AccelerationConfig
+
+Hardware acceleration and compute optimization settings.
+
+```yaml
+acceleration:
+  mixed_precision_mode: "bf16"                  # "no", "fp16", or "bf16"
+  quantization: null                            # Quantization options
+  load_text_encoder_in_8bit: false              # Load text encoder in 8-bit
+  offload_optimizer_during_validation: false    # Offload optimizer state to CPU during validation
+```
+
+**Key parameters:**
+
+| Parameter                             | Description                                                                                                                                                                              |
+|---------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `mixed_precision_mode`                | Precision mode - `"bf16"` recommended for modern GPUs                                                                                                                                    |
+| `quantization`                        | Model quantization: `null`, `"int8-quanto"`, `"int4-quanto"`, `"int2-quanto"`, `"fp8-quanto"`, or `"fp8uz-quanto"`                                                                       |
+| `load_text_encoder_in_8bit`           | Load the Gemma text encoder in 8-bit to save GPU memory                                                                                                                                  |
+| `offload_optimizer_during_validation` | Move optimizer state to CPU before validation video sampling and back afterwards. Useful when validation OOMs because VAE decoder + transformer + optimizer state can't coexist on the GPU (full fine-tune, high-rank LoRA). No effect for FSDP. |
+
+### DataConfig
+
+Data loading and processing configuration.
+
+```yaml
+data:
+  preprocessed_data_root: "/path/to/preprocessed/data"  # Path to precomputed dataset
+  num_dataloader_workers: 2                             # Background data loading workers
+```
+
+**Key parameters:**
+
+| Parameter                | Description                                                                                |
+|--------------------------|--------------------------------------------------------------------------------------------|
+| `preprocessed_data_root` | Path to your preprocessed dataset directory produced by `process_dataset.py` (contains `latents/`, `conditions/`, etc.) |
+| `num_dataloader_workers` | Number of parallel data loading processes (0 = synchronous loading, useful when debugging) |
+
+### ValidationConfig
+
+Validation and inference settings for monitoring training progress. Validation samples use a self-describing
+format where each sample specifies its own prompt and conditions.
+
+```yaml
+validation:
+  samples:
+    - prompt: "A cat playing with a ball"
+      conditions:
+        - type: first_frame
+          image_or_video: "/path/to/image.png"
+    - prompt: "A dog running in a field"
+  video_dims: [576, 576, 89]                # Output dimensions: [width, height, frames]
+  negative_prompt: "worst quality, inconsistent motion, blurry, jittery, distorted"  # Negative prompt for all samples
+  frame_rate: 25.0                          # Output video frame rate (fps)
+  seed: 42                                  # Random seed for reproducibility
+  inference_steps: 30                       # Number of denoising steps
+  interval: 100                             # Run validation every N steps (null to disable)
+  guidance_scale: 4.0                       # CFG scale (higher = stronger prompt adherence)
+  stg_scale: 1.0                            # STG scale (0.0 to disable)
+  stg_blocks: [29]                          # Transformer blocks to apply STG perturbation
+  stg_mode: "stg_av"                        # STG mode: "stg_av" (audio+video) or "stg_v" (video only)
+  generate_audio: true                      # Whether to generate audio during validation
+  generate_video: true                      # Whether to generate video during validation
+  skip_initial_validation: false            # Skip validation at step 0
+```
+
+**Key parameters:**
+
+| Parameter                | Description                                                                                                              |
+|--------------------------|--------------------------------------------------------------------------------------------------------------------------|
+| `samples`                | List of `ValidationSample` objects (see below). Replaces the legacy `prompts`/`images`/`reference_videos` fields.        |
+| `video_dims`             | Output dimensions `[width, height, frames]`. Width/height must be divisible by 32, frames must satisfy `frames % 8 == 1` |
+| `interval`               | Steps between validation runs (set to `null` to disable)                                                                 |
+| `guidance_scale`         | CFG (Classifier-Free Guidance) scale. Recommended: 4.0                                                                   |
+| `stg_scale`              | STG (Spatio-Temporal Guidance) scale. 0.0 disables STG. Recommended: 1.0                                                 |
+| `stg_blocks`             | Transformer blocks to perturb for STG. Recommended: `[29]` (single block)                                                |
+| `stg_mode`               | STG mode: `"stg_av"` perturbs both audio and video, `"stg_v"` perturbs video only                                        |
+| `generate_audio`         | Whether to generate audio in validation samples                                                                          |
+| `generate_video`         | Whether to generate video in validation samples. Set to `false` for V2A (video-to-audio) validation. Default: `true`     |
+| `skip_initial_validation`| Skip validation video sampling at step 0 (beginning of training)                                                         |
+
+#### ValidationSample
+
+Each sample in the `samples` list has:
+
+| Field        | Description                                                                                     |
+|--------------|-------------------------------------------------------------------------------------------------|
+| `prompt`     | Text prompt for this validation sample.                                                         |
+| `conditions` | List of validation conditions (see types below). Empty list = text-only generation.             |
+| `video_dims` | Optional per-sample override for `(width, height, frames)`. Inherits from `ValidationConfig` if not set. |
+| `seed`       | Optional per-sample override for random seed. Inherits from `ValidationConfig` if not set.      |
+
+#### Validation Condition Types
+
+| Type             | Parameters                                                 | Description                                                             |
+|------------------|------------------------------------------------------------|-------------------------------------------------------------------------|
+| `first_frame`    | `image_or_video` (path)                                    | Use the first frame of the image/video as conditioning.                 |
+| `prefix`         | `video` or `audio` (path), optional `num_frames`/`duration`| Use a video/audio clip as temporal prefix (for extension forward).      |
+| `suffix`         | `video` or `audio` (path), optional `num_frames`/`duration`| Use a video/audio clip as temporal suffix (for extension backward).     |
+| `spatial_crop`   | `video` (path), `spatial_region` (y1, x1, y2, x2)         | Provide spatial context for outpainting. Video only.                    |
+| `mask`           | `video` or `audio` (path), `mask` (path)                   | Mask-based inpainting with a binary mask file.                          |
+| `reference`      | `video` or `audio` (path), optional video-reference `downscale_factor`, `temporal_scale_factor`, `include_in_output` | IC-LoRA style reference conditioning. |
+| `video_to_audio` | `video` (path)                                             | Freeze video, generate audio. For Foley/V2A tasks.                      |
+| `audio_to_video` | `audio` (path)                                             | Freeze audio, generate video. For audio-driven generation.              |
+
+For video `reference` validation conditions, `downscale_factor` is the spatial reference scale and
+`temporal_scale_factor` is the temporal reference scale. Set both to match the factors used when
+preprocessing video reference latents for training; validation media is encoded on the fly and cannot infer
+those factors from the training dataset.
+
+> [!NOTE]
+> The legacy fields `prompts`, `images`, and `reference_videos` are deprecated but auto-converted to `samples`
+> internally. New configs should use the `samples` format.
+
+### CheckpointsConfig
+
+Model checkpointing configuration.
+
+```yaml
+checkpoints:
+  interval: 250       # Steps between checkpoint saves (null = disabled)
+  keep_last_n: 3      # Number of recent checkpoints to retain
+  precision: bfloat16 # Precision for saved weights (bfloat16 or float32)
+  no_resume: false            # Ignore saved state, start from step 0
+  save_training_state: "minimal"  # "full", "minimal", or "off"
+```
+
+**Key parameters:**
+
+| Parameter     | Description                                                                   |
+|---------------|-------------------------------------------------------------------------------|
+| `interval`    | Steps between intermediate checkpoint saves (set to `null` to disable)        |
+| `keep_last_n` | Number of most recent checkpoints to keep (-1 = keep all)                     |
+| `precision`   | Precision for saved checkpoint weights: `"bfloat16"` (default) or `"float32"` |
+| `no_resume` | When `true`, ignore saved training state and start from step 0. Model weights from `load_checkpoint` are still loaded. |
+| `save_training_state` | Save training state for resume: `"full"` (optimizer + scheduler + RNG), `"minimal"` (scheduler + RNG only, sufficient for LoRA), `"off"` (no resume). |
+
+### HubConfig
+
+Hugging Face Hub integration for automatic model uploads.
+
+```yaml
+hub:
+  push_to_hub: false                   # Enable Hub uploading
+  hub_model_id: "username/model-name"  # Hub repository ID
+```
+
+**Key parameters:**
+
+| Parameter      | Description                                                      |
+|----------------|------------------------------------------------------------------|
+| `push_to_hub`  | Whether to automatically push trained models to Hugging Face Hub |
+| `hub_model_id` | Repository ID in format `"username/repository-name"`             |
+
+### WandbConfig
+
+Weights & Biases logging configuration.
+
+```yaml
+wandb:
+  enabled: false               # Enable W&B logging
+  project: "ltx-2-trainer"     # W&B project name
+  entity: null                 # W&B username or team
+  tags: [ ]                    # Tags for the run
+  log_validation_videos: true  # Log validation videos to W&B
+```
+
+**Key parameters:**
+
+| Parameter               | Description                                      |
+|-------------------------|--------------------------------------------------|
+| `enabled`               | Whether to enable W&B logging                    |
+| `project`               | W&B project name                                 |
+| `entity`                | W&B username or team (null uses default account) |
+| `log_validation_videos` | Whether to log validation videos to W&B          |
+
+### FlowMatchingConfig
+
+Flow matching training configuration for timestep sampling.
+
+```yaml
+flow_matching:
+  timestep_sampling_mode: "shifted_logit_normal"  # Timestep sampling strategy
+  timestep_sampling_params: { }                   # Additional sampling parameters
+```
+
+**Key parameters:**
+
+| Parameter                  | Description                                                |
+|----------------------------|------------------------------------------------------------|
+| `timestep_sampling_mode`   | Sampling strategy: `"uniform"` or `"shifted_logit_normal"` |
+| `timestep_sampling_params` | Additional parameters for the sampling strategy            |
+
+### General Configuration
+
+Top-level settings for the training run.
+
+```yaml
+seed: 42                                  # Random seed for reproducibility
+output_dir: "outputs/my_training_run"     # Directory for outputs (checkpoints, validation videos, logs)
+```
+
+| Parameter    | Description                                              |
+|--------------|----------------------------------------------------------|
+| `seed`       | Random seed for reproducibility (default: `42`)          |
+| `output_dir` | Directory to save outputs (default: `"outputs"`)         |
+
+## 🚀 Next Steps
+
+Once you've configured your training parameters:
+
+- Set up your dataset using [Dataset Preparation](dataset-preparation.md)
+- Choose your training approach in [Training Modes](training-modes.md)
+- Start training with the [Training Guide](training-guide.md)

packages/ltx-trainer/docs/custom-training-strategies.md

@@ -0,0 +1,515 @@
+# Implementing Custom Training Strategies
+
+This guide explains how to implement your own training strategy for specialized recipes that cannot be expressed with
+the built-in `flexible` strategy.
+
+## 📋 Overview
+
+The trainer uses the **Strategy Pattern** to separate training logic from the core training loop. Each strategy defines:
+
+1. **What data is needed** - Which preprocessed data directories to load
+2. **How to prepare inputs** - Transform batch data into model inputs
+3. **How to compute loss** - Calculate the training objective
+
+This architecture lets you implement new training modes without modifying the core trainer code.
+
+### When You Need a Custom Strategy
+
+> [!NOTE]
+> The built-in `flexible` strategy already supports most conditioning scenarios out of the box:
+> first-frame conditioning, video extension (prefix/suffix), spatial crop (outpainting),
+> mask-based inpainting, IC-LoRA reference conditioning, and frozen modality cross-conditioning
+> (audio-to-video, video-to-audio). Only implement a custom strategy if your use case requires
+> fundamentally different training logic that cannot be expressed through the flexible strategy's
+> configuration.
+
+Consider implementing a custom strategy when you need:
+
+- **Custom loss computation** (e.g., weighted losses, auxiliary losses, perceptual losses)
+- **Non-standard noise application** (e.g., noise schedules different from flow matching)
+- **Novel conditioning mechanisms** not covered by the flexible strategy's condition types
+- **Additional model outputs** beyond the standard video/audio predictions
+
+## 🏗️ Architecture Overview
+
+### How Strategies Fit Into the Trainer
+
+The trainer delegates all training-mode-specific logic to the strategy:
+
+1. **Initialization** — The trainer calls `config.get_data_sources()` to determine which preprocessed data directories to load
+2. **Each training step:**
+    - Calls `prepare_training_inputs()` to transform the raw batch into model-ready inputs
+    - Runs the transformer forward pass
+    - Calls `compute_loss()` to compute the training objective
+
+The trainer handles everything else: optimization, checkpointing, validation, and distributed training.
+
+### Key Components
+
+| Component                                                                               | Purpose                                                      |
+|-----------------------------------------------------------------------------------------|--------------------------------------------------------------|
+| [`TrainingStrategyConfigBase`](../src/ltx_trainer/training_strategies/base_strategy.py) | Base class for strategy configuration (Pydantic model)       |
+| [`TrainingStrategy`](../src/ltx_trainer/training_strategies/base_strategy.py)           | Abstract base class defining the strategy interface          |
+| [`ModelInputs`](../src/ltx_trainer/training_strategies/base_strategy.py)                | Dataclass containing prepared inputs for the transformer     |
+| [`Modality`](../../ltx-core/src/ltx_core/model/transformer/modality.py)                 | ltx-core dataclass representing video or audio modality data |
+
+## 📝 Step-by-Step Implementation
+
+### Step 1: Plan Your Strategy
+
+Before writing code, answer these questions:
+
+1. **What additional data does your strategy need?**
+    - Example: A perceptual-loss strategy may need auxiliary feature targets
+    - Example: A novel conditioning mechanism may need an additional precomputed directory
+
+2. **What does conditioning look like?**
+    - Which tokens should be noised vs. kept clean?
+    - How should conditioning tokens be structured (e.g., first frame, reference video, mask)?
+
+3. **How should loss be computed?**
+    - Which tokens contribute to the loss?
+    - Are there multiple loss terms to combine?
+
+### Step 2: Extend Data Preprocessing (If Needed)
+
+If your strategy requires additional preprocessed data beyond video latents, audio latents, and text embeddings, you'll
+need to extend the preprocessing pipeline.
+
+#### Option A: Modify `process_dataset.py`
+
+For integrated preprocessing, add new arguments and processing steps to the main script. For example, to add mask
+preprocessing:
+
+```python
+# In process_dataset.py, add a new argument
+@app.command()
+def main(
+        # ... existing arguments ...
+        mask_column: str | None = typer.Option(
+            default=None,
+            help="Column name containing mask video paths (for inpainting)",
+        ),
+) -> None:
+    # ... existing processing ...
+
+    # Process masks if provided
+    if mask_column:
+        logger.info("Processing mask videos for inpainting training...")
+        mask_latents_dir = output_base / "mask_latents"
+
+        compute_latents(
+            dataset_file=dataset_path,
+            video_column=mask_column,
+            resolution_buckets=parsed_resolution_buckets,
+            output_dir=str(mask_latents_dir),
+            model_path=model_path,
+            # ... other args ...
+        )
+```
+
+#### Option B: Create a Standalone Script
+
+For complex preprocessing that doesn't fit naturally into the existing pipeline, create a dedicated script
+(e.g., `scripts/process_masks.py`). Use [`scripts/compute_reference.py`](../scripts/compute_reference.py) as a
+template - it shows how to process paired data and update the dataset JSON.
+
+#### Expected Output Structure
+
+Your preprocessing should create a directory structure that the strategy can reference:
+
+```
+preprocessed_data_root/
+├── latents/           # Video latents (standard)
+├── conditions/        # Text embeddings (standard)
+├── audio_latents/     # Audio latents (if with_audio)
+├── mask_latents/      # Your custom data directory
+└── reference_latents/ # Reference videos (for IC-LoRA)
+```
+
+### Step 3: Create the Strategy Configuration
+
+Create a new file for your strategy (e.g., `src/ltx_trainer/training_strategies/inpainting.py`):
+
+```python
+"""Inpainting training strategy.
+
+This strategy implements video inpainting training where:
+- Mask latents indicate which regions to inpaint
+- Loss is computed only on masked (inpainted) regions
+"""
+
+from typing import Any, Literal
+
+import torch
+from pydantic import Field
+from torch import Tensor
+
+from ltx_core.model.transformer.modality import Modality
+from ltx_trainer.timestep_samplers import TimestepSampler
+from ltx_trainer.training_strategies.base_strategy import (
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+
+
+class InpaintingConfig(TrainingStrategyConfigBase):
+    """Configuration for inpainting training strategy."""
+
+    # The 'name' field acts as a discriminator for the config union
+    name: Literal["inpainting"] = "inpainting"
+
+    mask_latents_dir: str = Field(
+        default="mask_latents",
+        description="Directory name for mask latents",
+    )
+
+    # Add any strategy-specific parameters
+    mask_threshold: float = Field(
+        default=0.5,
+        description="Threshold for binary mask conversion",
+        ge=0.0,
+        le=1.0,
+    )
+
+    def get_data_sources(self) -> dict[str, str]:
+        """Define which data directories to load.
+
+        Returns a mapping of directory names (under preprocessed_data_root) to
+        batch keys. The trainer loads .pt files from each directory and exposes
+        them in the batch under the specified key. The trainer also uses this
+        mapping to validate that all required directories exist.
+        """
+        return {
+            "latents": "latents",  # -> batch["latents"]
+            "conditions": "conditions",  # -> batch["conditions"]
+            self.mask_latents_dir: "masks",  # -> batch["masks"]
+        }
+```
+
+**Key points:**
+
+- Inherit from `TrainingStrategyConfigBase`
+- Use `Literal["your_strategy_name"]` for the `name` field - this enables automatic strategy selection
+- Use Pydantic `Field` for validation and documentation
+- Implement `get_data_sources()` on the config — it's the single source of truth for data directories (used for both dataset wiring and existence validation)
+
+### Step 4: Implement the Strategy Class
+
+```python
+class InpaintingStrategy(TrainingStrategy):
+    """Inpainting training strategy.
+
+    Trains the model to fill in masked regions of videos while
+    keeping unmasked regions as conditioning.
+    """
+
+    config: InpaintingConfig
+
+    def __init__(self, config: InpaintingConfig):
+        super().__init__(config)
+
+    def prepare_training_inputs(
+            self,
+            batch: dict[str, Any],
+            timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Transform batch data into model inputs.
+
+        This is where the core training logic lives:
+        1. Extract and patchify latents
+        2. Sample noise and apply it appropriately
+        3. Create conditioning masks
+        4. Build Modality objects for the transformer
+        """
+        # Get video latents [B, C, F, H, W]
+        latents_data = batch["latents"]
+        video_latents = latents_data["latents"]
+
+        # Get dimensions
+        num_frames = latents_data["num_frames"][0].item()
+        height = latents_data["height"][0].item()
+        width = latents_data["width"][0].item()
+
+        # Patchify: [B, C, F, H, W] -> [B, seq_len, C]
+        video_latents = self._video_patchifier.patchify(video_latents)
+
+        batch_size, seq_len, _ = video_latents.shape
+        device = video_latents.device
+        dtype = video_latents.dtype
+
+        # Get mask latents and process them
+        mask_data = batch["masks"]
+        mask_latents = mask_data["latents"]
+        mask_latents = self._video_patchifier.patchify(mask_latents)
+
+        # Create binary mask: True = inpaint this region, False = keep original
+        inpaint_mask = mask_latents.mean(dim=-1) > self.config.mask_threshold
+
+        # Sample noise and sigmas
+        sigmas = timestep_sampler.sample_for(video_latents)
+        noise = torch.randn_like(video_latents)
+
+        # Apply noise only to inpaint regions
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+        noisy_latents = (1 - sigmas_expanded) * video_latents + sigmas_expanded * noise
+
+        # Keep original latents for non-inpaint regions (conditioning)
+        inpaint_mask_expanded = inpaint_mask.unsqueeze(-1)
+        noisy_latents = torch.where(inpaint_mask_expanded, noisy_latents, video_latents)
+
+        # Create per-token timesteps
+        # Conditioning tokens (non-inpaint) get timestep=0
+        # Inpaint tokens get the sampled sigma
+        timesteps = self._create_per_token_timesteps(~inpaint_mask, sigmas.squeeze())
+
+        # Compute targets (velocity prediction: noise - clean)
+        targets = noise - video_latents
+
+        # Get text embeddings
+        conditions = batch["conditions"]
+        video_prompt_embeds = conditions["video_prompt_embeds"]
+        prompt_attention_mask = conditions["prompt_attention_mask"]
+
+        # Generate position embeddings
+        positions = self._get_video_positions(
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            batch_size=batch_size,
+            fps=24.0,  # Or get from latents_data
+            device=device,
+        )
+
+        # Create video Modality
+        video_modality = Modality(
+            enabled=True,
+            latent=noisy_latents,
+            sigma=sigmas,
+            timesteps=timesteps,
+            positions=positions,
+            context=video_prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        # Loss mask: only compute loss on inpaint regions
+        loss_mask = inpaint_mask
+
+        return ModelInputs(
+            video=video_modality,
+            audio=None,
+            video_targets=targets,
+            audio_targets=None,
+            video_loss_mask=loss_mask,
+            audio_loss_mask=None,
+        )
+
+    def compute_loss(
+            self,
+            video_pred: Tensor,
+            audio_pred: Tensor | None,
+            inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute training loss on inpaint regions only. Returns [B,]."""
+        # MSE loss
+        loss = (video_pred - inputs.video_targets).pow(2)
+
+        # Apply loss mask and reduce to per-element [B,]
+        loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
+        masked = loss.mul(loss_mask)
+        return masked.mean(dim=[-2, -1]) / loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
+```
+
+### Step 5: Register the Strategy
+
+You need to register your strategy in two places:
+
+**1. Update [`src/ltx_trainer/training_strategies/__init__.py`](../src/ltx_trainer/training_strategies/__init__.py):**
+
+```python
+# Add import for your strategy
+from ltx_trainer.training_strategies.inpainting import InpaintingConfig, InpaintingStrategy
+
+# Add to the TrainingStrategyConfig type alias
+TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig | FlexibleStrategyConfig | InpaintingConfig
+
+# Add to __all__
+__all__ = [
+    # ... existing exports ...
+    "InpaintingConfig",
+    "InpaintingStrategy",
+]
+
+
+# Add case in get_training_strategy()
+def get_training_strategy(config: TrainingStrategyConfig) -> TrainingStrategy:
+    match config:
+        # ... existing cases ...
+        case InpaintingConfig():
+            strategy = InpaintingStrategy(config)
+```
+
+**2. Update [`src/ltx_trainer/config.py`](../src/ltx_trainer/config.py):**
+
+```python
+# Add import
+from ltx_trainer.training_strategies.inpainting import InpaintingConfig
+
+# Add to the TrainingStrategyConfig union with a Tag matching your strategy name
+TrainingStrategyConfig = Annotated[
+    Annotated[TextToVideoConfig, Tag("text_to_video")]
+    | Annotated[VideoToVideoConfig, Tag("video_to_video")]
+    | Annotated[FlexibleStrategyConfig, Tag("flexible")]
+    | Annotated[InpaintingConfig, Tag("inpainting")],
+    Discriminator(_get_strategy_discriminator),
+]
+```
+
+### Step 6: Create a Configuration File
+
+Create an example config in `configs/`:
+
+```yaml
+# configs/custom_inpainting_lora.yaml
+
+model:
+  model_path: "/path/to/ltx2.safetensors"
+  text_encoder_path: "/path/to/gemma"
+  training_mode: "lora"
+
+training_strategy:
+  name: "inpainting"  # Must match your Literal type
+  mask_latents_dir: "mask_latents"
+  mask_threshold: 0.5
+
+lora:
+  rank: 32
+  alpha: 32
+  target_modules:
+    - "to_k"
+    - "to_q"
+    - "to_v"
+    - "to_out.0"
+
+data:
+  preprocessed_data_root: "/path/to/preprocessed/dataset"
+
+optimization:
+  learning_rate: 1e-4
+  steps: 2000
+  batch_size: 1
+
+# ... other config sections ...
+```
+
+## 🔧 Helper Methods Reference
+
+The base `TrainingStrategy` class provides these helper methods:
+
+| Method                                       | Purpose                                         |
+|----------------------------------------------|-------------------------------------------------|
+| `_video_patchifier.patchify(latents)`        | Convert `[B, C, F, H, W]` → `[B, seq_len, C]`   |
+| `_audio_patchifier.patchify(latents)`        | Convert `[B, C, T, F]` → `[B, T, C*F]`          |
+| `_get_video_positions(...)`                  | Generate position embeddings for video          |
+| `_get_audio_positions(...)`                  | Generate position embeddings for audio          |
+| `_create_per_token_timesteps(conditioning_mask, sampled_sigma)` | Create timesteps with 0 for conditioning tokens |
+| `_create_first_frame_conditioning_mask(...)`                    | Create mask for first-frame conditioning        |
+
+## 📊 Understanding ModelInputs
+
+The `ModelInputs` dataclass contains everything needed for the forward pass and loss computation:
+
+```python
+@dataclass
+class ModelInputs:
+    video: Modality | None  # Video modality data
+    audio: Modality | None  # Audio modality data
+
+    video_targets: Tensor | None  # Target values for video loss (velocity)
+    audio_targets: Tensor | None  # Target values for audio loss (velocity)
+
+    video_loss_mask: Tensor | None  # Boolean loss mask for video tokens
+    audio_loss_mask: Tensor | None  # Boolean loss mask for audio tokens
+```
+
+## 📊 Understanding Modality
+
+The `Modality` dataclass (from ltx-core) represents a single modality's data:
+
+```python
+@dataclass(frozen=True)
+class Modality:
+    latent: Tensor     # [B, T, D] — patchified latent tokens
+    sigma: Tensor      # [B,] — per-batch noise level (for cross-attn conditioning)
+    timesteps: Tensor  # [B, T] — per-token timestep embeddings
+    positions: Tensor  # [B, 3, T, 2] for video, [B, 1, T, 2] for audio — positional bounds
+    context: Tensor    # text conditioning embeddings
+    enabled: bool = True
+    context_mask: Tensor | None = None     # attention mask for text context
+    attention_mask: Tensor | None = None   # optional 2D self-attention mask [B, T, T]
+```
+
+> [!NOTE]
+> **Per-token timesteps:** Each token in the sequence has its own timestep. Conditioning tokens—those that should remain
+> un-noised—must have `timestep=0`. This is how the model distinguishes clean reference tokens from tokens to denoise. Use
+> `_create_per_token_timesteps(conditioning_mask, sampled_sigma)` to set this up correctly.
+
+> [!NOTE]
+> `Modality` is immutable (frozen dataclass). Use `dataclasses.replace()` to create modified copies.
+
+## ✅ Testing Your Strategy
+
+1. **Verify your training configuration is valid:**
+   ```bash
+   uv run python -c "
+   from ltx_trainer.config import LtxTrainerConfig
+   import yaml
+
+   with open('configs/custom_inpainting_lora.yaml') as f:
+       config = LtxTrainerConfig(**yaml.safe_load(f))
+   print(f'Strategy: {config.training_strategy.name}')
+   "
+   ```
+
+2. **Test strategy instantiation:**
+   ```bash
+   uv run python -c "
+   from ltx_trainer.training_strategies import get_training_strategy
+   from ltx_trainer.training_strategies.inpainting import InpaintingConfig
+
+   config = InpaintingConfig()
+   strategy = get_training_strategy(config)
+   print(f'Data sources: {config.get_data_sources()}')
+   "
+   ```
+
+3. **Run a short training test:**
+   ```bash
+   uv run python scripts/train.py configs/custom_inpainting_lora.yaml
+   ```
+
+## 💡 Tips and Best Practices
+
+### Debugging
+
+- Set `data.num_dataloader_workers: 0` to get clearer error messages
+- Use a small dataset and few steps for initial testing
+- Check tensor shapes at each step with print statements
+
+## 🔗 Related Documentation
+
+- [Training Modes](training-modes.md) - Overview of built-in training modes
+- [Configuration Reference](configuration-reference.md) - All configuration options
+- [Dataset Preparation](dataset-preparation.md) - Preprocessing workflow
+- [ltx-core Documentation](../../ltx-core/README.md) - Core model components
+
+## 📚 Reference: Existing Strategies
+
+Study these implementations for guidance:
+
+| Strategy | Complexity | Key Features |
+|----------|------------|--------------|
+| [`FlexibleStrategy`](../src/ltx_trainer/training_strategies/flexible.py) | Medium | Unified conditioning framework — supports all built-in modes |
+| [`TextToVideoStrategy`](../src/ltx_trainer/training_strategies/text_to_video.py) | Simple | First-frame conditioning, optional audio (deprecated) |
+| [`VideoToVideoStrategy`](../src/ltx_trainer/training_strategies/video_to_video.py) | Medium | Reference video concatenation, split loss mask (deprecated) |

packages/ltx-trainer/docs/dataset-preparation.md

@@ -0,0 +1,489 @@
+# Dataset Preparation Guide
+
+This guide covers the complete workflow for preparing and preprocessing your dataset for training.
+
+## 📋 Overview
+
+The general dataset preparation workflow is:
+
+1. **(Optional)** Split long videos into scenes using `split_scenes.py`
+2. **(Optional)** Generate captions for your videos using `caption_videos.py`
+3. **Preprocess your dataset** using `process_dataset.py` to compute and cache video/audio latents and text embeddings
+4. **Run the trainer** with your preprocessed dataset
+
+## 🎬 Step 1: Split Scenes
+
+If you're starting with raw, long-form videos (e.g., downloaded from YouTube), you should first split them into shorter, coherent scenes.
+
+```bash
+uv run python scripts/split_scenes.py input.mp4 scenes_output_dir/ \
+    --filter-shorter-than 5s
+```
+
+This will create multiple video clips in `scenes_output_dir`.
+These clips will be the input for the captioning step, if you choose to use it.
+
+The script supports many configuration options for scene detection (detector algorithms, thresholds, minimum scene lengths, etc.):
+
+```bash
+uv run python scripts/split_scenes.py --help
+```
+
+## 📝 Step 2: Caption Videos
+
+If your dataset doesn't include captions, you can automatically generate them using multimodal models that understand both video and audio.
+
+The default `qwen_omni` backend talks to a local vLLM server, which you launch once in a separate terminal:
+
+```bash
+# Terminal 1: start the captioner server (stays running)
+uv run python scripts/serve_captioner.py
+```
+
+```bash
+# Terminal 2: caption your videos
+uv run python scripts/caption_videos.py scenes_output_dir/ \
+    --output scenes_output_dir/dataset.json
+```
+
+This will create a `dataset.json` file containing video paths and their captions.
+
+**Captioning options:**
+
+| Option             | Description                                                      |
+| ------------------ | --------------------------------------------------------------- |
+| `--captioner-type` | `qwen_omni` (default, local vLLM server) or `gemini_flash` (API) |
+| `--vllm-url`       | Base URL of the vLLM server (default `http://127.0.0.1:8001/v1`) |
+| `--override`       | Re-caption files that already have captions                     |
+| `--api-key`        | Gemini API key (else `GEMINI_API_KEY`/`GOOGLE_API_KEY`; with no key, uses gcloud/Vertex AI auth) |
+
+**Caption format:**
+
+Each caption is a single, detailed paragraph describing both the visual content and the audio (speech, music, ambient sounds) of the clip. See the [Utility Scripts Reference](utility-scripts.md#automatic-video-captioning) for backend setup and the full list of options.
+
+> [!NOTE]
+> The automatically generated captions may contain inaccuracies or hallucinated content.
+> We recommend reviewing and correcting the generated captions in your `dataset.json` file before proceeding to preprocessing.
+
+## ⚡ Step 3: Dataset Preprocessing
+
+This step preprocesses your video dataset by:
+
+1. Resizing and cropping videos to fit specified resolution buckets
+2. Computing and caching video latent representations
+3. Computing and caching text embeddings for captions
+4. Extracting and caching audio latents from videos (automatic, use `--skip-audio` to disable)
+
+> [!WARNING]
+> Very large videos (especially high spatial resolution and/or many frames) can cause GPU out-of-memory (OOM)
+> during preprocessing/encoding.
+> The simplest fix is to reduce the target resolution (spatially: width/height) and/or the number of frames
+> (temporally) by using `--resolution-buckets` with smaller dimensions (lower width/height and/or fewer frames).
+
+### Basic Usage
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+Audio latents are automatically extracted from video files — no extra flag is needed. Use `--skip-audio`
+to disable this. For standalone audio files (`.wav`), use the `audio` column in your dataset instead
+(see [Convention-Based Column Detection](#convention-based-column-detection) below).
+
+### 🚀 Multi-GPU Preprocessing
+
+Preprocessing large datasets can take a while. To run it across multiple GPUs in parallel, wrap the command with
+`accelerate launch` (for example `--num_processes 4`). Each process handles an interleaved slice of the dataset.
+The same approach applies to `process_videos.py` and `process_captions.py` when you run them standalone.
+
+```bash
+uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+Outputs are written atomically (via a per-process temporary file, then renamed), so an interrupted run leaves no
+corrupt files. By default a rerun **resumes** — items whose output `.pt` already exists are skipped.
+
+> [!IMPORTANT]
+> Pass **`--overwrite`** when rerunning with changed parameters (different model checkpoint, resolution buckets,
+> text encoder, `--lora-trigger`, etc.). Without it the script keeps the stale outputs from the previous run.
+>
+> ```bash
+> uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+>     --resolution-buckets "960x544x49" \
+>     --model-path /path/to/ltx-2.3-model.safetensors \
+>     --text-encoder-path /path/to/gemma-model \
+>     --overwrite
+> ```
+
+### 📊 Dataset Format
+
+The trainer supports videos, single images, or a mix of both in the same dataset.
+
+> [!TIP]
+> **Image Datasets:** When using images, follow the same preprocessing steps and format requirements as with videos,
+> but use `1` for the frame count in the resolution bucket (e.g., `960x544x1`).
+
+> [!NOTE]
+> **Mixed image + video datasets:** Mixing stills and videos in a single dataset is supported, but requires some care:
+>
+> - Preprocess with **multiple resolution buckets** covering both frame counts — e.g.
+> `--resolution-buckets "960x544x1;960x544x49"`. Images are automatically assigned to the `F=1` bucket and
+> videos to an `F>1` bucket.
+> - You **must** set `optimization.batch_size: 1` in your training config (see the warning under
+> [Resolution Buckets](#resolution-buckets)), since samples with different shapes cannot be collated into a
+> single batch. Use `gradient_accumulation_steps` if you need a larger effective batch.
+> - Per-step cost differs substantially between a single-frame sample and a many-frame sample, which can lead to
+> uneven gradient magnitudes across steps. Consider weighting the two subsets or tuning the learning rate if
+> you observe instability.
+> - If you prefer a fully officially-supported path, train two separate LoRAs (one on stills, one on video) and
+> stack them at inference.
+
+The dataset must be a CSV, JSON, or JSONL metadata file with columns for captions and media paths.
+
+#### Convention-Based Column Detection
+
+The preprocessing script automatically detects and processes columns based on their names. The following columns are recognized:
+
+| Column | Output Dir | Description |
+|--------|-----------|-------------|
+| `video` (or legacy `media_path`) | `latents/` | Target video to encode |
+| `audio` | `audio_latents/` | Explicit audio file (overrides auto-extraction from video) |
+| `caption` | `conditions/` | Text caption for the sample |
+| `reference_video` (or legacy `ref_media_path`) | `reference_latents/` | IC-LoRA reference video |
+| `reference_audio` | `reference_audio_latents/` | IC-LoRA reference audio |
+| `video_mask` | `video_masks/` | Binary mask for video inpainting |
+| `audio_mask` | `audio_masks/` | Binary mask for audio inpainting |
+
+> [!NOTE]
+> **Legacy column names:** `media_path` and `ref_media_path` are accepted as aliases for `video` and `reference_video` respectively. Existing datasets using these names will continue to work without modification.
+
+**JSON format example:**
+
+```json
+[
+  {
+    "caption": "A cat playing with a ball of yarn",
+    "video": "videos/cat_playing.mp4"
+  },
+  {
+    "caption": "A dog running in the park",
+    "video": "videos/dog_running.mp4"
+  }
+]
+```
+
+**JSONL format example:**
+
+```jsonl
+{"caption": "A cat playing with a ball of yarn", "video": "videos/cat_playing.mp4"}
+{"caption": "A dog running in the park", "video": "videos/dog_running.mp4"}
+```
+
+**CSV format example:**
+
+```csv
+caption,video
+"A cat playing with a ball of yarn","videos/cat_playing.mp4"
+"A dog running in the park","videos/dog_running.mp4"
+```
+
+**Additional dataset format examples:**
+
+Audio-only dataset:
+```json
+{"audio": "song.wav", "caption": "piano melody"}
+```
+
+V2V IC-LoRA with reference video:
+```json
+{"video": "clip.mp4", "reference_video": "depth.mp4", "caption": "depth to video"}
+```
+
+A2A IC-LoRA with reference audio:
+```json
+{"video": "clip.mp4", "reference_audio": "ref.wav", "caption": "match this style"}
+```
+This form auto-extracts the target audio from `clip.mp4`. For pure audio datasets, use `audio` plus
+`reference_audio` columns and preprocess with `--audio-durations`.
+
+Video inpainting with mask:
+```json
+{"video": "clip.mp4", "video_mask": "mask.mp4", "caption": "fill the sky"}
+```
+
+### 📐 Resolution Buckets
+
+Videos are organized into "buckets" of specific dimensions (width × height × frames).
+Each video is assigned to the nearest matching bucket.
+You can preprocess with one or multiple resolution buckets.
+When training with multiple resolution buckets, you must use a batch size of 1.
+
+The dimensions of each bucket must follow these constraints due to LTX-2's VAE architecture:
+
+- **Spatial dimensions** (width and height) must be multiples of 32
+- **Number of frames** must satisfy `frames % 8 == 1` (e.g., 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, 97, 121, etc.)
+
+**Guidelines for choosing training resolution:**
+
+- For high-quality, detailed videos: use larger spatial dimensions (e.g. 768x448) with fewer frames (e.g. 89)
+- For longer, motion-focused videos: use smaller spatial dimensions (512×512) with more frames (121)
+- Memory usage increases with both spatial and temporal dimensions
+
+**Example usage:**
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+Multiple buckets are supported by separating entries with `;`:
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49;512x512x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+**Video processing workflow:**
+
+1. Videos are **resized** maintaining aspect ratio until either width or height matches the target
+2. The larger dimension is **center cropped** to match the bucket's dimensions
+3. Only the **first X frames are taken** to match the bucket's frame count, remaining frames are ignored
+
+> [!NOTE]
+> The sequence length processed by the transformer model can be calculated as:
+>
+> ```
+> sequence_length = (H/32) * (W/32) * ((F-1)/8 + 1)
+> ```
+>
+> Where:
+>
+> - H = Height of video
+> - W = Width of video
+> - F = Number of frames
+> - 32 = VAE's spatial downsampling factor
+> - 8 = VAE's temporal downsampling factor
+>
+> For example, a 768×448×89 video would have sequence length:
+>
+> ```
+> (768/32) * (448/32) * ((89-1)/8 + 1) = 24 * 14 * 12 = 4,032
+> ```
+>
+> Keep this in mind when choosing video dimensions, as longer sequences require more GPU memory.
+
+> [!WARNING]
+> When training with multiple resolution buckets, you must use a batch size of 1
+> (i.e., set `optimization.batch_size: 1` in your training config).
+
+### 📁 Output Structure
+
+The preprocessed data is saved in a `.precomputed` directory:
+
+```
+dataset/
+└── .precomputed/
+    ├── latents/                  # Video latents
+    ├── conditions/               # Text embeddings
+    ├── audio_latents/            # Audio latents (auto-extracted or explicit)
+    ├── reference_latents/        # Reference video latents (IC-LoRA)
+    ├── reference_audio_latents/  # Reference audio latents (audio IC-LoRA)
+    ├── video_masks/              # Video masks (inpainting)
+    └── audio_masks/              # Audio masks (audio inpainting)
+```
+
+Set `data.preprocessed_data_root` in your training config to this `.precomputed` directory — the parent directory that
+contains `latents/`, `conditions/`, and any mode-specific audio/reference/mask directories.
+
+## 🔊 Audio-Only Dataset Preprocessing
+
+For datasets containing only audio files (no `video` column), use `--audio-durations` to specify duration buckets:
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --audio-durations "2.0;4.0;8.0" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+The `--audio-durations` flag provides duration buckets (in seconds) for audio-only datasets. Since there is no video column to derive timing from, explicit duration buckets are required.
+
+## 🪄 IC-LoRA Reference Video Preprocessing
+
+For IC-LoRA training, you need to preprocess datasets that include reference videos.
+Reference videos provide the conditioning input while target videos represent the desired transformed output.
+
+### Dataset Format with Reference Videos
+
+The `reference_video` column is automatically detected by convention — no extra CLI flags are needed.
+
+**JSON format:**
+
+```json
+[
+  {
+    "caption": "A cat playing with a ball of yarn",
+    "video": "videos/cat_playing.mp4",
+    "reference_video": "references/cat_playing_depth.mp4"
+  }
+]
+```
+
+**JSONL format:**
+
+```jsonl
+{"caption": "A cat playing with a ball of yarn", "video": "videos/cat_playing.mp4", "reference_video": "references/cat_playing_depth.mp4"}
+{"caption": "A dog running in the park", "video": "videos/dog_running.mp4", "reference_video": "references/dog_running_depth.mp4"}
+```
+
+### Preprocessing with Reference Videos
+
+Convention-based detection means you just need the `reference_video` column in your dataset, and `process_dataset.py` will automatically detect and process it. No `--reference-column` flag is needed:
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model \
+    --reference-downscale-factor 2 \
+    --reference-temporal-scale-factor 1
+```
+
+This will create an additional `reference_latents/` directory containing the preprocessed reference video latents.
+Use `--reference-downscale-factor` for spatial subsampling and `--reference-temporal-scale-factor` for temporal
+subsampling. Validation reference conditions should use matching `downscale_factor` and `temporal_scale_factor` values.
+
+> [!NOTE]
+> **Legacy column names:** If your dataset uses `ref_media_path`, it is accepted as an alias for `reference_video`.
+
+### Generating Reference Videos
+
+**Dataset Requirements for IC-LoRA:**
+
+- Your dataset must contain paired videos where each target video has a corresponding reference video
+- Reference and target videos should cover the same content. Reference videos can optionally be lower spatial
+  resolution or temporally subsampled (see Scaled Reference Conditioning in [Training Modes](training-modes.md)).
+- Both reference and target videos should be preprocessed together using the same target resolution buckets, plus any
+  reference scale factors you choose.
+
+We provide an example script, `[scripts/compute_reference.py](../scripts/compute_reference.py)`, to generate reference
+videos for a given dataset. The default implementation generates Canny edge reference videos.
+
+```bash
+uv run python scripts/compute_reference.py scenes_output_dir/ \
+    --output scenes_output_dir/dataset.json
+```
+
+The script accepts a JSON file as the dataset configuration and updates it in-place by adding the filenames of the generated reference videos.
+
+> [!NOTE]
+> `compute_reference.py` writes generated references to the `reference_video` column, which `process_dataset.py`
+> detects automatically. The legacy `ref_media_path` column is also accepted.
+
+If you want to generate a different type of condition (depth maps, pose skeletons, etc.), modify or replace the `compute_reference()` function within this script.
+
+### Example Dataset
+
+For reference, see our **[Canny Control Dataset](https://huggingface.co/datasets/Lightricks/Canny-Control-Dataset)** which demonstrates proper IC-LoRA dataset structure with paired videos and Canny edge maps.
+
+## 🎭 Mask Preprocessing for Inpainting
+
+For inpainting training with the `mask` condition type, provide `video_mask` or `audio_mask` columns in your dataset
+metadata. These columns point to mask media files (for example a mask image/video for video inpainting, or a waveform or
+`.pt` tensor for audio inpainting). `process_dataset.py` downsamples and thresholds them into per-sample `.pt` tensors
+under `video_masks/` or `audio_masks/`.
+
+### Processed Video Mask Format
+
+If you create masks manually instead of using `process_dataset.py`, save them as `.pt` files with the key `"mask"`
+containing a tensor of shape `[F, H, W]` where:
+
+- `F` = number of latent frames (temporal dimension)
+- `H` = latent height (pixel height / 32)
+- `W` = latent width (pixel width / 32)
+- Values are thresholded at `0.5`: values `> 0.5` are conditioning tokens (clean, excluded from loss),
+  and values `<= 0.5` are generated tokens (noised, contributes to loss).
+
+### Audio Mask Format
+
+Audio masks follow the same thresholding pattern as video masks but with shape `[T]` (temporal dimension only), where `T` is the number of audio latent frames. They are stored in `audio_masks/`.
+
+### Directory Structure
+
+Place masks in a directory within your preprocessed data root:
+
+```
+preprocessed_data_root/
+├── latents/          # Video latents
+├── conditions/       # Text embeddings
+├── video_masks/      # Video masks (one .pt per sample, matching latent filenames)
+└── audio_masks/      # Audio masks (one .pt per sample, matching latent filenames)
+```
+
+Then reference the mask directory in your training config:
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: mask
+        mask_dir: "video_masks"
+```
+
+## 🎯 LoRA Trigger Words
+
+When training a LoRA, you can specify a trigger token that will be prepended to all captions:
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model \
+    --lora-trigger "MYTRIGGER"
+```
+
+This acts as a trigger word that activates the LoRA during inference when you include the same token in your prompts.
+
+> [!NOTE]
+> There is no need to manually insert the trigger word into your dataset JSON/JSONL/CSV file.
+> The trigger word specified with `--lora-trigger` is automatically prepended to each caption during preprocessing.
+
+## 🔍 Decoding Videos for Verification
+
+If you add the `--decode` flag, the script will VAE-decode the precomputed video latents and save the resulting videos
+in `.precomputed/decoded_videos`. Reference video latents are decoded to `.precomputed/decoded_reference_videos` when
+present. To inspect audio latents, run `scripts/decode_latents.py` with `--with-audio`.
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model \
+    --decode
+```
+
+For single-frame images, the decoded latents will be saved as PNG files rather than MP4 videos.
+
+## 🚀 Next Steps
+
+Once your dataset is preprocessed, you can proceed to:
+
+- Configure your training parameters in [Configuration Reference](configuration-reference.md)
+- Choose your training approach in [Training Modes](training-modes.md)
+- Start training with the [Training Guide](training-guide.md)
+
+> [!TIP]
+> The `flexible` strategy supports masks for inpainting (`mask` condition type) and spatial crop regions for outpainting (`spatial_crop` condition type) out of the box. For other custom preprocessing needs, see [Custom Training Strategies](custom-training-strategies.md).

packages/ltx-trainer/docs/quick-start.md

@@ -0,0 +1,160 @@
+# Quick Start Guide
+
+Get up and running with LTX-2 training in just a few steps!
+
+## 📋 Prerequisites
+
+Before you begin, ensure you have:
+
+1. **LTX-2 Model Checkpoint** - A local `.safetensors` file containing the LTX-2 model weights.
+   Download `ltx-2.3-22b-dev.safetensors` from: [HuggingFace Hub](https://huggingface.co/Lightricks/LTX-2.3)
+   The trainer supports LTX-2 and LTX-2.3 checkpoints through the same configuration API; version-specific components
+   are detected from the checkpoint.
+2. **Gemma Text Encoder** - A local directory containing the Gemma model (required for LTX-2).
+   Download from: [HuggingFace Hub](https://huggingface.co/google/gemma-3-12b-it-qat-q4_0-unquantized/)
+3. **Linux with CUDA** - The trainer requires `triton` which is Linux-only; CUDA 13+ is recommended
+4. **GPU with sufficient VRAM** - 80GB recommended for the standard config. For GPUs with 32GB VRAM (e.g., RTX 5090),
+   use the [low VRAM config](../configs/t2v_lora_low_vram.yaml) which enables INT8 quantization and other
+   memory optimizations
+
+## ⚡ Installation
+
+First, install [uv](https://docs.astral.sh/uv/getting-started/installation/) if you haven't already.
+Then clone the repository and install the dependencies:
+
+```bash
+git clone https://github.com/Lightricks/LTX-2
+```
+
+The `ltx-trainer` package is part of the `LTX-2` monorepo. Install the dependencies from the repository root,
+then navigate to the trainer package:
+
+```bash
+# From the repository root
+uv sync
+cd packages/ltx-trainer
+```
+
+> [!NOTE]
+> The trainer depends on [`ltx-core`](../../ltx-core/) and [`ltx-pipelines`](../../ltx-pipelines/)
+> packages which are automatically installed from the monorepo.
+
+## 🏋 Training Workflow
+
+If you are using an agent-enabled environment with repository skills, you can ask for the
+[`train-model`](../../../.claude/skills/train-model/SKILL.md) skill to run this workflow with you.
+It creates a run workspace, confirms the training mode, prepares data, preprocesses latents,
+launches training, and monitors the run while stopping for approval before expensive steps.
+
+### 1. Choose a Training Mode
+
+Start with [`t2v_lora.yaml`](../configs/t2v_lora.yaml) for a first run with videos and captions. For modes such as
+IC-LoRA, inpainting, or outpainting, check [Training Modes](training-modes.md) first because your metadata needs extra
+columns such as `reference_video`, `video_mask`, or `audio_mask` before preprocessing.
+
+### 2. Prepare Your Dataset
+
+Organize your videos and captions, then preprocess them:
+
+```bash
+# Split long videos into scenes (optional)
+uv run python scripts/split_scenes.py input.mp4 scenes_output_dir/ --filter-shorter-than 5s
+
+# Generate captions for videos (optional)
+uv run python scripts/caption_videos.py scenes_output_dir/ --output dataset.json
+
+# Preprocess the dataset (compute latents and embeddings)
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+By default, preprocessing writes to `.precomputed/`. Use that directory as `data.preprocessed_data_root`
+in your training config.
+
+See [Dataset Preparation](dataset-preparation.md) for detailed instructions.
+
+### 3. Configure Training
+
+Create or modify a configuration YAML file. Start with one of the example configs:
+
+- [`configs/t2v_lora.yaml`](../configs/t2v_lora.yaml) - Text-to-video LoRA
+- [`configs/t2v_lora_low_vram.yaml`](../configs/t2v_lora_low_vram.yaml) - Same as above, tuned for ~32GB VRAM (INT8 quantization and memory optimizations)
+- [`configs/v2v_ic_lora.yaml`](../configs/v2v_ic_lora.yaml) - IC-LoRA video-to-video
+
+Key settings to update:
+
+```yaml
+model:
+  model_path: "/path/to/ltx-2-model.safetensors"
+  text_encoder_path: "/path/to/gemma-model"
+
+data:
+  preprocessed_data_root: "/path/to/preprocessed/data"
+
+output_dir: "outputs/my_training_run"
+```
+
+See [Configuration Reference](configuration-reference.md) for all available options.
+
+### 4. Start Training
+
+```bash
+uv run python scripts/train.py configs/t2v_lora.yaml
+```
+
+For multi-GPU training:
+
+```bash
+uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
+```
+
+See [Training Guide](training-guide.md) for distributed training and advanced options.
+
+## 🎯 Training Modes
+
+> [!TIP]
+> **First time?** Start with [`t2v_lora.yaml`](../configs/t2v_lora.yaml) — it's the simplest mode
+> and only requires videos with captions. You can explore other modes once you've confirmed your
+> setup works.
+
+The trainer supports several training modes:
+
+| Mode                  | Description                                | Example Config                                                    |
+|-----------------------|--------------------------------------------|-------------------------------------------------------------------|
+| **Text-to-Video**     | Generate video+audio from text prompts     | [`t2v_lora.yaml`](../configs/t2v_lora.yaml)                      |
+| **Image-to-Video**    | Animate from a starting image              | [`i2v_lora.yaml`](../configs/i2v_lora.yaml)                      |
+| **Video Extension**   | Extend videos temporally (forward/backward)| [`video_extend_lora.yaml`](../configs/video_extend_lora.yaml), [`video_suffix_lora.yaml`](../configs/video_suffix_lora.yaml) |
+| **IC-LoRA (V2V)**     | Video-to-video transformations             | [`v2v_ic_lora.yaml`](../configs/v2v_ic_lora.yaml)                |
+| **Audio-to-Video**    | Generate video conditioned on audio        | [`a2v_lora.yaml`](../configs/a2v_lora.yaml)                      |
+| **Video-to-Audio**    | Generate audio/foley from video            | [`v2a_lora.yaml`](../configs/v2a_lora.yaml)                      |
+| **Video Inpainting**  | Fill in masked regions of video            | [`video_inpainting_lora.yaml`](../configs/video_inpainting_lora.yaml) |
+| **Video Outpainting** | Extend video spatially                     | [`video_outpainting_lora.yaml`](../configs/video_outpainting_lora.yaml) |
+| **Text-to-Audio**     | Generate audio from text prompts           | [`t2a_lora.yaml`](../configs/t2a_lora.yaml)                      |
+| **Audio Extension**   | Extend audio temporally                    | [`audio_extend_lora.yaml`](../configs/audio_extend_lora.yaml), [`audio_suffix_lora.yaml`](../configs/audio_suffix_lora.yaml) |
+| **Audio Inpainting**  | Fill in masked regions of audio            | [`audio_inpainting_lora.yaml`](../configs/audio_inpainting_lora.yaml) |
+| **IC-LoRA (A2A)**     | Audio-to-audio transformations             | [`a2a_ic_lora.yaml`](../configs/a2a_ic_lora.yaml)                |
+| **AV2AV IC-LoRA**     | Audio+video IC-LoRA transformations        | [`av2av_ic_lora.yaml`](../configs/av2av_ic_lora.yaml)            |
+| **Full Fine-tuning**  | Full model training (any mode above)       | Set `model.training_mode: "full"`                                 |
+
+See [Training Modes](training-modes.md) for detailed explanations of each mode.
+
+## Next Steps
+
+Once you've completed your first training run, you can:
+
+- **Use your trained LoRA for inference** - The [`ltx-pipelines`](../../ltx-pipelines/) package provides
+  production-ready inference
+  pipelines for various use cases (T2V, I2V, IC-LoRA, etc.). See the package documentation for details.
+- Learn more about [Dataset Preparation](dataset-preparation.md) for advanced preprocessing
+- Explore different [Training Modes](training-modes.md)
+- Dive deeper into [Training Configuration](configuration-reference.md)
+- Understand the model architecture in [LTX-Core Documentation](../../ltx-core/README.md)
+
+## Need Help?
+
+If you run into issues at any step, see the [Troubleshooting Guide](troubleshooting.md) for solutions to common
+problems.
+
+Join our [Discord community](https://discord.gg/ltxplatform) for real-time help and discussion!

packages/ltx-trainer/docs/training-guide.md

@@ -0,0 +1,215 @@
+# Training Guide
+
+This guide covers how to run training jobs, from basic single-GPU training to advanced distributed setups and automatic
+model uploads.
+
+## ⚡ Basic Training (Single GPU)
+
+After preprocessing your dataset and preparing a configuration file, you can start training using the trainer script:
+
+```bash
+uv run python scripts/train.py configs/t2v_lora.yaml
+```
+
+The trainer will:
+
+1. **Load your configuration** and validate all parameters
+2. **Initialize models** and apply optimizations
+3. **Run the training loop** with progress tracking
+4. **Generate validation videos** (if configured)
+5. **Save the trained weights** in your output directory
+
+### Agent-Assisted Training
+
+If your environment supports repository skills, the
+[`train-model`](../../../.claude/skills/train-model/SKILL.md) skill provides an end-to-end
+orchestrator for this package. It asks what you want the model to learn, maps that intent to
+one of the documented [training modes](training-modes.md), probes your filesystem and GPU,
+prepares/preprocesses the dataset, writes a run-specific config, launches training, and
+monitors the job. It uses the trainer docs as its source of truth and stops for approval before
+captioning, preprocessing, or starting expensive training work.
+
+### Output Files
+
+**For LoRA training:**
+
+- `checkpoints/lora_weights_step_00000.safetensors` - LoRA checkpoint weights, with the current step in the filename
+- `training_config.yaml` - Copy of training configuration
+- `samples/` - Generated validation samples (if enabled)
+- `checkpoints/training_state_step_00000.pt` - Optional resume state, depending on `checkpoints.save_training_state`
+
+**For full model fine-tuning:**
+
+- `checkpoints/model_weights_step_00000.safetensors` - Full model checkpoint weights, with the current step in the filename
+- `training_config.yaml` - Copy of training configuration
+- `samples/` - Generated validation samples (if enabled)
+- `checkpoints/training_state_step_00000.pt` - Optional resume state, depending on `checkpoints.save_training_state`
+
+## 🖥️ Distributed / Multi-GPU Training
+
+We use Hugging Face 🤗 [Accelerate](https://huggingface.co/docs/accelerate/index) for multi-GPU DDP and FSDP.
+
+### Configure Accelerate
+
+Run the interactive wizard once to set up your environment (DDP / FSDP, GPU count, etc.):
+
+```bash
+uv run accelerate config
+```
+
+This stores your preferences in `~/.cache/huggingface/accelerate/default_config.yaml`.
+
+### Use the Provided Accelerate Configs (Recommended)
+
+We include ready-to-use Accelerate config files in `configs/accelerate/`:
+
+- [ddp.yaml](../configs/accelerate/ddp.yaml) — Standard DDP
+- [ddp_compile.yaml](../configs/accelerate/ddp_compile.yaml) — DDP with `torch.compile` (Inductor)
+- [fsdp.yaml](../configs/accelerate/fsdp.yaml) — Standard FSDP (auto-wraps `BasicAVTransformerBlock`)
+- [fsdp_compile.yaml](../configs/accelerate/fsdp_compile.yaml) — FSDP with `torch.compile` (Inductor)
+
+Launch with a specific config using `--config_file`:
+
+```bash
+# DDP (2 GPUs shown as example)
+CUDA_VISIBLE_DEVICES=0,1 \
+uv run accelerate launch --config_file configs/accelerate/ddp.yaml \
+  scripts/train.py configs/t2v_lora.yaml
+
+# DDP + torch.compile
+CUDA_VISIBLE_DEVICES=0,1 \
+uv run accelerate launch --config_file configs/accelerate/ddp_compile.yaml \
+  scripts/train.py configs/t2v_lora.yaml
+
+# FSDP (4 GPUs shown as example)
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+uv run accelerate launch --config_file configs/accelerate/fsdp.yaml \
+  scripts/train.py configs/t2v_lora.yaml
+
+# FSDP + torch.compile
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+uv run accelerate launch --config_file configs/accelerate/fsdp_compile.yaml \
+  scripts/train.py configs/t2v_lora.yaml
+```
+
+**Notes:**
+
+- The number of processes is taken from the Accelerate config (`num_processes`). Override with `--num_processes X` or
+  restrict GPUs with `CUDA_VISIBLE_DEVICES`.
+- The compile variants enable `torch.compile` with the Inductor backend via Accelerate's `dynamo_config`.
+- FSDP configs auto-wrap the transformer blocks (`fsdp_transformer_layer_cls_to_wrap: BasicAVTransformerBlock`).
+
+### Launch with Your Default Accelerate Config
+
+If you prefer to use your default Accelerate profile:
+
+```bash
+# Use settings from your default accelerate config
+uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
+
+# Override number of processes on the fly (e.g., 2 GPUs)
+uv run accelerate launch --num_processes 2 scripts/train.py configs/t2v_lora.yaml
+
+# Select specific GPUs
+CUDA_VISIBLE_DEVICES=0,1 uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
+```
+
+> [!TIP]
+> You can disable the in-terminal progress bars with `--disable-progress-bars` flag in the trainer CLI if desired.
+
+### Benefits of Distributed Training
+
+- **Faster training**: Distribute workload across multiple GPUs
+- **Larger effective batch sizes**: Combine gradients from multiple GPUs
+- **Memory efficiency**: Each GPU handles a portion of the batch
+
+> [!NOTE]
+> Distributed training requires that all GPUs have sufficient memory for the model and batch size. The effective batch
+> size becomes `batch_size × num_processes`.
+
+## 🤗 Pushing Models to Hugging Face Hub
+
+You can automatically push your trained models to the Hugging Face Hub by adding the following to your configuration:
+
+```yaml
+hub:
+  push_to_hub: true
+  hub_model_id: "your-username/your-model-name"
+```
+
+### Prerequisites
+
+Before pushing, make sure you:
+
+1. **Have a Hugging Face account** - Sign up at [huggingface.co](https://huggingface.co)
+2. **Are logged in** via `huggingface-cli login` or have set the `HUGGING_FACE_HUB_TOKEN` environment variable
+3. **Have write access** to the specified repository (it will be created if it doesn't exist)
+
+### Login Options
+
+**Option 1: Interactive login**
+
+```bash
+uv run huggingface-cli login
+```
+
+**Option 2: Environment variable**
+
+```bash
+export HUGGING_FACE_HUB_TOKEN="your_token_here"
+```
+
+### What Gets Uploaded
+
+The trainer will automatically:
+
+- **Create a model card** with training details and sample outputs
+- **Upload model weights**
+- **Push sample videos as GIFs** in the model card
+- **Include training configuration and prompts**
+
+## 📊 Weights & Biases Logging
+
+Enable experiment tracking with W&B by adding to your configuration:
+
+```yaml
+wandb:
+  enabled: true
+  project: "ltx-2-trainer"
+  entity: null  # Your W&B username or team
+  tags: [ "ltx2", "lora" ]
+  log_validation_videos: true
+```
+
+This will log:
+
+- Training loss and learning rate
+- Validation videos
+- Model configuration
+- Training progress
+
+## 🚀 Next Steps
+
+After training completes:
+
+- **Run inference with your trained LoRA** - The [`ltx-pipelines`](../../ltx-pipelines/) package provides
+  production-ready inference
+  pipelines that support loading custom LoRAs. Available pipelines include text-to-video, image-to-video,
+  IC-LoRA video-to-video, and more. See the [`ltx-pipelines`](../../ltx-pipelines/) package for usage details.
+- **Test your model** with validation prompts
+- **Iterate and improve** based on validation results
+- **Share your results** by pushing to Hugging Face Hub
+
+## 💡 Tips for Successful Training
+
+- **Start small**: Begin with a small dataset and a few hundred steps to verify everything works
+- **Monitor validation**: Keep an eye on validation samples to catch overfitting
+- **Adjust learning rate**: Lower learning rates often produce better results
+- **Use gradient checkpointing**: Essential for training with limited GPU memory
+- **Save checkpoints**: Regular checkpoints help recover from interruptions
+
+## Need Help?
+
+If you encounter issues during training, see the [Troubleshooting Guide](troubleshooting.md).
+
+Join our [Discord community](https://discord.gg/ltxplatform) for real-time help!

packages/ltx-trainer/docs/training-modes.md

@@ -0,0 +1,600 @@
+# Training Modes Guide
+
+The trainer uses the **flexible** training strategy (`name: "flexible"`) — a unified conditioning framework that
+supports all training modes through configuration. Every scenario is expressed by setting `is_generated` on each
+modality and adding optional conditions, rather than choosing a separate strategy class.
+
+## Key Concepts
+
+Before diving into individual modes, here are the core ideas behind the flexible strategy:
+
+- **`is_generated: true`** — the modality is denoised during training and contributes to the loss. This is the
+  modality the model learns to generate.
+- **`is_generated: false`** — the modality is frozen (sigma=0, no noise, no loss). It passes through the transformer
+  clean and acts as cross-modal conditioning for the generated modality.
+- **At least one modality must have `is_generated: true`.**
+- **Conditions** are per-modality and can be composed (e.g., `reference` + `first_frame` together on the video
+  modality).
+- Audio does **not** support `first_frame` or `spatial_crop` conditions — only `prefix`, `suffix`, `mask`,
+  and `reference`.
+
+> [!TIP]
+> If you are using an agent-enabled environment with repository skills and are unsure which mode to choose,
+> ask for the [`train-model`](../../../.claude/skills/train-model/SKILL.md) skill. It maps your intent to one of
+> these configs and walks through dataset preparation, preprocessing, launch, and monitoring.
+
+## 📊 Quick Reference
+
+| Mode                  | Video     | Audio     | Conditions          | Config |
+|-----------------------|-----------|-----------|---------------------|--------|
+| **T2V**               | Generated | Generated | —                   | [`t2v_lora`](../configs/t2v_lora.yaml) |
+| **I2V**               | Generated | Generated | `first_frame`       | [`i2v_lora`](../configs/i2v_lora.yaml) |
+| **Video Extension**   | Generated | Generated | `prefix`/`suffix`   | [`video_extend_lora`](../configs/video_extend_lora.yaml) |
+| **V2V IC-LoRA**       | Generated | —         | `reference`         | [`v2v_ic_lora`](../configs/v2v_ic_lora.yaml) |
+| **A2V**               | Generated | Frozen    | —                   | [`a2v_lora`](../configs/a2v_lora.yaml) |
+| **V2A (Foley)**       | Frozen    | Generated | —                   | [`v2a_lora`](../configs/v2a_lora.yaml) |
+| **Video Inpainting**  | Generated | —         | `mask`              | [`video_inpainting_lora`](../configs/video_inpainting_lora.yaml) |
+| **Video Outpainting** | Generated | —         | `spatial_crop`      | [`video_outpainting_lora`](../configs/video_outpainting_lora.yaml) |
+| **T2A**               | —         | Generated | —                   | [`t2a_lora`](../configs/t2a_lora.yaml) |
+| **Audio Extension**   | —         | Generated | `prefix`/`suffix`   | [`audio_extend_lora`](../configs/audio_extend_lora.yaml) |
+| **Audio Inpainting**  | —         | Generated | `mask`              | [`audio_inpainting_lora`](../configs/audio_inpainting_lora.yaml) |
+| **A2A IC-LoRA**       | —         | Generated | `reference`         | [`a2a_ic_lora`](../configs/a2a_ic_lora.yaml) |
+| **AV2AV IC-LoRA**     | Generated | Generated | `reference` (both)  | [`av2av_ic_lora`](../configs/av2av_ic_lora.yaml) |
+
+---
+
+## 🎯 Text-to-Video (T2V)
+
+Generate video and audio from text prompts. Both modalities are denoised with no additional conditions.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+**Example config:** 📄 [t2v_lora.yaml](../configs/t2v_lora.yaml)
+
+---
+
+## 🖼️ Image-to-Video (I2V)
+
+Generate video conditioned on a starting image. The first frame is provided as a clean conditioning signal — no noise,
+timestep=0, excluded from loss. The `probability` parameter controls how often first-frame conditioning is applied;
+remaining samples train in pure T2V mode.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: first_frame
+        probability: 0.5
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+**Example config:** 📄 [i2v_lora.yaml](../configs/i2v_lora.yaml)
+
+---
+
+## ⏩ Video Extension
+
+Extend a video forward (or backward) in time. Prefix or suffix conditioning provides a span of existing latent frames
+as clean conditioning. The `temporal_boundary` sets the number of **latent frames** used as context (each latent frame
+= 8 pixel frames due to temporal compression).
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: prefix            # or "suffix" for backward extension
+        temporal_boundary: 8    # 8 latent frames = 64 pixel frames
+        probability: 1.0
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+> [!NOTE]
+> The `prefix` and `suffix` conditions also work on the audio modality for audio extension.
+> Set `temporal_boundary` on the audio modality's conditions list to condition on a prefix or suffix
+> of the audio latents.
+
+**Example configs:** 📄 [video_extend_lora.yaml](../configs/video_extend_lora.yaml) (forward), 📄 [video_suffix_lora.yaml](../configs/video_suffix_lora.yaml) (backward)
+
+---
+
+## 🔄 IC-LoRA / Video-to-Video (V2V)
+
+In-Context LoRA learns transformations from paired videos. Pre-encoded reference latents are concatenated to the target
+sequence — reference tokens participate in bidirectional self-attention but receive no noise and are excluded from loss.
+This enables control adapters (depth, pose), style transfer, deblurring, colorization, and more.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_latents"
+        probability: 1.0
+      - type: first_frame       # optional — composable with reference
+        probability: 0.2
+```
+
+> [!NOTE]
+> IC-LoRA is video-only by default (no audio modality block). Conditions can be composed — the example above also
+> applies first-frame conditioning with 20% probability alongside the reference.
+> Use [AV2AV IC-LoRA](#av2av-ic-lora) when both video and audio references should be trained jointly.
+
+**Example config:** 📄 [v2v_ic_lora.yaml](../configs/v2v_ic_lora.yaml)
+
+### Dataset Requirements
+
+- **Paired videos** — each target video has a corresponding reference video
+- **Same frame count** between reference and target
+- Reference videos can optionally be at **lower spatial resolution** (see [Scaled Reference](#scaled-reference-conditioning) below)
+- Both must be **preprocessed** before training
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── latents/            # Target video latents
+├── conditions/         # Text embeddings
+└── reference_latents/  # Reference video latents (conditioning input)
+```
+
+### Generating Reference Videos
+
+Use the `compute_reference.py` script to generate reference videos (e.g., Canny edge maps) for a dataset:
+
+```bash
+uv run python scripts/compute_reference.py scenes_output_dir/ \
+    --output scenes_output_dir/dataset.json
+```
+
+To compute a different condition (depth maps, pose skeletons, etc.), modify the `compute_reference()` function in the
+script.
+
+> [!NOTE]
+> `compute_reference.py` writes generated references to the `reference_video` column, which
+> `process_dataset.py` detects automatically. The legacy `ref_media_path` column is also accepted.
+
+### Scaled Reference Conditioning
+
+For more efficient training and inference, use **downscaled reference videos** while keeping targets at full
+resolution. During training, the strategy infers the spatial and temporal scale factors from the preprocessed
+reference and target latents and adjusts positional encodings accordingly. This reduces conditioning tokens, leading to:
+
+- **Faster training** — shorter sequence lengths
+- **Faster inference** — reduced memory usage
+- **Same aspect ratio** maintained between reference and target
+
+Preprocess with the `--reference-downscale-factor` option:
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets 768x768x25 \
+    --model-path /path/to/ltx2.safetensors \
+    --text-encoder-path /path/to/gemma \
+    --reference-downscale-factor 2
+```
+
+> [!NOTE]
+> The `reference_video` column is auto-detected by convention — no `--reference-column` flag needed.
+
+Validation encodes reference media on the fly, so set `downscale_factor` and `temporal_scale_factor`
+on each `reference` validation condition to match the preprocessing factors:
+
+```yaml
+validation:
+  samples:
+    - prompt: "..."
+      conditions:
+        - type: reference
+          video: "/path/to/reference.mp4"
+          downscale_factor: 2
+          temporal_scale_factor: 1
+          include_in_output: true
+```
+
+> [!NOTE]
+> The scale factor must be a positive integer, and all dimensions must be divisible by 32.
+> Common values are 1 (no scaling), 2 (half resolution), or 4 (quarter resolution).
+
+---
+
+## 🔊 Audio-to-Video (A2V)
+
+Generate video conditioned on frozen audio. Audio passes through the transformer clean (sigma=0) and influences video
+via the built-in cross-modal attention. Only video is denoised.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+  audio:
+    is_generated: false
+    latents_dir: "audio_latents"
+```
+
+**Example config:** 📄 [a2v_lora.yaml](../configs/a2v_lora.yaml)
+
+---
+
+## 🎵 Video-to-Audio / Foley (V2A)
+
+Generate audio (Foley) conditioned on frozen video. Video passes through the transformer clean (sigma=0) and
+conditions audio via cross-modal attention. Only audio is denoised.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: false
+    latents_dir: "latents"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+**Example config:** 📄 [v2a_lora.yaml](../configs/v2a_lora.yaml)
+
+---
+
+## 🎭 Video Inpainting
+
+Fill in masked regions of a video. Per-sample masks loaded from disk define which tokens are conditioning and which
+must be generated. Masks are thresholded at `0.5` to match validation/inference: tokens with `mask > 0.5` receive clean
+latents and timestep=0 and are excluded from loss; tokens with `mask <= 0.5` are denoised normally and contribute to
+loss.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: mask
+        mask_dir: "video_masks"
+        probability: 1.0
+```
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── latents/        # Video latents
+├── conditions/     # Text embeddings
+└── video_masks/    # Per-sample binary masks (1 → conditioning, 0 → generate)
+```
+
+In dataset metadata, provide mask media via the `video_mask` column; preprocessing converts it into `video_masks/`.
+
+**Example config:** 📄 [video_inpainting_lora.yaml](../configs/video_inpainting_lora.yaml)
+
+---
+
+## 🌅 Video Outpainting
+
+Extend a video spatially beyond its original boundaries. A rectangular pixel region is provided as clean conditioning
+(no noise, timestep=0, excluded from loss) — the model learns to generate the surrounding content. The `spatial_region`
+is specified in pixel coordinates `[y1, x1, y2, x2]` and automatically converted to latent space.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: spatial_crop
+        spatial_region: [0, 0, 288, 576]  # y1, x1, y2, x2 in pixels
+        probability: 1.0
+```
+
+> [!NOTE]
+> `spatial_crop` is a video-only condition — it is not supported on the audio modality.
+
+**Example config:** 📄 [video_outpainting_lora.yaml](../configs/video_outpainting_lora.yaml)
+
+---
+
+## 🔈 Text-to-Audio (T2A)
+
+Generate audio from text prompts with no video modality. Only the audio branch of the transformer is denoised. Since
+no video modality is configured, this mode uses **audio-only LoRA targets** — explicitly targeting `audio_attn1`,
+`audio_attn2`, and `audio_ff` modules.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+> [!NOTE]
+> With no `video` block in the strategy, the trainer only loads audio latents and text embeddings. LoRA adapters
+> should explicitly target audio modules (e.g., `audio_attn1.to_k`) rather than short patterns like `to_k` which
+> would also match video modules. See [LoRA Target Modules Guidance](#lora-target-modules-guidance) below.
+
+**Example config:** 📄 [t2a_lora.yaml](../configs/t2a_lora.yaml)
+
+---
+
+## 🔊 Audio Extension
+
+Extend audio forward (prefix) or backward (suffix) in time — the audio equivalent of Video Extension. A span of
+existing audio latent frames is provided as clean conditioning, and the model generates the continuation. The
+`temporal_boundary` sets the number of latent frames used as context. This mode uses **audio-only LoRA targets**.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: prefix            # or "suffix" for backward extension
+        temporal_boundary: 8
+        probability: 1.0
+```
+
+**Example configs:** 📄 [audio_extend_lora.yaml](../configs/audio_extend_lora.yaml), 📄 [audio_suffix_lora.yaml](../configs/audio_suffix_lora.yaml)
+
+---
+
+## 🎭 Audio Inpainting
+
+Fill in masked regions of audio. Per-sample masks loaded from disk define which audio tokens are conditioning and
+which must be generated — the audio equivalent of Video Inpainting. Masks are thresholded at `0.5` with the same
+binary semantics as video inpainting. This mode uses **audio-only LoRA targets**.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: mask
+        mask_dir: "audio_masks"
+        probability: 1.0
+```
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── conditions/      # Text embeddings
+├── audio_latents/   # Audio latents
+└── audio_masks/     # Per-sample binary masks (1 → conditioning, 0 → generate)
+```
+
+In dataset metadata, provide mask media via the `audio_mask` column; preprocessing converts it into `audio_masks/`.
+
+**Example config:** 📄 [audio_inpainting_lora.yaml](../configs/audio_inpainting_lora.yaml)
+
+---
+
+## 🔄 IC-LoRA / Audio-to-Audio (A2A)
+
+In-Context LoRA for audio-to-audio transformations. Pre-encoded reference audio latents are concatenated to the target
+sequence — reference tokens participate in bidirectional self-attention but receive no noise and are excluded from loss.
+This enables audio style transfer, voice conversion, sound effect transformation, and more. This mode uses
+**audio-only LoRA targets**.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_audio_latents"
+        probability: 1.0
+```
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── conditions/              # Text embeddings
+├── audio_latents/           # Target audio latents
+└── reference_audio_latents/ # Reference audio latents (conditioning input)
+```
+
+**Example config:** 📄 [a2a_ic_lora.yaml](../configs/a2a_ic_lora.yaml)
+
+---
+
+## 🔄 AV2AV IC-LoRA
+
+Joint audio-video In-Context LoRA — both modalities have reference conditioning. Pre-encoded reference latents are
+concatenated to each modality's target sequence independently. This enables joint audiovisual transformations such as
+synchronized style transfer across both video and audio.
+
+```yaml
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_latents"
+        probability: 1.0
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+    conditions:
+      - type: reference
+        latents_dir: "reference_audio_latents"
+        probability: 1.0
+```
+
+> [!NOTE]
+> Unlike audio-only IC-LoRA (A2A), AV2AV uses short LoRA target patterns like `"to_k"` to match all branches
+> (video, audio, and cross-modal attention), since both modalities are trained.
+
+**Dataset structure:**
+
+```
+preprocessed_data_root/
+├── latents/                 # Target video latents
+├── audio_latents/           # Target audio latents
+├── conditions/              # Text embeddings
+├── reference_latents/       # Reference video latents (conditioning input)
+└── reference_audio_latents/ # Reference audio latents (conditioning input)
+```
+
+**Example config:** 📄 [av2av_ic_lora.yaml](../configs/av2av_ic_lora.yaml)
+
+---
+
+## 🔥 Full Model Fine-tuning
+
+All modes above default to `training_mode: "lora"`. For full fine-tuning, set `training_mode: "full"` — this updates
+all model parameters rather than adding LoRA adapters.
+
+```yaml
+model:
+  training_mode: "full"
+
+training_strategy:
+  name: "flexible"
+  video:
+    is_generated: true
+    latents_dir: "latents"
+  audio:
+    is_generated: true
+    latents_dir: "audio_latents"
+```
+
+> [!IMPORTANT]
+> Full fine-tuning requires multiple high-end GPUs (e.g., 4-8× H100 80GB) and distributed training with FSDP.
+> See [Training Guide](training-guide.md) for multi-GPU setup instructions.
+
+---
+
+## 🎛️ LoRA Target Modules Guidance
+
+The `target_modules` configuration determines which transformer modules receive LoRA adapters. The right choice depends
+on whether your training involves cross-modal (audio ↔ video) interaction.
+
+**For T2V, I2V, A2V, V2A, or any mode involving both modalities** — use short patterns to match all branches
+(video, audio, and cross-modal attention):
+
+```yaml
+target_modules:
+  - "to_k"
+  - "to_q"
+  - "to_v"
+  - "to_out.0"
+```
+
+> [!IMPORTANT]
+> Short patterns like `"to_k"` match video modules (`attn1.to_k`, `attn2.to_k`), audio modules
+> (`audio_attn1.to_k`, `audio_attn2.to_k`), and cross-modal modules (`audio_to_video_attn.to_k`,
+> `video_to_audio_attn.to_k`). The cross-modal attention modules enable bidirectional information flow between
+> audio and video, which is critical for synchronized audiovisual generation.
+> See [Understanding Target Modules](configuration-reference.md#understanding-target-modules) for detailed guidance.
+
+**For video-only IC-LoRA** — explicitly target video modules (including FFN layers for better transformation quality):
+
+```yaml
+target_modules:
+  - "attn1.to_k"
+  - "attn1.to_q"
+  - "attn1.to_v"
+  - "attn1.to_out.0"
+  - "attn2.to_k"
+  - "attn2.to_q"
+  - "attn2.to_v"
+  - "attn2.to_out.0"
+  - "ff.net.0.proj"
+  - "ff.net.2"
+```
+
+**For audio-only modes (T2A, Audio Extension, Audio Inpainting, A2A IC-LoRA)** — explicitly target audio modules:
+
+```yaml
+target_modules:
+  - "audio_attn1.to_k"
+  - "audio_attn1.to_q"
+  - "audio_attn1.to_v"
+  - "audio_attn1.to_out.0"
+  - "audio_attn2.to_k"
+  - "audio_attn2.to_q"
+  - "audio_attn2.to_v"
+  - "audio_attn2.to_out.0"
+  - "audio_ff.net.0.proj"
+  - "audio_ff.net.2"
+```
+
+> [!NOTE]
+> Audio-only modes have no `video` block in the strategy, so there is no need to train video or cross-modal
+> attention modules. Targeting only `audio_*` modules keeps the LoRA small and focused.
+
+---
+
+## 🎬 Using Trained Models for Inference
+
+After training, use the [`ltx-pipelines`](../../ltx-pipelines/) package for production inference with your trained
+LoRAs:
+
+| Training Mode           | Recommended Pipeline                                  |
+|-------------------------|-------------------------------------------------------|
+| T2V / I2V / A2V / Extension / Inpainting / Outpainting | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` |
+| IC-LoRA (V2V / A2A / AV2AV) | `ICLoraPipeline`                                 |
+| V2A (Foley) / T2A / Audio Extension / Audio Inpainting | `TI2VidOneStagePipeline` or `TI2VidTwoStagesPipeline` |
+
+All pipelines support loading custom LoRAs via the `loras` parameter. See the [`ltx-pipelines`](../../ltx-pipelines/)
+package documentation for detailed usage instructions.
+
+> [!NOTE]
+> You can generate audio during validation even if you're not training the audio branch.
+> Set `validation.generate_audio: true` independently of whether audio has `is_generated: true`.
+
+---
+
+## 🔄 Migration from Legacy Strategies
+
+Legacy `text_to_video` and `video_to_video` strategy configs are forward-compatible and will continue to work (with a
+deprecation warning). We recommend migrating to `flexible` for access to all conditioning modes.
+
+---
+
+## 🚀 Next Steps
+
+Once you've chosen your training mode:
+
+- Set up your dataset using [Dataset Preparation](dataset-preparation.md)
+- Configure your training parameters in [Configuration Reference](configuration-reference.md)
+- Start training with the [Training Guide](training-guide.md)
+
+> [!TIP]
+> Need a training mode that's not covered here?
+> First check whether it can be expressed by composing existing `flexible` conditions. Use
+> [Implementing Custom Training Strategies](custom-training-strategies.md) only for custom losses,
+> noising rules, model outputs, or preprocessing that cannot be represented by configuration.

packages/ltx-trainer/docs/troubleshooting.md

@@ -0,0 +1,314 @@
+# Troubleshooting Guide
+
+This guide covers common issues and solutions when training with the LTX-2 trainer.
+
+## 🔧 VRAM and Memory Issues
+
+Memory management is crucial for successful training with LTX-2.
+
+> [!TIP]
+> For GPUs with 32GB VRAM, use the pre-configured low VRAM config:
+> [`configs/t2v_lora_low_vram.yaml`](../configs/t2v_lora_low_vram.yaml)
+> which combines 8-bit optimizer, INT8 quantization, and reduced LoRA rank.
+
+### Memory Optimization Techniques
+
+#### 1. Enable Gradient Checkpointing
+
+Gradient checkpointing trades training speed for memory savings. **Highly recommended** for most training runs:
+
+```yaml
+optimization:
+  enable_gradient_checkpointing: true
+```
+
+#### 2. Enable 8-bit Text Encoder
+
+Load the Gemma text encoder in 8-bit precision to save GPU memory:
+
+```yaml
+acceleration:
+  load_text_encoder_in_8bit: true
+```
+
+#### 3. Reduce Batch Size
+
+Lower the batch size if you encounter out-of-memory errors:
+
+```yaml
+optimization:
+  batch_size: 1  # Start with 1 and increase gradually
+```
+
+Use gradient accumulation to maintain a larger effective batch size:
+
+```yaml
+optimization:
+  batch_size: 1
+  gradient_accumulation_steps: 4  # Effective batch size = 4
+```
+
+#### 4. Use Lower Resolution
+
+Reduce spatial or temporal dimensions to save memory:
+
+```bash
+# Smaller spatial resolution
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "512x512x49" \
+    --model-path /path/to/model.safetensors \
+    --text-encoder-path /path/to/gemma
+
+# Fewer frames
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x25" \
+    --model-path /path/to/model.safetensors \
+    --text-encoder-path /path/to/gemma
+```
+
+#### 5. Enable Model Quantization
+
+Use quantization to reduce memory usage:
+
+```yaml
+acceleration:
+  quantization: "int8-quanto"  # Options: int8-quanto, int4-quanto, fp8-quanto
+```
+
+#### 6. Use 8-bit Optimizer
+
+The 8-bit AdamW optimizer uses less memory:
+
+```yaml
+optimization:
+  optimizer_type: "adamw8bit"
+```
+
+#### 7. Offload Optimizer State During Validation
+
+If you OOM specifically during validation video sampling — typically in
+full fine-tunes or high-rank LoRA runs where AdamW state and the VAE decoder
+can't coexist on the GPU — offload optimizer state to CPU during sampling:
+
+```yaml
+acceleration:
+  offload_optimizer_during_validation: true
+```
+
+The offload + reload happens once per validation interval, not per step.
+No effect for FSDP (sharded state).
+
+---
+
+## ⚠️ Common Usage Issues
+
+### Issue: "No module named 'ltx_trainer'" Error
+
+**Solution:**
+Ensure you've installed the dependencies and are using `uv run` to execute scripts:
+
+```bash
+# From the repository root
+uv sync
+cd packages/ltx-trainer
+uv run python scripts/train.py configs/t2v_lora.yaml
+```
+
+> [!TIP]
+> Always use `uv run` to execute Python scripts. This automatically uses the correct virtual environment
+> without requiring manual activation.
+
+### Issue: "Gemma model path is not a directory" Error
+
+**Solution:**
+The `text_encoder_path` must point to a directory containing the Gemma model, not a file:
+
+```yaml
+model:
+  model_path: "/path/to/ltx-2-model.safetensors"  # File path
+  text_encoder_path: "/path/to/gemma-model/"      # Directory path
+```
+
+### Issue: "Model path does not exist" Error
+
+**Solution:**
+LTX-2 requires local model paths. URLs are not supported:
+
+```yaml
+# ✅ Correct - local path
+model:
+  model_path: "/path/to/ltx-2-model.safetensors"
+
+# ❌ Wrong - URL not supported
+model:
+  model_path: "https://huggingface.co/..."
+```
+
+### Issue: "Frames must satisfy frames % 8 == 1" Error
+
+**Solution:**
+LTX-2 requires the number of frames to satisfy `frames % 8 == 1`:
+
+- ✅ Valid: 1, 9, 17, 25, 33, 41, 49, 57, 65, 73, 81, 89, 97, 121
+- ❌ Invalid: 24, 32, 48, 64, 100
+
+### Issue: Slow Training Speed
+
+**Optimizations:**
+
+1. **Disable gradient checkpointing** (if you have enough VRAM):
+
+   ```yaml
+   optimization:
+     enable_gradient_checkpointing: false
+   ```
+
+
+2. **Use torch.compile** via Accelerate:
+
+   ```bash
+   uv run accelerate launch --config_file configs/accelerate/ddp_compile.yaml \
+     scripts/train.py configs/t2v_lora.yaml
+   ```
+
+### Issue: Poor Quality Validation Outputs
+
+**Solutions:**
+
+1. **Use conditioned validation:** For more reliable validation, use image-to-video (first-frame conditioning) rather than pure text-to-video:
+
+   ```yaml
+   validation:
+     samples:
+       - prompt: "a professional portrait video of a person"
+         conditions:
+           - type: first_frame
+             image_or_video: "/path/to/first_frame.png"
+   ```
+
+2. **Increase inference steps:**
+
+   ```yaml
+   validation:
+     inference_steps: 30
+   ```
+
+3. **Adjust guidance settings:**
+
+   ```yaml
+   validation:
+     guidance_scale: 4.0  # CFG scale (recommended: 4.0)
+     stg_scale: 1.0       # STG scale for temporal coherence (recommended: 1.0)
+     stg_blocks: [29]     # Transformer block to perturb
+   ```
+
+4. **Check caption quality:**
+   Review and manually edit captions for accuracy if using auto-generated captions.
+   LTX-2 prefers long, detailed captions that describe both visual content and audio (e.g., ambient sounds, speech,
+   music).
+
+5. **Check target modules:**
+   Ensure your `target_modules` configuration matches your training goals. For audio-video training,
+   use patterns that match both branches (e.g., `"to_k"` instead of `"attn1.to_k"`).
+   See [Understanding Target Modules](configuration-reference.md#understanding-target-modules) for details.
+
+6. **Adjust LoRA rank:**
+   Try higher values for more capacity:
+
+   ```yaml
+   lora:
+     rank: 64  # Or 128 for more capacity
+   ```
+
+7. **Increase training steps:**
+
+   ```yaml
+   optimization:
+     steps: 3000
+   ```
+
+---
+
+## 🔍 Debugging Tools
+
+### Monitor GPU Memory Usage
+
+Track memory usage during training:
+
+```bash
+# Watch GPU memory in real-time
+watch -n 1 nvidia-smi
+
+# Log memory usage to file
+nvidia-smi --query-gpu=memory.used,memory.total --format=csv --loop=5 > memory_log.csv
+```
+
+### Verify Preprocessed Data
+
+Decode latents to visualize the preprocessed videos:
+
+```bash
+uv run python scripts/decode_latents.py dataset/.precomputed/latents debug_output \
+    --model-path /path/to/model.safetensors
+```
+
+To also decode audio latents, add the `--with-audio` flag:
+
+```bash
+uv run python scripts/decode_latents.py dataset/.precomputed/latents debug_output \
+    --model-path /path/to/model.safetensors \
+    --with-audio
+```
+
+Compare decoded videos and audio with originals to ensure quality.
+
+---
+
+## 💡 Best Practices
+
+### Before Training
+
+- [ ] Test preprocessing with a small subset first
+- [ ] Verify all video files are accessible
+- [ ] Check available GPU memory
+- [ ] Review configuration against hardware capabilities
+- [ ] Ensure model and text encoder paths are correct
+
+### During Training
+
+- [ ] Monitor GPU memory usage
+- [ ] Check loss convergence regularly
+- [ ] Review validation samples periodically
+- [ ] Save checkpoints frequently
+
+### After Training
+
+- [ ] Test trained model with diverse prompts
+- [ ] Document training parameters and results
+- [ ] Archive training data and configs
+
+## 🆘 Getting Help
+
+If you're still experiencing issues:
+
+1. **Check logs:** Review console output for error details
+2. **Search issues:** Look through GitHub issues for similar problems
+3. **Provide details:** When reporting issues, include:
+    - Hardware specifications (GPU model, VRAM)
+    - Configuration file used
+    - Complete error message
+    - Steps to reproduce the issue
+
+---
+
+## 🤝 Join the Community
+
+Have questions, want to share your results, or need real-time help?
+Join our [community Discord server](https://discord.gg/ltxplatform)
+to connect with other users and the development team!
+
+- Get troubleshooting help
+- Share your training results and workflows
+- Stay up to date with announcements and updates
+
+We look forward to seeing you there!

packages/ltx-trainer/docs/utility-scripts.md

@@ -0,0 +1,239 @@
+# Utility Scripts Reference
+
+This guide covers the various utility scripts available for preprocessing, conversion, and debugging tasks.
+
+## 🎬 Dataset Processing Scripts
+
+### Video Scene Splitting
+
+The `scripts/split_scenes.py` script automatically splits long videos into shorter, coherent scenes.
+
+```bash
+# Basic scene splitting
+uv run python scripts/split_scenes.py input.mp4 output_dir/ --filter-shorter-than 5s
+```
+
+**Key features:**
+
+- **Automatic scene detection**: Uses PySceneDetect for intelligent splitting
+- **Multiple algorithms**: Content-based, adaptive, threshold, and histogram detection
+- **Filtering options**: Remove scenes shorter than specified duration
+- **Customizable parameters**: Thresholds, window sizes, and detection modes
+
+**Common options:**
+
+```bash
+# See all available options
+uv run python scripts/split_scenes.py --help
+
+# Use adaptive detection with custom threshold
+uv run python scripts/split_scenes.py video.mp4 scenes/ --detector adaptive --threshold 30.0
+
+# Limit to maximum number of scenes
+uv run python scripts/split_scenes.py video.mp4 scenes/ --max-scenes 50
+```
+
+### Automatic Video Captioning
+
+The `scripts/caption_videos.py` script generates a single, detailed combined audio-visual
+caption per video as a continuous paragraph of prose. Two backends are available:
+
+- **`qwen_omni` (default)** — Qwen3-Omni-30B-A3B-Thinking served via a local
+  [vLLM](https://docs.vllm.ai/) HTTP server (~1-3 s/video on H100). Highest quality, runs
+  fully offline once the model is downloaded.
+- **`gemini_flash`** — Google Gemini (cloud, `gemini-3.5-flash`). No GPU required. Auth is
+  automatic: set `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) for the Developer API, or just have
+  Google Cloud credentials available (`gcloud auth` / an attached service account) and it
+  uses Vertex AI with no extra setup.
+
+**Step 1 — launch the captioner server** (`qwen_omni` only, one-time).
+
+`scripts/serve_captioner.py` runs vLLM in an isolated environment via `uvx`, so vLLM's heavy
+CUDA dependencies never touch the trainer's venv. It defaults to dynamic FP8 quantization
+(~31 GiB weights, fits on 40 GB GPUs, same speed as BF16 on H100):
+
+```bash
+# Terminal 1 - stays running
+uv run python packages/ltx-trainer/scripts/serve_captioner.py
+
+# Useful variants:
+#   --print-cmd           show the vLLM command without running it
+#   --quantization bf16   use BF16 instead (needs ~66 GiB free VRAM)
+#   --hf-home /mnt/disk   override where the ~65 GB model is downloaded
+```
+
+**Step 2 — caption your videos.**
+
+```bash
+# Terminal 2 - default backend talks to the server above
+uv run python packages/ltx-trainer/scripts/caption_videos.py videos_dir/ --output dataset.json
+
+# Remote server:           --vllm-url http://other-host:8001/v1
+# Gemini (gemini-3.5-flash): --captioner-type gemini_flash   (uses GEMINI_API_KEY, else gcloud/Vertex)
+# Gemini, parallel calls:  --captioner-type gemini_flash --num-workers 5
+# Re-caption everything:   --override
+```
+
+Captioning is incremental (already-captioned files are skipped, progress saves every 5 videos)
+and writes JSON, JSONL, CSV, or TXT based on the output extension.
+
+Qwen3-Omni-Thinking can optionally emit a `<think>...</think>` chain-of-thought before the
+caption (`--enable-thinking`). It is off by default, which is recommended for bulk captioning
+(thinking is slower as it generates the reasoning trace first).
+
+For Gemini, keep `--num-workers` at 3-5 (higher values may hit API rate limits).
+
+### Dataset Preprocessing
+
+The `scripts/process_dataset.py` script processes videos and caches latents for training.
+
+```bash
+# Basic preprocessing
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+
+# With video decoding for verification
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model \
+    --decode
+```
+
+Multiple resolution buckets can be specified, separated by `;`:
+
+```bash
+uv run python scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49;512x512x81" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+```
+
+> [!NOTE]
+> When training with multiple resolution buckets, set `optimization.batch_size: 1`.
+
+**Multi-GPU preprocessing.** Launch with `accelerate launch` to shard the dataset across processes. Reruns resume
+by default (existing `.pt` outputs are skipped); writes are atomic so interrupted runs are safe. Pass `--overwrite`
+when rerunning with changed parameters (different model, resolution buckets, text encoder, `--lora-trigger`, etc.)
+so stale outputs are replaced. Use the same `accelerate launch` pattern (and `--overwrite` when needed) with
+`process_videos.py` or `process_captions.py` when you run those scripts standalone.
+
+```bash
+# Multi-GPU preprocessing
+uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --text-encoder-path /path/to/gemma-model
+
+# Force re-encoding of all items (e.g. after switching model or resolution)
+uv run accelerate launch --num_processes 4 scripts/process_dataset.py dataset.json \
+    --resolution-buckets "960x544x49" \
+    --model-path /path/to/ltx-2.3-model.safetensors \
+    --text-encoder-path /path/to/gemma-model \
+    --overwrite
+```
+
+For detailed usage, see the [Dataset Preparation Guide](dataset-preparation.md).
+
+### Reference Video Generation
+
+The `scripts/compute_reference.py` script provides a template for creating reference videos needed for IC-LoRA training.
+The default implementation generates Canny edge reference videos.
+
+```bash
+# Generate Canny edge reference videos
+uv run python scripts/compute_reference.py videos_dir/ --output dataset.json
+```
+
+**Key features:**
+
+- **Canny edge detection**: Creates edge-based reference videos
+- **In-place editing**: Updates existing dataset JSON files
+- **Customizable**: Modify the `compute_reference()` function for different conditions (depth, pose, etc.)
+
+> [!TIP]
+> You can edit this script to generate other types of reference videos for IC-LoRA training,
+> such as depth maps, segmentation masks, or any custom video transformation.
+
+> [!NOTE]
+> `compute_reference.py` writes generated references to the `reference_video` column, which
+> `process_dataset.py` detects automatically.
+
+## 🔍 Debugging and Verification Scripts
+
+### Latents Decoding
+
+The `scripts/decode_latents.py` script decodes precomputed video latents back into video files for visual inspection.
+
+```bash
+# Basic usage
+uv run python scripts/decode_latents.py /path/to/latents/dir \
+    --output-dir /path/to/output \
+    --model-path /path/to/ltx-2-model.safetensors
+
+# With VAE tiling for large videos
+uv run python scripts/decode_latents.py /path/to/latents/dir \
+    --output-dir /path/to/output \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --vae-tiling
+
+# Decode both video and audio latents
+uv run python scripts/decode_latents.py /path/to/latents/dir \
+    --output-dir /path/to/output \
+    --model-path /path/to/ltx-2-model.safetensors \
+    --with-audio
+```
+
+**The script will:**
+
+1. **Load the VAE model** from the specified path
+2. **Process all `.pt` latent files** in the input directory
+3. **Decode each latent** back into a video using the VAE
+4. **Save resulting videos** as MP4 files in the output directory
+
+**When to use:**
+
+- **Verify preprocessing quality**: Check that your videos were encoded correctly
+- **Debug training data**: Visualize what the model actually sees during training
+- **Quality assessment**: Ensure latent encoding preserves important visual details
+
+### Inference with Trained Models
+
+For inference with trained LoRAs, use the [`ltx-pipelines`](../../ltx-pipelines/) package which provides
+production-ready pipelines:
+
+- **Text/Image-to-Video**: `TI2VidOneStagePipeline`, `TI2VidTwoStagesPipeline`
+- **Distilled (fast) inference**: `DistilledPipeline`
+- **IC-LoRA video-to-video**: `ICLoraPipeline`
+- **Keyframe interpolation**: `KeyframeInterpolationPipeline`
+
+All pipelines support loading custom LoRAs trained with this trainer.
+
+## 🚀 Training Scripts
+
+### Basic and Distributed Training
+
+Use `scripts/train.py` for both single GPU and multi-GPU runs:
+
+```bash
+# Single-GPU training
+uv run python scripts/train.py configs/t2v_lora.yaml
+
+# Multi-GPU (uses your accelerate config)
+uv run accelerate launch scripts/train.py configs/t2v_lora.yaml
+
+# Override number of processes
+uv run accelerate launch --num_processes 4 scripts/train.py configs/t2v_lora.yaml
+```
+
+For detailed usage, see the [Training Guide](training-guide.md).
+
+## 💡 Tips for Using Utility Scripts
+
+- **Start with `--help`**: Always check available options for each script
+- **Test on small datasets**: Verify workflows with a few files before processing large datasets
+- **Use decode verification**: Always decode a few samples to verify preprocessing quality
+- **Monitor VRAM usage**: Reach for quantization or lower-memory settings (e.g. FP8 for the captioner server) when running into memory issues
+- **Keep backups**: Make copies of important dataset files before running conversion scripts

packages/ltx-trainer/pyproject.toml

@@ -0,0 +1,98 @@
+[project]
+name = "ltx-trainer"
+version = "1.1.7"
+description = "LTX-2 training, democratized."
+readme = "README.md"
+authors = [
+    { name = "Matan Ben-Yosef", email = "mbyosef@lightricks.com" }
+]
+requires-python = ">=3.10"
+dependencies = [
+    "ltx-core",
+    "accelerate>=1.2.1",
+    "av>=14.2.1",
+    "bitsandbytes >=0.45.2; sys_platform == 'linux'",
+    "google-genai>=2.0",
+    "huggingface-hub[hf-xet]>=0.31.4",
+    "imageio>=2.37.0",
+    "imageio-ffmpeg>=0.6.0",
+    "openai>=2.0",
+    "opencv-python>=4.11.0.86",
+    "optimum-quanto>=0.2.6",
+    "pandas>=2.2.3",
+    "peft>=0.14.0",
+    "pillow-heif>=0.21.0",
+    "pydantic>=2.10.4",
+    "rich>=13.9.4",
+    "safetensors>=0.5.0",
+    "scenedetect>=0.6.5.2",
+    "sentencepiece>=0.2.0",
+    "soundfile>=0.12.1",
+    "torch>=2.6.0",
+    "torchaudio>=2.7.0",
+    # torchcodec must match the torch version (it ships a torch-ABI C++ extension and declares
+    # no torch pin of its own); the 0.9 line matches torch 2.9. torchaudio>=2.9 routes
+    # torchaudio.load() through torchcodec, so audio preprocessing needs it installed.
+    "torchcodec>=0.8.1,<0.10",
+    "torchvision>=0.21.0",
+    "typer>=0.15.1",
+    "wandb>=0.27.0",
+    "setuptools>=79.0.0",
+]
+
+[dependency-groups]
+dev = [
+    "pre-commit>=4.0.1",
+    "ruff>=0.8.6",
+]
+
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+
+
+[tool.ruff]
+target-version = "1.1.7"
+line-length = 120
+# Restrict isort first-party detection to src/ so stray dirs (e.g. wandb/ run output)
+# next to pyproject.toml don't get classified as first-party packages. See ruff#10519.
+src = ["src"]
+
+[tool.ruff.lint]
+select = [
+    "E", # pycodestyle
+    "F", # pyflakes
+    "W", # pycodestyle (warnings)
+    "I", # isort
+    "N", # pep8-naming
+    "ANN", # flake8-annotations
+    "B", # flake8-bugbear
+    "A", # flake8-builtins
+    "COM", # flake8-commas
+    "C4", # flake8-comprehensions
+    "DTZ", # flake8-datetimez
+    "EXE", # flake8-executable
+    "PIE", # flake8-pie
+    "T20", # flake8-print
+    "PT", # flake8-pytest
+    "SIM", # flake8-simplify
+    "ARG", # flake8-unused-arguments
+    "PTH", # flake8--use-pathlib
+    "ERA", # flake8-eradicate
+    "RUF", # ruff specific rules
+    "PL", # pylint
+]
+ignore = [
+    "ANN002", # Missing type annotation for *args
+    "ANN003", # Missing type annotation for **kwargs
+    "ANN204", # Missing type annotation for special method
+    "COM812", # Missing trailing comma
+    "PTH123", # `open()` should be replaced by `Path.open()`
+    "PLR2004", # Magic value used in comparison, consider replacing with a constant variable
+]
+[tool.ruff.lint.pylint]
+max-args = 10
+[tool.ruff.lint.isort]
+known-first-party = ["ltx_trainer", "ltx_core", "ltx_pipelines"]

packages/ltx-trainer/scripts/caption_videos.py

@@ -0,0 +1,536 @@
+#!/usr/bin/env python3
+
+"""
+Auto-caption videos with audio using multimodal models.
+Backends:
+- Qwen3-Omni-30B-A3B-Thinking via a local vLLM HTTP server (default,
+  ``qwen_omni``). Launch the server once with ``scripts/serve_captioner.py``.
+- Gemini Flash 3.5 via Google's API (``gemini_flash``).
+The paths in the output file are RELATIVE to the output file's directory,
+making the dataset portable.
+Basic usage:
+    # Launch the captioner server once (separate terminal)
+    uv run python scripts/serve_captioner.py
+    # Caption a directory
+    caption_videos.py videos_dir/ --output captions.json
+    # Caption a single video with a custom prompt
+    caption_videos.py video.mp4 --output cap.json --instruction "Describe in detail."
+Advanced usage:
+    # Use Gemini Flash 3.5 (cloud, requires GEMINI_API_KEY)
+    caption_videos.py videos_dir/ --captioner-type gemini_flash
+    # Gemini with parallel workers
+    caption_videos.py videos_dir/ --captioner-type gemini_flash --num-workers 5
+    # Talk to a remote vLLM server
+    caption_videos.py videos_dir/ --vllm-url http://192.168.1.10:8001/v1
+    # Enable Qwen3 chain-of-thought (slower, more detail)
+    caption_videos.py videos_dir/ --enable-thinking
+"""
+
+import csv
+import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from enum import Enum
+from pathlib import Path
+
+import typer
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+
+from ltx_trainer.captioning import (
+    DEFAULT_QWEN_MODEL,
+    DEFAULT_VLLM_BASE_URL,
+    CaptionerType,
+    MediaCaptioningModel,
+    create_captioner,
+)
+
+VIDEO_EXTENSIONS = ["mp4", "avi", "mov", "mkv", "webm"]
+IMAGE_EXTENSIONS = ["jpg", "jpeg", "png"]
+MEDIA_EXTENSIONS = VIDEO_EXTENSIONS + IMAGE_EXTENSIONS
+SAVE_INTERVAL = 5
+
+console = Console()
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=True,
+    help="Auto-caption videos with audio using multimodal models.",
+)
+
+
+class OutputFormat(str, Enum):
+    """Available output formats for captions."""
+
+    TXT = "txt"  # Separate files for captions and video paths, one caption / video path per line
+    CSV = "csv"  # CSV file with video path and caption columns
+    JSON = "json"  # JSON file with video paths as keys and captions as values
+    JSONL = "jsonl"  # JSON Lines file with one JSON object per line
+
+
+def caption_media(
+    input_path: Path,
+    output_path: Path,
+    captioner: MediaCaptioningModel,
+    extensions: list[str],
+    recursive: bool,
+    fps: int,
+    output_format: OutputFormat,
+    override: bool,
+    num_workers: int = 1,
+) -> None:
+    """Caption videos and images using the provided captioning model.
+    Args:
+        input_path: Path to input video file or directory
+        output_path: Path to output caption file
+        captioner: Media captioning model
+        extensions: List of media file extensions to include
+        recursive: Whether to search subdirectories recursively
+        fps: Frames per second to sample from videos (ignored for images)
+        output_format: Format to save the captions in
+        override: Whether to override existing captions
+        num_workers: Number of parallel workers (only for cloud-based captioners like Gemini)
+    """
+
+    # Get list of media files to process
+    media_files = _get_media_files(input_path, extensions, recursive)
+
+    if not media_files:
+        console.print("[bold yellow]No media files found to process.[/]")
+        return
+
+    console.print(f"Found [bold]{len(media_files)}[/] media files to process.")
+
+    # Load existing captions and determine which files need processing
+    base_dir = output_path.parent.resolve()
+    existing_captions = _load_existing_captions(output_path, output_format)
+    existing_abs_paths = {str((base_dir / p).resolve()) for p in existing_captions}
+
+    if override:
+        media_to_process = media_files
+    else:
+        media_to_process = [f for f in media_files if str(f.resolve()) not in existing_abs_paths]
+        if skipped := len(media_files) - len(media_to_process):
+            console.print(f"[bold yellow]Skipping {skipped} media that already have captions.[/]")
+
+    if not media_to_process:
+        console.print("[bold yellow]All media already have captions. Use --override to recaption.[/]")
+        return
+
+    if num_workers > 1:
+        console.print(f"Running with [bold cyan]{num_workers}[/] parallel workers.")
+
+    captions = existing_captions.copy()
+    successfully_captioned = 0
+    completed_since_save = 0
+
+    progress = Progress(
+        SpinnerColumn(),
+        TextColumn("{task.description}"),
+        BarColumn(bar_width=40),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TextColumn("•"),
+        TimeRemainingColumn(),
+        console=console,
+    )
+
+    def process_one(media_file: Path) -> tuple[str, str]:
+        """Caption a single media file and return (relative_path, caption)."""
+        caption = captioner.caption(
+            path=media_file,
+            fps=fps,
+        )
+        # Don't resolve the file itself, so a symlinked clip keeps its logical path under the
+        # dataset dir instead of jumping to its (possibly external) link target.
+        rel_path = str((media_file.parent.resolve() / media_file.name).relative_to(base_dir))
+        return rel_path, caption
+
+    with progress:
+        task = progress.add_task(
+            f"Captioning (workers: {num_workers})" if num_workers > 1 else "Captioning",
+            total=len(media_to_process),
+        )
+
+        with ThreadPoolExecutor(max_workers=num_workers) as executor:
+            futures = {executor.submit(process_one, f): f for f in media_to_process}
+
+            for future in as_completed(futures):
+                media_file = futures[future]
+                progress.update(task, description=f"Captioning [bold blue]{media_file.name}[/]")
+
+                try:
+                    rel_path, caption = future.result()
+
+                    captions[rel_path] = caption
+                    successfully_captioned += 1
+                    completed_since_save += 1
+
+                    if completed_since_save >= SAVE_INTERVAL:
+                        _save_captions(captions, output_path, output_format)
+                        completed_since_save = 0
+
+                except Exception as e:
+                    console.print(f"[bold red]Error captioning {media_file.name}: {e}[/]")
+
+                progress.advance(task)
+
+    # Final save with everything accumulated
+    _save_captions(captions, output_path, output_format)
+
+    # Print summary
+    console.print(
+        f"[bold green]✓[/] Captioned [bold]{successfully_captioned}/{len(media_to_process)}[/] media successfully.",
+    )
+
+
+def _get_media_files(
+    input_path: Path,
+    extensions: list[str] = MEDIA_EXTENSIONS,
+    recursive: bool = False,
+) -> list[Path]:
+    """Get all media files from the input path."""
+    input_path = Path(input_path)
+    # Normalize extensions to lowercase without dots
+    extensions_set = {ext.lower().lstrip(".") for ext in extensions}
+
+    if input_path.is_file():
+        # If input is a file, check if it has a valid extension
+        if input_path.suffix.lstrip(".").lower() in extensions_set:
+            return [input_path]
+        else:
+            typer.echo(f"Warning: {input_path} is not a recognized media file. Skipping.")
+            return []
+    elif input_path.is_dir():
+        # Find all files and filter by extension case-insensitively
+        glob_pattern = "**/*" if recursive else "*"
+        media_files = [
+            f for f in input_path.glob(glob_pattern) if f.is_file() and f.suffix.lstrip(".").lower() in extensions_set
+        ]
+        return sorted(media_files)
+    else:
+        typer.echo(f"Error: {input_path} does not exist.")
+        raise typer.Exit(code=1)
+
+
+def _save_captions(
+    captions: dict[str, str],
+    output_path: Path,
+    format_type: OutputFormat,
+) -> None:
+    """Save captions to a file in the specified format.
+    Args:
+        captions: Dictionary mapping media paths to captions
+        output_path: Path to save the output file
+        format_type: Format to save the captions in
+    """
+    # Create parent directories if they don't exist
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    console.print("[bold blue]Saving captions...[/]")
+
+    match format_type:
+        case OutputFormat.TXT:
+            # Create two separate files for captions and media paths
+            captions_file = output_path.with_stem(f"{output_path.stem}_captions")
+            paths_file = output_path.with_stem(f"{output_path.stem}_paths")
+
+            with captions_file.open("w", encoding="utf-8") as f:
+                for caption in captions.values():
+                    f.write(f"{caption}\n")
+
+            with paths_file.open("w", encoding="utf-8") as f:
+                for media_path in captions:
+                    f.write(f"{media_path}\n")
+
+            console.print(f"[bold green]✓[/] Captions saved to [cyan]{captions_file}[/]")
+            console.print(f"[bold green]✓[/] Media paths saved to [cyan]{paths_file}[/]")
+
+        case OutputFormat.CSV:
+            with output_path.open("w", encoding="utf-8", newline="") as f:
+                writer = csv.writer(f)
+                writer.writerow(["caption", "media_path"])
+                for media_path, caption in captions.items():
+                    writer.writerow([caption, media_path])
+
+            console.print(f"[bold green]✓[/] Captions saved to [cyan]{output_path}[/]")
+
+        case OutputFormat.JSON:
+            # Format as list of dictionaries with caption and media_path keys
+            json_data = [{"caption": caption, "media_path": media_path} for media_path, caption in captions.items()]
+
+            with output_path.open("w", encoding="utf-8") as f:
+                json.dump(json_data, f, indent=2, ensure_ascii=False)
+
+            console.print(f"[bold green]✓[/] Captions saved to [cyan]{output_path}[/]")
+
+        case OutputFormat.JSONL:
+            with output_path.open("w", encoding="utf-8") as f:
+                for media_path, caption in captions.items():
+                    f.write(json.dumps({"caption": caption, "media_path": media_path}, ensure_ascii=False) + "\n")
+
+            console.print(f"[bold green]✓[/] Captions saved to [cyan]{output_path}[/]")
+
+        case _:
+            raise ValueError(f"Unsupported output format: {format_type}")
+
+
+def _load_existing_captions(  # noqa: PLR0912
+    output_path: Path,
+    format_type: OutputFormat,
+) -> dict[str, str]:
+    """Load existing captions from a file.
+    Args:
+        output_path: Path to the captions file
+        format_type: Format of the captions file
+    Returns:
+        Dictionary mapping media paths to captions, or empty dict if file doesn't exist
+    """
+    if not output_path.exists():
+        return {}
+
+    console.print(f"[bold blue]Loading existing captions from [cyan]{output_path}[/]...[/]")
+
+    existing_captions = {}
+
+    try:
+        match format_type:
+            case OutputFormat.TXT:
+                # For TXT format, we have two separate files
+                captions_file = output_path.with_stem(f"{output_path.stem}_captions")
+                paths_file = output_path.with_stem(f"{output_path.stem}_paths")
+
+                if captions_file.exists() and paths_file.exists():
+                    captions = captions_file.read_text(encoding="utf-8").splitlines()
+                    paths = paths_file.read_text(encoding="utf-8").splitlines()
+
+                    if len(captions) == len(paths):
+                        existing_captions = dict(zip(paths, captions, strict=False))
+
+            case OutputFormat.CSV:
+                with output_path.open("r", encoding="utf-8", newline="") as f:
+                    reader = csv.reader(f)
+                    # Skip header
+                    next(reader, None)
+                    for row in reader:
+                        if len(row) >= 2:
+                            caption, media_path = row[0], row[1]
+                            existing_captions[media_path] = caption
+
+            case OutputFormat.JSON:
+                with output_path.open("r", encoding="utf-8") as f:
+                    json_data = json.load(f)
+                    for item in json_data:
+                        if "caption" in item and "media_path" in item:
+                            existing_captions[item["media_path"]] = item["caption"]
+
+            case OutputFormat.JSONL:
+                with output_path.open("r", encoding="utf-8") as f:
+                    for line in f:
+                        item = json.loads(line)
+                        if "caption" in item and "media_path" in item:
+                            existing_captions[item["media_path"]] = item["caption"]
+
+            case _:
+                raise ValueError(f"Unsupported output format: {format_type}")
+
+        console.print(f"[bold green]✓[/] Loaded [bold]{len(existing_captions)}[/] existing captions")
+        return existing_captions
+
+    except Exception as e:
+        console.print(f"[bold yellow]Warning: Could not load existing captions: {e}[/]")
+        return {}
+
+
+@app.command()
+def main(  # noqa: PLR0913
+    input_path: Path = typer.Argument(  # noqa: B008
+        ...,
+        help="Path to input video/image file or directory containing media files",
+        exists=True,
+    ),
+    output: Path | None = typer.Option(  # noqa: B008
+        None,
+        "--output",
+        "-o",
+        help="Path to output file for captions. Format determined by file extension.",
+    ),
+    captioner_type: CaptionerType = typer.Option(  # noqa: B008
+        CaptionerType.QWEN_OMNI,
+        "--captioner-type",
+        "-c",
+        help="Type of captioner to use. Valid values: 'qwen_omni' (local), 'gemini_flash' (API)",
+        case_sensitive=False,
+    ),
+    vllm_url: str = typer.Option(
+        DEFAULT_VLLM_BASE_URL,
+        "--vllm-url",
+        help=(
+            "Base URL of the vLLM OpenAI-compatible server (qwen_omni only). "
+            "Launch the server with `uv run python scripts/serve_captioner.py`."
+        ),
+    ),
+    vllm_model: str = typer.Option(
+        DEFAULT_QWEN_MODEL,
+        "--vllm-model",
+        help="Served model identifier on the vLLM server (qwen_omni only).",
+    ),
+    enable_thinking: bool = typer.Option(
+        False,
+        "--enable-thinking/--no-thinking",
+        help=(
+            "Let Qwen3-Omni produce a <think>...</think> chain-of-thought before the caption. "
+            "Off by default: ~5x slower with marginal quality benefit and occasional hallucinations."
+        ),
+    ),
+    max_tokens: int = typer.Option(
+        4096,
+        "--max-tokens",
+        help="Maximum new tokens to generate per caption (qwen_omni only).",
+    ),
+    instruction: str | None = typer.Option(
+        None,
+        "--instruction",
+        "-i",
+        help="Custom instruction for the captioning model. If not provided, uses an appropriate default.",
+    ),
+    extensions: str = typer.Option(
+        ",".join(MEDIA_EXTENSIONS),
+        "--extensions",
+        "-e",
+        help="Comma-separated list of media file extensions to process",
+    ),
+    recursive: bool = typer.Option(
+        False,
+        "--recursive",
+        "-r",
+        help="Search for media files in subdirectories recursively",
+    ),
+    fps: int = typer.Option(
+        2,
+        "--fps",
+        "-f",
+        help=(
+            "Frames per second to sample from videos. 2 is a typical default; "
+            "lower values use less compute per video. Ignored for images and for the "
+            "Gemini backend (which decides its own sampling rate)."
+        ),
+    ),
+    override: bool = typer.Option(
+        False,
+        "--override",
+        help="Whether to override existing captions for media",
+    ),
+    api_key: str | None = typer.Option(
+        None,
+        "--api-key",
+        envvar=["GOOGLE_API_KEY", "GEMINI_API_KEY"],
+        help="API key for Gemini Flash (can also use GOOGLE_API_KEY or GEMINI_API_KEY env var)",
+    ),
+    num_workers: int = typer.Option(
+        1,
+        "--num-workers",
+        "-w",
+        min=1,
+        max=10,
+        help=(
+            "Number of parallel workers for captioning (1-10). "
+            "Values above 1 are only supported for cloud-based captioners (gemini_flash). "
+            "Using multiple workers with a local model will raise an error."
+        ),
+    ),
+) -> None:
+    """Auto-caption videos with audio using multimodal models.
+    Backends:
+    - ``qwen_omni`` (default): Qwen3-Omni-30B-A3B-Thinking via a local vLLM
+      HTTP server. Launch the server once in a separate terminal with
+      ``uv run python scripts/serve_captioner.py``. The server stays loaded
+      across script invocations.
+    - ``gemini_flash``: Google Gemini (``gemini-3.5-flash``) via the google-genai SDK.
+      Auth is automatic -- ``GEMINI_API_KEY``/``GOOGLE_API_KEY`` for the Developer API,
+      or Google Cloud credentials (gcloud / service account) for Vertex AI with no env vars.
+    The paths in the output file will be relative to the output file's directory.
+    Examples:
+        # Caption videos using the local vLLM server (default)
+        caption_videos.py videos_dir/ -o captions.json
+        # Point at a remote vLLM server
+        caption_videos.py videos_dir/ -o captions.json --vllm-url http://other-host:8001/v1
+        # Caption using Gemini Flash 3.5
+        caption_videos.py videos_dir/ -o captions.json -c gemini_flash
+        # Caption with custom instruction
+        caption_videos.py video.mp4 -o captions.json -i "Describe this video in detail"
+    """
+
+    # Parallel workers are only supported for the cloud Gemini backend; qwen_omni
+    # drives a single shared vLLM server and is captioned serially from here.
+    if num_workers > 1 and captioner_type != CaptionerType.GEMINI_FLASH:
+        console.print(
+            "[bold red]Error:[/] --num-workers > 1 is only supported with "
+            "[bold]--captioner-type gemini_flash[/]. Use --num-workers 1 (default) "
+            "for the qwen_omni backend."
+        )
+        raise typer.Exit(code=1)
+
+    # Parse extensions
+    ext_list = [ext.strip() for ext in extensions.split(",")]
+
+    # Determine output path and format
+    if output is None:
+        output_format = OutputFormat.JSON
+        if input_path.is_file():  # noqa: SIM108
+            # Default to a JSON file with the same name as the input media
+            output = input_path.with_suffix(".dataset.json")
+        else:
+            # Default to a JSON file in the input directory
+            output = input_path / "dataset.json"
+    else:
+        # Determine format from file extension
+        output_format = OutputFormat(Path(output).suffix.lstrip(".").lower())
+
+    # Ensure output path is absolute
+    output = Path(output).resolve()
+    console.print(f"Output will be saved to [bold blue]{output}[/]")
+
+    with console.status("Initializing captioner...", spinner="dots"):
+        if captioner_type == CaptionerType.QWEN_OMNI:
+            captioner = create_captioner(
+                captioner_type=captioner_type,
+                base_url=vllm_url,
+                model=vllm_model,
+                instruction=instruction,
+                max_tokens=max_tokens,
+                enable_thinking=enable_thinking,
+            )
+        elif captioner_type == CaptionerType.GEMINI_FLASH:
+            captioner = create_captioner(
+                captioner_type=captioner_type,
+                api_key=api_key,
+                instruction=instruction,
+            )
+        else:
+            raise ValueError(f"Unsupported captioner type: {captioner_type}")
+
+        console.print(f"[bold green]✓[/] {captioner_type.value} captioner ready")
+
+    # Caption media files
+    caption_media(
+        input_path=input_path,
+        output_path=output,
+        captioner=captioner,
+        extensions=ext_list,
+        recursive=recursive,
+        fps=fps,
+        output_format=output_format,
+        override=override,
+        num_workers=num_workers,
+    )
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/scripts/compute_reference.py

@@ -0,0 +1,304 @@
+"""
+Compute reference videos for IC-LoRA training.
+This script provides a command-line interface for generating reference videos to be used for IC-LoRA training.
+Note that it reads and writes to the same file (the output of caption_videos.py),
+where it adds the "reference_video" field to the JSON.
+Basic usage:
+    # Compute reference videos for all videos in a directory
+    compute_reference.py videos_dir/ --output videos_dir/captions.json
+"""
+
+# Standard library imports
+import json
+from pathlib import Path
+from typing import Any
+
+# Third-party imports
+import cv2
+import torch
+import torchvision.transforms.functional as TF  # noqa: N812
+import typer
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from transformers.utils.logging import disable_progress_bar
+
+# Local imports
+from ltx_trainer.video_utils import read_video, save_video
+
+# Initialize console and disable progress bars
+console = Console()
+disable_progress_bar()
+
+VIDEO_COLUMNS = ("video", "media_path")
+REFERENCE_VIDEO_COLUMN = "reference_video"
+LEGACY_REFERENCE_COLUMN = "reference_path"
+
+
+def compute_reference(
+    images: torch.Tensor,
+) -> torch.Tensor:
+    """Compute Canny edge detection on a batch of images.
+    Args:
+        images: Batch of images tensor of shape [B, C, H, W]
+    Returns:
+        Binary edge masks tensor of shape [B, H, W]
+    """
+    # Convert to grayscale if needed
+    if images.shape[1] == 3:
+        images = TF.rgb_to_grayscale(images)
+
+    # Ensure images are in [0, 1] range
+    if images.max() > 1.0:
+        images = images / 255.0
+
+    # Compute Canny edges
+    edge_masks = []
+    for image in images:
+        # Convert to numpy for OpenCV
+        image_np = (image.squeeze().cpu().numpy() * 255).astype("uint8")
+
+        # Apply Canny edge detection
+        edges = cv2.Canny(
+            image_np,
+            threshold1=100,
+            threshold2=200,
+        )
+
+        # Convert back to tensor
+        edge_mask = torch.from_numpy(edges).float()
+        edge_masks.append(edge_mask)
+
+    edges = torch.stack(edge_masks)
+    edges = torch.stack([edges] * 3, dim=1)  # Convert to 3-channel
+    return edges
+
+
+def _get_meta_data(
+    output_path: Path,
+) -> list[dict[str, Any]]:
+    """Get set of existing reference video paths without loading the actual files.
+    Args:
+        output_path: Path to the reference video paths file
+    Returns:
+        Dataset rows with media paths and captions
+    """
+    if not output_path.exists():
+        return []
+
+    console.print(f"[bold blue]Reading meta data from [cyan]{output_path}[/]...[/]")
+
+    try:
+        with output_path.open("r", encoding="utf-8") as f:
+            json_data = json.load(f)
+        return json_data
+
+    except Exception as e:
+        console.print(f"[bold yellow]Warning: Could not check meta data: {e}[/]")
+        return []
+
+
+def _get_media_path(item: dict[str, Any]) -> str:
+    for column in VIDEO_COLUMNS:
+        if column in item:
+            return item[column]
+    raise KeyError(f"Dataset row must contain one of {VIDEO_COLUMNS}")
+
+
+def _save_dataset_json(
+    reference_paths: dict[str, str],
+    output_path: Path,
+) -> None:
+    """Save dataset json with reference video paths.
+    Args:
+        reference_paths: Dictionary mapping media paths to reference video paths
+        output_path: Path to save the output file
+    """
+
+    with output_path.open("r", encoding="utf-8") as f:
+        json_data = json.load(f)
+        new_json_data = json_data.copy()
+        for i, item in enumerate(json_data):
+            media_path = _get_media_path(item)
+            reference_path = reference_paths[media_path]
+            new_json_data[i].pop(LEGACY_REFERENCE_COLUMN, None)
+            new_json_data[i][REFERENCE_VIDEO_COLUMN] = reference_path
+
+    with output_path.open("w", encoding="utf-8") as f:
+        json.dump(new_json_data, f, indent=2, ensure_ascii=False)
+
+    console.print(f"[bold green]✓[/] Reference video paths saved to [cyan]{output_path}[/]")
+    console.print("[bold yellow]Note:[/] Reference videos were written to the '[cyan]reference_video[/]' column.")
+    console.print("  [cyan]process_dataset.py[/] detects this column automatically for IC-LoRA preprocessing.")
+
+
+def process_media(
+    input_path: Path,
+    output_path: Path,
+    override: bool,
+    batch_size: int = 100,
+) -> None:
+    """Process videos and images to compute condition on videos.
+    Args:
+        input_path: Path to input video/image file or directory
+        output_path: Path to output reference video file
+        override: Whether to override existing reference video files
+    """
+    if not output_path.exists():
+        raise FileNotFoundError(
+            f"Output file does not exist: {output_path}. This is also the input file for the dataset."
+        )
+
+    # Check for existing reference video files
+    meta_data = _get_meta_data(output_path)
+
+    base_dir = input_path.resolve()
+    console.print(f"Using [bold blue]{base_dir}[/] as base directory for relative paths")
+
+    # Filter media files
+    media_to_process = []
+    skipped_media = []
+
+    def media_path_to_reference_path(media_file: Path) -> Path:
+        return media_file.parent / (media_file.stem + "_reference" + media_file.suffix)
+
+    media_files = [base_dir / Path(_get_media_path(sample)) for sample in meta_data]
+    for media_file in media_files:
+        reference_path = media_path_to_reference_path(media_file)
+        media_to_process.append(media_file)
+
+    console.print(f"Processing [bold]{len(media_to_process)}[/] media.")
+
+    # Initialize progress tracking
+    progress = Progress(
+        SpinnerColumn(),
+        TextColumn("{task.description}"),
+        BarColumn(bar_width=40),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TextColumn("•"),
+        TimeRemainingColumn(),
+        console=console,
+    )
+
+    # Process media files
+    media_paths = [_get_media_path(item) for item in meta_data]
+    reference_paths = {rel_path: str(media_path_to_reference_path(Path(rel_path))) for rel_path in media_paths}
+
+    with progress:
+        task = progress.add_task("Computing condition on videos", total=len(media_to_process))
+
+        for media_file, rel_path in zip(media_to_process, media_paths, strict=True):
+            progress.update(task, description=f"Processing [bold blue]{media_file.name}[/]")
+
+            # Key by the original media-path string (matches the dict seeded above). Avoid
+            # resolve()/relative_to here — they crash on symlinked or absolute media paths.
+            reference_path = media_path_to_reference_path(media_file)
+            try:
+                ref_stored = str(reference_path.relative_to(base_dir))
+            except ValueError:
+                ref_stored = str(reference_path)  # absolute/out-of-tree: keep it next to the source
+            reference_paths[rel_path] = ref_stored
+
+            if not reference_path.resolve().exists() or override:
+                try:
+                    video, fps = read_video(media_file)
+
+                    # Process frames in batches
+                    condition_frames = []
+
+                    for i in range(0, len(video), batch_size):
+                        batch = video[i : i + batch_size]
+                        condition_batch = compute_reference(batch)
+                        condition_frames.append(condition_batch)
+
+                    # Concatenate all edge frames
+                    all_condition = torch.cat(condition_frames, dim=0)
+
+                    # Save the edge video
+                    save_video(all_condition, reference_path.resolve(), fps=fps)
+
+                except Exception as e:
+                    console.print(f"[bold red]Error processing [bold blue]{media_file}[/]: {e}[/]")
+                    reference_paths.pop(rel_path)
+            else:
+                skipped_media.append(media_file)
+
+            progress.advance(task)
+
+    # Save results
+    _save_dataset_json(reference_paths, output_path)
+
+    # Print summary
+    total_to_process = len(media_files) - len(skipped_media)
+    console.print(
+        f"[bold green]✓[/] Processed [bold]{total_to_process}/{len(media_files)}[/] media successfully.",
+    )
+
+
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=True,
+    help="Compute reference videos for IC-LoRA training.",
+)
+
+
+@app.command()
+def main(
+    input_path: Path = typer.Argument(  # noqa: B008
+        ...,
+        help="Path to input video/image file or directory containing media files",
+        exists=True,
+    ),
+    output: Path = typer.Option(  # noqa: B008
+        ...,
+        "--output",
+        "-o",
+        help="Path to json output file for reference video paths. "
+        "This is also the input file for the dataset, the output of compute_captions.py.",
+    ),
+    override: bool = typer.Option(
+        False,
+        "--override",
+        help="Whether to override existing reference video files",
+    ),
+    batch_size: int = typer.Option(
+        100,
+        "--batch-size",
+        help="Batch size for processing videos",
+    ),
+) -> None:
+    """Compute reference videos for IC-LoRA training.
+    This script generates reference videos (e.g., Canny edge maps) for given videos.
+    The paths in the output file will be relative to the output file's directory.
+    Examples:
+        # Process all videos in a directory
+        compute_reference.py videos_dir/ -o videos_dir/captions.json
+    """
+
+    # Ensure output path is absolute
+    output = Path(output).resolve()
+    console.print(f"Output will be saved to [bold blue]{output}[/]")
+
+    # Verify output path exists
+    if not output.exists():
+        raise FileNotFoundError(f"Output file does not exist: {output}. This is also the input file for the dataset.")
+
+    # Process media files
+    process_media(
+        input_path=input_path,
+        output_path=output,
+        override=override,
+        batch_size=batch_size,
+    )
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/scripts/decode_latents.py

@@ -0,0 +1,369 @@
+#!/usr/bin/env python3
+
+"""
+Decode precomputed video latents back into videos using the VAE.
+This script loads latent files saved during preprocessing and decodes them
+back into video clips using the same VAE model.
+Basic usage:
+    python scripts/decode_latents.py /path/to/latents/dir /path/to/output \
+        --model-source /path/to/ltx2.safetensors
+"""
+
+from pathlib import Path
+
+import torch
+import torchaudio
+import torchvision.utils
+import typer
+from einops import rearrange
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from transformers.utils.logging import disable_progress_bar
+
+from ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
+from ltx_trainer import logger
+from ltx_trainer.model_loader import load_audio_vae_decoder, load_video_vae_decoder, load_vocoder
+from ltx_trainer.video_utils import save_video
+
+DEFAULT_TILE_SIZE_PIXELS = 512  # Spatial tile size in pixels (must be ≥64 and divisible by 32)
+DEFAULT_TILE_OVERLAP_PIXELS = 128  # Spatial tile overlap in pixels (must be divisible by 32)
+DEFAULT_TILE_SIZE_FRAMES = 128  # Temporal tile size in frames (must be ≥16 and divisible by 8)
+DEFAULT_TILE_OVERLAP_FRAMES = 24  # Temporal tile overlap in frames (must be divisible by 8)
+
+disable_progress_bar()
+console = Console()
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=True,
+    help="Decode precomputed video latents back into videos using the VAE.",
+)
+
+
+class LatentsDecoder:
+    def __init__(
+        self,
+        model_path: str,
+        device: str = "cuda",
+        vae_tiling: bool = False,
+        with_audio: bool = False,
+    ):
+        """Initialize the decoder with model configuration.
+        Args:
+            model_path: Path to LTX-2 checkpoint (.safetensors)
+            device: Device to use for computation
+            vae_tiling: Whether to enable VAE tiling for larger video resolutions
+            with_audio: Whether to load audio VAE for audio decoding
+        """
+        self.device = torch.device(device)
+        self.model_path = model_path
+        self.vae = None
+        self.audio_vae = None
+        self.vocoder = None
+        self.vae_tiling = vae_tiling
+
+        self._load_model(model_path, with_audio)
+
+    def _load_model(self, model_path: str, with_audio: bool = False) -> None:
+        """Initialize and load the VAE model(s)."""
+        with console.status(f"[bold]Loading video VAE decoder from {model_path}...", spinner="dots"):
+            self.vae = load_video_vae_decoder(model_path, device=self.device, dtype=torch.bfloat16)
+
+        if with_audio:
+            with console.status(f"[bold]Loading audio VAE decoder from {model_path}...", spinner="dots"):
+                self.audio_vae = load_audio_vae_decoder(model_path, device=self.device, dtype=torch.bfloat16)
+
+            with console.status(f"[bold]Loading vocoder from {model_path}...", spinner="dots"):
+                self.vocoder = load_vocoder(model_path, device=self.device)
+
+    @torch.inference_mode()
+    def decode(self, latents_dir: Path, output_dir: Path, seed: int | None = None) -> None:
+        """Decode all latent files in the directory recursively.
+        Args:
+            latents_dir: Directory containing latent files (.pt)
+            output_dir: Directory to save decoded videos
+            seed: Optional random seed for noise generation
+        """
+        # Find all .pt files recursively
+        latent_files = list(latents_dir.rglob("*.pt"))
+
+        if not latent_files:
+            logger.warning(f"No .pt files found in {latents_dir}")
+            return
+
+        logger.info(f"Found {len(latent_files):,} latent files to decode")
+
+        # Process files with progress bar
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            MofNCompleteColumn(),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            console=console,
+        ) as progress:
+            task = progress.add_task("Decoding latents", total=len(latent_files))
+
+            for latent_file in latent_files:
+                # Calculate relative path to maintain directory structure
+                rel_path = latent_file.relative_to(latents_dir)
+                output_subdir = output_dir / rel_path.parent
+                output_subdir.mkdir(parents=True, exist_ok=True)
+
+                try:
+                    self._process_file(latent_file, output_subdir, seed)
+                except Exception as e:
+                    logger.error(f"Error processing {latent_file}: {e}")
+                    continue
+
+                progress.advance(task)
+
+        logger.info(f"Decoding complete! Videos saved to {output_dir}")
+
+    @torch.inference_mode()
+    def decode_audio(self, latents_dir: Path, output_dir: Path) -> None:
+        """Decode all audio latent files in the directory recursively.
+        Args:
+            latents_dir: Directory containing audio latent files (.pt)
+            output_dir: Directory to save decoded audio files
+        """
+        # Check if audio VAE is loaded
+        if self.audio_vae is None or self.vocoder is None:
+            logger.warning("Audio VAE or vocoder not loaded. Skipping audio decoding.")
+            return
+
+        # Find all .pt files recursively
+        latent_files = list(latents_dir.rglob("*.pt"))
+
+        if not latent_files:
+            logger.warning(f"No .pt files found in {latents_dir}")
+            return
+
+        logger.info(f"Found {len(latent_files):,} audio latent files to decode")
+
+        # Process files with progress bar
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            MofNCompleteColumn(),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            console=console,
+        ) as progress:
+            task = progress.add_task("Decoding audio latents", total=len(latent_files))
+
+            for latent_file in latent_files:
+                # Calculate relative path to maintain directory structure
+                rel_path = latent_file.relative_to(latents_dir)
+                output_subdir = output_dir / rel_path.parent
+                output_subdir.mkdir(parents=True, exist_ok=True)
+
+                try:
+                    self._process_audio_file(latent_file, output_subdir)
+                except Exception as e:
+                    logger.error(f"Error processing audio {latent_file}: {e}")
+                    continue
+
+                progress.advance(task)
+
+        logger.info(f"Audio decoding complete! Audio files saved to {output_dir}")
+
+    def _process_file(self, latent_file: Path, output_dir: Path, seed: int | None) -> None:
+        """Process a single latent file."""
+        # Load the latent data
+        data = torch.load(latent_file, map_location=self.device, weights_only=False)
+
+        # Get latents - handle both old patchified [seq_len, C] and new [C, F, H, W] formats
+        latents = data["latents"]
+        num_frames = data["num_frames"]
+        height = data["height"]
+        width = data["width"]
+
+        # Check if latents need reshaping (old patchified format)
+        if latents.dim() == 2:
+            # Old format: [seq_len, C] -> reshape to [C, F, H, W]
+            latents = rearrange(latents, "(f h w) c -> c f h w", f=num_frames, h=height, w=width)
+
+        # Add batch dimension: [C, F, H, W] -> [1, C, F, H, W]
+        latents = latents.unsqueeze(0).to(device=self.device, dtype=torch.bfloat16)
+
+        # Create generator only if seed is provided
+        generator = None
+        if seed is not None:
+            generator = torch.Generator(device=self.device)
+            generator.manual_seed(seed)
+
+        # Decode the video
+        video = self._decode_video(latents, generator)
+
+        # Determine output format and save
+        is_image = video.shape[0] == 1
+        if is_image:
+            # Save as PNG for single frame
+            output_path = output_dir / f"{latent_file.stem}.png"
+            torchvision.utils.save_image(
+                video[0],  # [C, H, W] in [0, 1]
+                str(output_path),
+            )
+        else:
+            # Save as MP4 for video using PyAV-based save_video
+            output_path = output_dir / f"{latent_file.stem}.mp4"
+            fps = data.get("fps", 24)  # Use stored FPS or default to 24
+            save_video(
+                video_tensor=video,  # [F, C, H, W] in [0, 1]
+                output_path=output_path,
+                fps=fps,
+            )
+
+    def _decode_video(self, latents: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+        """Decode latents to video frames."""
+        if self.vae_tiling:
+            # Use tiled decoding for reduced VRAM
+            tiling_config = TilingConfig(
+                spatial_config=SpatialTilingConfig(
+                    tile_size_in_pixels=DEFAULT_TILE_SIZE_PIXELS,
+                    tile_overlap_in_pixels=DEFAULT_TILE_OVERLAP_PIXELS,
+                ),
+                temporal_config=TemporalTilingConfig(
+                    tile_size_in_frames=DEFAULT_TILE_SIZE_FRAMES,
+                    tile_overlap_in_frames=DEFAULT_TILE_OVERLAP_FRAMES,
+                ),
+            )
+            chunks = list(
+                self.vae.tiled_decode(
+                    latents,
+                    tiling_config=tiling_config,
+                    generator=generator,
+                )
+            )
+            # Concatenate along temporal dimension
+            video = torch.cat(chunks, dim=2)  # [B, C, F, H, W]
+        else:
+            # Standard full decoding
+            video = self.vae(latents, generator=generator)  # [B, C, F, H, W]
+
+        # Convert to [F, C, H, W] format and normalize to [0, 1]
+        video = rearrange(video, "1 c f h w -> f c h w")
+        video = (video + 1) / 2  # Denormalize from [-1, 1] to [0, 1]
+        video = video.clamp(0, 1)
+
+        return video
+
+    def _process_audio_file(self, latent_file: Path, output_dir: Path) -> None:
+        """Process a single audio latent file."""
+        # Load the latent data
+        data = torch.load(latent_file, map_location=self.device, weights_only=False)
+
+        latents = data["latents"].to(device=self.device, dtype=torch.float32)
+        num_time_steps = data["num_time_steps"]
+        freq_bins = data["frequency_bins"]
+
+        # Handle both old patchified [seq_len, C] and new [C, T, F] formats
+        if latents.dim() == 2:
+            # Old format: [seq_len, channels] where seq_len = time * freq
+            # Reshape to [C, T, F]
+            latents = rearrange(latents, "(t f) c -> c t f", t=num_time_steps, f=freq_bins)
+
+        # Add batch dimension: [C, T, F] -> [1, C, T, F]
+        latents = latents.unsqueeze(0)
+
+        # Set correct dtype for audio VAE
+        latents = latents.to(dtype=torch.bfloat16)
+
+        # Decode audio using audio VAE decoder (produces mel spectrogram)
+        mel_spectrogram = self.audio_vae(latents)
+
+        # Convert mel spectrogram to waveform using vocoder
+        waveform = self.vocoder(mel_spectrogram)
+
+        # Save as WAV
+        output_path = output_dir / f"{latent_file.stem}.wav"
+        sample_rate = self.vocoder.output_sampling_rate
+        torchaudio.save(str(output_path), waveform[0].cpu(), sample_rate)
+
+
+@app.command()
+def main(
+    latents_dir: str = typer.Argument(
+        ...,
+        help="Directory containing the precomputed latent files (searched recursively)",
+    ),
+    output_dir: str = typer.Argument(
+        ...,
+        help="Directory to save the decoded videos (maintains same folder hierarchy as input)",
+    ),
+    model_path: str = typer.Option(
+        ...,
+        help="Path to LTX-2 checkpoint (.safetensors file)",
+    ),
+    device: str = typer.Option(
+        default="cuda",
+        help="Device to use for computation",
+    ),
+    vae_tiling: bool = typer.Option(
+        default=True,
+        help="Enable VAE tiling for larger video resolutions",
+    ),
+    seed: int | None = typer.Option(
+        default=None,
+        help="Random seed for noise generation during decoding",
+    ),
+    with_audio: bool = typer.Option(
+        default=False,
+        help="Also decode audio latents (requires audio_latents directory)",
+    ),
+    audio_latents_dir: str | None = typer.Option(
+        default=None,
+        help="Directory containing audio latent files (defaults to 'audio_latents' sibling of latents_dir)",
+    ),
+) -> None:
+    """Decode precomputed video latents back into videos using the VAE.
+    This script recursively searches for .pt latent files in the input directory
+    and decodes them to videos, maintaining the same folder hierarchy in the output.
+    Examples:
+        # Basic usage
+        python scripts/decode_latents.py /path/to/latents /path/to/videos \\
+            --model-path /path/to/ltx2.safetensors
+        # With VAE tiling for large videos
+        python scripts/decode_latents.py /path/to/latents /path/to/videos \\
+            --model-path /path/to/ltx2.safetensors --vae-tiling
+        # With audio decoding
+        python scripts/decode_latents.py /path/to/latents /path/to/videos \\
+            --model-path /path/to/ltx2.safetensors --with-audio
+    """
+    latents_path = Path(latents_dir)
+    output_path = Path(output_dir)
+
+    if not latents_path.exists() or not latents_path.is_dir():
+        raise typer.BadParameter(f"Latents directory does not exist: {latents_path}")
+
+    decoder = LatentsDecoder(
+        model_path=model_path,
+        device=device,
+        vae_tiling=vae_tiling,
+        with_audio=with_audio,
+    )
+    decoder.decode(latents_path, output_path, seed=seed)
+
+    # Decode audio if requested
+    if with_audio:
+        audio_path = Path(audio_latents_dir) if audio_latents_dir else latents_path.parent / "audio_latents"
+
+        if audio_path.exists():
+            audio_output_path = output_path.parent / "decoded_audio"
+            decoder.decode_audio(audio_path, audio_output_path)
+        else:
+            logger.warning(f"Audio latents directory not found: {audio_path}")
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/scripts/process_captions.py

@@ -0,0 +1,497 @@
+#!/usr/bin/env python
+
+"""
+Compute text embeddings for video generation training.
+This module provides functionality for processing text captions, including:
+- Loading captions from various file formats (CSV, JSON, JSONL)
+- Cleaning and preprocessing text (removing LLM prefixes, adding ID tokens)
+- CaptionsDataset for caption-only preprocessing workflows
+Can be used as a standalone script:
+    python scripts/process_captions.py dataset.json --output-dir /path/to/output \
+        --model-source /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma
+"""
+
+import json
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+import torch
+import typer
+from accelerate import PartialState
+from rich.console import Console
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+from torch.utils.data import DataLoader, Dataset, Subset
+from transformers.utils.logging import disable_progress_bar
+
+from ltx_trainer import logger
+from ltx_trainer.model_loader import load_embeddings_processor, load_text_encoder
+
+# Disable tokenizers parallelism to avoid warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+disable_progress_bar()
+
+# Common phrases that LLMs often add to captions that we might want to remove
+COMMON_BEGINNING_PHRASES: tuple[str, ...] = (
+    "This video",
+    "The video",
+    "This clip",
+    "The clip",
+    "The animation",
+    "This image",
+    "The image",
+    "This picture",
+    "The picture",
+)
+
+COMMON_CONTINUATION_WORDS: tuple[str, ...] = (
+    "shows",
+    "depicts",
+    "features",
+    "captures",
+    "highlights",
+    "introduces",
+    "presents",
+)
+
+COMMON_LLM_START_PHRASES: tuple[str, ...] = (
+    "In the video,",
+    "In this video,",
+    "In this video clip,",
+    "In the clip,",
+    "Caption:",
+    *(
+        f"{beginning} {continuation}"
+        for beginning in COMMON_BEGINNING_PHRASES
+        for continuation in COMMON_CONTINUATION_WORDS
+    ),
+)
+
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=True,
+    help="Process text captions and save embeddings for video generation training.",
+)
+
+
+class CaptionsDataset(Dataset):
+    """
+    Dataset for processing text captions only.
+    This dataset is designed for caption preprocessing workflows where you only need
+    to process text without loading videos. Useful for:
+    - Precomputing text embeddings
+    - Caption cleaning and preprocessing
+    - Text-only preprocessing pipelines
+    """
+
+    def __init__(
+        self,
+        dataset_file: str | Path,
+        caption_column: str,
+        media_column: str = "media_path",
+        lora_trigger: str | None = None,
+        remove_llm_prefixes: bool = False,
+    ) -> None:
+        """
+        Initialize the captions dataset.
+        Args:
+            dataset_file: Path to CSV/JSON/JSONL metadata file
+            caption_column: Column name for captions in the metadata file
+            media_column: Column name for media paths (used for output naming)
+            lora_trigger: Optional trigger word to prepend to each caption
+            remove_llm_prefixes: Whether to remove common LLM-generated prefixes
+        """
+        super().__init__()
+
+        self.dataset_file = Path(dataset_file)
+        self.caption_column = caption_column
+        self.media_column = media_column
+        self.lora_trigger = f"{lora_trigger.strip()} " if lora_trigger else ""
+
+        # Load captions with their corresponding output embedding paths
+        self.caption_data = self._load_caption_data()
+
+        # Convert to lists for indexing
+        self.output_paths = list(self.caption_data.keys())
+        self.prompts = list(self.caption_data.values())
+
+        # Clean LLM start phrases if requested
+        if remove_llm_prefixes:
+            self._clean_llm_prefixes()
+
+    def __len__(self) -> int:
+        return len(self.prompts)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        """Get a single caption with optional trigger word prepended and output path."""
+        prompt = self.lora_trigger + self.prompts[index]
+        return {
+            "prompt": prompt,
+            "output_path": self.output_paths[index],
+            "index": index,
+        }
+
+    def _load_caption_data(self) -> dict[str, str]:
+        """Load captions and compute their output embedding paths."""
+        if self.dataset_file.suffix == ".csv":
+            return self._load_caption_data_from_csv()
+        elif self.dataset_file.suffix == ".json":
+            return self._load_caption_data_from_json()
+        elif self.dataset_file.suffix == ".jsonl":
+            return self._load_caption_data_from_jsonl()
+        else:
+            raise ValueError("Expected `dataset_file` to be a path to a CSV, JSON, or JSONL file.")
+
+    def _embedding_output_path(self, media_path: Path) -> str:
+        """Output `.pt` path relative to the dataset dir; mirrors `process_videos._output_relative`
+        so caption keys match video/audio latent keys (and absolute paths don't escape output_dir)."""
+        data_root = self.dataset_file.parent
+        resolved = data_root / media_path  # pathlib: an absolute media_path overrides data_root
+        try:
+            rel = resolved.relative_to(data_root)
+        except ValueError:
+            rel = Path(*resolved.parts[1:]) if resolved.is_absolute() else resolved
+        return str(rel.with_suffix(".pt"))
+
+    def _load_caption_data_from_csv(self) -> dict[str, str]:
+        """Load captions from a CSV file and compute output embedding paths."""
+        df = pd.read_csv(self.dataset_file)
+
+        if self.caption_column not in df.columns:
+            raise ValueError(f"Column '{self.caption_column}' not found in CSV file")
+        if self.media_column not in df.columns:
+            raise ValueError(f"Column '{self.media_column}' not found in CSV file")
+
+        caption_data = {}
+        for _, row in df.iterrows():
+            media_path = Path(row[self.media_column].strip())
+            output_path = self._embedding_output_path(media_path)
+            caption_data[output_path] = row[self.caption_column]
+
+        return caption_data
+
+    def _load_caption_data_from_json(self) -> dict[str, str]:
+        """Load captions from a JSON file and compute output embedding paths."""
+        with open(self.dataset_file, "r", encoding="utf-8") as file:
+            data = json.load(file)
+
+        if not isinstance(data, list):
+            raise ValueError("JSON file must contain a list of objects")
+
+        caption_data = {}
+        for entry in data:
+            if self.caption_column not in entry:
+                raise ValueError(f"Key '{self.caption_column}' not found in JSON entry: {entry}")
+            if self.media_column not in entry:
+                raise ValueError(f"Key '{self.media_column}' not found in JSON entry: {entry}")
+
+            media_path = Path(entry[self.media_column].strip())
+            output_path = self._embedding_output_path(media_path)
+            caption_data[output_path] = entry[self.caption_column]
+
+        return caption_data
+
+    def _load_caption_data_from_jsonl(self) -> dict[str, str]:
+        """Load captions from a JSONL file and compute output embedding paths."""
+        caption_data = {}
+        with open(self.dataset_file, "r", encoding="utf-8") as file:
+            for line in file:
+                entry = json.loads(line)
+                if self.caption_column not in entry:
+                    raise ValueError(f"Key '{self.caption_column}' not found in JSONL entry: {entry}")
+                if self.media_column not in entry:
+                    raise ValueError(f"Key '{self.media_column}' not found in JSONL entry: {entry}")
+
+                media_path = Path(entry[self.media_column].strip())
+                output_path = self._embedding_output_path(media_path)
+                caption_data[output_path] = entry[self.caption_column]
+
+        return caption_data
+
+    def _clean_llm_prefixes(self) -> None:
+        """Remove common LLM-generated prefixes from captions."""
+        for i in range(len(self.prompts)):
+            self.prompts[i] = self.prompts[i].strip()
+            for phrase in COMMON_LLM_START_PHRASES:
+                if self.prompts[i].startswith(phrase):
+                    self.prompts[i] = self.prompts[i].removeprefix(phrase).strip()
+                    break
+
+
+def compute_captions_embeddings(  # noqa: PLR0913
+    dataset_file: str | Path,
+    output_dir: str,
+    model_path: str,
+    text_encoder_path: str,
+    caption_column: str = "caption",
+    media_column: str = "media_path",
+    lora_trigger: str | None = None,
+    remove_llm_prefixes: bool = False,
+    batch_size: int = 8,
+    device: str = "cuda",
+    load_in_8bit: bool = False,
+    overwrite: bool = False,
+) -> None:
+    """
+    Process captions and save text embeddings.
+    Under ``accelerate launch``, each process handles an interleaved shard of
+    the dataset (rank/world read from ``accelerate.PartialState``). Already-
+    computed ``.pt`` outputs are skipped unless ``overwrite=True``; writes are
+    atomic so an interrupted run is safe to resume.
+    Args:
+        dataset_file: Path to metadata file (CSV/JSON/JSONL) containing captions and media paths
+        output_dir: Directory to save embeddings
+        model_path: Path to LTX-2 checkpoint (.safetensors)
+        text_encoder_path: Path to Gemma text encoder directory
+        caption_column: Column name containing captions in the metadata file
+        media_column: Column name containing media paths (used for output naming)
+        lora_trigger: Optional trigger word to prepend to each caption
+        remove_llm_prefixes: Whether to remove common LLM-generated prefixes
+        batch_size: Batch size for processing
+        device: Device to use for computation
+        load_in_8bit: Whether to load the Gemma text encoder in 8-bit precision
+        overwrite: Re-encode every item even if its output exists. Use when rerunning with
+            changed parameters (different text encoder, lora_trigger, etc.) so stale
+            outputs are replaced.
+    """
+    console = Console()
+
+    dataset = CaptionsDataset(
+        dataset_file=dataset_file,
+        caption_column=caption_column,
+        media_column=media_column,
+        lora_trigger=lora_trigger,
+        remove_llm_prefixes=remove_llm_prefixes,
+    )
+    logger.info(f"Loaded {len(dataset):,} captions")
+
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    # TODO(batch-tokenization): The current Gemma tokenizer doesn't support batched tokenization.
+    if batch_size > 1:
+        logger.warning(
+            "Batch size greater than 1 is not currently supported with the Gemma tokenizer. "
+            "Overriding batch_size to 1. This will be fixed in a future update."
+        )
+        batch_size = 1
+
+    dataloader = _build_sharded_dataloader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=2,
+        is_done=lambda idx: (output_path / dataset.output_paths[idx]).is_file(),
+        overwrite=overwrite,
+    )
+    if dataloader is None:
+        return
+
+    # Load text encoder and embeddings processor
+    with console.status("[bold]Loading Gemma text encoder...", spinner="dots"):
+        text_encoder = load_text_encoder(
+            text_encoder_path,
+            device=device,
+            dtype=torch.bfloat16,
+            load_in_8bit=load_in_8bit,
+        )
+        embeddings_processor = load_embeddings_processor(
+            model_path,
+            device=device,
+            dtype=torch.bfloat16,
+        )
+
+    logger.info("Text encoder and embeddings processor loaded successfully")
+    logger.info(f"Processing captions in {len(dataloader):,} batches...")
+
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        TimeRemainingColumn(),
+        console=console,
+    ) as progress:
+        task = progress.add_task("Processing captions", total=len(dataloader))
+        for batch in dataloader:
+            # Encode prompts using text_encoder.encode() + feature_extractor
+            # (returns video/audio features before connector).
+            # The connector is applied during training via embeddings_processor
+            with torch.inference_mode():
+                # TODO(batch-tokenization): When tokenizer supports batching, encode all prompts at once.
+                # For now, process one at a time:
+                for i in range(len(batch["prompt"])):
+                    encoded = text_encoder.encode([batch["prompt"][i]], padding_side="left")
+                    hidden_states, prompt_attention_mask = encoded[0]
+                    video_prompt_embeds, audio_prompt_embeds = embeddings_processor.feature_extractor(
+                        hidden_states, prompt_attention_mask, "left"
+                    )
+
+                    output_rel_path = Path(batch["output_path"][i])
+
+                    # Create output directory maintaining structure
+                    output_dir_path = output_path / output_rel_path.parent
+                    output_dir_path.mkdir(parents=True, exist_ok=True)
+
+                    embedding_data = {
+                        "video_prompt_embeds": video_prompt_embeds[0].cpu().contiguous(),
+                        "prompt_attention_mask": prompt_attention_mask[0].cpu().contiguous(),
+                    }
+                    if audio_prompt_embeds is not None:
+                        embedding_data["audio_prompt_embeds"] = audio_prompt_embeds[0].cpu().contiguous()
+
+                    output_file = output_path / output_rel_path
+                    _atomic_save(embedding_data, output_file)
+
+            progress.advance(task)
+
+    logger.info(f"Processed {len(dataloader.dataset):,} captions -> {output_path}")  # type: ignore[arg-type]
+
+
+def _atomic_save(data: Any, out: Path) -> None:  # noqa: ANN401
+    """Save to ``out`` atomically via per-PID temp file + replace.
+    Crash mid-write leaves an orphan ``.tmp.<pid>`` file that the skip logic
+    ignores. The per-PID suffix makes concurrent writes from multiple ranks
+    collision-free.
+    """
+    tmp = out.with_suffix(f"{out.suffix}.tmp.{os.getpid()}")
+    torch.save(data, tmp)
+    tmp.replace(out)
+
+
+def _build_sharded_dataloader(
+    dataset: Dataset,
+    *,
+    batch_size: int,
+    num_workers: int,
+    is_done: Callable[[int], bool],
+    overwrite: bool,
+) -> DataLoader | None:
+    """Return a DataLoader over this rank's interleaved shard of ``dataset``.
+    When ``overwrite`` is False, items whose outputs already exist (per
+    ``is_done``) are filtered out. Returns ``None`` if this rank has nothing
+    to do, so the caller can early-return without loading any models.
+    """
+    state = PartialState()
+    todo = [i for i in range(state.process_index, len(dataset), state.num_processes) if overwrite or not is_done(i)]
+    if not todo:
+        logger.info(f"Rank {state.process_index}/{state.num_processes}: nothing to do")
+        return None
+    logger.info(f"Rank {state.process_index}/{state.num_processes}: processing {len(todo):,} of {len(dataset):,} items")
+    return DataLoader(Subset(dataset, todo), batch_size=batch_size, shuffle=False, num_workers=num_workers)
+
+
+@app.command()
+def main(  # noqa: PLR0913
+    dataset_file: str = typer.Argument(
+        ...,
+        help="Path to metadata file (CSV/JSON/JSONL) containing captions and media paths",
+    ),
+    output_dir: str = typer.Option(
+        ...,
+        help="Output directory to save text embeddings",
+    ),
+    model_path: str = typer.Option(
+        ...,
+        help="Path to LTX-2 checkpoint (.safetensors file)",
+    ),
+    text_encoder_path: str = typer.Option(
+        ...,
+        help="Path to Gemma text encoder directory",
+    ),
+    caption_column: str = typer.Option(
+        default="caption",
+        help="Column name containing captions in the dataset JSON/JSONL/CSV file",
+    ),
+    media_column: str = typer.Option(
+        default="media_path",
+        help="Column name in the dataset JSON/JSONL/CSV file containing media paths "
+        "(used for output file naming and folder structure)",
+    ),
+    batch_size: int = typer.Option(
+        default=8,
+        help="Batch size for processing",
+    ),
+    device: str = typer.Option(
+        default="cuda",
+        help="Device to use for computation",
+    ),
+    lora_trigger: str | None = typer.Option(
+        default=None,
+        help="Optional trigger word to prepend to each caption (activates the LoRA during inference)",
+    ),
+    remove_llm_prefixes: bool = typer.Option(
+        default=False,
+        help="Remove common LLM-generated prefixes from captions",
+    ),
+    load_text_encoder_in_8bit: bool = typer.Option(
+        default=False,
+        help="Load the Gemma text encoder in 8-bit precision to save GPU memory (requires bitsandbytes)",
+    ),
+    overwrite: bool = typer.Option(
+        default=False,
+        help="Re-encode every caption even if its output exists. Use when rerunning with "
+        "changed parameters (different text encoder, lora_trigger, etc.) so stale outputs are replaced.",
+    ),
+) -> None:
+    """Process text captions and save embeddings for video generation training.
+    For multi-GPU preprocessing, invoke under ``accelerate launch`` - each process
+    will handle an interleaved shard of the dataset.
+    This script processes captions from metadata files and saves text embeddings
+    that can be used for training video generation models. The output embeddings
+    will maintain the same folder structure and naming as the corresponding media files.
+    Note: This script is designed for LTX-2 models which use the Gemma text encoder.
+    Examples:
+        # Process captions with LTX-2 model
+        python scripts/process_captions.py dataset.json --output-dir ./embeddings \\
+            --model-path /path/to/ltx2_checkpoint.safetensors \\
+            --text-encoder-path /path/to/gemma
+        # Add a trigger word for LoRA training
+        python scripts/process_captions.py dataset.json --output-dir ./embeddings \\
+            --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
+            --lora-trigger "mytoken"
+        # Remove LLM-generated prefixes from captions
+        python scripts/process_captions.py dataset.json --output-dir ./embeddings \\
+            --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma \\
+            --remove-llm-prefixes
+    """
+
+    # Validate dataset file
+    if not Path(dataset_file).is_file():
+        raise typer.BadParameter(f"Dataset file not found: {dataset_file}")
+
+    if lora_trigger:
+        logger.info(f'LoRA trigger word "{lora_trigger}" will be prepended to all captions')
+
+    # Process embeddings
+    compute_captions_embeddings(
+        dataset_file=dataset_file,
+        output_dir=output_dir,
+        model_path=model_path,
+        text_encoder_path=text_encoder_path,
+        caption_column=caption_column,
+        media_column=media_column,
+        lora_trigger=lora_trigger,
+        remove_llm_prefixes=remove_llm_prefixes,
+        batch_size=batch_size,
+        device=device,
+        load_in_8bit=load_text_encoder_in_8bit,
+        overwrite=overwrite,
+    )
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/scripts/process_dataset.py

@@ -0,0 +1,432 @@
+#!/usr/bin/env python3
+
+"""
+Preprocess a media dataset for LTX-2 training.
+Automatically detects dataset columns and processes each according to a convention table.
+Column names determine what gets encoded and where outputs go — no per-role CLI flags needed.
+Convention table:
+    video           → Video VAE    → latents/
+    audio           → Audio VAE    → audio_latents/
+    reference_video → Video VAE    → reference_latents/
+    reference_audio → Audio VAE    → reference_audio_latents/
+    video_mask      → (downsample) → video_masks/
+    audio_mask      → (downsample) → audio_masks/
+    caption         → Text encoder → conditions/
+Legacy aliases: media_path → video, ref_media_path → reference_video
+Basic usage:
+    python scripts/process_dataset.py /path/to/dataset.json --resolution-buckets 768x768x49 \\
+        --model-path /path/to/ltx2.safetensors --text-encoder-path /path/to/gemma
+"""
+
+from pathlib import Path
+
+import typer
+from decode_latents import LatentsDecoder
+from process_captions import compute_captions_embeddings
+from process_videos import (
+    compute_audio_latents,
+    compute_audio_masks,
+    compute_latents,
+    compute_scaled_resolution_buckets,
+    compute_video_masks,
+    detect_dataset_columns,
+    parse_resolution_buckets,
+)
+from rich.console import Console
+
+from ltx_trainer import logger
+from ltx_trainer.gpu_utils import free_gpu_memory_context
+
+console = Console()
+
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=True,
+    help="Preprocess a media dataset for LTX-2 training. "
+    "Automatically detects columns (video, audio, reference_video, reference_audio, caption) "
+    "and processes each with the appropriate encoder.",
+)
+
+_KNOWN_ROLES = {"video", "audio", "reference_video", "reference_audio", "video_mask", "audio_mask", "caption"}
+_LEGACY_ALIASES = {"media_path": "video", "ref_media_path": "reference_video"}
+
+
+def preprocess_dataset(  # noqa: PLR0912, PLR0913, PLR0915
+    dataset_file: str,
+    resolution_buckets: list[tuple[int, int, int]] | None,
+    model_path: str,
+    text_encoder_path: str,
+    device: str,
+    output_dir: str | None = None,
+    video_column: str | None = None,
+    caption_column: str | None = None,
+    batch_size: int = 1,
+    lora_trigger: str | None = None,
+    vae_tiling: bool = False,
+    decode: bool = False,
+    remove_llm_prefixes: bool = False,
+    reference_downscale_factor: int = 1,
+    reference_temporal_scale_factor: int = 1,
+    skip_audio: bool = False,
+    audio_durations: list[float] | None = None,
+    load_text_encoder_in_8bit: bool = False,
+    overwrite: bool = False,
+) -> None:
+    """Run the preprocessing pipeline with convention-based column detection."""
+    _validate_dataset_file(dataset_file)
+
+    # Detect columns and resolve roles
+    dataset_columns = detect_dataset_columns(dataset_file)
+    roles = _resolve_columns(dataset_columns, video_column, caption_column)
+
+    # Log detected roles
+    for role, col in sorted(roles.items()):
+        alias_note = f" (alias for '{role}')" if col != role else ""
+        logger.info(f"Detected column '{col}'{alias_note} → {role}")
+
+    # Validate: need at least caption
+    if "caption" not in roles:
+        raise ValueError(
+            f"No caption column found. Dataset has columns: {dataset_columns}. "
+            f"Expected 'caption' or use --caption-column to specify."
+        )
+
+    # Validate: need video or audio
+    has_video = "video" in roles
+    has_audio = "audio" in roles
+    if not has_video and not has_audio:
+        raise ValueError(
+            f"No media column found. Dataset has columns: {dataset_columns}. "
+            f"Expected 'video', 'audio', or 'media_path' (legacy)."
+        )
+
+    # Validate: video modes need resolution buckets
+    if has_video and not resolution_buckets:
+        raise ValueError("--resolution-buckets is required when the dataset has a video column.")
+
+    output_base = Path(output_dir) if output_dir else Path(dataset_file).parent / ".precomputed"
+
+    if lora_trigger:
+        logger.info(f'LoRA trigger word "{lora_trigger}" will be prepended to all captions')
+
+    # --- Phase 1: Text encoder ---
+    with free_gpu_memory_context():
+        compute_captions_embeddings(
+            dataset_file=dataset_file,
+            output_dir=str(output_base / "conditions"),
+            model_path=model_path,
+            text_encoder_path=text_encoder_path,
+            caption_column=roles["caption"],
+            media_column=roles.get("video") or roles.get("audio") or roles["caption"],
+            lora_trigger=lora_trigger,
+            remove_llm_prefixes=remove_llm_prefixes,
+            batch_size=batch_size,
+            device=device,
+            load_in_8bit=load_text_encoder_in_8bit,
+            overwrite=overwrite,
+        )
+
+    # --- Phase 2: Video VAE (video, reference_video) ---
+    if has_video and resolution_buckets:
+        # Determine if audio should be auto-extracted from video files
+        auto_audio = not skip_audio and "audio" not in roles
+
+        audio_latents_dir = str(output_base / "audio_latents") if auto_audio else None
+        if auto_audio:
+            logger.info("Audio will be auto-extracted from video files (use --skip-audio to disable)")
+
+        with free_gpu_memory_context():
+            compute_latents(
+                dataset_file=dataset_file,
+                video_column=roles["video"],
+                resolution_buckets=resolution_buckets,
+                output_dir=str(output_base / "latents"),
+                model_path=model_path,
+                batch_size=batch_size,
+                device=device,
+                vae_tiling=vae_tiling,
+                with_audio=auto_audio,
+                audio_output_dir=audio_latents_dir,
+                overwrite=overwrite,
+            )
+
+        # Process reference video if present
+        if "reference_video" in roles:
+            if reference_downscale_factor > 1 and len(resolution_buckets) > 1:
+                raise ValueError(
+                    "When using --reference-downscale-factor > 1, only a single resolution bucket is supported."
+                )
+            if reference_temporal_scale_factor > 1 and len(resolution_buckets) > 1:
+                raise ValueError(
+                    "When using --reference-temporal-scale-factor > 1, only a single resolution bucket is supported."
+                )
+
+            reference_buckets = compute_scaled_resolution_buckets(resolution_buckets, reference_downscale_factor)
+            if reference_downscale_factor > 1:
+                logger.info(f"Processing reference videos at 1/{reference_downscale_factor} resolution...")
+            if reference_temporal_scale_factor > 1:
+                logger.info(
+                    f"Temporally subsampling reference videos by {reference_temporal_scale_factor}x "
+                    f"(VAE-aligned pattern)..."
+                )
+
+            with free_gpu_memory_context():
+                compute_latents(
+                    dataset_file=dataset_file,
+                    main_media_column=roles["video"],
+                    video_column=roles["reference_video"],
+                    resolution_buckets=reference_buckets,
+                    output_dir=str(output_base / "reference_latents"),
+                    model_path=model_path,
+                    batch_size=batch_size,
+                    device=device,
+                    vae_tiling=vae_tiling,
+                    overwrite=overwrite,
+                    temporal_subsample_factor=reference_temporal_scale_factor,
+                )
+
+    # --- Phase 2b: Masks (video_mask, audio_mask) — processed after video latents for alignment ---
+    if "video_mask" in roles and has_video:
+        compute_video_masks(
+            dataset_file=dataset_file,
+            mask_column=roles["video_mask"],
+            latents_dir=str(output_base / "latents"),
+            output_dir=str(output_base / "video_masks"),
+            main_media_column=roles["video"],
+        )
+
+    # --- Phase 3: Audio VAE (audio, reference_audio) ---
+    audio_roles_to_process = [
+        ("audio", "audio_latents"),
+        ("reference_audio", "reference_audio_latents"),
+    ]
+    active_audio_roles = [(role, subdir) for role, subdir in audio_roles_to_process if role in roles]
+
+    if active_audio_roles:
+        # Determine audio duration constraint: video bucket → max_duration, or explicit buckets
+        max_audio_duration = None
+        audio_duration_buckets = None
+        if has_video and resolution_buckets:
+            max_audio_duration = max(f for f, _h, _w in resolution_buckets) / 25.0
+        elif audio_durations:
+            audio_duration_buckets = audio_durations
+
+        for role, output_subdir in active_audio_roles:
+            with free_gpu_memory_context():
+                compute_audio_latents(
+                    dataset_file=dataset_file,
+                    audio_column=roles[role],
+                    output_dir=str(output_base / output_subdir),
+                    model_path=model_path,
+                    main_media_column=roles.get("video"),
+                    max_duration=max_audio_duration,
+                    duration_buckets=audio_duration_buckets,
+                    device=device,
+                    overwrite=overwrite,
+                )
+
+    # --- Phase 4: Audio masks (after audio latents exist for temporal alignment) ---
+    if "audio_mask" in roles:
+        audio_latents_source = output_base / "audio_latents"
+        if audio_latents_source.exists():
+            compute_audio_masks(
+                dataset_file=dataset_file,
+                mask_column=roles["audio_mask"],
+                audio_latents_dir=str(audio_latents_source),
+                output_dir=str(output_base / "audio_masks"),
+                main_media_column=roles.get("video") or roles.get("audio"),
+            )
+        else:
+            logger.warning("audio_mask column found but no audio_latents/ — run with audio first")
+
+    # --- Decode for verification ---
+    if decode:
+        logger.info("Decoding latents for verification...")
+        decoder = LatentsDecoder(model_path=model_path, device=device, vae_tiling=vae_tiling, with_audio=has_audio)
+        if has_video:
+            decoder.decode(output_base / "latents", output_base / "decoded_videos")
+        if "reference_video" in roles and (output_base / "reference_latents").exists():
+            decoder.decode(output_base / "reference_latents", output_base / "decoded_reference_videos")
+
+    # --- Summary ---
+    logger.info(f"Dataset preprocessing complete! Results saved to {output_base}")
+    produced = [d.name for d in output_base.iterdir() if d.is_dir() and not d.name.startswith("decoded")]
+    logger.info(f"Output directories: {', '.join(sorted(produced))}")
+
+
+def _validate_dataset_file(dataset_path: str) -> None:
+    """Validate that the dataset file exists and has the correct format."""
+    dataset_file = Path(dataset_path)
+    if not dataset_file.exists():
+        raise FileNotFoundError(f"Dataset file does not exist: {dataset_file}")
+    if not dataset_file.is_file():
+        raise ValueError(f"Dataset path must be a file, not a directory: {dataset_file}")
+    if dataset_file.suffix.lower() not in [".csv", ".json", ".jsonl"]:
+        raise ValueError(f"Dataset file must be CSV, JSON, or JSONL format: {dataset_file}")
+
+
+def _resolve_columns(
+    dataset_columns: set[str],
+    video_column_override: str | None = None,
+    caption_column_override: str | None = None,
+) -> dict[str, str]:
+    """Map canonical role names to actual dataset column names.
+    Returns a dict of role → column_name for recognized roles found in the dataset.
+    """
+    roles: dict[str, str] = {}
+    for col in dataset_columns:
+        role = _LEGACY_ALIASES.get(col, col)
+        if role in _KNOWN_ROLES:
+            roles[role] = col
+
+    if video_column_override and video_column_override in dataset_columns:
+        roles["video"] = video_column_override
+    if caption_column_override and caption_column_override in dataset_columns:
+        roles["caption"] = caption_column_override
+
+    return roles
+
+
+@app.command()
+def main(  # noqa: PLR0913
+    dataset_path: str = typer.Argument(
+        ...,
+        help="Path to metadata file (CSV/JSON/JSONL) with columns matching the convention table",
+    ),
+    resolution_buckets: str | None = typer.Option(
+        default=None,
+        help='Resolution buckets in format "WxHxF;WxHxF;..." (e.g. "768x768x25"). '
+        "Required when dataset has a video column.",
+    ),
+    model_path: str = typer.Option(
+        ...,
+        help="Path to LTX-2 checkpoint (.safetensors file)",
+    ),
+    text_encoder_path: str = typer.Option(
+        ...,
+        help="Path to Gemma text encoder directory",
+    ),
+    caption_column: str | None = typer.Option(
+        default=None,
+        help="Override: treat this column as 'caption' (default: auto-detect 'caption')",
+    ),
+    video_column: str | None = typer.Option(
+        default=None,
+        help="Override: treat this column as 'video' (default: auto-detect 'video' or 'media_path')",
+    ),
+    batch_size: int = typer.Option(
+        default=1,
+        help="Batch size for preprocessing",
+    ),
+    device: str = typer.Option(
+        default="cuda",
+        help="Device to use for computation",
+    ),
+    vae_tiling: bool = typer.Option(
+        default=False,
+        help="Enable VAE tiling for larger video resolutions",
+    ),
+    output_dir: str | None = typer.Option(
+        default=None,
+        help="Output directory (defaults to .precomputed in dataset directory)",
+    ),
+    lora_trigger: str | None = typer.Option(
+        default=None,
+        help="Optional trigger word to prepend to each caption",
+    ),
+    decode: bool = typer.Option(
+        default=False,
+        help="Decode and save latents after encoding for verification",
+    ),
+    remove_llm_prefixes: bool = typer.Option(
+        default=False,
+        help="Remove LLM prefixes from captions",
+    ),
+    skip_audio: bool = typer.Option(
+        default=False,
+        help="Don't extract audio from video files (audio extraction is on by default)",
+    ),
+    audio_durations: str | None = typer.Option(
+        default=None,
+        help='Audio duration buckets in seconds for audio-only datasets (e.g. "2.0;4.0;8.0"). '
+        "When set, audio files are trimmed to the best matching duration. "
+        "Not needed when a video column is present (audio duration derived from video bucket).",
+    ),
+    with_audio: bool = typer.Option(
+        default=False,
+        hidden=True,
+        help="[DEPRECATED: audio is now on by default, use --skip-audio to disable]",
+    ),
+    load_text_encoder_in_8bit: bool = typer.Option(
+        default=False,
+        help="Load the Gemma text encoder in 8-bit precision to save GPU memory",
+    ),
+    reference_downscale_factor: int = typer.Option(
+        default=1,
+        help="Downscale factor for reference video resolution (e.g., 2 = half resolution for IC-LoRA)",
+    ),
+    reference_temporal_scale_factor: int = typer.Option(
+        default=1,
+        help="Temporal subsampling factor for reference videos (e.g., 2 = half frame rate, "
+        "VAE-aligned: keeps frame 0, then every Nth frame from frame 1 onwards)",
+    ),
+    overwrite: bool = typer.Option(
+        default=False,
+        help="Re-compute every item even if its output exists. Use when rerunning with "
+        "changed parameters (different model, resolution, etc.) so stale outputs are replaced.",
+    ),
+) -> None:
+    """Preprocess a media dataset for LTX-2 training.
+    See module docstring for the convention table. Audio is auto-extracted from
+    video files by default — use --skip-audio to disable.
+    For multi-GPU preprocessing, invoke under ``accelerate launch`` -- each process
+    will handle an interleaved shard of the dataset.
+    """
+    # Handle deprecated --with-audio flag
+    if with_audio:
+        logger.warning(
+            "--with-audio is deprecated. Audio extraction is now on by default. Use --skip-audio to disable."
+        )
+
+    parsed_buckets = parse_resolution_buckets(resolution_buckets) if resolution_buckets else None
+
+    if parsed_buckets and len(parsed_buckets) > 1:
+        logger.warning("Using multiple resolution buckets. Training batch size must be 1.")
+
+    if reference_downscale_factor < 1:
+        raise typer.BadParameter("--reference-downscale-factor must be >= 1")
+
+    if reference_temporal_scale_factor < 1:
+        raise typer.BadParameter("--reference-temporal-scale-factor must be >= 1")
+
+    parsed_audio_durations = None
+    if audio_durations:
+        parsed_audio_durations = [float(d) for d in audio_durations.split(";")]
+        if any(d <= 0 for d in parsed_audio_durations):
+            raise typer.BadParameter("All audio durations must be positive")
+
+    preprocess_dataset(
+        dataset_file=dataset_path,
+        resolution_buckets=parsed_buckets,
+        model_path=model_path,
+        text_encoder_path=text_encoder_path,
+        device=device,
+        output_dir=output_dir,
+        video_column=video_column,
+        caption_column=caption_column,
+        batch_size=batch_size,
+        lora_trigger=lora_trigger,
+        vae_tiling=vae_tiling,
+        decode=decode,
+        remove_llm_prefixes=remove_llm_prefixes,
+        reference_downscale_factor=reference_downscale_factor,
+        reference_temporal_scale_factor=reference_temporal_scale_factor,
+        skip_audio=skip_audio,
+        audio_durations=parsed_audio_durations,
+        load_text_encoder_in_8bit=load_text_encoder_in_8bit,
+        overwrite=overwrite,
+    )
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/scripts/process_videos.py

@@ -0,0 +1,1474 @@
+#!/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()

packages/ltx-trainer/scripts/serve_captioner.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+
+"""Launch a vLLM server for Qwen3-Omni captioning.
+Runs the actual server via ``uvx`` so that vLLM and its CUDA-tied
+dependencies live in their own isolated environment (no impact on this
+package's dependency tree).
+The captioning script (``caption_videos.py``) talks to the server over its
+OpenAI-compatible HTTP API. Once the server is up it stays loaded across
+captioning runs; no per-script model warmup cost.
+Typical usage::
+    # Default: dynamic FP8 quantization, listen on 127.0.0.1:8001
+    uv run python scripts/serve_captioner.py
+    # Just print the chosen `uvx vllm serve ...` command without running it
+    uv run python scripts/serve_captioner.py --print-cmd
+    # Use full bf16 on a GPU with >= 66 GiB free VRAM (slightly more reliable
+    # numerics but 2x the weight memory)
+    uv run python scripts/serve_captioner.py --quantization bf16
+    # Use a different port or expose on all interfaces
+    uv run python scripts/serve_captioner.py --port 9000 --host 0.0.0.0
+"""
+
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+console = Console()
+
+# Model identifier we serve. The captioner client must use the same string.
+DEFAULT_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
+
+# Pinned vLLM version known to support Qwen3-Omni on CUDA 12.x.
+# vLLM 0.20+ requires CUDA 13. Update both as the environment evolves.
+# The ``[audio]`` extra is required for Qwen3-Omni to decode audio at all.
+DEFAULT_VLLM_SPEC = "vllm[audio]==0.11.2"
+
+# Approximate disk needed for the model download (HF cache structure).
+MODEL_DISK_GIB = 65.0
+
+
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=False,
+    help="Launch a local vLLM server for Qwen3-Omni captioning.",
+)
+
+
+def _query_disk_free_gib(path: Path) -> float:
+    return shutil.disk_usage(str(path)).free / 1024**3
+
+
+def _build_vllm_args(
+    *,
+    model: str,
+    host: str,
+    port: int,
+    quantization: str,
+    max_model_len: int,
+    gpu_memory_utilization: float,
+    extra_args: list[str],
+) -> list[str]:
+    """Construct the `vllm serve ...` argv."""
+    args = [
+        "vllm",
+        "serve",
+        model,
+        "--host",
+        host,
+        "--port",
+        str(port),
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        str(max_model_len),
+        "--gpu-memory-utilization",
+        str(gpu_memory_utilization),
+        # Let the server accept ``file://`` URLs pointing at local videos.
+        "--allowed-local-media-path",
+        "/",
+        # The model is a multimodal MoE; cap each input to one of each
+        # modality to match what our captioner sends.
+        "--limit-mm-per-prompt",
+        '{"image": 1, "video": 1, "audio": 1}',
+        # Small concurrent-sequence cap so KV cache headroom isn't fragmented.
+        "--max-num-seqs",
+        "4",
+    ]
+    if quantization == "fp8":
+        args += ["--quantization", "fp8"]
+    args += extra_args
+    return args
+
+
+@app.command()
+def main(
+    model: str = typer.Option(DEFAULT_MODEL, "--model", help="Model identifier to serve."),
+    host: str = typer.Option("127.0.0.1", "--host", help="Listen address. Use 0.0.0.0 for remote access."),
+    port: int = typer.Option(8001, "--port", help="HTTP port."),
+    quantization: str = typer.Option(
+        "fp8",
+        "--quantization",
+        "-q",
+        help=(
+            "Weight precision. 'fp8' (default, dynamic FP8 -- ~31 GiB weights) is "
+            "the recommended choice; it fits on 40 GiB GPUs and runs at the same "
+            "speed as bf16 on H100. 'bf16' uses ~60 GiB of weights -- pick it if "
+            "you have abundant VRAM and want minimal numerical drift."
+        ),
+    ),
+    max_model_len: int = typer.Option(
+        32768,
+        "--max-model-len",
+        help="Maximum context length the server accepts (must fit input video tokens + max_tokens).",
+    ),
+    gpu_memory_utilization: float = typer.Option(
+        0.9,
+        "--gpu-memory-utilization",
+        help="Fraction of GPU memory vLLM may reserve (model + KV cache).",
+    ),
+    hf_home: Path | None = typer.Option(  # noqa: B008
+        None,
+        "--hf-home",
+        help=(
+            "Override HF_HOME (where the model is downloaded). The model is ~65 GB; "
+            "by default this follows your environment's HF_HOME or HuggingFace's default."
+        ),
+    ),
+    vllm_spec: str = typer.Option(
+        DEFAULT_VLLM_SPEC,
+        "--vllm-spec",
+        help="pip-style spec passed to `uvx --from`. Pin a version that matches your CUDA.",
+    ),
+    print_cmd: bool = typer.Option(
+        False,
+        "--print-cmd",
+        help="Print the chosen command without running it.",
+    ),
+    extra_args: list[str] | None = typer.Argument(  # noqa: B008
+        None,
+        help="Additional args passed through to `vllm serve` after `--`.",
+    ),
+) -> None:
+    """Launch the vLLM server for Qwen3-Omni."""
+    extra = extra_args or []
+
+    if quantization not in ("bf16", "fp8"):
+        console.print(f"[red]--quantization must be 'bf16' or 'fp8'; got {quantization!r}.[/]")
+        raise typer.Exit(code=1)
+
+    # Disk check (only meaningful before first download).
+    cache_root = hf_home or Path(os.environ.get("HF_HOME", str(Path.home() / ".cache" / "huggingface")))
+    cache_root.mkdir(parents=True, exist_ok=True)
+    free_disk = _query_disk_free_gib(cache_root)
+    if free_disk < MODEL_DISK_GIB:
+        console.print(
+            f"[yellow]\u26a0 Only {free_disk:.1f} GiB free on disk under {cache_root} but the "
+            f"model needs ~{MODEL_DISK_GIB:.0f} GiB. Either free up space, set --hf-home "
+            f"to a larger volume, or expect the download to fail mid-way.[/]"
+        )
+
+    vllm_args = _build_vllm_args(
+        model=model,
+        host=host,
+        port=port,
+        quantization=quantization,
+        max_model_len=max_model_len,
+        gpu_memory_utilization=gpu_memory_utilization,
+        extra_args=extra,
+    )
+
+    # Use ``uvx --from vllm==...`` so vLLM lives in its own throwaway venv
+    # (or a cached tool venv). The `--` separates uvx args from the command's.
+    uvx_cmd = ["uvx", "--from", vllm_spec, *vllm_args]
+
+    env = os.environ.copy()
+    # vLLM 0.11.x requires the V0 engine for Qwen3-Omni's multimodal pipeline.
+    env.setdefault("VLLM_USE_V1", "0")
+    if hf_home is not None:
+        env["HF_HOME"] = str(hf_home)
+
+    console.print("\n[bold]Command:[/]")
+    console.print("  " + " ".join(uvx_cmd))
+    if hf_home is not None:
+        console.print(f"  [dim](with HF_HOME={hf_home})[/]")
+
+    if print_cmd:
+        return
+
+    console.print("\n[dim]Launching... (first run downloads the model -- ~5 min on a fast link)[/]\n")
+    try:
+        completed = subprocess.run(uvx_cmd, env=env, check=False)
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Interrupted.[/]")
+        return
+    sys.exit(completed.returncode)
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/scripts/split_scenes.py

@@ -0,0 +1,417 @@
+#!/usr/bin/env python3
+
+"""
+Split video into scenes using PySceneDetect.
+This script provides a command-line interface for splitting videos into scenes using various detection algorithms.
+It supports multiple detection methods, preview image generation, and customizable parameters for fine-tuning
+the scene detection process.
+Basic usage:
+    # Split video using default content-based detection
+    scenes_split.py input.mp4 output_dir/
+    # Save 3 preview images per scene
+    scenes_split.py input.mp4 output_dir/ --save-images 3
+    # Process specific duration and filter short scenes
+    scenes_split.py input.mp4 output_dir/ --duration 60s --filter-shorter-than 2s
+Advanced usage:
+    # Content detection with minimum scene length and frame skip
+    scenes_split.py input.mp4 output_dir/ --detector content --min-scene-length 30 --frame-skip 2
+    # Use adaptive detection with custom detector and detector parameters
+    scenes_split.py input.mp4 output_dir/ --detector adaptive --threshold 3.0 --adaptive-window 10
+"""
+
+from enum import Enum
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import typer
+from scenedetect import (
+    AdaptiveDetector,
+    ContentDetector,
+    HistogramDetector,
+    SceneManager,
+    ThresholdDetector,
+    open_video,
+)
+from scenedetect.frame_timecode import FrameTimecode
+from scenedetect.scene_manager import SceneDetector, write_scene_list_html
+from scenedetect.scene_manager import save_images as save_scene_images
+from scenedetect.stats_manager import StatsManager
+from scenedetect.video_splitter import split_video_ffmpeg
+
+app = typer.Typer(no_args_is_help=True, help="Split video into scenes using PySceneDetect.")
+
+
+class DetectorType(str, Enum):
+    """Available scene detection algorithms."""
+
+    CONTENT = "content"  # Detects fast cuts using HSV color space
+    ADAPTIVE = "adaptive"  # Detects fast two-phase cuts
+    THRESHOLD = "threshold"  # Detects fast cuts/slow fades in from and out to a given threshold level
+    HISTOGRAM = "histogram"  # Detects based on YUV histogram differences in adjacent frames
+
+
+def create_detector(
+    detector_type: DetectorType,
+    threshold: Optional[float] = None,
+    min_scene_len: Optional[int] = None,
+    luma_only: Optional[bool] = None,
+    adaptive_window: Optional[int] = None,
+    fade_bias: Optional[float] = None,
+) -> SceneDetector:
+    """Create a scene detector based on the specified type and parameters.
+    Args:
+        detector_type: Type of detector to create
+        threshold: Detection threshold (meaning varies by detector)
+        min_scene_len: Minimum scene length in frames
+        luma_only: If True, only use brightness for content detection
+        adaptive_window: Window size for adaptive detection
+        fade_bias: Bias for fade in/out detection (-1.0 to 1.0)
+    Note: Parameters set to None will use the detector's built-in default values.
+    Returns:
+        Configured scene detector instance
+    """
+    # Set common arguments
+    kwargs = {}
+    if threshold is not None:
+        kwargs["threshold"] = threshold
+
+    if min_scene_len is not None:
+        kwargs["min_scene_len"] = min_scene_len
+
+    match detector_type:
+        case DetectorType.CONTENT:
+            if luma_only is not None:
+                kwargs["luma_only"] = luma_only
+            return ContentDetector(**kwargs)
+        case DetectorType.ADAPTIVE:
+            if adaptive_window is not None:
+                kwargs["window_width"] = adaptive_window
+            if luma_only is not None:
+                kwargs["luma_only"] = luma_only
+            if "threshold" in kwargs:
+                # Special case for adaptive detector which uses different param name
+                kwargs["adaptive_threshold"] = kwargs.pop("threshold")
+            return AdaptiveDetector(**kwargs)
+        case DetectorType.THRESHOLD:
+            if fade_bias is not None:
+                kwargs["fade_bias"] = fade_bias
+            return ThresholdDetector(**kwargs)
+        case DetectorType.HISTOGRAM:
+            return HistogramDetector(**kwargs)
+        case _:
+            raise ValueError(f"Unknown detector type: {detector_type}")
+
+
+def validate_output_dir(output_dir: str) -> Path:
+    """Validate and create output directory if it doesn't exist.
+    Args:
+        output_dir: Path to the output directory
+    Returns:
+        Path object of the validated output directory
+    """
+    path = Path(output_dir)
+
+    if path.exists() and not path.is_dir():
+        raise typer.BadParameter(f"{output_dir} exists but is not a directory")
+
+    return path
+
+
+def parse_timecode(video: any, time_str: Optional[str]) -> Optional[FrameTimecode]:
+    """Parse a timecode string into a FrameTimecode object.
+    Supports formats:
+    - Frames: '123'
+    - Seconds: '123s' or '123.45s'
+    - Timecode: '00:02:03' or '00:02:03.456'
+    Args:
+        video: Video object to get framerate from
+        time_str: String to parse, or None
+    Returns:
+        FrameTimecode object or None if input is None
+    """
+    if time_str is None:
+        return None
+
+    try:
+        if time_str.endswith("s"):
+            # Seconds format
+            seconds = float(time_str[:-1])
+            return FrameTimecode(timecode=seconds, fps=video.frame_rate)
+        elif ":" in time_str:
+            # Timecode format
+            return FrameTimecode(timecode=time_str, fps=video.frame_rate)
+        else:
+            # Frame number format
+            return FrameTimecode(timecode=int(time_str), fps=video.frame_rate)
+    except ValueError as e:
+        raise typer.BadParameter(
+            f"Invalid timecode format: {time_str}. Use frames (123), "
+            f"seconds (123s/123.45s), or timecode (HH:MM:SS[.nnn])",
+        ) from e
+
+
+def detect_and_split_scenes(  # noqa: PLR0913
+    video_path: str,
+    output_dir: Path,
+    detector_type: DetectorType,
+    threshold: Optional[float] = None,
+    min_scene_len: Optional[int] = None,
+    max_scenes: Optional[int] = None,
+    filter_shorter_than: Optional[str] = None,
+    skip_start: Optional[int] = None,  # noqa: ARG001
+    skip_end: Optional[int] = None,  # noqa: ARG001
+    save_images_per_scene: int = 0,
+    stats_file: Optional[str] = None,
+    luma_only: bool = False,
+    adaptive_window: Optional[int] = None,
+    fade_bias: Optional[float] = None,
+    downscale_factor: Optional[int] = None,
+    frame_skip: int = 0,
+    duration: Optional[str] = None,
+) -> List[Tuple[FrameTimecode, FrameTimecode]]:
+    """Detect and split scenes in a video using the specified parameters.
+    Args:
+        video_path: Path to input video.
+        output_dir: Directory to save output split scenes.
+        detector_type: Type of scene detector to use.
+        threshold: Detection threshold.
+        min_scene_len: Minimum scene length in frames.
+        max_scenes: Maximum number of scenes to detect.
+        filter_shorter_than: Filter out scenes shorter than this duration (frames/seconds/timecode)
+        skip_start: Number of frames to skip at start.
+        skip_end: Number of frames to skip at end.
+        save_images_per_scene: Number of images to save per scene (0 to disable).
+        stats_file: Path to save detection statistics (optional).
+        luma_only: Only use brightness for content detection.
+        adaptive_window: Window size for adaptive detection.
+        fade_bias: Bias for fade detection (-1.0 to 1.0).
+        downscale_factor: Factor to downscale frames by during detection.
+        frame_skip: Number of frames to skip (i.e. process every 1 in N+1 frames,
+            where N is frame_skip, processing only 1/N+1 percent of the video,
+            speeding up the detection time at the expense of accuracy).
+            frame_skip must be 0 (the default) when using a StatsManager.
+        duration: How much of the video to process from start position.
+            Can be specified as frames (123), seconds (123s/123.45s),
+            or timecode (HH:MM:SS[.nnn]).
+    Returns:
+        List of detected scenes as (start, end) FrameTimecode pairs.
+    """
+    # Create video stream
+    video = open_video(video_path, backend="opencv")
+
+    # Parse duration if specified
+    duration_tc = parse_timecode(video, duration)
+
+    # Parse filter_shorter_than if specified
+    filter_shorter_than_tc = parse_timecode(video, filter_shorter_than)
+
+    # Initialize scene manager with optional stats manager
+    stats_manager = StatsManager() if stats_file else None
+    scene_manager = SceneManager(stats_manager)
+
+    # Configure scene manager
+    if downscale_factor:
+        scene_manager.auto_downscale = False
+        scene_manager.downscale = downscale_factor
+
+    # Create and add detector
+    detector = create_detector(
+        detector_type=detector_type,
+        threshold=threshold,
+        min_scene_len=min_scene_len,
+        luma_only=luma_only,
+        adaptive_window=adaptive_window,
+        fade_bias=fade_bias,
+    )
+    scene_manager.add_detector(detector)
+
+    # Detect scenes
+    typer.echo("Detecting scenes...")
+    scene_manager.detect_scenes(
+        video=video,
+        show_progress=True,
+        frame_skip=frame_skip,
+        duration=duration_tc,
+    )
+
+    # Get scene list
+    scenes = scene_manager.get_scene_list()
+
+    # Filter out scenes that are too short if filter_shorter_than is specified
+    if filter_shorter_than_tc:
+        original_count = len(scenes)
+        scenes = [
+            (start, end)
+            for start, end in scenes
+            if (end.get_frames() - start.get_frames()) >= filter_shorter_than_tc.get_frames()
+        ]
+        if len(scenes) < original_count:
+            typer.echo(
+                f"Filtered out {original_count - len(scenes)} scenes shorter "
+                f"than {filter_shorter_than_tc.get_seconds():.1f} seconds "
+                f"({filter_shorter_than_tc.get_frames()} frames)",
+            )
+
+    # Apply max scenes limit if specified
+    if max_scenes and len(scenes) > max_scenes:
+        typer.echo(f"Dropping last {len(scenes) - max_scenes} scenes to meet max_scenes ({max_scenes}) limit")
+        scenes = scenes[:max_scenes]
+
+    # Print scene information
+    typer.echo(f"Found {len(scenes)} scenes:")
+    for i, (start, end) in enumerate(scenes, 1):
+        typer.echo(
+            f"Scene {i}: {start.get_timecode()} to {end.get_timecode()} "
+            f"({end.get_frames() - start.get_frames()} frames)",
+        )
+
+    # Save stats if requested
+    if stats_file:
+        typer.echo(f"Saving detection stats to {stats_file}")
+        stats_manager.save_to_csv(stats_file)
+
+    # Split video into scenes
+    typer.echo("Splitting video into scenes...")
+    try:
+        split_video_ffmpeg(
+            input_video_path=video_path,
+            scene_list=scenes,
+            output_dir=output_dir,
+            show_progress=True,
+        )
+        typer.echo(f"Scenes have been saved to: {output_dir}")
+    except Exception as e:
+        raise typer.BadParameter(f"Error splitting video: {e}") from e
+
+    # Save preview images if requested
+    if save_images_per_scene > 0:
+        typer.echo(f"Saving {save_images_per_scene} preview images per scene...")
+        image_filenames = save_scene_images(
+            scene_list=scenes,
+            video=video,
+            num_images=save_images_per_scene,
+            output_dir=str(output_dir),
+            show_progress=True,
+        )
+
+        # Generate HTML report with scene information and previews
+        html_path = output_dir / "scene_report.html"
+        write_scene_list_html(
+            output_html_filename=str(html_path),
+            scene_list=scenes,
+            image_filenames=image_filenames,
+        )
+        typer.echo(f"Scene report saved to: {html_path}")
+
+    return scenes
+
+
+@app.command()
+def main(  # noqa: PLR0913
+    video_path: Path = typer.Argument(  # noqa: B008
+        ...,
+        help="Path to the input video file",
+        exists=True,
+        dir_okay=False,
+    ),
+    output_dir: str = typer.Argument(
+        ...,
+        help="Directory where split scenes will be saved",
+    ),
+    detector: DetectorType = typer.Option(  # noqa: B008
+        DetectorType.CONTENT,
+        help="Scene detection algorithm to use",
+    ),
+    threshold: Optional[float] = typer.Option(
+        None,
+        help="Detection threshold (meaning varies by detector)",
+    ),
+    max_scenes: Optional[int] = typer.Option(
+        None,
+        help="Maximum number of scenes to produce",
+    ),
+    min_scene_length: Optional[int] = typer.Option(
+        None,
+        help="Minimum scene length during detection. Forces the detector to make scenes at least this many frames. "
+        "This affects scene detection behavior but does not filter out short scenes.",
+    ),
+    filter_shorter_than: Optional[str] = typer.Option(
+        None,
+        help="Filter out scenes shorter than this duration. Can be specified as frames (123), "
+        "seconds (123s/123.45s), or timecode (HH:MM:SS[.nnn]). These scenes will be detected but not saved.",
+    ),
+    skip_start: Optional[int] = typer.Option(
+        None,
+        help="Number of frames to skip at the start of the video",
+    ),
+    skip_end: Optional[int] = typer.Option(
+        None,
+        help="Number of frames to skip at the end of the video",
+    ),
+    duration: Optional[str] = typer.Option(
+        None,
+        "-d",
+        help="How much of the video to process. Can be specified as frames (123), "
+        "seconds (123s/123.45s), or timecode (HH:MM:SS[.nnn])",
+    ),
+    save_images: int = typer.Option(
+        0,
+        help="Number of preview images to save per scene (0 to disable)",
+    ),
+    stats_file: Optional[str] = typer.Option(
+        None,
+        help="Path to save detection statistics CSV",
+    ),
+    luma_only: bool = typer.Option(
+        False,
+        help="Only use brightness for content detection",
+    ),
+    adaptive_window: Optional[int] = typer.Option(
+        None,
+        help="Window size for adaptive detection",
+    ),
+    fade_bias: Optional[float] = typer.Option(
+        None,
+        help="Bias for fade detection (-1.0 to 1.0)",
+    ),
+    downscale: Optional[int] = typer.Option(
+        None,
+        help="Factor to downscale frames by during detection",
+    ),
+    frame_skip: int = typer.Option(
+        0,
+        help="Number of frames to skip during processing",
+    ),
+) -> None:
+    """Split video into scenes using PySceneDetect."""
+    if skip_start or skip_end:
+        typer.echo("Skipping start and end frames is not supported yet.")
+        return
+
+    # Validate output directory
+    output_path = validate_output_dir(output_dir)
+
+    # Detect and split scenes
+    detect_and_split_scenes(
+        video_path=str(video_path),
+        output_dir=output_path,
+        detector_type=detector,
+        threshold=threshold,
+        min_scene_len=min_scene_length,
+        max_scenes=max_scenes,
+        filter_shorter_than=filter_shorter_than,
+        skip_start=skip_start,
+        skip_end=skip_end,
+        duration=duration,
+        save_images_per_scene=save_images,
+        stats_file=stats_file,
+        luma_only=luma_only,
+        adaptive_window=adaptive_window,
+        fade_bias=fade_bias,
+        downscale_factor=downscale,
+        frame_skip=frame_skip,
+    )
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/scripts/train.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+"""
+Train LTXV models using configuration from YAML files.
+This script provides a command-line interface for training LTXV models using
+either LoRA fine-tuning or full model fine-tuning. It loads configuration from
+a YAML file and passes it to the trainer.
+Basic usage:
+    python scripts/train.py CONFIG_PATH [--disable-progress-bars]
+Resume is automatic when a training state file exists next to the loaded checkpoint.
+To start fresh, set `checkpoints.no_resume: true` in the YAML config.
+For multi-GPU/FSDP training, configure and launch via Accelerate:
+    accelerate config
+    accelerate launch scripts/train.py CONFIG_PATH
+"""
+
+from pathlib import Path
+
+import typer
+import yaml
+from rich.console import Console
+
+from ltx_trainer.config import LtxTrainerConfig
+from ltx_trainer.trainer import LtxvTrainer
+
+console = Console()
+app = typer.Typer(
+    pretty_exceptions_enable=False,
+    no_args_is_help=True,
+    help="Train LTXV models using configuration from YAML files.",
+)
+
+
+@app.command()
+def main(
+    config_path: str = typer.Argument(..., help="Path to YAML configuration file"),
+    disable_progress_bars: bool = typer.Option(
+        False,
+        "--disable-progress-bars",
+        help="Disable progress bars (useful for multi-process runs)",
+    ),
+) -> None:
+    """Train the model using the provided configuration file."""
+    config_path = Path(config_path)
+    if not config_path.exists():
+        typer.echo(f"Error: Configuration file {config_path} does not exist.")
+        raise typer.Exit(code=1)
+
+    with open(config_path, "r") as file:
+        config_data = yaml.safe_load(file)
+
+    try:
+        trainer_config = LtxTrainerConfig(**config_data)
+    except Exception as e:
+        typer.echo(f"Error: Invalid configuration data: {e}")
+        raise typer.Exit(code=1) from e
+
+    trainer = LtxvTrainer(trainer_config)
+    trainer.train(disable_progress_bars=disable_progress_bars)
+
+
+if __name__ == "__main__":
+    app()

packages/ltx-trainer/src/ltx_trainer/__init__.py

@@ -0,0 +1,44 @@
+import logging
+import os
+import sys
+from logging import getLogger
+from pathlib import Path
+
+from rich.logging import RichHandler
+
+# Get the process rank
+IS_MULTI_GPU = os.environ.get("LOCAL_RANK") is not None
+RANK = int(os.environ.get("LOCAL_RANK", "0"))
+
+# Configure with Rich
+logging.basicConfig(
+    level="INFO",
+    format=f"\\[rank {RANK}] %(message)s" if IS_MULTI_GPU else "%(message)s",
+    handlers=[
+        RichHandler(
+            rich_tracebacks=True,
+            show_time=False,
+            markup=True,
+        )
+    ],
+)
+
+# Get the logger and configure it
+logger = getLogger("ltxv_trainer")
+logger.setLevel(logging.DEBUG)
+logger.propagate = True
+
+# Set level based on process
+if RANK != 0:
+    logger.setLevel(logging.WARNING)
+
+# Expose common logging functions directly
+debug = logger.debug
+info = logger.info
+warning = logger.warning
+error = logger.error
+critical = logger.critical
+
+
+# Add the root directory to the Python path so we can import from scripts.
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))

packages/ltx-trainer/src/ltx_trainer/captioning.py

@@ -0,0 +1,455 @@
+"""
+Audio-visual media captioning using multimodal models.
+This module provides captioning capabilities for videos with audio using:
+- Qwen3-Omni via a local vLLM server (default)
+- Gemini Flash 3.5 (cloud API)
+Both produce a single combined English caption per video as a single
+continuous paragraph of prose.
+The Qwen3-Omni backend runs in a separately-launched vLLM server rather than
+in-process, so vLLM's heavy CUDA dependencies stay out of this package. The
+captioner talks to it over the OpenAI-compatible HTTP API.
+Launch the server once (in an isolated environment) with:
+.. code-block:: bash
+    uv run python scripts/serve_captioner.py
+That helper picks BF16 vs FP8 dynamic quantization based on the GPU's free
+memory and forwards everything else to ``vllm serve``. To check the recommended
+command without running it, pass ``--print-cmd``.
+To use Gemini instead, install ``google-genai`` and either set ``GEMINI_API_KEY``
+(Gemini Developer API) or have Google Cloud credentials available (gcloud / an
+attached service account), in which case it uses Vertex AI automatically.
+"""
+
+import json
+import os
+import re
+import subprocess
+import tempfile
+from abc import ABC, abstractmethod
+from enum import Enum
+from pathlib import Path
+from typing import ClassVar
+
+DEFAULT_VIDEO_CAPTION_INSTRUCTION = """\
+Analyze this video and produce a single detailed caption covering both its visual content and its audio. Be \
+detailed enough that someone reading the caption could form an accurate mental picture of what happens on screen \
+and what can be heard. Be exhaustive: include every meaningful detail you can see and hear, including small \
+objects, textures, secondary movements, and minor background sounds.
+
+Begin the caption directly with the action or visual detail; do not preface it with phrases like \
+"The video opens with...", "The scene shows...", "We see...", or "There is...".
+
+For every shot, include:
+- The shot type and framing (extreme wide / wide / medium / medium close-up / close-up / extreme close-up) and any \
+camera motion.
+- Characters' clothing, appearance, posture, and movement (direction, speed, quality).
+- The environment's materials, textures, lighting, and colors.
+- All audio: spoken dialogue (quoted exactly in the original language), tone of voice, music (style, mood, \
+volume changes), and environmental sounds. If a category is absent -- for example no music is playing, or no one is \
+speaking -- state that explicitly. Do not invent specific instruments, music genres, moods, or ambient sounds \
+that are not actually present.
+- Any on-screen text (signs, titles, labels).
+
+Describe only what is visible or audible. Do not infer emotions, intentions, or anything outside the segment. \
+Refer to people descriptively (e.g., "the man in the blue jacket"). Narrate strictly in chronological order; if \
+the video contains multiple shots, describe each one in turn.
+
+Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or labels \
+like "Audio:" / "Visual:" / "Shot:". Integrate visual and audio details naturally within the same sentences.
+
+Return a JSON object with exactly one key:
+
+{"combined_caption_english": "<your caption here>"}"""
+
+
+DEFAULT_IMAGE_CAPTION_INSTRUCTION = """\
+Analyze this image and produce a single detailed caption of its visual content. Be detailed enough that \
+someone reading the caption could form an accurate mental picture of the image. Be thorough: include every meaningful \
+detail that is actually present, including small objects, textures, and background elements.
+
+Begin the caption directly with the main subject or a visual detail; do not preface it with phrases like \
+"The image shows...", "This is a photo of...", "We see...", or "There is...".
+
+Include:
+- The framing and composition (close-up / medium / wide / overhead, etc.) and the vantage point.
+- The medium or style if distinctive (photograph, illustration, 3D render, painting).
+- People's clothing, appearance, and posture, and what they are doing.
+- The setting's materials, textures, lighting, and colors.
+- Transcribe any visible text verbatim (signs, labels, titles, captions).
+
+Describe only what is visible. Do not infer emotions or intentions, and do not describe sounds, motion, or \
+events before or after the moment shown -- this is a single still image. When something is ambiguous, describe \
+the visible cue (e.g., "warm low-angle light") rather than guessing the underlying fact (e.g., "sunrise"). \
+Refer to people descriptively (e.g., "the man in the blue jacket").
+
+Only describe what is present. Never state that something is absent or missing -- do not write phrases like \
+"there is no text", "no people are present", or "no other objects". If a category such as people or text does \
+not appear, simply leave it out.
+
+Write everything as a single continuous paragraph of prose. Do not use section headers, bullet points, or \
+labels.
+
+Return a JSON object with exactly one key:
+
+{"combined_caption_english": "<your caption here>"}"""
+
+
+# Default model served by ``scripts/serve_captioner.py``. The captioner does not
+# download or load this model itself -- it just sends requests to the vLLM
+# server, which already has the model loaded.
+DEFAULT_QWEN_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
+DEFAULT_VLLM_BASE_URL = "http://127.0.0.1:8001/v1"
+
+# Key the combined-caption prompt asks the model to return its caption under.
+_CAPTION_JSON_KEY = "combined_caption_english"
+
+
+class CaptionerType(str, Enum):
+    """Enum for different types of media captioners."""
+
+    QWEN_OMNI = "qwen_omni"  # Qwen3-Omni via local vLLM HTTP server
+    GEMINI_FLASH = "gemini_flash"  # Gemini Flash 3.5 cloud API
+
+
+def create_captioner(captioner_type: CaptionerType, **kwargs) -> "MediaCaptioningModel":
+    """Factory function to create a media captioner."""
+    match captioner_type:
+        case CaptionerType.QWEN_OMNI:
+            return QwenOmniCaptioner(**kwargs)
+        case CaptionerType.GEMINI_FLASH:
+            return GeminiFlashCaptioner(**kwargs)
+        case _:
+            raise ValueError(f"Unsupported captioner type: {captioner_type}")
+
+
+class MediaCaptioningModel(ABC):
+    """Abstract base class for audio-visual media captioning models."""
+
+    instruction: str | None = None
+
+    @abstractmethod
+    def caption(self, path: str | Path, **kwargs) -> str:
+        """Generate a caption for the given video or image."""
+
+    def _resolve_instruction(self, path: str | Path) -> str:
+        """Return the custom instruction, or the image/video default for this input."""
+        if self.instruction is not None:
+            return self.instruction
+        return DEFAULT_IMAGE_CAPTION_INSTRUCTION if self._is_image_file(path) else DEFAULT_VIDEO_CAPTION_INSTRUCTION
+
+    @staticmethod
+    def _is_image_file(path: str | Path) -> bool:
+        return str(path).lower().endswith((".png", ".jpg", ".jpeg", ".heic", ".heif", ".webp"))
+
+    @staticmethod
+    def _is_video_file(path: str | Path) -> bool:
+        return str(path).lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm"))
+
+
+class QwenOmniCaptioner(MediaCaptioningModel):
+    """Audio-visual captioning via a local vLLM server running Qwen3-Omni.
+    The vLLM server must already be running. See ``scripts/serve_captioner.py``
+    for a helper that launches one in an isolated environment (no impact on
+    this package's dependency tree).
+    The captioner uses the OpenAI-compatible chat completions API. It sends
+    a ``file://`` URL pointing at the local video, the default combined-caption
+    prompt, and parses the JSON-wrapped response.
+    """
+
+    def __init__(
+        self,
+        base_url: str = DEFAULT_VLLM_BASE_URL,
+        model: str = DEFAULT_QWEN_MODEL,
+        api_key: str = "EMPTY",
+        instruction: str | None = None,
+        max_tokens: int = 4096,
+        enable_thinking: bool = False,
+        timeout_s: float = 600.0,
+    ):
+        """Initialize the Qwen3-Omni captioner.
+        Args:
+            base_url: Base URL of the vLLM OpenAI-compatible server (default
+                ``http://127.0.0.1:8001/v1``).
+            model: Model identifier the server is serving. Must match the
+                server's ``--served-model-name`` (defaults to the HuggingFace
+                model ID).
+            api_key: Token sent in the ``Authorization`` header. vLLM accepts
+                any value by default.
+            instruction: Custom instruction prompt. If ``None``, uses the
+                default combined-caption prompt.
+            max_tokens: Maximum new tokens to generate per caption. 4096 leaves
+                comfortable headroom for both ``enable_thinking`` modes.
+            enable_thinking: Whether to let the Thinking model produce a
+                ``<think>...</think>`` chain-of-thought before the caption.
+                Off by default: it makes captioning ~5x slower with little
+                quality benefit and occasionally introduces hallucinations
+                (e.g., inventing dialogue or background music).
+            timeout_s: Per-request HTTP timeout.
+        """
+        from openai import OpenAI  # noqa: PLC0415
+
+        self.model = model
+        self.instruction = instruction
+        self.max_tokens = max_tokens
+        self.enable_thinking = enable_thinking
+        self._client = OpenAI(base_url=base_url, api_key=api_key, timeout=timeout_s)
+
+    def caption(
+        self,
+        path: str | Path,
+        fps: int = 2,
+    ) -> str:
+        """Generate a caption for the given video or image.
+        Args:
+            path: Path to the video/image file to caption.
+            fps: Frames per second to sample from the video. Passed through to
+                vLLM's multimodal processor (``mm_processor_kwargs.fps``).
+                Default 2 is a typical choice for video MLLMs at this resolution.
+                Ignored for image inputs.
+        Returns:
+            The extracted caption string.
+        """
+        path = Path(path)
+        is_image = self._is_image_file(path)
+        is_video = self._is_video_file(path)
+        if not (is_image or is_video):
+            raise ValueError(f"Unsupported media file: {path}")
+
+        instruction = self._resolve_instruction(path)
+
+        if is_image:
+            content = [
+                {"type": "image_url", "image_url": {"url": f"file://{path.resolve()}"}},
+                {"type": "text", "text": instruction},
+            ]
+            return _parse_caption_response(self._chat(content)).strip()
+
+        return self._caption_video(path, instruction, fps)
+
+    def _chat(self, content: list[dict], mm_kwargs: dict | None = None) -> str:
+        """Send one chat-completions request and return the raw response text."""
+        extra_body: dict = {
+            "repetition_penalty": 1.05,
+            "chat_template_kwargs": {"enable_thinking": self.enable_thinking},
+        }
+        if mm_kwargs:
+            extra_body["mm_processor_kwargs"] = mm_kwargs
+        response = self._client.chat.completions.create(
+            model=self.model,
+            messages=[{"role": "user", "content": content}],
+            max_tokens=self.max_tokens,
+            temperature=0.0,
+            extra_body=extra_body,
+        )
+        return response.choices[0].message.content or ""
+
+    def _caption_video(self, path: Path, instruction: str, fps: int) -> str:
+        """Caption a video, sending its audio track as a separate modality.
+        vLLM does not extract a video's audio on its own (and its
+        ``use_audio_in_video`` path is broken server-side), so we pull the audio
+        into a 16 kHz mono WAV and send it alongside the video -- otherwise the
+        model only sees frames and fabricates any spoken content.
+        """
+        with tempfile.TemporaryDirectory(prefix="qwencap_") as tmp:
+            work = Path(tmp)
+
+            # Best-effort: ffmpeg fails (and we send video only) if there's no audio.
+            audio_url: str | None = None
+            try:
+                wav = work / "audio.wav"
+                _extract_audio_wav(path, wav)
+                audio_url = f"file://{wav.resolve()}"
+            except subprocess.CalledProcessError:
+                pass
+
+            def content(video: Path) -> list[dict]:
+                parts: list[dict] = [{"type": "video_url", "video_url": {"url": f"file://{video.resolve()}"}}]
+                if audio_url:
+                    parts.append({"type": "audio_url", "audio_url": {"url": audio_url}})
+                parts.append({"type": "text", "text": instruction})
+                return parts
+
+            mm_kwargs = {"fps": fps}
+            try:
+                raw = self._chat(content(path), mm_kwargs)
+            except Exception as e:
+                # Raw / variable-frame-rate videos over-report their frame count, which
+                # breaks the server's frame sampler ("... frames from video"). Re-encode
+                # to a constant frame rate and retry once.
+                if "frames from video" not in str(e):
+                    raise
+                cfr = work / "video_cfr.mp4"
+                _transcode_cfr(path, cfr)
+                raw = self._chat(content(cfr), mm_kwargs)
+
+            return _parse_caption_response(raw).strip()
+
+
+class GeminiFlashCaptioner(MediaCaptioningModel):
+    """Audio-visual captioning using Google's Gemini via the Google Gen AI SDK.
+    Uses the ``google-genai`` package (the current SDK; ``google-generativeai``
+    is deprecated). Auth is resolved automatically:
+    1. If an API key is given (``api_key`` argument, or ``GEMINI_API_KEY`` /
+       ``GOOGLE_API_KEY`` in the environment) -> the Gemini Developer API (AI Studio).
+    2. Otherwise, if Google Cloud Application Default Credentials are available
+       (an attached service account or ``gcloud auth application-default login``)
+       -> Vertex AI. The project comes from ADC (or ``GOOGLE_CLOUD_PROJECT``) and
+       the location defaults to ``global`` (override with ``GOOGLE_CLOUD_LOCATION``).
+       This means it "just works" on a gcloud-authed GCP VM with no env vars.
+    If neither is available, a clear error explains how to authenticate.
+    Media is sent inline (``Part.from_bytes``), which works on both backends.
+    """
+
+    MODEL_ID = "gemini-3.5-flash"
+
+    _MIME_TYPES: ClassVar[dict[str, str]] = {
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+        ".png": "image/png",
+        ".webp": "image/webp",
+        ".heic": "image/heic",
+        ".heif": "image/heif",
+        ".mp4": "video/mp4",
+        ".mov": "video/quicktime",
+        ".avi": "video/x-msvideo",
+        ".mkv": "video/x-matroska",
+        ".webm": "video/webm",
+    }
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        instruction: str | None = None,
+        model: str | None = None,
+    ):
+        """Initialize the Gemini captioner.
+        Args:
+            api_key: Gemini Developer API key. If ``None``, falls back to
+                ``GEMINI_API_KEY`` / ``GOOGLE_API_KEY``; if no key is set at all,
+                uses Vertex AI via Application Default Credentials.
+            instruction: Custom instruction prompt. If ``None``, uses the default
+                image or video prompt depending on the input.
+            model: Override the served model id (defaults to ``MODEL_ID``).
+        """
+        self.instruction = instruction
+        self.model = model or self.MODEL_ID
+        self._client = self._make_client(api_key)
+
+    def caption(
+        self,
+        path: str | Path,
+        fps: int = 2,  # noqa: ARG002 - kept for API compatibility
+    ) -> str:
+        from google.genai import types  # noqa: PLC0415
+
+        path = Path(path)
+        instruction = self._resolve_instruction(path)
+        media = types.Part.from_bytes(data=path.read_bytes(), mime_type=self._mime_type(path))
+        response = self._client.models.generate_content(
+            model=self.model,
+            contents=[media, instruction],
+            config=types.GenerateContentConfig(temperature=0.0),
+        )
+
+        # Gemini may also return JSON if it followed our prompt format.
+        return _parse_caption_response(response.text or "").strip()
+
+    @classmethod
+    def _mime_type(cls, path: Path) -> str:
+        try:
+            return cls._MIME_TYPES[path.suffix.lower()]
+        except KeyError:
+            raise ValueError(f"Unsupported media type for Gemini: {path.suffix}") from None
+
+    def _make_client(self, api_key: str | None):  # noqa: ANN202 - genai.Client type is lazy-imported
+        from google import genai  # noqa: PLC0415
+
+        # 1. API key (explicit arg or env) -> Gemini Developer API.
+        key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+        if key:
+            return genai.Client(api_key=key)
+
+        # 2. No key -> Vertex AI via Application Default Credentials (gcloud / service account).
+        import google.auth  # noqa: PLC0415
+
+        try:
+            _, adc_project = google.auth.default()
+        except Exception as e:
+            raise ValueError(
+                "No Gemini credentials found. Provide an API key (--api-key, or "
+                "GEMINI_API_KEY / GOOGLE_API_KEY), or set up Google Cloud credentials "
+                "for Vertex AI (e.g. `gcloud auth application-default login` or an "
+                "attached service account)."
+            ) from e
+
+        project = os.environ.get("GOOGLE_CLOUD_PROJECT") or adc_project
+        location = os.environ.get("GOOGLE_CLOUD_LOCATION", "global")
+        return genai.Client(vertexai=True, project=project, location=location)
+
+
+def _parse_caption_response(raw: str) -> str:
+    """Extract the caption text from a model response.
+    Backend-agnostic: works for any model that follows the combined-caption
+    prompt. Handles the formats a model may produce:
+    - Plain caption text
+    - JSON ``{"combined_caption_english": "..."}``
+    - ``<think>...</think>`` chain-of-thought followed by either of the above
+    - Truncated JSON (when generation hits a token limit mid-string)
+    """
+    text = re.sub(r"<think>[\s\S]*?</think>", "", raw).strip()
+
+    # Thinking models (e.g. Qwen3-Omni-*-Thinking) emit the reasoning trace
+    # without an opening ``<think>`` tag, because the chat template injects it
+    # for them -- so the response starts mid-thought and is terminated by a lone
+    # ``</think>`` before the real answer. Drop everything up to that closer.
+    if "</think>" in text:
+        text = text.rsplit("</think>", 1)[1].strip()
+
+    if not text:
+        return raw.strip()
+
+    try:
+        parsed = json.loads(text)
+        if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
+            return parsed[_CAPTION_JSON_KEY]
+    except (json.JSONDecodeError, ValueError):
+        pass
+
+    match = re.search(rf"\{{[^{{}}]*\"{_CAPTION_JSON_KEY}\"[^{{}}]*\}}", text)
+    if match:
+        try:
+            parsed = json.loads(match.group())
+            if isinstance(parsed, dict) and _CAPTION_JSON_KEY in parsed:
+                return parsed[_CAPTION_JSON_KEY]
+        except (json.JSONDecodeError, ValueError):
+            pass
+
+    # Truncated JSON: extract the string value even if the closing quote/brace is missing.
+    match = re.search(rf'"{_CAPTION_JSON_KEY}"\s*:\s*"((?:[^"\\]|\\.)*)', text)
+    if match:
+        try:
+            return json.loads('"' + match.group(1) + '"')
+        except (json.JSONDecodeError, ValueError):
+            return match.group(1)
+
+    return text
+
+
+def _run_ffmpeg(args: list[str]) -> None:
+    """Run the ffmpeg binary bundled with ``imageio-ffmpeg`` (a dependency)."""
+    import imageio_ffmpeg  # noqa: PLC0415
+
+    cmd = [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args]
+    subprocess.run(cmd, check=True, capture_output=True)
+
+
+def _extract_audio_wav(src: Path, dest: Path) -> None:
+    """Extract the audio track to a 16 kHz mono PCM WAV (matches pretraining).
+    Raises ``CalledProcessError`` when the video has no audio stream.
+    """
+    _run_ffmpeg(["-i", str(src), "-vn", "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", str(dest)])
+
+
+def _transcode_cfr(src: Path, dest: Path) -> None:
+    """Re-encode the video to a constant frame rate so the server's frame sampler can
+    read every requested index (raw / variable-frame-rate videos over-report frames)."""
+    _run_ffmpeg(["-i", str(src), "-fps_mode", "cfr", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-an", str(dest)])

packages/ltx-trainer/src/ltx_trainer/config.py

@@ -0,0 +1,833 @@
+from pathlib import Path
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, ValidationInfo, field_validator, model_validator
+
+from ltx_trainer.quantization import QuantizationOptions
+from ltx_trainer.training_strategies.base_strategy import TrainingStrategyConfigBase
+from ltx_trainer.training_strategies.flexible import FlexibleStrategyConfig
+from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig
+from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig
+
+
+class ConfigBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+# =============================================================================
+# Validation Condition Types
+# =============================================================================
+
+
+class FirstFrameConditionConfig(ConfigBaseModel):
+    """First-frame conditioning (intrinsic, latent_idx=0). Always targets video.
+    If image_or_video points to a video file, the first frame is automatically extracted.
+    """
+
+    type: Literal["first_frame"] = "first_frame"
+    image_or_video: str | Path
+
+
+class PrefixConditionConfig(ConfigBaseModel):
+    """Prefix conditioning for temporal extension (intrinsic). Exactly one of video/audio must be set."""
+
+    type: Literal["prefix"] = "prefix"
+    video: str | None = None
+    audio: str | None = None
+    num_frames: int | None = Field(
+        default=None,
+        ge=1,
+        description="Number of pixel frames for video prefix. Must satisfy num_frames %% 8 == 1.",
+    )
+    duration: float | None = Field(default=None, gt=0, description="Duration in seconds for audio prefix")
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "PrefixConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for prefix condition")
+        return self
+
+    @model_validator(mode="after")
+    def validate_num_frames_constraint(self) -> "PrefixConditionConfig":
+        if self.video is not None and self.num_frames is not None and self.num_frames % 8 != 1:
+            raise ValueError(
+                f"num_frames ({self.num_frames}) must satisfy num_frames % 8 == 1 "
+                f"for video prefix (e.g., 1, 9, 17, 25, ...)"
+            )
+        return self
+
+
+class SuffixConditionConfig(ConfigBaseModel):
+    """Suffix conditioning for temporal extension (intrinsic). Exactly one of video/audio must be set."""
+
+    type: Literal["suffix"] = "suffix"
+    video: str | None = None
+    audio: str | None = None
+    num_frames: int | None = Field(
+        default=None,
+        ge=1,
+        description="Number of pixel frames for video suffix. Must satisfy num_frames %% 8 == 0.",
+    )
+    duration: float | None = Field(default=None, gt=0, description="Duration in seconds for audio suffix")
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "SuffixConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for suffix condition")
+        return self
+
+    @model_validator(mode="after")
+    def validate_num_frames_constraint(self) -> "SuffixConditionConfig":
+        if self.video is not None and self.num_frames is not None and self.num_frames % 8 != 0:
+            raise ValueError(
+                f"num_frames ({self.num_frames}) must satisfy num_frames % 8 == 0 "
+                f"for video suffix (e.g., 8, 16, 24, 32, ...)"
+            )
+        return self
+
+
+class SpatialCropConditionConfig(ConfigBaseModel):
+    """Spatial crop conditioning for outpainting (intrinsic, video only)."""
+
+    type: Literal["spatial_crop"] = "spatial_crop"
+    video: str
+    spatial_region: tuple[int, int, int, int] = Field(
+        ..., description="Spatial crop region as (y1, x1, y2, x2) in pixel coordinates"
+    )
+
+
+class MaskConditionConfig(ConfigBaseModel):
+    """Mask-based conditioning for inpainting (intrinsic). Exactly one of video/audio must be set."""
+
+    type: Literal["mask"] = "mask"
+    video: str | None = None
+    audio: str | None = None
+    mask: str
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "MaskConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for mask condition")
+        return self
+
+
+class ReferenceConditionConfig(ConfigBaseModel):
+    """Reference conditioning (IC-LoRA style concatenation). Exactly one of video/audio must be set."""
+
+    type: Literal["reference"] = "reference"
+    video: str | None = None
+    audio: str | None = None
+    downscale_factor: int = Field(default=1, ge=1)
+    temporal_scale_factor: int = Field(default=1, ge=1)
+    include_in_output: bool = False
+
+    @model_validator(mode="after")
+    def validate_exactly_one_modality(self) -> "ReferenceConditionConfig":
+        if (self.video is None) == (self.audio is None):
+            raise ValueError("Exactly one of 'video' or 'audio' must be set for reference condition")
+        return self
+
+
+class VideoToAudioConditionConfig(ConfigBaseModel):
+    """Video-to-audio — video is provided as frozen cross-modal conditioning.
+    The video is kept clean (sigma=0) and influences audio generation via cross-modal attention.
+    """
+
+    type: Literal["video_to_audio"] = "video_to_audio"
+    video: str
+
+
+class AudioToVideoConditionConfig(ConfigBaseModel):
+    """Audio-to-video — audio is provided as frozen cross-modal conditioning.
+    The audio is kept clean (sigma=0) and influences video generation via cross-modal attention.
+    """
+
+    type: Literal["audio_to_video"] = "audio_to_video"
+    audio: str
+
+
+ValidationCondition = Annotated[
+    Union[
+        FirstFrameConditionConfig,
+        PrefixConditionConfig,
+        SuffixConditionConfig,
+        SpatialCropConditionConfig,
+        MaskConditionConfig,
+        ReferenceConditionConfig,
+        VideoToAudioConditionConfig,
+        AudioToVideoConditionConfig,
+    ],
+    Field(discriminator="type"),
+]
+
+
+def _condition_targets_video(cond: ValidationCondition) -> bool:
+    """Check if a validation condition targets the video modality."""
+    if cond.type in ("first_frame", "spatial_crop", "video_to_audio"):
+        return True
+    if cond.type in ("prefix", "suffix", "mask", "reference"):
+        return getattr(cond, "video", None) is not None
+    return False
+
+
+def _condition_targets_audio(cond: ValidationCondition) -> bool:
+    """Check if a validation condition targets the audio modality."""
+    if cond.type == "audio_to_video":
+        return True
+    if cond.type in ("prefix", "suffix", "mask", "reference"):
+        return getattr(cond, "audio", None) is not None
+    return False
+
+
+class ValidationSample(ConfigBaseModel):
+    """Configuration for a single validation sample — fully self-describing."""
+
+    prompt: str
+    conditions: list[ValidationCondition] = Field(default_factory=list)
+
+    video_dims: tuple[int, int, int] | None = Field(
+        default=None,
+        description="Per-sample override for (width, height, frames). None = inherit from ValidationConfig.",
+    )
+    seed: int | None = Field(
+        default=None,
+        description="Per-sample override for random seed. None = inherit from ValidationConfig.",
+    )
+
+    @field_validator("video_dims")
+    @classmethod
+    def validate_video_dims(cls, v: tuple[int, int, int] | None) -> tuple[int, int, int] | None:
+        if v is None:
+            return v
+        width, height, frames = v
+        if width % 32 != 0:
+            raise ValueError(f"Width ({width}) must be divisible by 32")
+        if height % 32 != 0:
+            raise ValueError(f"Height ({height}) must be divisible by 32")
+        if frames % 8 != 1:
+            raise ValueError(f"Frames ({frames}) must satisfy frames % 8 == 1 for LTX-2 (e.g., 1, 9, 17, 25, ...)")
+        return v
+
+    @model_validator(mode="after")
+    def validate_frozen_modality_conflicts(self) -> "ValidationSample":
+        frozen_types = {c.type for c in self.conditions if c.type in ("video_to_audio", "audio_to_video")}
+
+        if "video_to_audio" in frozen_types and "audio_to_video" in frozen_types:
+            raise ValueError(
+                "Cannot have both video_to_audio and audio_to_video conditions — nothing would be generated"
+            )
+
+        if "video_to_audio" in frozen_types:
+            for c in self.conditions:
+                if c.type != "video_to_audio" and _condition_targets_video(c):
+                    raise ValueError(
+                        f"Cannot use video-targeting '{c.type}' condition when video is frozen (video_to_audio)"
+                    )
+
+        if "audio_to_video" in frozen_types:
+            for c in self.conditions:
+                if c.type != "audio_to_video" and _condition_targets_audio(c):
+                    raise ValueError(
+                        f"Cannot use audio-targeting '{c.type}' condition when audio is frozen (audio_to_video)"
+                    )
+
+        return self
+
+
+class ModelConfig(ConfigBaseModel):
+    """Configuration for the base model and training mode"""
+
+    model_path: str | Path = Field(
+        ...,
+        description="Model path - local path to safetensors checkpoint file",
+    )
+
+    text_encoder_path: str | Path | None = Field(
+        default=None,
+        description="Path to text encoder (required for LTX-2/Gemma models, optional for LTXV/T5 models)",
+    )
+
+    training_mode: Literal["lora", "full"] = Field(
+        default="lora",
+        description="Training mode - either LoRA fine-tuning or full model fine-tuning",
+    )
+
+    load_checkpoint: str | Path | None = Field(
+        default=None,
+        description="Path to a checkpoint file or directory to load from. "
+        "If a directory is provided, the latest checkpoint will be used.",
+    )
+
+    @field_validator("model_path")
+    @classmethod
+    def validate_model_path(cls, v: str | Path) -> str | Path:
+        """Validate that model_path is either a valid URL or an existing local path."""
+        is_url = str(v).startswith(("http://", "https://"))
+
+        if is_url:
+            raise ValueError(f"Model path cannot be a URL: {v}")
+
+        if not Path(v).exists():
+            raise ValueError(f"Model path does not exist: {v}")
+
+        return v
+
+
+class LoraConfig(ConfigBaseModel):
+    """Configuration for LoRA fine-tuning"""
+
+    rank: int = Field(
+        default=64,
+        description="Rank of LoRA adaptation",
+        ge=2,
+    )
+
+    alpha: int = Field(
+        default=64,
+        description="Alpha scaling factor for LoRA",
+        ge=1,
+    )
+
+    dropout: float = Field(
+        default=0.0,
+        description="Dropout probability for LoRA layers",
+        ge=0.0,
+        le=1.0,
+    )
+
+    target_modules: list[str] = Field(
+        default=["to_k", "to_q", "to_v", "to_out.0"],
+        description="List of modules to target with LoRA",
+    )
+
+
+def _get_strategy_discriminator(v: dict | TrainingStrategyConfigBase) -> str:
+    """Discriminator function for strategy config union."""
+    if isinstance(v, dict):
+        return v.get("name", "text_to_video")
+    return v.name
+
+
+# Union type for all strategy configs with discriminator
+TrainingStrategyConfig = Annotated[
+    Annotated[TextToVideoConfig, Tag("text_to_video")]
+    | Annotated[VideoToVideoConfig, Tag("video_to_video")]
+    | Annotated[FlexibleStrategyConfig, Tag("flexible")],
+    Discriminator(_get_strategy_discriminator),
+]
+
+
+class OptimizationConfig(ConfigBaseModel):
+    """Configuration for optimization parameters"""
+
+    learning_rate: float = Field(
+        default=5e-4,
+        description="Learning rate for optimization",
+    )
+
+    steps: int = Field(
+        default=3000,
+        description="Number of training steps",
+    )
+
+    batch_size: int = Field(
+        default=2,
+        description="Batch size for training",
+    )
+
+    gradient_accumulation_steps: int = Field(
+        default=1,
+        description="Number of steps to accumulate gradients",
+    )
+
+    max_grad_norm: float = Field(
+        default=1.0,
+        description="Maximum gradient norm for clipping",
+    )
+
+    optimizer_type: Literal["adamw", "adamw8bit"] = Field(
+        default="adamw",
+        description="Type of optimizer to use for training",
+    )
+
+    scheduler_type: Literal[
+        "constant",
+        "linear",
+        "cosine",
+        "cosine_with_restarts",
+        "polynomial",
+        "step",
+    ] = Field(
+        default="linear",
+        description="Type of scheduler to use for training",
+    )
+
+    scheduler_params: dict = Field(
+        default_factory=dict,
+        description="Parameters for the scheduler",
+    )
+
+    enable_gradient_checkpointing: bool = Field(
+        default=False,
+        description="Enable gradient checkpointing to save memory at the cost of slower training",
+    )
+
+
+class AccelerationConfig(ConfigBaseModel):
+    """Configuration for hardware acceleration and compute optimization"""
+
+    mixed_precision_mode: Literal["no", "fp16", "bf16"] | None = Field(
+        default="bf16",
+        description="Mixed precision training mode",
+    )
+
+    quantization: QuantizationOptions | None = Field(
+        default=None,
+        description="Quantization precision to use",
+    )
+
+    load_text_encoder_in_8bit: bool = Field(
+        default=False,
+        description="Whether to load the text encoder in 8-bit precision to save memory",
+    )
+
+    offload_optimizer_during_validation: bool = Field(
+        default=False,
+        description="Offload optimizer state to CPU before validation video sampling and reload "
+        "it afterwards, to free VRAM for inference. Useful when optimizer state is large "
+        "(e.g. AdamW for full fine-tuning or high-rank LoRA) and validation OOMs because the "
+        "VAE decoder + transformer + optimizer state cannot coexist on the GPU. Has no effect "
+        "for FSDP (sharded state). Disabled by default.",
+    )
+
+
+class DataConfig(ConfigBaseModel):
+    """Configuration for data loading and processing"""
+
+    preprocessed_data_root: str = Field(
+        description="Path to folder containing preprocessed training data",
+    )
+
+    num_dataloader_workers: int = Field(
+        default=2,
+        description="Number of background processes for data loading (0 means synchronous loading)",
+        ge=0,
+    )
+
+    @field_validator("preprocessed_data_root")
+    @classmethod
+    def validate_preprocessed_data_root(cls, v: str) -> str:
+        """Validate that preprocessed_data_root exists."""
+        path = Path(v).expanduser().resolve()
+        if not path.exists():
+            raise ValueError(f"Dataset path does not exist: {v}")
+        if not path.is_dir():
+            raise ValueError(f"Dataset path is not a directory: {v}")
+        return str(path)
+
+
+class ValidationConfig(ConfigBaseModel):
+    """Configuration for validation during training"""
+
+    # Per-sample configuration (new format — preferred)
+    samples: list[ValidationSample] = Field(
+        default_factory=list,
+        description="List of validation samples. Each sample is fully self-describing with its own "
+        "prompt, conditions, and optional overrides. Replaces prompts/images/reference_videos.",
+    )
+
+    # Legacy fields (deprecated — converted to samples internally via convert_legacy_format)
+    prompts: list[str] = Field(
+        default_factory=list,
+        description="[DEPRECATED: use 'samples' instead] List of prompts to use for validation",
+    )
+
+    negative_prompt: str = Field(
+        default="worst quality, inconsistent motion, blurry, jittery, distorted",
+        description="Negative prompt to use for validation examples",
+    )
+
+    images: list[str] | None = Field(
+        default=None,
+        description="[DEPRECATED: use 'samples' with first_frame conditions] "
+        "List of image paths to use for validation. "
+        "One image path must be provided for each validation prompt",
+    )
+
+    reference_videos: list[str] | None = Field(
+        default=None,
+        description="[DEPRECATED: use 'samples' with reference conditions] "
+        "List of reference video paths to use for validation. "
+        "One video path must be provided for each validation prompt",
+    )
+
+    reference_downscale_factor: int = Field(
+        default=1,
+        description="[DEPRECATED: use downscale_factor on ReferenceCondition] "
+        "Downscale factor for reference videos in IC-LoRA validation. "
+        "When > 1, reference videos are processed at 1/n resolution (e.g., 2 means half resolution). "
+        "Must match the factor used during dataset preprocessing.",
+        ge=1,
+    )
+
+    video_dims: tuple[int, int, int] = Field(
+        default=(960, 544, 97),
+        description="Dimensions of validation videos (width, height, frames). "
+        "Width and height must be divisible by 32. Frames must satisfy frames % 8 == 1 for LTX-2.",
+    )
+
+    @field_validator("video_dims")
+    @classmethod
+    def validate_video_dims(cls, v: tuple[int, int, int]) -> tuple[int, int, int]:
+        """Validate video dimensions for LTX-2 compatibility."""
+        width, height, frames = v
+
+        if width % 32 != 0:
+            raise ValueError(f"Width ({width}) must be divisible by 32")
+        if height % 32 != 0:
+            raise ValueError(f"Height ({height}) must be divisible by 32")
+        if frames % 8 != 1:
+            raise ValueError(f"Frames ({frames}) must satisfy frames % 8 == 1 for LTX-2 (e.g., 1, 9, 17, 25, ...)")
+
+        return v
+
+    frame_rate: float = Field(
+        default=25.0,
+        description="Frame rate for validation videos",
+        gt=0,
+    )
+
+    seed: int = Field(
+        default=42,
+        description="Random seed used when sampling validation videos",
+    )
+
+    inference_steps: int = Field(
+        default=50,
+        description="Number of inference steps for validation",
+        gt=0,
+    )
+
+    interval: int | None = Field(
+        default=100,
+        description="Number of steps between validation runs. If None, validation is disabled.",
+        gt=0,
+    )
+
+    guidance_scale: float = Field(
+        default=4.0,
+        description="CFG guidance scale to use during validation",
+        ge=1.0,
+    )
+
+    stg_scale: float = Field(
+        default=1.0,
+        description="STG (Spatio-Temporal Guidance) scale. 0.0 disables STG. "
+        "Recommended value is 1.0. STG is combined with CFG for improved video quality.",
+        ge=0.0,
+    )
+
+    stg_blocks: list[int] | None = Field(
+        default=[29],
+        description="Which transformer blocks to perturb for STG. "
+        "None means all blocks are perturbed. Recommended for LTX-2: [29].",
+    )
+
+    stg_mode: Literal["stg_av", "stg_v"] = Field(
+        default="stg_av",
+        description="STG mode: 'stg_av' skips both audio and video self-attention, "
+        "'stg_v' skips only video self-attention.",
+    )
+
+    generate_audio: bool = Field(
+        default=True,
+        description="Whether to generate audio in validation samples. "
+        "Independent of training strategy setting - you can generate audio "
+        "in validation even when not training the audio branch.",
+    )
+
+    generate_video: bool = Field(
+        default=True,
+        description="Whether to generate video in validation samples. "
+        "Set to False for audio-only or v2a validation to save VRAM by skipping video VAE decoder loading. "
+        "When False, validation will only generate audio (requires generate_audio=True).",
+    )
+
+    skip_initial_validation: bool = Field(
+        default=False,
+        description="Skip validation video sampling at step 0 (beginning of training)",
+    )
+
+    include_reference_in_output: bool = Field(
+        default=False,
+        description="[DEPRECATED: use include_in_output on ReferenceCondition] "
+        "For video-to-video training: concatenate the original reference video side-by-side "
+        "with the generated output. The reference comes from the input video, not from the model's output.",
+    )
+
+    @field_validator("images")
+    @classmethod
+    def validate_images(cls, v: list[str] | None, info: ValidationInfo) -> list[str] | None:
+        """Validate that number of images (if provided) matches number of prompts."""
+        if v is None:
+            return None
+
+        num_prompts = len(info.data.get("prompts", []))
+        if v is not None and len(v) != num_prompts:
+            raise ValueError(f"Number of images ({len(v)}) must match number of prompts ({num_prompts})")
+
+        for image_path in v:
+            if not Path(image_path).exists():
+                raise ValueError(f"Image path '{image_path}' does not exist")
+
+        return v
+
+    @field_validator("reference_videos")
+    @classmethod
+    def validate_reference_videos(cls, v: list[str] | None, info: ValidationInfo) -> list[str] | None:
+        """Validate that number of reference videos (if provided) matches number of prompts."""
+        if v is None:
+            return None
+
+        num_prompts = len(info.data.get("prompts", []))
+        if v is not None and len(v) != num_prompts:
+            raise ValueError(f"Number of reference videos ({len(v)}) must match number of prompts ({num_prompts})")
+
+        for video_path in v:
+            if not Path(video_path).exists():
+                raise ValueError(f"Reference video path '{video_path}' does not exist")
+
+        return v
+
+    @model_validator(mode="after")
+    def convert_legacy_format(self) -> "ValidationConfig":
+        """Convert deprecated prompts/images/reference_videos to the new samples format."""
+        if self.prompts and not self.samples:
+            samples = []
+            for i, prompt in enumerate(self.prompts):
+                conditions: list[ValidationCondition] = []
+                if self.images and i < len(self.images):
+                    conditions.append(FirstFrameConditionConfig(image_or_video=self.images[i]))
+                if self.reference_videos and i < len(self.reference_videos):
+                    conditions.append(
+                        ReferenceConditionConfig(
+                            video=self.reference_videos[i],
+                            downscale_factor=self.reference_downscale_factor,
+                            include_in_output=self.include_reference_in_output,
+                        )
+                    )
+                samples.append(ValidationSample(prompt=prompt, conditions=conditions))
+            self.samples = samples
+        return self
+
+    @model_validator(mode="after")
+    def validate_scaled_reference_dimensions(self) -> "ValidationConfig":
+        """Validate that scaled reference dimensions are valid when reference_downscale_factor > 1."""
+        if self.reference_downscale_factor > 1:
+            width, height, _frames = self.video_dims
+
+            if width % self.reference_downscale_factor != 0:
+                raise ValueError(
+                    f"Width {width} is not evenly divisible by reference_downscale_factor "
+                    f"{self.reference_downscale_factor}. Choose a downscale factor that divides {width} evenly."
+                )
+            if height % self.reference_downscale_factor != 0:
+                raise ValueError(
+                    f"Height {height} is not evenly divisible by reference_downscale_factor "
+                    f"{self.reference_downscale_factor}. Choose a downscale factor that divides {height} evenly."
+                )
+
+            scaled_width = width // self.reference_downscale_factor
+            scaled_height = height // self.reference_downscale_factor
+
+            if scaled_width % 32 != 0:
+                raise ValueError(
+                    f"Scaled reference width {scaled_width} (from {width} / {self.reference_downscale_factor}) "
+                    f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
+                )
+            if scaled_height % 32 != 0:
+                raise ValueError(
+                    f"Scaled reference height {scaled_height} (from {height} / {self.reference_downscale_factor}) "
+                    f"is not divisible by 32. Choose a different downscale factor or adjust video_dims."
+                )
+
+        return self
+
+    @model_validator(mode="after")
+    def validate_output_modality_requirements(self) -> "ValidationConfig":
+        """Validate output modality settings when validation is configured."""
+        has_validation = bool(self.prompts) or bool(self.samples)
+        if has_validation and not self.generate_video and not self.generate_audio:
+            raise ValueError(
+                "At least one of generate_video or generate_audio must be True when validation is configured."
+            )
+        return self
+
+
+class CheckpointsConfig(ConfigBaseModel):
+    """Configuration for model checkpointing during training"""
+
+    interval: int | None = Field(
+        default=None,
+        description="Number of steps between checkpoint saves. If None, intermediate checkpoints are disabled.",
+        gt=0,
+    )
+
+    keep_last_n: int = Field(
+        default=1,
+        description="Number of most recent checkpoints to keep. Set to -1 to keep all checkpoints.",
+        ge=-1,
+    )
+
+    precision: Literal["bfloat16", "float32"] = Field(
+        default="bfloat16",
+        description="Precision to use when saving checkpoint weights. Options: 'bfloat16' or 'float32'.",
+    )
+
+    no_resume: bool = Field(
+        default=False,
+        description="When True, ignore any saved training state and start from step 0. "
+        "Model weights from load_checkpoint are still loaded, but optimizer/scheduler "
+        "state and step counter are reset.",
+    )
+
+    save_training_state: Literal["full", "minimal", "off"] = Field(
+        default="minimal",
+        description="Save training state alongside checkpoints for resume. "
+        "'full': optimizer + scheduler + RNG + step (~800MB for LoRA, much larger for full fine-tuning). "
+        "'minimal': scheduler + RNG + step only (~few KB, sufficient for LoRA). "
+        "'off': nothing saved, resume not possible.",
+    )
+
+
+class HubConfig(ConfigBaseModel):
+    """Configuration for Hugging Face Hub integration"""
+
+    push_to_hub: bool = Field(default=False, description="Whether to push the model weights to the Hugging Face Hub")
+    hub_model_id: str | None = Field(
+        default=None, description="Hugging Face Hub repository ID (e.g., 'username/repo-name')"
+    )
+
+    @model_validator(mode="after")
+    def validate_hub_config(self) -> "HubConfig":
+        """Validate that hub_model_id is not None when push_to_hub is True."""
+        if self.push_to_hub and not self.hub_model_id:
+            raise ValueError("hub_model_id must be specified when push_to_hub is True")
+        return self
+
+
+class WandbConfig(ConfigBaseModel):
+    """Configuration for Weights & Biases logging"""
+
+    enabled: bool = Field(
+        default=False,
+        description="Whether to enable W&B logging",
+    )
+
+    project: str = Field(
+        default="ltxv-trainer",
+        description="W&B project name",
+    )
+
+    entity: str | None = Field(
+        default=None,
+        description="W&B username or team",
+    )
+
+    tags: list[str] = Field(
+        default_factory=list,
+        description="Tags to add to the W&B run",
+    )
+
+    log_validation_videos: bool = Field(
+        default=True,
+        description="Whether to log validation videos to W&B",
+    )
+
+
+class FlowMatchingConfig(ConfigBaseModel):
+    """Configuration for flow matching training"""
+
+    timestep_sampling_mode: Literal["uniform", "shifted_logit_normal"] = Field(
+        default="shifted_logit_normal",
+        description="Mode to use for timestep sampling",
+    )
+
+    timestep_sampling_params: dict = Field(
+        default_factory=dict,
+        description="Parameters for timestep sampling",
+    )
+
+
+class LtxTrainerConfig(ConfigBaseModel):
+    """Unified configuration for LTXV training"""
+
+    # Sub-configurations
+    model: ModelConfig = Field(default_factory=ModelConfig)
+    lora: LoraConfig | None = Field(default=None)
+    training_strategy: TrainingStrategyConfig = Field(
+        default_factory=TextToVideoConfig,
+        description="Training strategy configuration. Determines the training mode and its parameters.",
+    )
+    optimization: OptimizationConfig = Field(default_factory=OptimizationConfig)
+    acceleration: AccelerationConfig = Field(default_factory=AccelerationConfig)
+    data: DataConfig
+    validation: ValidationConfig = Field(default_factory=ValidationConfig)
+    checkpoints: CheckpointsConfig = Field(default_factory=CheckpointsConfig)
+    hub: HubConfig = Field(default_factory=HubConfig)
+    flow_matching: FlowMatchingConfig = Field(default_factory=FlowMatchingConfig)
+    wandb: WandbConfig = Field(default_factory=WandbConfig)
+
+    # General configuration
+    seed: int = Field(
+        default=42,
+        description="Random seed for reproducibility",
+    )
+
+    output_dir: str = Field(
+        default="outputs",
+        description="Directory to save model outputs",
+    )
+
+    # noinspection PyNestedDecorators
+    @field_validator("output_dir")
+    @classmethod
+    def expand_output_path(cls, v: str) -> str:
+        """Expand user home directory in output path."""
+        return str(Path(v).expanduser().resolve())
+
+    def _validate_data_dirs_exist(self) -> None:
+        """Verify that every directory declared by the training strategy exists under the data root."""
+        data_root = Path(self.data.preprocessed_data_root)
+        for dir_name in self.training_strategy.get_data_sources():
+            dir_path = data_root / dir_name
+            if not dir_path.is_dir():
+                raise ValueError(
+                    f"Required data directory '{dir_name}' does not exist under preprocessed_data_root: {dir_path}"
+                )
+
+    @model_validator(mode="after")
+    def validate_strategy_compatibility(self) -> "LtxTrainerConfig":
+        """Validate that training strategy and other configurations are compatible."""
+        self._validate_data_dirs_exist()
+
+        # Check that reference videos are provided when using video_to_video strategy
+        if self.training_strategy.name == "video_to_video" and self.validation.interval:
+            has_reference = bool(self.validation.reference_videos) or any(
+                cond.type == "reference" for sample in self.validation.samples for cond in sample.conditions
+            )
+            if not has_reference:
+                raise ValueError(
+                    "reference_videos or samples with reference conditions must be provided "
+                    "in validation config when using video_to_video strategy"
+                )
+
+        # Check that LoRA config is provided when training mode is lora
+        if self.model.training_mode == "lora" and self.lora is None:
+            raise ValueError("LoRA configuration must be provided when training_mode is 'lora'")
+
+        # Check that LoRA config is provided when using video_to_video strategy
+        if self.training_strategy.name == "video_to_video" and self.model.training_mode != "lora":
+            raise ValueError("Training mode must be 'lora' when using video_to_video strategy")
+
+        return self

packages/ltx-trainer/src/ltx_trainer/config_display.py

@@ -0,0 +1,156 @@
+"""Display utilities for training configuration.
+This module provides formatted console output for LtxTrainerConfig.
+"""
+
+from rich import box
+from rich.console import Console
+from rich.table import Table
+
+from ltx_trainer.config import LtxTrainerConfig
+
+
+def print_config(config: LtxTrainerConfig) -> None:
+    """Print configuration as a nicely formatted table with sections."""
+
+    def fmt(v: object, max_len: int = 55) -> str:
+        """Format any value for display."""
+        if v is None:
+            return "[dim]—[/]"
+        if isinstance(v, bool):
+            return "[green]✓[/]" if v else "[dim]✗[/]"
+        if isinstance(v, (list, tuple)):
+            if not v:
+                return "[dim]—[/]"
+            return ", ".join(str(x) for x in v)
+        s = str(v)
+        return s[: max_len - 3] + "..." if len(s) > max_len else s
+
+    cfg = config
+    opt = cfg.optimization
+    val = cfg.validation
+    accel = cfg.acceleration
+
+    # Build sections: list of (section_title, [(key, value), ...])
+    sections: list[tuple[str, list[tuple[str, str]]]] = [
+        (
+            "🎬 Model",
+            [
+                ("Base", fmt(cfg.model.model_path)),
+                ("Text Encoder", fmt(cfg.model.text_encoder_path) or "[dim]Built-in[/]"),
+                ("Training Mode", f"[bold green]{cfg.model.training_mode.upper()}[/]"),
+                ("Load Checkpoint", fmt(cfg.model.load_checkpoint) if cfg.model.load_checkpoint else "[dim]—[/]"),
+            ],
+        ),
+    ]
+
+    if cfg.lora:
+        sections.append(
+            (
+                "🔗 LoRA",
+                [
+                    ("Rank / Alpha", f"{cfg.lora.rank} / {cfg.lora.alpha}"),
+                    ("Dropout", str(cfg.lora.dropout)),
+                    ("Target Modules", fmt(cfg.lora.target_modules)),
+                ],
+            )
+        )
+
+    # Strategy section - include strategy-specific fields
+    strategy_items: list[tuple[str, str]] = [("Name", cfg.training_strategy.name)]
+    if hasattr(cfg.training_strategy, "with_audio"):
+        strategy_items.append(("Audio", fmt(cfg.training_strategy.with_audio)))
+    if hasattr(cfg.training_strategy, "first_frame_conditioning_p"):
+        strategy_items.append(("First Frame Cond P", str(cfg.training_strategy.first_frame_conditioning_p)))
+
+    sections.append(("🎯 Strategy", strategy_items))
+
+    sections.extend(
+        [
+            (
+                "⚡ Optimization",
+                [
+                    ("Steps", f"[bold]{opt.steps:,}[/]"),
+                    ("Learning Rate", f"{opt.learning_rate:.2e}"),
+                    ("Batch Size", str(opt.batch_size)),
+                    ("Grad Accumulation", str(opt.gradient_accumulation_steps)),
+                    ("Optimizer", opt.optimizer_type),
+                    ("Scheduler", opt.scheduler_type),
+                    ("Max Grad Norm", str(opt.max_grad_norm)),
+                    ("Grad Checkpointing", fmt(opt.enable_gradient_checkpointing)),
+                ],
+            ),
+            (
+                "🚀 Acceleration",
+                [
+                    ("Mixed Precision", accel.mixed_precision_mode or "[dim]—[/]"),
+                    ("Quantization", str(accel.quantization) if accel.quantization else "[dim]—[/]"),
+                    ("Text Encoder 8bit", fmt(accel.load_text_encoder_in_8bit)),
+                    ("Optimizer CPU Offload", fmt(accel.offload_optimizer_during_validation)),
+                ],
+            ),
+            (
+                "🎥 Validation",
+                [
+                    ("Prompts", f"{len(val.prompts)} prompt(s)" if val.prompts else "[dim]—[/]"),
+                    ("Interval", f"Every {val.interval} steps" if val.interval else "[dim]Disabled[/]"),
+                    ("Video Dims", f"{val.video_dims[0]}x{val.video_dims[1]}, {val.video_dims[2]} frames"),
+                    ("Frame Rate", f"{val.frame_rate} fps"),
+                    ("Inference Steps", str(val.inference_steps)),
+                    ("CFG Scale", str(val.guidance_scale)),
+                    (
+                        "STG",
+                        f"scale={val.stg_scale}; blocks={fmt(val.stg_blocks)}; mode={val.stg_mode}"
+                        if val.stg_scale > 0
+                        else "[dim]Disabled[/]",
+                    ),
+                    ("Seed", str(val.seed)),
+                ],
+            ),
+            (
+                "📂 Data & Output",
+                [
+                    ("Dataset", fmt(cfg.data.preprocessed_data_root)),
+                    ("Dataloader Workers", str(cfg.data.num_dataloader_workers)),
+                    ("Output Dir", fmt(cfg.output_dir)),
+                    ("Seed", str(cfg.seed)),
+                ],
+            ),
+            (
+                "🔌 Integrations",
+                [
+                    (
+                        "Checkpoints",
+                        f"Every {cfg.checkpoints.interval} steps (keep {cfg.checkpoints.keep_last_n})"
+                        if cfg.checkpoints.interval
+                        else "[dim]Disabled[/]",
+                    ),
+                    ("W&B", f"{cfg.wandb.project}" if cfg.wandb.enabled else "[dim]Disabled[/]"),
+                    ("HF Hub", cfg.hub.hub_model_id if cfg.hub.push_to_hub else "[dim]Disabled[/]"),
+                ],
+            ),
+        ]
+    )
+
+    # Build table with section headers
+    table = Table(
+        title="[bold]⚙️  Training Configuration[/]",
+        show_header=False,
+        box=box.ROUNDED,
+        border_style="bright_blue",
+        padding=(0, 1),
+        title_style="bold bright_blue",
+    )
+    table.add_column("Key", style="white", width=20)
+    table.add_column("Value", style="cyan")
+
+    for i, (section_title, items) in enumerate(sections):
+        if i > 0:
+            table.add_row("", "")  # Blank line between sections
+        table.add_row(f"[bold yellow]{section_title}[/]", "")
+        for key, value in items:
+            table.add_row(f"  {key}", value)
+
+    console = Console()
+    console.print()
+    console.print(table)
+    console.print()

packages/ltx-trainer/src/ltx_trainer/datasets.py

@@ -0,0 +1,315 @@
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import torch
+from einops import rearrange
+from torch import Tensor
+from torch.utils.data import Dataset
+
+from ltx_trainer import logger
+
+# Constants for precomputed data directories
+PRECOMPUTED_DIR_NAME = ".precomputed"
+
+
+class DummyDataset(Dataset):
+    """Produce random latents and prompt embeddings. For minimal demonstration and benchmarking purposes"""
+
+    def __init__(
+        self,
+        width: int = 1024,
+        height: int = 1024,
+        num_frames: int = 25,
+        fps: int = 24,
+        dataset_length: int = 200,
+        latent_dim: int = 128,
+        latent_spatial_compression_ratio: int = 32,
+        latent_temporal_compression_ratio: int = 8,
+        prompt_embed_dim: int = 4096,
+        prompt_sequence_length: int = 256,
+    ) -> None:
+        if width % 32 != 0:
+            raise ValueError(f"Width must be divisible by 32, got {width=}")
+
+        if height % 32 != 0:
+            raise ValueError(f"Height must be divisible by 32, got {height=}")
+
+        if num_frames % 8 != 1:
+            raise ValueError(f"Number of frames must have a remainder of 1 when divided by 8, got {num_frames=}")
+
+        self.width = width
+        self.height = height
+        self.num_frames = num_frames
+        self.fps = fps
+        self.dataset_length = dataset_length
+        self.latent_dim = latent_dim
+        self.num_latent_frames = (num_frames - 1) // latent_temporal_compression_ratio + 1
+        self.latent_height = height // latent_spatial_compression_ratio
+        self.latent_width = width // latent_spatial_compression_ratio
+        self.latent_sequence_length = self.num_latent_frames * self.latent_height * self.latent_width
+        self.prompt_embed_dim = prompt_embed_dim
+        self.prompt_sequence_length = prompt_sequence_length
+
+    def __len__(self) -> int:
+        return self.dataset_length
+
+    def __getitem__(self, idx: int) -> dict[str, dict[str, Tensor]]:
+        return {
+            "latent_conditions": {
+                "latents": torch.randn(
+                    self.latent_dim,
+                    self.num_latent_frames,
+                    self.latent_height,
+                    self.latent_width,
+                ),
+                "num_frames": self.num_latent_frames,
+                "height": self.latent_height,
+                "width": self.latent_width,
+                "fps": self.fps,
+            },
+            "text_conditions": {
+                "video_prompt_embeds": torch.randn(
+                    self.prompt_sequence_length,
+                    self.prompt_embed_dim,
+                ),
+                "audio_prompt_embeds": torch.randn(
+                    self.prompt_sequence_length,
+                    self.prompt_embed_dim,
+                ),
+                "prompt_attention_mask": torch.ones(
+                    self.prompt_sequence_length,
+                    dtype=torch.bool,
+                ),
+            },
+        }
+
+
+class PrecomputedDataset(Dataset):
+    def __init__(self, data_root: str, data_sources: dict[str, str] | list[str] | None = None) -> None:
+        """
+        Generic dataset for loading precomputed data from multiple sources.
+        Args:
+            data_root: Root directory containing preprocessed data
+            data_sources: Either:
+              - Dict mapping directory names to output keys
+              - List of directory names (keys will equal values)
+              - None (defaults to ["latents", "conditions"])
+        Example:
+            # Standard mode (list)
+            dataset = PrecomputedDataset("data/", ["latents", "conditions"])
+            # Standard mode (dict)
+            dataset = PrecomputedDataset("data/", {"latents": "latent_conditions", "conditions": "text_conditions"})
+            # IC-LoRA mode
+            dataset = PrecomputedDataset("data/", ["latents", "conditions", "reference_latents"])
+        Note:
+            Latents are always returned in non-patchified format [C, F, H, W].
+            Legacy patchified format [seq_len, C] is automatically converted.
+        """
+        super().__init__()
+
+        self.data_root = self._setup_data_root(data_root)
+        self.data_sources = self._normalize_data_sources(data_sources)
+        self.source_paths = self._setup_source_paths()
+        self.sample_files = self._discover_samples()
+        self._validate_setup()
+
+    @staticmethod
+    def _setup_data_root(data_root: str) -> Path:
+        """Setup and validate the data root directory."""
+        data_root = Path(data_root).expanduser().resolve()
+
+        if not data_root.exists():
+            raise FileNotFoundError(f"Data root directory does not exist: {data_root}")
+
+        # If the given path is the dataset root, use the precomputed subdirectory
+        if (data_root / PRECOMPUTED_DIR_NAME).exists():
+            data_root = data_root / PRECOMPUTED_DIR_NAME
+
+        return data_root
+
+    @staticmethod
+    def _normalize_data_sources(data_sources: dict[str, str] | list[str] | None) -> dict[str, str]:
+        """Normalize data_sources input to a consistent dict format."""
+        if data_sources is None:
+            # Default sources
+            return {"latents": "latent_conditions", "conditions": "text_conditions"}
+        elif isinstance(data_sources, list):
+            # Convert list to dict where keys equal values
+            return {source: source for source in data_sources}
+        elif isinstance(data_sources, dict):
+            return data_sources.copy()
+        else:
+            raise TypeError(f"data_sources must be dict, list, or None, got {type(data_sources)}")
+
+    def _setup_source_paths(self) -> dict[str, Path]:
+        """Map data source names to their actual directory paths."""
+        source_paths = {}
+
+        for dir_name in self.data_sources:
+            source_path = self.data_root / dir_name
+            source_paths[dir_name] = source_path
+
+            # Check that all sources exist.
+            if not source_path.exists():
+                raise FileNotFoundError(f"Required {dir_name} directory does not exist: {source_path}")
+
+        return source_paths
+
+    def _discover_samples(self) -> dict[str, list[Path]]:
+        """Discover all valid sample files across all data sources.
+        Uses a fast two-pass approach: first globs all sources in parallel to build
+        full-path sets in memory, then checks expected paths via set membership.
+        This avoids O(N * num_sources) stat calls on networked filesystems while
+        correctly handling path remapping (e.g. latent_X.pt -> condition_X.pt).
+        """
+        if not self.data_sources:
+            raise ValueError("No data sources configured")
+
+        data_key = "latents" if "latents" in self.data_sources else next(iter(self.data_sources.keys()))
+        data_path = self.source_paths[data_key]
+
+        # Pass 1: Glob all sources in parallel, build full-path sets
+        def _glob_source(dir_name: str) -> tuple[list[Path], set[str]]:
+            source_path = self.source_paths[dir_name]
+            paths = list(source_path.glob("**/*.pt"))
+            path_set = {str(p) for p in paths}
+            return paths, path_set
+
+        with ThreadPoolExecutor(max_workers=len(self.data_sources)) as executor:
+            glob_results = dict(
+                zip(
+                    self.data_sources.keys(),
+                    executor.map(_glob_source, self.data_sources.keys()),
+                    strict=True,
+                )
+            )
+
+        # Get primary source files (cached from glob, no second scan)
+        data_files, _ = glob_results[data_key]
+        if not data_files:
+            raise ValueError(f"No data files found in {data_path}")
+        data_files.sort()
+
+        # Log source sizes
+        for dir_name, (paths, _) in glob_results.items():
+            logger.debug(f"Source {dir_name}: {len(paths)} files")
+
+        # Build path sets for non-primary sources
+        other_path_sets = {
+            dir_name: path_set for dir_name, (_, path_set) in glob_results.items() if dir_name != data_key
+        }
+
+        # Pass 2: For each primary file, check if expected paths exist in other sources' sets
+        sample_files: dict[str, list[Path]] = {output_key: [] for output_key in self.data_sources.values()}
+        valid_count = 0
+
+        for data_file in data_files:
+            rel_path = data_file.relative_to(data_path)
+
+            # Check all other sources via set lookup (O(1) per source, no stat calls)
+            all_exist = True
+            for dir_name, path_set in other_path_sets.items():
+                expected = self._get_expected_file_path(dir_name, data_file, rel_path)
+                if str(expected) not in path_set:
+                    logger.debug(f"Skipping {data_file.name}: no matching {dir_name} file at {expected}")
+                    all_exist = False
+                    break
+
+            if all_exist:
+                self._fill_sample_data_files(data_file, rel_path, sample_files)
+                valid_count += 1
+
+        skipped = len(data_files) - valid_count
+        if skipped > 0:
+            logger.info(f"Fast index: {valid_count} valid samples from {len(data_files)} total ({skipped} skipped)")
+        else:
+            logger.debug(f"Fast index: {valid_count} valid samples from {len(data_files)} total")
+
+        return sample_files
+
+    def _get_expected_file_path(self, dir_name: str, data_file: Path, rel_path: Path) -> Path:
+        """Get the expected file path for a given data source."""
+        source_path = self.source_paths[dir_name]
+
+        # For conditions, handle legacy naming where latent_X.pt maps to condition_X.pt
+        if dir_name == "conditions" and data_file.name.startswith("latent_"):
+            return source_path / f"condition_{data_file.stem[7:]}.pt"
+
+        return source_path / rel_path
+
+    def _fill_sample_data_files(self, data_file: Path, rel_path: Path, sample_files: dict[str, list[Path]]) -> None:
+        """Add a valid sample to the sample_files tracking."""
+        for dir_name, output_key in self.data_sources.items():
+            expected_path = self._get_expected_file_path(dir_name, data_file, rel_path)
+            sample_files[output_key].append(expected_path.relative_to(self.source_paths[dir_name]))
+
+    def _validate_setup(self) -> None:
+        """Validate that the dataset setup is correct."""
+        sample_counts = {key: len(files) for key, files in self.sample_files.items()}
+        if not sample_counts or all(count == 0 for count in sample_counts.values()):
+            raise ValueError(
+                f"No valid samples found in {self.data_root} - all configured data sources "
+                f"({list(self.data_sources)}) must have matching files (per-source counts: {sample_counts})"
+            )
+
+        # Verify all output keys have the same number of samples
+        if len(set(sample_counts.values())) > 1:
+            raise ValueError(f"Mismatched sample counts across sources: {sample_counts}")
+
+    def __len__(self) -> int:
+        # Use the first output key as reference count
+        first_key = next(iter(self.sample_files.keys()))
+        return len(self.sample_files[first_key])
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        result = {}
+
+        for dir_name, output_key in self.data_sources.items():
+            source_path = self.source_paths[dir_name]
+            file_rel_path = self.sample_files[output_key][index]
+            file_path = source_path / file_rel_path
+
+            try:
+                data = torch.load(file_path, map_location="cpu", weights_only=True)
+
+                # Normalize video latent format if this is a latent source
+                if "latent" in dir_name.lower():
+                    data = self._normalize_video_latents(data)
+
+                result[output_key] = data
+            except Exception as e:
+                raise RuntimeError(f"Failed to load {output_key} from {file_path}: {e}") from e
+
+        # Add index for debugging
+        result["idx"] = index
+        return result
+
+    @staticmethod
+    def _normalize_video_latents(data: dict) -> dict:
+        """
+        Normalize video latents to non-patchified format [C, F, H, W].
+        Used for keeping backward compatibility with legacy datasets.
+        """
+        latents = data["latents"]
+
+        # Check if latents are in legacy patchified format [seq_len, C]
+        if latents.dim() == 2:
+            # Legacy format: [seq_len, C] where seq_len = F * H * W
+            num_frames = data["num_frames"]
+            height = data["height"]
+            width = data["width"]
+
+            # Unpatchify: [seq_len, C] -> [C, F, H, W]
+            latents = rearrange(
+                latents,
+                "(f h w) c -> c f h w",
+                f=num_frames,
+                h=height,
+                w=width,
+            )
+
+            # Update the data dict with unpatchified latents
+            data = data.copy()
+            data["latents"] = latents
+
+        return data

packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py

@@ -0,0 +1,104 @@
+# ruff: noqa: PLC0415
+
+"""
+8-bit Gemma text encoder loading utilities.
+This module provides functionality for loading the Gemma text encoder in 8-bit precision
+using bitsandbytes, which significantly reduces GPU memory usage.
+Example usage:
+    from ltx_trainer.gemma_8bit import load_8bit_gemma
+    text_encoder = load_8bit_gemma(gemma_model_path="/path/to/gemma")
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Generator
+from contextlib import contextmanager
+from pathlib import Path
+
+import torch
+
+from ltx_core.text_encoders.gemma.encoders.base_encoder import GemmaTextEncoder
+from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer
+
+
+def load_8bit_gemma(
+    gemma_model_path: str | Path,
+    dtype: torch.dtype = torch.bfloat16,
+    device: torch.device | str | int | None = None,
+) -> GemmaTextEncoder:
+    """Load the Gemma text encoder in 8-bit precision using bitsandbytes.
+    Only the Gemma LLM backbone is loaded here.  The embeddings processor
+    (feature extractor + connectors) should be loaded separately via
+    :func:`ltx_trainer.model_loader.load_embeddings_processor`.
+    Args:
+        gemma_model_path: Path to Gemma model directory
+        dtype: Data type for non-quantized model weights
+        device: Device to place the quantized model on. When ``None`` (default),
+            the device is inferred from ``LOCAL_RANK`` if CUDA is available, so
+            multi-process launches put each rank's encoder on its own GPU
+            instead of all colliding on ``cuda:0``.
+    Returns:
+        GemmaTextEncoder with 8-bit quantized Gemma backbone
+    Raises:
+        ImportError: If bitsandbytes is not installed
+        FileNotFoundError: If required model files are not found
+    """
+    try:
+        from transformers import BitsAndBytesConfig, Gemma3ForConditionalGeneration
+    except ImportError as e:
+        raise ImportError(
+            "8-bit text encoder loading requires bitsandbytes. Install it with: uv pip install bitsandbytes"
+        ) from e
+
+    gemma_path = _find_gemma_subpath(gemma_model_path, "model*.safetensors")
+    tokenizer_path = _find_gemma_subpath(gemma_model_path, "tokenizer.model")
+
+    # Pin the entire model to a single device. `device_map="auto"` collides on cuda:0
+    # in multi-process launches because every rank picks the same default device.
+    device_map: str | dict[str, int | str | torch.device]
+    if device is not None:
+        device_map = {"": device}
+    elif torch.cuda.is_available():
+        device_map = {"": int(os.environ.get("LOCAL_RANK", "0"))}
+    else:
+        device_map = "auto"
+
+    quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+    with _suppress_accelerate_memory_warnings():
+        gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
+            gemma_path,
+            quantization_config=quantization_config,
+            torch_dtype=torch.bfloat16,
+            device_map=device_map,
+            local_files_only=True,
+        )
+
+    tokenizer = LTXVGemmaTokenizer(tokenizer_path, 1024)
+
+    return GemmaTextEncoder(
+        tokenizer=tokenizer,
+        model=gemma_model,
+        dtype=dtype,
+    )
+
+
+def _find_gemma_subpath(root_path: str | Path, pattern: str) -> str:
+    """Find a file matching a glob pattern and return its parent directory."""
+    matches = list(Path(root_path).rglob(pattern))
+    if not matches:
+        raise FileNotFoundError(f"No files matching pattern '{pattern}' found under {root_path}")
+    return str(matches[0].parent)
+
+
+@contextmanager
+def _suppress_accelerate_memory_warnings() -> Generator[None, None, None]:
+    """Temporarily suppress INFO warnings from accelerate about memory allocation."""
+    accelerate_logger = logging.getLogger("accelerate.utils.modeling")
+    old_level = accelerate_logger.level
+    accelerate_logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        accelerate_logger.setLevel(old_level)

packages/ltx-trainer/src/ltx_trainer/gpu_utils.py

@@ -0,0 +1,90 @@
+"""GPU memory management utilities for training and inference."""
+
+import functools
+import gc
+import subprocess
+from typing import Callable, TypeVar
+
+import torch
+
+from ltx_trainer import logger
+
+F = TypeVar("F", bound=Callable)
+
+
+def free_gpu_memory(log: bool = False) -> None:
+    """Free GPU memory by running garbage collection and emptying CUDA cache.
+    Args:
+        log: If True, log memory stats after clearing
+    """
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        if log:
+            allocated = torch.cuda.memory_allocated() / 1024**3
+            reserved = torch.cuda.memory_reserved() / 1024**3
+            logger.debug(f"GPU memory freed. Allocated: {allocated:.2f}GB, Reserved: {reserved:.2f}GB")
+
+
+class free_gpu_memory_context:  # noqa: N801
+    """Context manager and decorator to free GPU memory before and/or after execution.
+    Can be used as a decorator:
+        @free_gpu_memory_context(after=True)
+        def my_function():
+            ...
+    Or as a context manager:
+        with free_gpu_memory_context():
+            heavy_operation()
+    Args:
+        before: Free memory before execution (default: False)
+        after: Free memory after execution (default: True)
+        log: Log memory stats when freeing (default: False)
+    """
+
+    def __init__(self, *, before: bool = False, after: bool = True, log: bool = False) -> None:
+        self.before = before
+        self.after = after
+        self.log = log
+
+    def __enter__(self) -> "free_gpu_memory_context":
+        if self.before:
+            free_gpu_memory(log=self.log)
+        return self
+
+    def __exit__(self, exc_type: type | None, exc_val: Exception | None, exc_tb: object) -> None:
+        if self.after:
+            free_gpu_memory(log=self.log)
+
+    def __call__(self, func: F) -> F:
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs) -> object:
+            with self:
+                return func(*args, **kwargs)
+
+        return wrapper  # type: ignore
+
+
+def get_gpu_memory_gb(device: torch.device) -> float:
+    """Get current GPU memory usage in GB using nvidia-smi.
+    Args:
+        device: torch.device to get memory usage for
+    Returns:
+        Current GPU memory usage in GB
+    """
+    try:
+        device_id = device.index if device.index is not None else 0
+        result = subprocess.check_output(
+            [
+                "nvidia-smi",
+                "--query-gpu=memory.used",
+                "--format=csv,nounits,noheader",
+                "-i",
+                str(device_id),
+            ],
+            encoding="utf-8",
+        )
+        return float(result.strip()) / 1024  # Convert MB to GB
+    except (subprocess.CalledProcessError, FileNotFoundError, ValueError) as e:
+        logger.error(f"Failed to get GPU memory from nvidia-smi: {e}")
+        # Fallback to torch
+        return torch.cuda.memory_allocated(device) / 1024**3

packages/ltx-trainer/src/ltx_trainer/hf_hub_utils.py

@@ -0,0 +1,212 @@
+import shutil
+import tempfile
+from pathlib import Path
+from typing import List, Optional, Union
+
+import imageio
+from huggingface_hub import HfApi, create_repo
+from huggingface_hub.utils import are_progress_bars_disabled, disable_progress_bars, enable_progress_bars
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+from ltx_trainer import logger
+from ltx_trainer.config import LtxTrainerConfig
+
+
+def push_to_hub(
+    weights_path: Path,
+    sampled_videos_paths: Optional[List[Path]],
+    config: LtxTrainerConfig,
+) -> None:
+    """Push the trained LoRA weights to HuggingFace Hub."""
+    if not config.hub.hub_model_id:
+        logger.warning("⚠️ HuggingFace hub_model_id not specified, skipping push to hub")
+        return
+
+    api = HfApi()
+
+    # Save original progress bar state
+    original_progress_state = are_progress_bars_disabled()
+    disable_progress_bars()  # Disable during our custom progress tracking
+
+    try:
+        # Try to create repo if it doesn't exist
+        try:
+            repo = create_repo(
+                repo_id=config.hub.hub_model_id,
+                repo_type="model",
+                exist_ok=True,  # Don't raise error if repo exists
+            )
+            repo_id = repo.repo_id
+            logger.info(f"🤗 Successfully created HuggingFace model repository at: {repo.url}")
+        except Exception as e:
+            logger.error(f"❌ Failed to create HuggingFace model repository: {e}")
+            return
+
+        # Create a single temporary directory for all files
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_path = Path(temp_dir)
+
+            with Progress(
+                SpinnerColumn(),
+                TextColumn("[progress.description]{task.description}"),
+                transient=True,
+            ) as progress:
+                try:
+                    # Copy weights
+                    task_copy = progress.add_task("Copying weights...", total=None)
+                    weights_dest = temp_path / weights_path.name
+                    shutil.copy2(weights_path, weights_dest)
+                    progress.update(task_copy, description="✓ Weights copied")
+
+                    # Create model card and save samples
+                    task_card = progress.add_task("Creating model card and samples...", total=None)
+                    _create_model_card(
+                        output_dir=temp_path,
+                        videos=sampled_videos_paths,
+                        config=config,
+                    )
+                    progress.update(task_card, description="✓ Model card and samples created")
+
+                    # Upload everything at once
+                    task_upload = progress.add_task("Pushing files to HuggingFace Hub...", total=None)
+                    api.upload_folder(
+                        folder_path=str(temp_path),
+                        repo_id=repo_id,
+                        repo_type="model",
+                    )
+                    progress.update(task_upload, description="✓ Files pushed to HuggingFace Hub")
+                    logger.info("✅ Successfully pushed files to HuggingFace Hub")
+
+                except Exception as e:
+                    logger.error(f"❌ Failed to process and push files to HuggingFace Hub: {e}")
+                    raise  # Re-raise to handle in outer try block
+
+    finally:
+        # Restore original progress bar state
+        if not original_progress_state:
+            enable_progress_bars()
+
+
+def convert_video_to_gif(video_path: Path, output_path: Path) -> None:
+    """Convert a video file to GIF format."""
+    try:
+        # Read the video file
+        reader = imageio.get_reader(str(video_path))
+        fps = reader.get_meta_data()["fps"]
+
+        # Write GIF file with infinite loop
+        writer = imageio.get_writer(
+            str(output_path),
+            fps=min(fps, 15),  # Cap FPS at 15 for reasonable file size
+            loop=0,  # 0 means infinite loop
+        )
+
+        for frame in reader:
+            writer.append_data(frame)
+
+        writer.close()
+        reader.close()
+    except Exception as e:
+        logger.error(f"Failed to convert video to GIF: {e}")
+
+
+def _create_model_card(
+    output_dir: Union[str, Path],
+    videos: Optional[List[Path]],
+    config: LtxTrainerConfig,
+) -> Path:
+    """Generate and save a model card for the trained model."""
+
+    repo_id = config.hub.hub_model_id
+    pretrained_model_name_or_path = config.model.model_path
+    validation_prompts = config.validation.prompts
+    output_dir = Path(output_dir)
+    template_path = Path(__file__).parent.parent.parent / "templates" / "model_card.md"
+
+    # Read the template
+    template = template_path.read_text()
+
+    # Get model name from repo_id
+    model_name = repo_id.split("/")[-1]
+
+    # Get base model information
+    base_model_link = str(pretrained_model_name_or_path)
+    model_path_str = str(pretrained_model_name_or_path)
+    is_url = model_path_str.startswith(("http://", "https://"))
+
+    # For URLs, extract the filename from the URL. For local paths, use the filename stem
+    base_model_name = model_path_str.split("/")[-1] if is_url else Path(pretrained_model_name_or_path).name
+
+    # Format validation prompts and create grid layout
+    prompts_text = ""
+    sample_grid = []
+
+    if validation_prompts and videos:
+        prompts_text = "Example prompts used during validation:\n\n"
+
+        # Create samples directory
+        samples_dir = output_dir / "samples"
+        samples_dir.mkdir(exist_ok=True, parents=True)
+
+        # Process videos and create cells
+        cells = []
+        for i, (prompt, video) in enumerate(zip(validation_prompts, videos, strict=False)):
+            if video.exists():
+                # Add prompt to text section
+                prompts_text += f"- `{prompt}`\n"
+
+                # Convert video to GIF
+                gif_path = samples_dir / f"sample_{i}.gif"
+                try:
+                    convert_video_to_gif(video, gif_path)
+
+                    # Create grid cell with collapsible description
+                    cell = (
+                        f"![example{i + 1}](./samples/sample_{i}.gif)"
+                        "<br>"
+                        '<details style="max-width: 300px; margin: auto;">'
+                        f"<summary>Prompt</summary>"
+                        f"{prompt}"
+                        "</details>"
+                    )
+                    cells.append(cell)
+                except Exception as e:
+                    logger.error(f"Failed to process video {video}: {e}")
+
+        # Calculate optimal grid dimensions
+        num_cells = len(cells)
+        if num_cells > 0:
+            # Aim for a roughly square grid, with max 4 columns
+            num_cols = min(4, num_cells)
+            num_rows = (num_cells + num_cols - 1) // num_cols  # Ceiling division
+
+            # Create grid rows
+            for row in range(num_rows):
+                start_idx = row * num_cols
+                end_idx = min(start_idx + num_cols, num_cells)
+                row_cells = cells[start_idx:end_idx]
+                # Properly format the row with table markers and exact number of cells
+                formatted_row = "| " + " | ".join(row_cells) + " |"
+                sample_grid.append(formatted_row)
+
+    # Join grid rows with just the content, no headers needed
+    grid_text = "\n".join(sample_grid) if sample_grid else ""
+
+    # Fill in the template
+    model_card_content = template.format(
+        base_model=base_model_name,
+        base_model_link=base_model_link,
+        model_name=model_name,
+        training_type="LoRA fine-tuning" if config.model.training_mode == "lora" else "Full model fine-tuning",
+        training_steps=config.optimization.steps,
+        learning_rate=config.optimization.learning_rate,
+        batch_size=config.optimization.batch_size,
+        validation_prompts=prompts_text,
+        sample_grid=grid_text,
+    )
+
+    # Save the model card directly
+    model_card_path = output_dir / "README.md"
+    model_card_path.write_text(model_card_content)
+
+    return model_card_path

packages/ltx-trainer/src/ltx_trainer/model_loader.py

@@ -0,0 +1,378 @@
+# ruff: noqa: PLC0415
+
+"""
+Model loader for LTX-2 trainer using the new ltx-core package.
+This module provides a unified interface for loading LTX-2 model components
+for training, using SingleGPUModelBuilder from ltx-core.
+Example usage:
+    # Load individual components
+    vae_encoder = load_video_vae_encoder("/path/to/checkpoint.safetensors", device="cuda")
+    vae_decoder = load_video_vae_decoder("/path/to/checkpoint.safetensors", device="cuda")
+    text_encoder = load_text_encoder("/path/to/gemma", device="cuda")
+    # Load all components at once
+    components = load_model("/path/to/checkpoint.safetensors", text_encoder_path="/path/to/gemma")
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+
+from ltx_trainer import logger
+
+# Type alias for device specification
+Device = str | torch.device
+
+# Type checking imports (not loaded at runtime)
+if TYPE_CHECKING:
+    from ltx_core.components.schedulers import LTX2Scheduler
+    from ltx_core.model.audio_vae import AudioDecoder, AudioEncoder, Vocoder
+    from ltx_core.model.transformer import LTXModel
+    from ltx_core.model.video_vae import VideoDecoder, VideoEncoder
+    from ltx_core.text_encoders.gemma import GemmaTextEncoder
+    from ltx_core.text_encoders.gemma.embeddings_processor import EmbeddingsProcessor
+
+
+def _to_torch_device(device: Device) -> torch.device:
+    """Convert device specification to torch.device."""
+    return torch.device(device) if isinstance(device, str) else device
+
+
+# =============================================================================
+# Individual Component Loaders
+# =============================================================================
+
+
+def load_transformer(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "LTXModel":
+    """Load the LTX transformer model.
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded LTXModel transformer
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.model.transformer.model_configurator import (
+        LTXV_MODEL_COMFY_RENAMING_MAP,
+        LTXModelConfigurator,
+    )
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=LTXModelConfigurator,
+        model_sd_ops=LTXV_MODEL_COMFY_RENAMING_MAP,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_video_vae_encoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "VideoEncoder":
+    """Load the video VAE encoder (for preprocessing).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded VideoEncoder
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.model.video_vae import VAE_ENCODER_COMFY_KEYS_FILTER, VideoEncoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=VideoEncoderConfigurator,
+        model_sd_ops=VAE_ENCODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_video_vae_decoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "VideoDecoder":
+    """Load the video VAE decoder (for inference/validation).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded VideoDecoder
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.model.video_vae import VAE_DECODER_COMFY_KEYS_FILTER, VideoDecoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=VideoDecoderConfigurator,
+        model_sd_ops=VAE_DECODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_audio_vae_encoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "AudioEncoder":
+    """Load the audio VAE encoder (for preprocessing).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights (default bfloat16, but float32 recommended for quality)
+    Returns:
+        Loaded AudioEncoder
+    """
+    from ltx_core.loader import SingleGPUModelBuilder
+    from ltx_core.model.audio_vae import AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER, AudioEncoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=AudioEncoderConfigurator,
+        model_sd_ops=AUDIO_VAE_ENCODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_audio_vae_decoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "AudioDecoder":
+    """Load the audio VAE decoder.
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded AudioDecoder
+    """
+    from ltx_core.loader import SingleGPUModelBuilder
+    from ltx_core.model.audio_vae import AUDIO_VAE_DECODER_COMFY_KEYS_FILTER, AudioDecoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=AudioDecoderConfigurator,
+        model_sd_ops=AUDIO_VAE_DECODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_vocoder(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "Vocoder":
+    """Load the vocoder (for audio waveform generation).
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded Vocoder
+    """
+    from ltx_core.loader import SingleGPUModelBuilder
+    from ltx_core.model.audio_vae import VOCODER_COMFY_KEYS_FILTER, VocoderConfigurator
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=VocoderConfigurator,
+        model_sd_ops=VOCODER_COMFY_KEYS_FILTER,
+    ).build(device=_to_torch_device(device), dtype=dtype)
+
+
+def load_text_encoder(
+    gemma_model_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+    load_in_8bit: bool = False,
+) -> "GemmaTextEncoder":
+    """Load the Gemma text encoder.
+    Args:
+        gemma_model_path: Path to Gemma model directory
+        device: Device to load model on
+        dtype: Data type for model weights
+        load_in_8bit: Whether to load the Gemma model in 8-bit precision using bitsandbytes.
+    Returns:
+        Loaded GemmaTextEncoder
+    """
+    if not Path(gemma_model_path).is_dir():
+        raise ValueError(f"Gemma model path is not a directory: {gemma_model_path}")
+
+    # Use 8-bit loading path if requested
+    if load_in_8bit:
+        from ltx_trainer.gemma_8bit import load_8bit_gemma
+
+        return load_8bit_gemma(gemma_model_path, dtype, device=device)
+
+    # Standard loading path
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.text_encoders.gemma import (
+        GEMMA_LLM_KEY_OPS,
+        GEMMA_MODEL_OPS,
+        GemmaTextEncoderConfigurator,
+        module_ops_from_gemma_root,
+    )
+    from ltx_core.utils import find_matching_file
+
+    torch_device = _to_torch_device(device)
+
+    gemma_model_folder = find_matching_file(str(gemma_model_path), "model*.safetensors").parent
+    gemma_weight_paths = [str(p) for p in gemma_model_folder.rglob("*.safetensors")]
+
+    text_encoder = SingleGPUModelBuilder(
+        model_path=tuple(gemma_weight_paths),
+        model_class_configurator=GemmaTextEncoderConfigurator,
+        model_sd_ops=GEMMA_LLM_KEY_OPS,
+        module_ops=(GEMMA_MODEL_OPS, *module_ops_from_gemma_root(str(gemma_model_path))),
+    ).build(device=torch_device, dtype=dtype)
+
+    return text_encoder
+
+
+def load_embeddings_processor(
+    checkpoint_path: str | Path,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+) -> "EmbeddingsProcessor":
+    """Load the embeddings processor (feature extractor + video/audio connectors).
+    Args:
+        checkpoint_path: Path to the LTX-2 safetensors checkpoint file
+        device: Device to load model on
+        dtype: Data type for model weights
+    Returns:
+        Loaded EmbeddingsProcessor with feature extractor and connectors
+    """
+    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder
+    from ltx_core.text_encoders.gemma import (
+        EMBEDDINGS_PROCESSOR_KEY_OPS,
+        EmbeddingsProcessorConfigurator,
+    )
+
+    torch_device = _to_torch_device(device)
+
+    return SingleGPUModelBuilder(
+        model_path=str(checkpoint_path),
+        model_class_configurator=EmbeddingsProcessorConfigurator,
+        model_sd_ops=EMBEDDINGS_PROCESSOR_KEY_OPS,
+    ).build(device=torch_device, dtype=dtype)
+
+
+# =============================================================================
+# Combined Component Loader
+# =============================================================================
+
+
+@dataclass
+class LtxModelComponents:
+    """Container for all LTX-2 model components."""
+
+    transformer: "LTXModel"
+    video_vae_encoder: "VideoEncoder | None" = None
+    video_vae_decoder: "VideoDecoder | None" = None
+    audio_vae_decoder: "AudioDecoder | None" = None
+    vocoder: "Vocoder | None" = None
+    text_encoder: "GemmaTextEncoder | None" = None
+    scheduler: "LTX2Scheduler | None" = None
+
+
+def load_model(
+    checkpoint_path: str | Path,
+    text_encoder_path: str | Path | None = None,
+    device: Device = "cpu",
+    dtype: torch.dtype = torch.bfloat16,
+    with_video_vae_encoder: bool = False,
+    with_video_vae_decoder: bool = True,
+    with_audio_vae_decoder: bool = True,
+    with_vocoder: bool = True,
+    with_text_encoder: bool = True,
+) -> LtxModelComponents:
+    """
+    Load LTX-2 model components from a safetensors checkpoint.
+    This is a convenience function that loads multiple components at once.
+    For loading individual components, use the dedicated functions:
+    - load_transformer()
+    - load_video_vae_encoder()
+    - load_video_vae_decoder()
+    - load_audio_vae_decoder()
+    - load_vocoder()
+    - load_text_encoder()
+    Args:
+        checkpoint_path: Path to the safetensors checkpoint file
+        text_encoder_path: Path to Gemma model directory (required if with_text_encoder=True)
+        device: Device to load models on ("cuda", "cpu", etc.)
+        dtype: Data type for model weights
+        with_video_vae_encoder: Whether to load the video VAE encoder (for preprocessing)
+        with_video_vae_decoder: Whether to load the video VAE decoder (for inference/validation)
+        with_audio_vae_decoder: Whether to load the audio VAE decoder
+        with_vocoder: Whether to load the vocoder
+        with_text_encoder: Whether to load the text encoder
+    Returns:
+        LtxModelComponents containing all loaded model components
+    """
+    from ltx_core.components.schedulers import LTX2Scheduler
+
+    checkpoint_path = Path(checkpoint_path)
+
+    # Validate checkpoint exists
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+    logger.info(f"Loading LTX-2 model from {checkpoint_path}")
+
+    torch_device = _to_torch_device(device)
+
+    # Load transformer
+    logger.debug("Loading transformer...")
+    transformer = load_transformer(checkpoint_path, torch_device, dtype)
+
+    # Load video VAE encoder
+    video_vae_encoder = None
+    if with_video_vae_encoder:
+        logger.debug("Loading video VAE encoder...")
+        video_vae_encoder = load_video_vae_encoder(checkpoint_path, torch_device, dtype)
+
+    # Load video VAE decoder
+    video_vae_decoder = None
+    if with_video_vae_decoder:
+        logger.debug("Loading video VAE decoder...")
+        video_vae_decoder = load_video_vae_decoder(checkpoint_path, torch_device, dtype)
+
+    # Load audio VAE decoder
+    audio_vae_decoder = None
+    if with_audio_vae_decoder:
+        logger.debug("Loading audio VAE decoder...")
+        audio_vae_decoder = load_audio_vae_decoder(checkpoint_path, torch_device, dtype)
+
+    # Load vocoder
+    vocoder = None
+    if with_vocoder:
+        logger.debug("Loading vocoder...")
+        vocoder = load_vocoder(checkpoint_path, torch_device, dtype)
+
+    # Load text encoder
+    text_encoder = None
+    if with_text_encoder:
+        if text_encoder_path is None:
+            raise ValueError("text_encoder_path must be provided when with_text_encoder=True")
+        logger.debug("Loading Gemma text encoder...")
+        text_encoder = load_text_encoder(text_encoder_path, torch_device, dtype)
+
+    # Create scheduler (stateless, no loading needed)
+    scheduler = LTX2Scheduler()
+
+    return LtxModelComponents(
+        transformer=transformer,
+        video_vae_encoder=video_vae_encoder,
+        video_vae_decoder=video_vae_decoder,
+        audio_vae_decoder=audio_vae_decoder,
+        vocoder=vocoder,
+        text_encoder=text_encoder,
+        scheduler=scheduler,
+    )

packages/ltx-trainer/src/ltx_trainer/progress.py

@@ -0,0 +1,182 @@
+"""Progress tracking for LTX training.
+This module provides a unified progress display for training and validation sampling,
+encapsulating all Rich progress bar logic in one place.
+"""
+
+from rich.progress import (
+    BarColumn,
+    Progress,
+    TaskID,
+    TextColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+)
+
+
+class SamplingContext:
+    """Context for validation sampling progress tracking.
+    Provides a unified progress display showing current video and denoising step.
+    Display format: "Sampling X/Y [████████████] step Z/W"
+    The progress bar shows the denoising progress for the current video.
+    """
+
+    def __init__(self, progress: Progress | None, task: TaskID | None, num_prompts: int, num_steps: int):
+        self._progress = progress
+        self._task = task
+        self._num_prompts = num_prompts
+        self._num_steps = num_steps
+
+    def start_video(self, video_idx: int) -> None:
+        """Start tracking a new video (resets step progress)."""
+        if self._progress is None or self._task is None:
+            return
+        # Reset task for new video: completed=0, total=num_steps
+        self._progress.reset(self._task, total=self._num_steps)
+        self._progress.update(
+            self._task,
+            completed=0,
+            video=f"{video_idx + 1}/{self._num_prompts}",
+            info=f"step 0/{self._num_steps}",
+        )
+
+    def advance_step(self) -> None:
+        """Advance the denoising step by one."""
+        if self._progress is None or self._task is None:
+            return
+        self._progress.advance(self._task)
+        completed = int(self._progress.tasks[self._task].completed)
+        self._progress.update(self._task, info=f"step {completed}/{self._num_steps}")
+
+    def cleanup(self) -> None:
+        """Hide sampling task when done."""
+        if self._progress is None or self._task is None:
+            return
+        self._progress.update(self._task, visible=False)
+
+
+class TrainingProgress:
+    """Manages Rich progress display for training and validation.
+    This class encapsulates all progress bar logic, providing a clean interface
+    for the trainer to update progress without dealing with Rich internals.
+    Usage:
+        with TrainingProgress(enabled=True, total_steps=1000) as progress:
+            for step in range(1000):
+                # ... training step ...
+                progress.update_training(loss=0.1, lr=1e-4, step_time=0.5)
+                if should_validate:
+                    sampling_ctx = progress.start_sampling(num_prompts=3, num_steps=30)
+                    sampler = ValidationSampler(..., sampling_context=sampling_ctx)
+                    for prompt_idx, prompt in enumerate(prompts):
+                        sampling_ctx.start_video(prompt_idx)
+                        sampler.generate(...)
+                    sampling_ctx.cleanup()
+    """
+
+    def __init__(self, enabled: bool, total_steps: int):
+        """Initialize progress tracking.
+        Args:
+            enabled: Whether to display progress bars (False for non-main processes)
+            total_steps: Total number of training steps
+        """
+        self._enabled = enabled
+        self._total_steps = total_steps
+        self._train_task: TaskID | None = None
+
+        if not enabled:
+            self._progress = None
+            return
+
+        # Single Progress instance with flexible columns
+        self._progress = Progress(
+            TextColumn("[progress.description]{task.description}"),
+            TextColumn("{task.fields[video]}", style="magenta"),
+            BarColumn(bar_width=40, style="blue"),
+            TextColumn("{task.fields[info]}", style="cyan"),
+            TimeElapsedColumn(),
+            TextColumn("ETA:"),
+            TimeRemainingColumn(compact=True),
+        )
+
+    def __enter__(self) -> "TrainingProgress":
+        """Enter the progress context, starting the live display."""
+        if self._progress is not None:
+            self._progress.__enter__()
+            self._train_task = self._progress.add_task(
+                "Training",
+                total=self._total_steps,
+                video=f"0/{self._total_steps}",
+                info="Starting...",
+            )
+        return self
+
+    def __exit__(self, *args) -> None:
+        """Exit the progress context, stopping the live display."""
+        if self._progress is not None:
+            self._progress.__exit__(*args)
+
+    @property
+    def enabled(self) -> bool:
+        """Whether progress display is enabled."""
+        return self._enabled
+
+    def update_training(
+        self,
+        *,
+        loss: float,
+        lr: float,
+        step_time: float,
+        advance: bool = True,
+    ) -> None:
+        """Update the training progress display.
+        Args:
+            loss: Current training loss
+            lr: Current learning rate
+            step_time: Time taken for this step in seconds
+            advance: Whether to advance the progress by one step
+        """
+        if self._progress is None or self._train_task is None:
+            return
+
+        info = f"Loss: {loss:.4f} | LR: {lr:.2e} | {step_time:.2f}s/step"
+        self._progress.update(
+            self._train_task,
+            advance=1 if advance else 0,
+            info=info,
+        )
+        # Update step count in video column
+        completed = int(self._progress.tasks[self._train_task].completed)
+        self._progress.update(self._train_task, video=f"{completed}/{self._total_steps}")
+
+    def start_sampling(self, num_prompts: int, num_steps: int) -> SamplingContext:
+        """Start validation sampling progress tracking.
+        Creates a task that shows current video and denoising step progress.
+        Format: "Sampling X/Y [████████████] step Z/W"
+        Args:
+            num_prompts: Number of validation prompts to sample
+            num_steps: Number of denoising steps per sample
+        Returns:
+            SamplingContext for tracking progress (no-op if progress is disabled)
+        """
+        if self._progress is None:
+            # Return a no-op context when progress is disabled
+            return SamplingContext(
+                progress=None,
+                task=None,
+                num_prompts=num_prompts,
+                num_steps=num_steps,
+            )
+
+        task = self._progress.add_task(
+            "Sampling",
+            total=num_steps,
+            completed=0,
+            video=f"0/{num_prompts}",
+            info=f"step 0/{num_steps}",
+        )
+
+        return SamplingContext(
+            progress=self._progress,
+            task=task,
+            num_prompts=num_prompts,
+            num_steps=num_steps,
+        )

packages/ltx-trainer/src/ltx_trainer/quantization.py

@@ -0,0 +1,195 @@
+# Adapted from: https://github.com/bghira/SimpleTuner
+# With improvements from: https://github.com/ostris/ai-toolkit
+from typing import Literal
+
+import torch
+from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
+
+from ltx_trainer import logger
+
+QuantizationOptions = Literal[
+    "int8-quanto",
+    "int4-quanto",
+    "int2-quanto",
+    "fp8-quanto",
+    "fp8uz-quanto",
+]
+
+# Modules to exclude from quantization.
+# These are glob patterns passed to quanto's `exclude` parameter.
+# When quantizing the full model at once, these patterns match against full module paths.
+# When quantizing block-by-block, we also use SKIP_ROOT_MODULES for top-level modules.
+EXCLUDE_PATTERNS = [
+    # Input/output projection layers
+    "patchify_proj",
+    "audio_patchify_proj",
+    "proj_out",
+    "audio_proj_out",
+    # Timestep embedding layers - int4 tinygemm requires strict bfloat16 input
+    # and these receive float32 sinusoidal embeddings that are cast to bfloat16
+    "*adaln*",
+    "time_proj",
+    "timestep_embedder*",
+    # Caption/text projection layers
+    "caption_projection*",
+    "audio_caption_projection*",
+    # Normalization layers (usually excluded from quantization)
+    "*norm*",
+]
+
+# Top-level modules to skip entirely during block-by-block quantization.
+# These are exact matches against model.named_children() names.
+# (Needed because quanto's exclude patterns don't work when calling quantize() directly on a module)
+SKIP_ROOT_MODULES = {
+    "patchify_proj",
+    "audio_patchify_proj",
+    "proj_out",
+    "audio_proj_out",
+    "audio_caption_projection",
+}
+
+
+def quantize_model(
+    model: torch.nn.Module,
+    precision: QuantizationOptions,
+    quantize_activations: bool = False,
+    device: torch.device | str | None = None,
+) -> torch.nn.Module:
+    """
+    Quantize a model using optimum-quanto.
+    For large models with transformer_blocks, this function quantizes block-by-block
+    on GPU then moves back to CPU, which is much faster than quantizing on CPU and
+    uses less peak VRAM than loading the entire model to GPU at once.
+    Args:
+        model: The model to quantize.
+        precision: The quantization precision (e.g. "int8-quanto", "fp8-quanto").
+        quantize_activations: Whether to quantize activations in addition to weights.
+        device: Device to use for quantization. If None, uses CUDA if available, else CPU.
+    Returns:
+        The quantized model.
+    """
+    from optimum.quanto import freeze, quantize  # noqa: PLC0415
+
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    elif isinstance(device, str):
+        device = torch.device(device)
+
+    weight_quant = _get_quanto_dtype(precision)
+
+    if quantize_activations:
+        logger.debug("Quantizing model weights and activations")
+        activations_quant = weight_quant
+    else:
+        activations_quant = None
+
+    # Remember original device to restore after quantization
+    original_device = next(model.parameters()).device
+
+    # Check if model has transformer_blocks for block-by-block quantization
+    if hasattr(model, "transformer_blocks"):
+        logger.debug("Quantizing model using block-by-block approach for memory efficiency")
+        _quantize_blockwise(
+            model,
+            weight_quant=weight_quant,
+            activations_quant=activations_quant,
+            device=device,
+        )
+    else:
+        # Fallback: quantize entire model at once
+        model.to(device)
+        quantize(model, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
+        freeze(model)
+
+    # Restore model to original device
+    model.to(original_device)
+
+    return model
+
+
+def _quantize_blockwise(
+    model: torch.nn.Module,
+    weight_quant: torch.dtype,
+    activations_quant: torch.dtype | None,
+    device: torch.device,
+) -> None:
+    """Quantize a model block-by-block using optimum-quanto.
+    This approach:
+    1. Moves each transformer block to GPU
+    2. Quantizes on GPU (fast!)
+    3. Freezes the quantized weights
+    4. Moves back to CPU
+    This is much faster than quantizing on CPU and uses less peak VRAM
+    than loading the entire model to GPU.
+    """
+    from optimum.quanto import freeze, quantize  # noqa: PLC0415
+
+    original_dtype = next(model.parameters()).dtype
+    transformer_blocks = list(model.transformer_blocks)
+
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        transient=True,
+    ) as progress:
+        task = progress.add_task("Quantizing transformer blocks", total=len(transformer_blocks))
+
+        for block in transformer_blocks:
+            # Move block to GPU
+            block.to(device, dtype=original_dtype, non_blocking=True)
+
+            # Quantize on GPU
+            quantize(block, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
+            freeze(block)
+
+            # Move back to CPU to free up VRAM for next block
+            block.to("cpu", non_blocking=True)
+
+            progress.advance(task)
+
+    # Quantize remaining non-transformer-block modules (e.g., embeddings, timestep projections)
+    # Skip modules that should not be quantized (patchify_proj, proj_out, etc.)
+    logger.debug("Quantizing remaining model components")
+
+    for name, module in model.named_children():
+        if name == "transformer_blocks":
+            continue  # Already quantized
+
+        if name in SKIP_ROOT_MODULES:
+            logger.debug(f"Skipping quantization for module: {name}")
+            continue  # Don't quantize these modules
+
+        # Move to device, quantize, freeze, move back
+        module.to(device, dtype=original_dtype, non_blocking=True)
+        quantize(module, weights=weight_quant, activations=activations_quant, exclude=EXCLUDE_PATTERNS)
+        freeze(module)
+        module.to("cpu", non_blocking=True)
+
+
+def _get_quanto_dtype(precision: QuantizationOptions) -> torch.dtype:
+    """Map precision string to quanto dtype."""
+    from optimum.quanto import (  # noqa: PLC0415
+        qfloat8,
+        qfloat8_e4m3fnuz,
+        qint2,
+        qint4,
+        qint8,
+    )
+
+    if precision == "int2-quanto":
+        return qint2
+    elif precision == "int4-quanto":
+        return qint4
+    elif precision == "int8-quanto":
+        return qint8
+    elif precision in ("fp8-quanto", "fp8uz-quanto"):
+        if torch.backends.mps.is_available():
+            raise ValueError("FP8 quantization is not supported on MPS devices. Use int2, int4, or int8 instead.")
+        if precision == "fp8-quanto":
+            return qfloat8
+        elif precision == "fp8uz-quanto":
+            return qfloat8_e4m3fnuz
+
+    raise ValueError(f"Invalid quantization precision: {precision}")

packages/ltx-trainer/src/ltx_trainer/sigma_tracker.py

@@ -0,0 +1,59 @@
+"""Sigma-bucketed loss tracking.
+Maps each training step's per-element sigmas and losses to buckets.
+Smoothing is left to wandb's UI.
+"""
+
+import bisect
+from collections import defaultdict
+
+
+class SigmaBucketTracker:
+    """Map per-element sigma values to named buckets for per-bucket loss logging.
+    By default, partitions [0, 1] into four equal-width buckets.
+    Custom boundaries can be provided for non-uniform bucketing.
+    Each call to update() receives per-element sigmas and losses (both [B,]),
+    buckets each element, and computes the mean loss per bucket. This gives
+    accurate per-sigma loss tracking even for batch_size > 1.
+    """
+
+    def __init__(
+        self,
+        bucket_boundaries: list[float] | None = None,
+    ) -> None:
+        if bucket_boundaries is None:
+            bucket_boundaries = [0.0, 0.25, 0.5, 0.75, 1.0]
+        if len(bucket_boundaries) < 2:
+            raise ValueError("bucket_boundaries must have at least 2 elements")
+        if any(bucket_boundaries[i] >= bucket_boundaries[i + 1] for i in range(len(bucket_boundaries) - 1)):
+            raise ValueError("bucket_boundaries must be strictly increasing")
+        self._boundaries = list(bucket_boundaries)
+        self._num_buckets = len(bucket_boundaries) - 1
+        self._bucket_labels = [
+            f"{bucket_boundaries[i]:.2f}-{bucket_boundaries[i + 1]:.2f}" for i in range(self._num_buckets)
+        ]
+        self._last_metrics: dict[str, float] = {}
+
+    def _get_bucket_index(self, sigma: float) -> int:
+        """Map sigma value to bucket index."""
+        idx = bisect.bisect_right(self._boundaries, sigma) - 1
+        return max(0, min(idx, self._num_buckets - 1))
+
+    def update(self, sigmas: list[float], losses: list[float]) -> None:
+        """Record per-element losses into their sigma buckets.
+        Args:
+            sigmas: Per-element sigma values, one per batch element.
+            losses: Per-element losses, one per batch element.
+        """
+        if not sigmas:
+            self._last_metrics = {}
+            return
+        bucket_losses: dict[int, list[float]] = defaultdict(list)
+        for sigma, loss in zip(sigmas, losses, strict=True):
+            bucket_losses[self._get_bucket_index(sigma)].append(loss)
+        self._last_metrics = {self._bucket_labels[b]: sum(vals) / len(vals) for b, vals in bucket_losses.items()}
+
+    def get_metrics(self, prefix: str = "train") -> dict[str, float]:
+        """Return the mean loss for each bucket hit on the last update.
+        Wandb handles smoothing in the UI.
+        """
+        return {f"{prefix}/loss_sigma_{label}": loss for label, loss in self._last_metrics.items()}

packages/ltx-trainer/src/ltx_trainer/timestep_samplers.py

@@ -0,0 +1,160 @@
+import torch
+
+
+class TimestepSampler:
+    """Base class for timestep samplers.
+    Timestep samplers are used to sample timesteps for diffusion models.
+    They should implement both sample() and sample_for() methods.
+    """
+
+    def sample(self, batch_size: int, seq_length: int | None = None, device: torch.device = None) -> torch.Tensor:
+        """Sample timesteps for a batch.
+        Args:
+            batch_size: Number of timesteps to sample
+            seq_length: (optional) Length of the sequence being processed
+            device: Device to place the samples on
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps
+        """
+        raise NotImplementedError
+
+    def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
+        """Sample timesteps for a specific batch tensor.
+        Args:
+            batch: Input tensor of shape (batch_size, seq_length, ...)
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps
+        """
+        raise NotImplementedError
+
+
+class UniformTimestepSampler(TimestepSampler):
+    """Samples timesteps uniformly between min_value and max_value (default 0 and 1)."""
+
+    def __init__(self, min_value: float = 0.0, max_value: float = 1.0):
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def sample(self, batch_size: int, seq_length: int | None = None, device: torch.device = None) -> torch.Tensor:  # noqa: ARG002
+        return torch.rand(batch_size, device=device) * (self.max_value - self.min_value) + self.min_value
+
+    def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
+        if batch.ndim != 3:
+            raise ValueError(f"Batch should have 3 dimensions, got {batch.ndim}")
+
+        return self.sample(batch.shape[0], device=batch.device)
+
+
+class ShiftedLogitNormalTimestepSampler(TimestepSampler):
+    """
+    Samples timesteps from a stretched shifted logit-normal distribution,
+    where the shift is determined by the sequence length.
+    The stretching normalizes samples between percentile bounds to ensure
+    the distribution covers [0, 1] more evenly. A uniform fallback prevents
+    collapse at high token counts.
+    """
+
+    def __init__(self, std: float = 1.0, eps: float = 1e-3, uniform_prob: float = 0.1):
+        self.std = std
+        self.eps = eps
+        self.uniform_prob = uniform_prob
+        # Percentile values for stretching (scaled by std)
+        # 99.9th percentile of standard normal ≈ 3.0902
+        # 0.5th percentile of standard normal ≈ -2.5758
+        self.normal_999_percentile = 3.0902 * std
+        self.normal_005_percentile = -2.5758 * std
+
+    def sample(self, batch_size: int, seq_length: int, device: torch.device = None) -> torch.Tensor:
+        """Sample timesteps for a batch from a stretched shifted logit-normal distribution.
+        Args:
+            batch_size: Number of timesteps to sample
+            seq_length: Length of the sequence being processed, used to determine the shift
+            device: Device to place the samples on
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps sampled from a stretched
+            shifted logit-normal distribution, where the shift is determined by seq_length
+        """
+        mu = self._get_shift_for_sequence_length(seq_length)
+
+        # Sample from shifted logit-normal
+        normal_samples = torch.randn((batch_size,), device=device) * self.std + mu
+        logitnormal_samples = torch.sigmoid(normal_samples)
+
+        # Compute percentile bounds for stretching
+        percentile_999 = torch.sigmoid(torch.tensor(mu + self.normal_999_percentile, device=device))
+        percentile_005 = torch.sigmoid(torch.tensor(mu + self.normal_005_percentile, device=device))
+
+        # Stretch to [0, 1] range by normalizing between percentiles
+        zero_terminal_raw = (logitnormal_samples - percentile_005) / (percentile_999 - percentile_005)
+
+        # Reflect small values around eps for numerical stability
+        stretched_logit = torch.where(
+            zero_terminal_raw >= self.eps,
+            zero_terminal_raw,
+            2 * self.eps - zero_terminal_raw,
+        )
+        stretched_logit = torch.clamp(stretched_logit, 0, 1)
+
+        # Mix with uniform samples (uniform_prob of the time)
+        uniform = (1 - self.eps) * torch.rand((batch_size,), device=device) + self.eps
+        prob = torch.rand((batch_size,), device=device)
+
+        return torch.where(prob > self.uniform_prob, stretched_logit, uniform)
+
+    def sample_for(self, batch: torch.Tensor) -> torch.Tensor:
+        """Sample timesteps for a specific batch tensor.
+        Args:
+            batch: Input tensor of shape (batch_size, seq_length, ...)
+        Returns:
+            Tensor of shape (batch_size,) containing timesteps sampled from a shifted
+            logit-normal distribution, where the shift is determined by the sequence length
+            of the input batch
+        Raises:
+            ValueError: If the input batch does not have 3 dimensions
+        """
+        if batch.ndim != 3:
+            raise ValueError(f"Batch should have 3 dimensions, got {batch.ndim}")
+
+        batch_size, seq_length, _ = batch.shape
+        return self.sample(batch_size, seq_length, device=batch.device)
+
+    @staticmethod
+    def _get_shift_for_sequence_length(
+        seq_length: int,
+        min_tokens: int = 1024,
+        max_tokens: int = 4096,
+        min_shift: float = 0.95,
+        max_shift: float = 2.05,
+    ) -> float:
+        # Calculate the shift value for a given sequence length using linear interpolation
+        # between min_shift and max_shift based on sequence length.
+        m = (max_shift - min_shift) / (max_tokens - min_tokens)  # Calculate slope
+        b = min_shift - m * min_tokens  # Calculate y-intercept
+        shift = m * seq_length + b  # Apply linear equation y = mx + b
+        return shift
+
+
+SAMPLERS = {
+    "uniform": UniformTimestepSampler,
+    "shifted_logit_normal": ShiftedLogitNormalTimestepSampler,
+}
+
+
+def example() -> None:
+    # noinspection PyUnresolvedReferences
+    import matplotlib.pyplot as plt  # noqa: PLC0415
+
+    sampler = ShiftedLogitNormalTimestepSampler()
+    for seq_length in [1024, 2048, 4096, 8192]:
+        samples = sampler.sample(batch_size=1_000_000, seq_length=seq_length)
+
+        # plot the histogram of the samples
+        plt.hist(samples.numpy(), bins=100, density=True)
+        plt.title(f"Timestep Samples for Sequence Length {seq_length}")
+        plt.xlabel("Timestep")
+        plt.ylabel("Density")
+        plt.show()
+
+
+if __name__ == "__main__":
+    example()