diff --git a/packages/ltx-trainer/configs/a2v_lora.yaml b/packages/ltx-trainer/configs/a2v_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..45c91bfeb39df5549504d9d298753a1d867f524e
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/accelerate/ddp.yaml b/packages/ltx-trainer/configs/accelerate/ddp.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..2edb13252bf19959630a5d1d53186ff156c84c4f
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/configs/accelerate/ddp_compile.yaml b/packages/ltx-trainer/configs/accelerate/ddp_compile.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..9398f94928d2e769ca14908011cc3ad6949ca6af
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/configs/accelerate/fsdp.yaml b/packages/ltx-trainer/configs/accelerate/fsdp.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..6a8b0028c5d5d0f01766b927577eb91ceeb1b6f5
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/configs/accelerate/fsdp_compile.yaml b/packages/ltx-trainer/configs/accelerate/fsdp_compile.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..a6b7089d9a50e5590c5360c47fdbe8b6dd5f9560
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/configs/audio_extend_lora.yaml b/packages/ltx-trainer/configs/audio_extend_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..fd655882a0918f0cfe483f27e795185de2444050
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/audio_inpainting_lora.yaml b/packages/ltx-trainer/configs/audio_inpainting_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..e5c8361b769dc15c6a628e24ec423837366016fa
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/audio_suffix_lora.yaml b/packages/ltx-trainer/configs/audio_suffix_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..9ef48a1f6f667739d8de228cf76437921f49b42e
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/av2av_ic_lora.yaml b/packages/ltx-trainer/configs/av2av_ic_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..78d12c9b0fadc03041a5a20337bd0b2b7828ae22
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/i2v_lora.yaml b/packages/ltx-trainer/configs/i2v_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..373a7911327b0f59af7e2d13990735bbc66f3cad
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/t2a_lora.yaml b/packages/ltx-trainer/configs/t2a_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..641a9bcdcc26c15004639c310b69e772878d2eeb
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/t2v_lora.yaml b/packages/ltx-trainer/configs/t2v_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..5d8fdae6ab3d78d111319249c9ccdf7c571908bf
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/t2v_lora_low_vram.yaml b/packages/ltx-trainer/configs/t2v_lora_low_vram.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..2d2b434fd9d7fe5da52055e61563d588b2f58f33
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/v2a_lora.yaml b/packages/ltx-trainer/configs/v2a_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..cc7231f28ae5bad08299fc8c12adab2f255d7ec8
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/v2v_ic_lora.yaml b/packages/ltx-trainer/configs/v2v_ic_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..349e4b3e462c8979f6b71a9d2a0be2d071b11aeb
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/video_extend_lora.yaml b/packages/ltx-trainer/configs/video_extend_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..39a2d4f5ff7e3862582c60824dcaf67eb7b2bbdf
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/video_inpainting_lora.yaml b/packages/ltx-trainer/configs/video_inpainting_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..46829b38dff6b7a6799d659c8aa4d62abe0a6f00
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/video_outpainting_lora.yaml b/packages/ltx-trainer/configs/video_outpainting_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..e26518aab24f4c2aacc51db9100f548a0d789845
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/configs/video_suffix_lora.yaml b/packages/ltx-trainer/configs/video_suffix_lora.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..b9617f19caf3a9897c020907158dc8ef3d5e39a8
--- /dev/null
+++ b/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"
diff --git a/packages/ltx-trainer/docs/configuration-reference.md b/packages/ltx-trainer/docs/configuration-reference.md
new file mode 100644
index 0000000000000000000000000000000000000000..e9d17bec80c88143f8a923004fb95b8718dcc071
--- /dev/null
+++ b/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)
diff --git a/packages/ltx-trainer/docs/custom-training-strategies.md b/packages/ltx-trainer/docs/custom-training-strategies.md
new file mode 100644
index 0000000000000000000000000000000000000000..e24d9a03b90f7e9c55f9ab4cc27c13505949137a
--- /dev/null
+++ b/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) |
diff --git a/packages/ltx-trainer/docs/dataset-preparation.md b/packages/ltx-trainer/docs/dataset-preparation.md
new file mode 100644
index 0000000000000000000000000000000000000000..6cb0c4023b25c130d92bc9002e83ad9cf7d4f715
--- /dev/null
+++ b/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).
diff --git a/packages/ltx-trainer/docs/quick-start.md b/packages/ltx-trainer/docs/quick-start.md
new file mode 100644
index 0000000000000000000000000000000000000000..c81953a5f88d9886c272206f6dacbc8703e4485d
--- /dev/null
+++ b/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!
diff --git a/packages/ltx-trainer/docs/training-guide.md b/packages/ltx-trainer/docs/training-guide.md
new file mode 100644
index 0000000000000000000000000000000000000000..a61d246bd80b0935b6e51e68ec907989d644a628
--- /dev/null
+++ b/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!
diff --git a/packages/ltx-trainer/docs/training-modes.md b/packages/ltx-trainer/docs/training-modes.md
new file mode 100644
index 0000000000000000000000000000000000000000..2469c77e7cbf93c226f849bd54263c9e6ff6b8a9
--- /dev/null
+++ b/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.
diff --git a/packages/ltx-trainer/docs/troubleshooting.md b/packages/ltx-trainer/docs/troubleshooting.md
new file mode 100644
index 0000000000000000000000000000000000000000..ddf8c4ae4b4cf1578f5c7caefa14037b52650828
--- /dev/null
+++ b/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!
diff --git a/packages/ltx-trainer/docs/utility-scripts.md b/packages/ltx-trainer/docs/utility-scripts.md
new file mode 100644
index 0000000000000000000000000000000000000000..3364a83dd2d58a05eb6ff2a04ba06fab784348a1
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/pyproject.toml b/packages/ltx-trainer/pyproject.toml
new file mode 100644
index 0000000000000000000000000000000000000000..ce956d61f6dafa2c134f2769449261b157990f8f
--- /dev/null
+++ b/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"]
diff --git a/packages/ltx-trainer/scripts/caption_videos.py b/packages/ltx-trainer/scripts/caption_videos.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed3fa93d19d7454b5f968435dfbf0a98f71902b9
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/compute_reference.py b/packages/ltx-trainer/scripts/compute_reference.py
new file mode 100644
index 0000000000000000000000000000000000000000..20f515413e74143330b1e6ea388901f267c61cc2
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/decode_latents.py b/packages/ltx-trainer/scripts/decode_latents.py
new file mode 100644
index 0000000000000000000000000000000000000000..f27ff4fcecfb842adc905a827dcadc3b4ddd2f4a
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/process_captions.py b/packages/ltx-trainer/scripts/process_captions.py
new file mode 100644
index 0000000000000000000000000000000000000000..037848e93365c07b18d9644b78d70412f48a138a
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/process_dataset.py b/packages/ltx-trainer/scripts/process_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..74d3e3970655bf1a1ef27f0e1dc48bb4a477b7e3
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/process_videos.py b/packages/ltx-trainer/scripts/process_videos.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d2f96c5d65755c7538e97d9d5180ec1b45f76e0
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/serve_captioner.py b/packages/ltx-trainer/scripts/serve_captioner.py
new file mode 100644
index 0000000000000000000000000000000000000000..4cecc1070f70f01994ed79f0d5e1fe57d45f52a7
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/split_scenes.py b/packages/ltx-trainer/scripts/split_scenes.py
new file mode 100644
index 0000000000000000000000000000000000000000..0cfce5c6baedb5e26e22499ae9c804d1fa9e181b
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/scripts/train.py b/packages/ltx-trainer/scripts/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd58e410a4a43ed7288afd324540c277fe6eab75
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/src/ltx_trainer/__init__.py b/packages/ltx-trainer/src/ltx_trainer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..56efa403f9384f9dc4f631b4d63a52d2df6a7680
--- /dev/null
+++ b/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))
diff --git a/packages/ltx-trainer/src/ltx_trainer/captioning.py b/packages/ltx-trainer/src/ltx_trainer/captioning.py
new file mode 100644
index 0000000000000000000000000000000000000000..1aa55a8d26fb359481bacb377c26ecaa4a032d6a
--- /dev/null
+++ b/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)])
diff --git a/packages/ltx-trainer/src/ltx_trainer/config.py b/packages/ltx-trainer/src/ltx_trainer/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4d30092183e978eeeb251cb192216712ea9b820
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/src/ltx_trainer/config_display.py b/packages/ltx-trainer/src/ltx_trainer/config_display.py
new file mode 100644
index 0000000000000000000000000000000000000000..c9ddb2070e19625258a0e2bd91ee32fd977623b5
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/src/ltx_trainer/datasets.py b/packages/ltx-trainer/src/ltx_trainer/datasets.py
new file mode 100644
index 0000000000000000000000000000000000000000..2d8f59cb9d1142a967f01bc63ea6dcb96361c3ae
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py b/packages/ltx-trainer/src/ltx_trainer/gemma_8bit.py
new file mode 100644
index 0000000000000000000000000000000000000000..83727907cea391c4f6b7df638b9b02761764eb45
--- /dev/null
+++ b/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)
diff --git a/packages/ltx-trainer/src/ltx_trainer/gpu_utils.py b/packages/ltx-trainer/src/ltx_trainer/gpu_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..de385a1cefcadf504b537583fd2e180c5541f121
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/src/ltx_trainer/hf_hub_utils.py b/packages/ltx-trainer/src/ltx_trainer/hf_hub_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..32f9a998d9a5e8d13b7efbeb410a39781c537bb7
--- /dev/null
+++ b/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
diff --git a/packages/ltx-trainer/src/ltx_trainer/model_loader.py b/packages/ltx-trainer/src/ltx_trainer/model_loader.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6aeba62359144eb1d28372098df881b6688815e
--- /dev/null
+++ b/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,
+    )
diff --git a/packages/ltx-trainer/src/ltx_trainer/progress.py b/packages/ltx-trainer/src/ltx_trainer/progress.py
new file mode 100644
index 0000000000000000000000000000000000000000..114b2db57503890011b5335aa781d7609757657c
--- /dev/null
+++ b/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,
+        )
diff --git a/packages/ltx-trainer/src/ltx_trainer/quantization.py b/packages/ltx-trainer/src/ltx_trainer/quantization.py
new file mode 100644
index 0000000000000000000000000000000000000000..31bc9bc62f16c5c7a42729a8eba98a8de1461402
--- /dev/null
+++ b/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}")
diff --git a/packages/ltx-trainer/src/ltx_trainer/sigma_tracker.py b/packages/ltx-trainer/src/ltx_trainer/sigma_tracker.py
new file mode 100644
index 0000000000000000000000000000000000000000..499c0c044d520b5313c1cac7de53a7662dd69391
--- /dev/null
+++ b/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()}
diff --git a/packages/ltx-trainer/src/ltx_trainer/timestep_samplers.py b/packages/ltx-trainer/src/ltx_trainer/timestep_samplers.py
new file mode 100644
index 0000000000000000000000000000000000000000..30e849479bcc00ffeb5fa00ef2e77759112cc8e2
--- /dev/null
+++ b/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()
diff --git a/packages/ltx-trainer/src/ltx_trainer/trainer.py b/packages/ltx-trainer/src/ltx_trainer/trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..8eb3fbc5d5931a94360d2c81829a8dfcbe0c8bf2
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/trainer.py
@@ -0,0 +1,1117 @@
+import contextlib
+import math
+import os
+import re
+import time
+import warnings
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+import wandb
+import yaml
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import DistributedDataParallelKwargs, gather_object, set_seed
+from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
+from peft.tuners.tuners_utils import BaseTunerLayer
+from peft.utils import ModulesToSaveWrapper
+from pydantic import BaseModel
+from safetensors.torch import load_file, save_file
+from torch import Tensor
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import (
+    CosineAnnealingLR,
+    CosineAnnealingWarmRestarts,
+    LinearLR,
+    LRScheduler,
+    PolynomialLR,
+    StepLR,
+)
+from torch.utils.data import DataLoader
+
+from ltx_core.text_encoders.gemma import convert_to_additive_mask
+from ltx_trainer import logger
+from ltx_trainer.config import LtxTrainerConfig
+from ltx_trainer.config_display import print_config
+from ltx_trainer.datasets import PrecomputedDataset
+from ltx_trainer.gpu_utils import free_gpu_memory, get_gpu_memory_gb
+from ltx_trainer.hf_hub_utils import push_to_hub
+from ltx_trainer.model_loader import load_embeddings_processor, load_transformer
+from ltx_trainer.progress import TrainingProgress
+from ltx_trainer.quantization import quantize_model
+from ltx_trainer.sigma_tracker import SigmaBucketTracker
+from ltx_trainer.timestep_samplers import SAMPLERS
+from ltx_trainer.training_state import ConfigFingerprint, RngStates, TrainingState
+from ltx_trainer.training_strategies import get_training_strategy
+from ltx_trainer.validation_runner import ValidationRunner
+
+# Disable irrelevant warnings from transformers
+os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+# Silence bitsandbytes warnings about casting
+warnings.filterwarnings(
+    "ignore", message="MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantization"
+)
+
+# Disable progress bars if not main process
+IS_MAIN_PROCESS = os.environ.get("LOCAL_RANK", "0") == "0"
+if not IS_MAIN_PROCESS:
+    from transformers.utils.logging import disable_progress_bar
+
+    disable_progress_bar()
+
+StepCallback = Callable[[int, int, list[Path]], None]  # (step, total, list[sampled_video_path]) -> None
+
+MEMORY_CHECK_INTERVAL = 200
+
+
+class TrainingStats(BaseModel):
+    """Statistics collected during training"""
+
+    total_time_seconds: float
+    steps_per_second: float
+    samples_per_second: float
+    peak_gpu_memory_gb: float
+    global_batch_size: int
+    num_processes: int
+
+
+@dataclass(frozen=True)
+class TrainingStepOutput:
+    """Output from a single training step."""
+
+    loss: Tensor  # [B,] per-element loss (unreduced)
+    sigma: Tensor  # [B,] sampled sigma, detached from computational graph
+
+
+class LtxvTrainer:
+    def __init__(self, trainer_config: LtxTrainerConfig) -> None:
+        self._config = trainer_config
+        if IS_MAIN_PROCESS:
+            print_config(trainer_config)
+        self._training_strategy = get_training_strategy(self._config.training_strategy)
+
+        # ValidationRunner loads its own models (text encoder, VAE encoder/decoder, etc.),
+        # caches prompt embeddings and conditioning media, then unloads encoders.
+        self._validation_runner = ValidationRunner(
+            config=self._config.validation,
+            model_path=self._config.model.model_path,
+            text_encoder_path=self._config.model.text_encoder_path,
+            load_text_encoder_in_8bit=self._config.acceleration.load_text_encoder_in_8bit,
+        )
+
+        self._load_models()
+        self._setup_accelerator()
+        self._collect_trainable_params()
+        self._loaded_checkpoint_path: Path | None = None
+        self._load_checkpoint()
+        self._prepare_models_for_training()
+        self._dataset = None
+        self._global_step = -1
+        self._checkpoint_paths: list[Path] = []
+        self._training_state_paths: list[Path] = []
+        self._training_state_size_warned = False
+        self._sigma_tracker = SigmaBucketTracker()
+        self._wandb_run = None
+
+    def train(  # noqa: PLR0912, PLR0915
+        self,
+        disable_progress_bars: bool = False,
+        step_callback: StepCallback | None = None,
+    ) -> tuple[Path, TrainingStats]:
+        """
+        Start the training process.
+        Args:
+            disable_progress_bars: Disable Rich progress bars (useful for multi-process runs).
+            step_callback: Optional callback invoked after each optimization step.
+        Returns:
+            Tuple of (saved_model_path, training_stats)
+        """
+        device = self._accelerator.device
+        cfg = self._config
+        start_mem = get_gpu_memory_gb(device)
+
+        train_start_time = time.time()
+
+        initial_step, training_state = self._resume_state
+        resuming = training_state is not None
+
+        set_seed(cfg.seed)
+        logger.debug(f"Process {self._accelerator.process_index} using seed: {cfg.seed}")
+
+        self._init_optimizer()
+
+        if training_state is not None and not self._restore_training_state(training_state):
+            initial_step = 0
+            resuming = False
+
+        # Initialize W&B after restore so we only resume the run when state restore succeeds.
+        resume_run_id = training_state.wandb_run_id if resuming and training_state is not None else None
+        self._init_wandb(resume_run_id=resume_run_id)
+
+        self._init_dataloader()
+        data_iter = iter(self._dataloader)
+        self._init_timestep_sampler()
+
+        # Synchronize all processes after initialization
+        self._accelerator.wait_for_everyone()
+
+        Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
+
+        # Save the training configuration as YAML
+        self._save_config()
+
+        remaining_steps = cfg.optimization.steps - initial_step
+        if remaining_steps <= 0:
+            raise ValueError(
+                f"No remaining training steps: initial_step={initial_step} >= "
+                f"target_steps={cfg.optimization.steps}. Nothing to train."
+            )
+
+        if resuming:
+            logger.info(f"🚀 Resuming training from step {initial_step} → {cfg.optimization.steps}")
+        else:
+            logger.info("🚀 Starting training...")
+
+        # Create progress tracking (disabled for non-main processes or when explicitly disabled)
+        progress_enabled = IS_MAIN_PROCESS and not disable_progress_bars
+        progress = TrainingProgress(
+            enabled=progress_enabled,
+            total_steps=remaining_steps,
+        )
+
+        if IS_MAIN_PROCESS and disable_progress_bars:
+            logger.warning("Progress bars disabled. Intermediate status messages will be logged instead.")
+
+        self._transformer.train()
+        self._global_step = initial_step
+
+        peak_mem_during_training = start_mem
+
+        sampled_videos_paths = None
+
+        with progress:
+            if cfg.validation.interval and not cfg.validation.skip_initial_validation:
+                with self._offloaded_optimizer_state():
+                    sampled_videos_paths = self._run_validation(progress)
+
+            self._accelerator.wait_for_everyone()
+
+            for step in range(remaining_steps * cfg.optimization.gradient_accumulation_steps):
+                # Get next batch, reset the dataloader if needed
+                try:
+                    batch = next(data_iter)
+                except StopIteration:
+                    data_iter = iter(self._dataloader)
+                    batch = next(data_iter)
+
+                step_start_time = time.time()
+                with self._accelerator.accumulate(self._transformer):
+                    is_optimization_step = (step + 1) % cfg.optimization.gradient_accumulation_steps == 0
+                    if is_optimization_step:
+                        self._global_step += 1
+
+                    output = self._training_step(batch)
+                    self._accelerator.backward(output.loss.mean())
+
+                    if self._accelerator.sync_gradients and cfg.optimization.max_grad_norm > 0:
+                        self._accelerator.clip_grad_norm_(
+                            self._trainable_params,
+                            cfg.optimization.max_grad_norm,
+                        )
+
+                    self._optimizer.step()
+                    self._optimizer.zero_grad()
+
+                    if self._lr_scheduler is not None:
+                        self._lr_scheduler.step()
+
+                    # Run validation if needed (handles DDP/FSDP work distribution internally)
+                    if (
+                        cfg.validation.interval
+                        and self._global_step > 0
+                        and self._global_step % cfg.validation.interval == 0
+                        and is_optimization_step
+                    ):
+                        with self._offloaded_optimizer_state():
+                            sampled_videos_paths = self._run_validation(progress)
+
+                    # Save checkpoint if needed
+                    if (
+                        cfg.checkpoints.interval
+                        and self._global_step > 0
+                        and self._global_step % cfg.checkpoints.interval == 0
+                        and is_optimization_step
+                    ):
+                        self._save_checkpoint()
+
+                    self._accelerator.wait_for_everyone()
+
+                    # Call step callback if provided
+                    if step_callback and is_optimization_step:
+                        step_callback(self._global_step, cfg.optimization.steps, sampled_videos_paths)
+
+                    self._accelerator.wait_for_everyone()
+
+                    # Update progress and log metrics
+                    current_lr = self._optimizer.param_groups[0]["lr"]
+                    step_time = (time.time() - step_start_time) * cfg.optimization.gradient_accumulation_steps
+                    step_loss = output.loss.detach().mean().item()
+
+                    progress.update_training(
+                        loss=step_loss,
+                        lr=current_lr,
+                        step_time=step_time,
+                        advance=is_optimization_step,
+                    )
+
+                    # Log metrics to W&B (only on main process and optimization steps)
+                    if IS_MAIN_PROCESS and is_optimization_step:
+                        # Track per-element loss by sigma bucket
+                        self._sigma_tracker.update(output.sigma.cpu().tolist(), output.loss.detach().cpu().tolist())
+                        metrics = {
+                            "train/loss": step_loss,
+                            "train/learning_rate": current_lr,
+                            "train/step_time": step_time,
+                            "train/global_step": self._global_step,
+                        }
+                        metrics.update(self._sigma_tracker.get_metrics())
+                        self._log_metrics(metrics)
+
+                    # Fallback logging when progress bars are disabled
+                    if disable_progress_bars and IS_MAIN_PROCESS and self._global_step % 20 == 0:
+                        elapsed = time.time() - train_start_time
+                        steps_done = self._global_step - initial_step
+                        if steps_done > 0:
+                            total_estimated = elapsed / steps_done * remaining_steps
+                            total_time = f"{total_estimated // 3600:.0f}h {(total_estimated % 3600) // 60:.0f}m"
+                        else:
+                            total_time = "calculating..."
+                        logger.info(
+                            f"Step {self._global_step}/{cfg.optimization.steps} - "
+                            f"Loss: {step_loss:.4f}, LR: {current_lr:.2e}, "
+                            f"Time/Step: {step_time:.2f}s, Total Time: {total_time}",
+                        )
+
+                    # Sample GPU memory periodically
+                    if step % MEMORY_CHECK_INTERVAL == 0:
+                        current_mem = get_gpu_memory_gb(device)
+                        peak_mem_during_training = max(peak_mem_during_training, current_mem)
+
+        # Collect final stats
+        train_end_time = time.time()
+        end_mem = get_gpu_memory_gb(device)
+        peak_mem = max(start_mem, end_mem, peak_mem_during_training)
+
+        # Calculate steps/second over entire training
+        total_time_seconds = train_end_time - train_start_time
+        steps_per_second = remaining_steps / total_time_seconds
+
+        samples_per_second = steps_per_second * self._accelerator.num_processes * cfg.optimization.batch_size
+
+        stats = TrainingStats(
+            total_time_seconds=total_time_seconds,
+            steps_per_second=steps_per_second,
+            samples_per_second=samples_per_second,
+            peak_gpu_memory_gb=peak_mem,
+            num_processes=self._accelerator.num_processes,
+            global_batch_size=cfg.optimization.batch_size * self._accelerator.num_processes,
+        )
+
+        saved_path = self._save_checkpoint()
+
+        if IS_MAIN_PROCESS:
+            # Log the training statistics
+            self._log_training_stats(stats)
+
+            # Upload artifacts to hub if enabled
+            if cfg.hub.push_to_hub:
+                push_to_hub(saved_path, sampled_videos_paths, self._config)
+
+            # Log final stats to W&B
+            if self._wandb_run is not None:
+                self._log_metrics(
+                    {
+                        "stats/total_time_minutes": stats.total_time_seconds / 60,
+                        "stats/steps_per_second": stats.steps_per_second,
+                        "stats/samples_per_second": stats.samples_per_second,
+                        "stats/peak_gpu_memory_gb": stats.peak_gpu_memory_gb,
+                    }
+                )
+                self._wandb_run.finish()
+
+        self._accelerator.wait_for_everyone()
+        self._accelerator.end_training()
+
+        return saved_path, stats
+
+    def _training_step(self, batch: dict[str, dict[str, Tensor]]) -> TrainingStepOutput:
+        """Perform a single training step using the configured strategy."""
+        # Apply embedding connectors to transform pre-computed text embeddings
+        conditions = batch["conditions"]
+
+        if "video_prompt_embeds" in conditions:
+            # New format: separate video/audio features from precompute()
+            video_features = conditions["video_prompt_embeds"]
+            audio_features = conditions.get("audio_prompt_embeds")
+        else:
+            # Legacy format: single prompt_embeds tensor — duplicate for both modalities
+            video_features = conditions["prompt_embeds"]
+            audio_features = conditions["prompt_embeds"]
+
+        mask = conditions["prompt_attention_mask"]
+        additive_mask = convert_to_additive_mask(mask, video_features.dtype)
+        video_embeds, audio_embeds, attention_mask = self._embeddings_processor.create_embeddings(
+            video_features, audio_features, additive_mask
+        )
+
+        conditions["video_prompt_embeds"] = video_embeds
+        conditions["audio_prompt_embeds"] = audio_embeds
+        conditions["prompt_attention_mask"] = attention_mask
+
+        # Use strategy to prepare training inputs (returns ModelInputs with Modality objects)
+        model_inputs = self._training_strategy.prepare_training_inputs(batch, self._timestep_sampler)
+
+        # Run transformer forward pass with Modality-based interface
+        video_pred, audio_pred = self._transformer(
+            video=model_inputs.video,
+            audio=model_inputs.audio,
+            perturbations=None,
+        )
+
+        # Use strategy to compute loss (returns per-element [B,] for sigma-bucket tracking)
+        loss = self._training_strategy.compute_loss(video_pred, audio_pred, model_inputs)
+
+        # Sigma comes from whichever modality is generated (video preferred, else audio).
+        if model_inputs.video is not None and model_inputs.video.enabled:
+            sigma = model_inputs.video.sigma.detach()
+        else:
+            sigma = model_inputs.audio.sigma.detach()
+
+        return TrainingStepOutput(loss=loss, sigma=sigma)
+
+    def _load_models(self) -> None:
+        """Load the transformer and embeddings processor for training."""
+        logger.debug("Loading transformer...")
+        self._transformer = load_transformer(
+            checkpoint_path=self._config.model.model_path,
+            device="cpu",
+            dtype=torch.bfloat16,
+        )
+
+        # DDP-safe: LOCAL_RANK is set by accelerate before trainer init. Loading on bare
+        # "cuda" would resolve to cuda:0 on every rank and crash with a device mismatch.
+        local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+        init_device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
+
+        logger.debug("Loading embeddings processor...")
+        self._embeddings_processor = load_embeddings_processor(
+            checkpoint_path=self._config.model.model_path,
+            device=init_device,
+            dtype=torch.bfloat16,
+        )
+        self._embeddings_processor.feature_extractor = None
+
+        transformer_dtype = torch.bfloat16 if self._config.model.training_mode == "lora" else torch.float32
+        self._transformer = self._transformer.to(dtype=transformer_dtype)
+
+        if self._config.acceleration.quantization is not None:
+            if self._config.model.training_mode == "full":
+                raise ValueError("Quantization is not supported in full training mode.")
+
+            logger.info(f'Quantizing model with "{self._config.acceleration.quantization}". This may take a while...')
+            self._transformer = quantize_model(
+                self._transformer,
+                precision=self._config.acceleration.quantization,
+            )
+
+        self._transformer.requires_grad_(False)
+
+    def _collect_trainable_params(self) -> None:
+        """Collect trainable parameters based on training mode."""
+        if self._config.model.training_mode == "lora":
+            # For LoRA training, first set up LoRA layers
+            self._setup_lora()
+        elif self._config.model.training_mode == "full":
+            # For full training, unfreeze all transformer parameters
+            self._transformer.requires_grad_(True)
+        else:
+            raise ValueError(f"Unknown training mode: {self._config.model.training_mode}")
+
+        self._trainable_params = [p for p in self._transformer.parameters() if p.requires_grad]
+        logger.debug(f"Trainable params count: {sum(p.numel() for p in self._trainable_params):,}")
+
+    def _init_timestep_sampler(self) -> None:
+        """Initialize the timestep sampler based on the config."""
+        sampler_cls = SAMPLERS[self._config.flow_matching.timestep_sampling_mode]
+        self._timestep_sampler = sampler_cls(**self._config.flow_matching.timestep_sampling_params)
+
+    def _setup_lora(self) -> None:
+        """Configure LoRA adapters for the transformer. Only called in LoRA training mode."""
+        logger.debug(f"Adding LoRA adapter with rank {self._config.lora.rank}")
+        lora_config = LoraConfig(
+            r=self._config.lora.rank,
+            lora_alpha=self._config.lora.alpha,
+            target_modules=self._config.lora.target_modules,
+            lora_dropout=self._config.lora.dropout,
+            init_lora_weights=True,
+        )
+        # Wrap the transformer with PEFT to add LoRA layers
+        # noinspection PyTypeChecker
+        self._transformer = get_peft_model(self._transformer, lora_config)
+
+    def _load_checkpoint(self) -> None:
+        """Load checkpoint if specified in config, then resolve resume state."""
+        if not self._config.model.load_checkpoint:
+            self._resume_state: tuple[int, TrainingState | None] = (0, None)
+            return
+
+        checkpoint_path = self._find_checkpoint(self._config.model.load_checkpoint)
+        if not checkpoint_path:
+            logger.warning(f"⚠️ Could not find checkpoint at {self._config.model.load_checkpoint}")
+            self._resume_state = (0, None)
+            return
+
+        self._loaded_checkpoint_path = checkpoint_path
+        logger.info(f"📥 Loading checkpoint from {checkpoint_path}")
+
+        if self._config.model.training_mode == "full":
+            self._load_full_checkpoint(checkpoint_path)
+        else:  # LoRA mode
+            self._load_lora_checkpoint(checkpoint_path)
+
+        self._resume_state = self._resolve_resume_state()
+
+    def _load_full_checkpoint(self, checkpoint_path: Path) -> None:
+        """Load full model checkpoint."""
+        state_dict = load_file(checkpoint_path)
+        self._transformer.load_state_dict(state_dict, strict=True)
+
+        logger.info("✅ Full model checkpoint loaded successfully")
+
+    def _load_lora_checkpoint(self, checkpoint_path: Path) -> None:
+        """Load LoRA checkpoint with DDP/FSDP compatibility."""
+        state_dict = load_file(checkpoint_path)
+
+        # Adjust layer names to match internal format.
+        # (Weights are saved in ComfyUI-compatible format, with "diffusion_model." prefix)
+        state_dict = {k.replace("diffusion_model.", "", 1): v for k, v in state_dict.items()}
+
+        # Load LoRA weights and verify all weights were loaded
+        base_model = self._transformer.get_base_model()
+        set_peft_model_state_dict(base_model, state_dict)
+
+        logger.info("✅ LoRA checkpoint loaded successfully")
+
+    def _resolve_resume_state(self) -> tuple[int, TrainingState | None]:
+        """Determine resume state by looking for a training state file next to the loaded checkpoint.
+        Returns (initial_step, TrainingState or None).
+        If no_resume config is set, no checkpoint loaded, or no state file found: returns (0, None).
+        """
+        if self._config.checkpoints.no_resume or self._loaded_checkpoint_path is None:
+            return 0, None
+
+        state = self._load_training_state(self._loaded_checkpoint_path)
+        if state is None:
+            return 0, None
+
+        fp = state.config_fingerprint
+        cfg = self._config
+        mismatches: list[str] = []
+        if fp.optimizer_type != cfg.optimization.optimizer_type:
+            mismatches.append(f"optimizer_type: {fp.optimizer_type} → {cfg.optimization.optimizer_type}")
+        if fp.scheduler_type != cfg.optimization.scheduler_type:
+            mismatches.append(f"scheduler_type: {fp.scheduler_type} → {cfg.optimization.scheduler_type}")
+        if fp.training_mode != cfg.model.training_mode:
+            mismatches.append(f"training_mode: {fp.training_mode} → {cfg.model.training_mode}")
+        if (
+            cfg.model.training_mode == "lora"
+            and cfg.lora is not None
+            and fp.lora_rank is not None
+            and fp.lora_rank != cfg.lora.rank
+        ):
+            mismatches.append(f"lora_rank: {fp.lora_rank} → {cfg.lora.rank}")
+        if mismatches:
+            logger.warning(
+                f"⚠️ Training state config mismatch ({', '.join(mismatches)}). "
+                "Starting from step 0. Set checkpoints.no_resume=true to silence this warning."
+            )
+            return 0, None
+
+        if state.global_step < 0:
+            logger.warning(f"⚠️ Training state has invalid global_step={state.global_step!r}. Starting from step 0.")
+            return 0, None
+        logger.info(f"📌 Resuming from step {state.global_step}")
+        return state.global_step, state
+
+    @staticmethod
+    def _load_training_state(checkpoint_path: Path) -> TrainingState | None:
+        """Load training state file that corresponds to a checkpoint weights file."""
+        match = re.search(r"step_(\d+)", checkpoint_path.name)
+        if not match:
+            return None
+
+        step_str = match.group(1)
+        state_path = checkpoint_path.parent / f"training_state_step_{step_str}.pt"
+
+        if not state_path.exists():
+            return None
+
+        try:
+            raw: dict = torch.load(state_path, map_location="cpu", weights_only=False)
+            state = TrainingState.from_save_dict(raw)
+            logger.info(f"📥 Loaded training state from {state_path}")
+            return state
+        except Exception as e:
+            logger.warning(f"⚠️ Failed to load training state from {state_path}: {e}. Starting from step 0.")
+            return None
+
+    def _restore_training_state(self, training_state: TrainingState) -> bool:
+        """Restore optimizer, scheduler, and RNG states from a loaded TrainingState.
+        Must be called after _init_optimizer() (which calls accelerator.prepare).
+        Returns True if restore succeeded, False if it failed (caller should fall back to step 0).
+        """
+        try:
+            if training_state.optimizer_state_dict is not None:
+                self._optimizer.load_state_dict(training_state.optimizer_state_dict)
+                logger.debug("Restored optimizer state (full mode)")
+
+            if training_state.lr_scheduler_state_dict is not None and self._lr_scheduler is not None:
+                self._lr_scheduler.load_state_dict(training_state.lr_scheduler_state_dict)
+                logger.debug("Restored LR scheduler state")
+        except Exception as e:
+            logger.warning(f"⚠️ Failed to restore training state: {e}. Starting from step 0.")
+            return False
+
+        rng = training_state.rng_states
+        if self._accelerator.num_processes > 1:
+            logger.debug("Skipping RNG restore in multi-process mode (only main process state was saved)")
+        else:
+            if rng.torch_state is not None:
+                torch.random.set_rng_state(rng.torch_state)
+            if rng.cuda_state is not None and torch.cuda.is_available():
+                torch.cuda.set_rng_state(rng.cuda_state)
+            logger.debug("Restored RNG states")
+
+        return True
+
+    def _prepare_models_for_training(self) -> None:
+        """Prepare models for training with Accelerate."""
+
+        # For FSDP + LoRA: Cast entire model to FP32.
+        # FSDP requires uniform dtype across all parameters in wrapped modules.
+        # In LoRA mode, PEFT creates LoRA params in FP32 while base model is BF16.
+        # We cast the base model to FP32 to match the LoRA params.
+        if self._accelerator.distributed_type == DistributedType.FSDP and self._config.model.training_mode == "lora":
+            logger.debug("FSDP: casting transformer to FP32 for uniform dtype")
+            self._transformer = self._transformer.to(dtype=torch.float32)
+
+        # Enable gradient checkpointing if requested
+        # For PeftModel, we need to access the underlying base model
+        transformer = (
+            self._transformer.get_base_model() if hasattr(self._transformer, "get_base_model") else self._transformer
+        )
+
+        transformer.set_gradient_checkpointing(self._config.optimization.enable_gradient_checkpointing)
+
+        # noinspection PyTypeChecker
+        self._transformer = self._accelerator.prepare(self._transformer)
+
+        # Log GPU memory usage after model preparation
+        vram_usage_gb = torch.cuda.memory_allocated() / 1024**3
+        logger.debug(f"GPU memory usage after models preparation: {vram_usage_gb:.2f} GB")
+
+    @staticmethod
+    def _find_checkpoint(checkpoint_path: str | Path) -> Path | None:
+        """Find the checkpoint file to load, handling both file and directory paths."""
+        checkpoint_path = Path(checkpoint_path)
+
+        if checkpoint_path.is_file():
+            if not checkpoint_path.suffix == ".safetensors":
+                raise ValueError(f"Checkpoint file must have a .safetensors extension: {checkpoint_path}")
+            return checkpoint_path
+
+        if checkpoint_path.is_dir():
+            # Look for checkpoint files in the directory
+            checkpoints = list(checkpoint_path.rglob("*step_*.safetensors"))
+
+            if not checkpoints:
+                return None
+
+            # Sort by step number and return the latest
+            def _get_step_num(p: Path) -> int:
+                try:
+                    return int(p.stem.split("step_")[1])
+                except (IndexError, ValueError):
+                    return -1
+
+            latest = max(checkpoints, key=_get_step_num)
+            return latest
+
+        else:
+            raise ValueError(f"Invalid checkpoint path: {checkpoint_path}. Must be a file or directory.")
+
+    def _init_dataloader(self) -> None:
+        """Initialize the training data loader using the strategy's data sources."""
+        if self._dataset is None:
+            # Get data sources from the training strategy
+            data_sources = self._config.training_strategy.get_data_sources()
+
+            self._dataset = PrecomputedDataset(self._config.data.preprocessed_data_root, data_sources=data_sources)
+            logger.debug(f"Loaded dataset with {len(self._dataset):,} samples from sources: {list(data_sources)}")
+
+        num_workers = self._config.data.num_dataloader_workers
+        dataloader = DataLoader(
+            self._dataset,
+            batch_size=self._config.optimization.batch_size,
+            shuffle=True,
+            drop_last=True,
+            num_workers=num_workers,
+            pin_memory=num_workers > 0,
+            persistent_workers=num_workers > 0,
+        )
+
+        self._dataloader = self._accelerator.prepare(dataloader)
+
+    def _init_lora_weights(self) -> None:
+        """Initialize LoRA weights for the transformer."""
+        logger.debug("Initializing LoRA weights...")
+        for _, module in self._transformer.named_modules():
+            if isinstance(module, (BaseTunerLayer, ModulesToSaveWrapper)):
+                module.reset_lora_parameters(adapter_name="default", init_lora_weights=True)
+
+    def _init_optimizer(self) -> None:
+        """Initialize the optimizer and learning rate scheduler."""
+        opt_cfg = self._config.optimization
+
+        lr = opt_cfg.learning_rate
+        if opt_cfg.optimizer_type == "adamw":
+            optimizer = AdamW(self._trainable_params, lr=lr)
+        elif opt_cfg.optimizer_type == "adamw8bit":
+            # noinspection PyUnresolvedReferences
+            from bitsandbytes.optim import AdamW8bit  # noqa: PLC0415
+
+            optimizer = AdamW8bit(self._trainable_params, lr=lr)
+        else:
+            raise ValueError(f"Unknown optimizer type: {opt_cfg.optimizer_type}")
+
+        lr_scheduler = self._create_scheduler(optimizer)
+
+        # noinspection PyTypeChecker
+        self._optimizer, self._lr_scheduler = self._accelerator.prepare(optimizer, lr_scheduler)
+
+    def _create_scheduler(self, optimizer: torch.optim.Optimizer) -> LRScheduler | None:
+        """Create learning rate scheduler based on config."""
+        scheduler_type = self._config.optimization.scheduler_type
+        steps = self._config.optimization.steps
+        params = self._config.optimization.scheduler_params or {}
+
+        if scheduler_type is None:
+            return None
+
+        if scheduler_type == "linear":
+            scheduler = LinearLR(
+                optimizer,
+                start_factor=params.pop("start_factor", 1.0),
+                end_factor=params.pop("end_factor", 0.1),
+                total_iters=steps,
+                **params,
+            )
+        elif scheduler_type == "cosine":
+            scheduler = CosineAnnealingLR(
+                optimizer,
+                T_max=steps,
+                eta_min=params.pop("eta_min", 0),
+                **params,
+            )
+        elif scheduler_type == "cosine_with_restarts":
+            scheduler = CosineAnnealingWarmRestarts(
+                optimizer,
+                T_0=params.pop("T_0", steps // 4),
+                T_mult=params.pop("T_mult", 1),
+                eta_min=params.pop("eta_min", 5e-5),
+                **params,
+            )
+        elif scheduler_type == "polynomial":
+            scheduler = PolynomialLR(
+                optimizer,
+                total_iters=steps,
+                power=params.pop("power", 1.0),
+                **params,
+            )
+        elif scheduler_type == "step":
+            scheduler = StepLR(
+                optimizer,
+                step_size=params.pop("step_size", steps // 2),
+                gamma=params.pop("gamma", 0.1),
+                **params,
+            )
+        elif scheduler_type == "constant":
+            scheduler = None
+        else:
+            raise ValueError(f"Unknown scheduler type: {scheduler_type}")
+
+        return scheduler
+
+    def _setup_accelerator(self) -> None:
+        """Initialize the Accelerator with the appropriate settings."""
+
+        # find_unused_parameters=True keeps DDP happy when LoRA targets a branch the forward
+        # pass skips (e.g. audio LoRA with `with_audio: false`, or short module patterns like
+        # "to_k" that match the audio branch unintentionally). It's a no-op for FSDP and
+        # single-GPU runs. The probing cost is paid only on the first step.
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+
+        # All distributed setup (DDP/FSDP, number of processes, etc.) is controlled by
+        # the user's Accelerate configuration (accelerate config / accelerate launch).
+        self._accelerator = Accelerator(
+            mixed_precision=self._config.acceleration.mixed_precision_mode,
+            gradient_accumulation_steps=self._config.optimization.gradient_accumulation_steps,
+            kwargs_handlers=[ddp_kwargs],
+        )
+
+        if self._accelerator.num_processes > 1:
+            logger.info(
+                f"{self._accelerator.distributed_type.value} distributed training enabled "
+                f"with {self._accelerator.num_processes} processes"
+            )
+
+            local_batch = self._config.optimization.batch_size
+            global_batch = self._config.optimization.batch_size * self._accelerator.num_processes
+            logger.info(f"Local batch size: {local_batch}, global batch size: {global_batch}")
+
+        # Log torch.compile status from Accelerate's dynamo plugin
+        is_compile_enabled = (
+            hasattr(self._accelerator.state, "dynamo_plugin") and self._accelerator.state.dynamo_plugin.backend != "NO"
+        )
+        if is_compile_enabled:
+            plugin = self._accelerator.state.dynamo_plugin
+            logger.info(f"🔥 torch.compile enabled via Accelerate: backend={plugin.backend}, mode={plugin.mode}")
+
+            if self._accelerator.distributed_type == DistributedType.FSDP:
+                logger.warning(
+                    "⚠️ FSDP + torch.compile is experimental and may hang on the first training iteration. "
+                    "If this occurs, disable torch.compile by removing dynamo_config from your Accelerate config."
+                )
+
+        if self._accelerator.distributed_type == DistributedType.FSDP and self._config.acceleration.quantization:
+            logger.warning(
+                f"FSDP with quantization ({self._config.acceleration.quantization}) may have compatibility issues."
+                "Monitor training stability and consider disabling quantization if issues arise."
+            )
+
+    @contextlib.contextmanager
+    def _offloaded_optimizer_state(self) -> Iterator[None]:
+        """Context manager that offloads optimizer state to CPU during validation.
+        Opt-in via `acceleration.offload_optimizer_during_validation`. Frees VRAM for
+        validation video generation when optimizer state is large (e.g. full fine-tune
+        AdamW, high-rank LoRA). No-op for FSDP (sharded state -- manual `.cpu()` breaks
+        metadata).
+        """
+        enabled = (
+            self._config.acceleration.offload_optimizer_during_validation
+            and self._accelerator.distributed_type != DistributedType.FSDP
+        )
+
+        # Track exactly which tensors we move so we don't promote ones that were
+        # intentionally on CPU (e.g. AdamW's `step` scalar on recent PyTorch).
+        offloaded: list[tuple[dict, str]] = []
+        if enabled:
+            offloaded_bytes = 0
+            for state in self._optimizer.state.values():
+                for k, v in state.items():
+                    if isinstance(v, torch.Tensor) and v.is_cuda:
+                        offloaded.append((state, k))
+                        offloaded_bytes += v.nbytes
+            if offloaded:
+                logger.info(f"Offloading optimizer state to CPU ({offloaded_bytes / 1e9:.1f} GB)")
+                for state, k in offloaded:
+                    state[k] = state[k].cpu()
+
+        try:
+            yield
+        finally:
+            device = self._accelerator.device
+            for state, k in offloaded:
+                state[k] = state[k].to(device)
+
+    def _run_validation(self, progress: TrainingProgress) -> list[Path]:
+        """Run distributed validation by delegating to the ValidationRunner.
+        Each rank generates its assigned subset of validation samples (round-robin by
+        `process_index`/`num_processes`), so all GPUs stay busy and no rank idles long
+        enough to trigger NCCL timeouts. Paths are gathered across ranks so rank 0 has
+        the full list for W&B logging.
+        Under FSDP with multiple processes, ranks pad with extra generate passes
+        (same sample, no disk write) so every rank runs the same number of forwards --
+        avoids collective mismatch.
+        Note: Multi-node training requires a shared filesystem so rank 0 can read
+        videos written by other ranks.
+        """
+        self._optimizer.zero_grad(set_to_none=True)
+        free_gpu_memory()
+
+        num_samples = len(self._config.validation.samples)
+        if num_samples == 0:
+            return []
+
+        rank = self._accelerator.process_index
+        world_size = self._accelerator.num_processes
+
+        rank_indices = list(range(rank, num_samples, world_size))
+        work_items: list[tuple[int, bool]] = [(i, True) for i in rank_indices]
+        if self._accelerator.distributed_type == DistributedType.FSDP and world_size > 1:
+            # FSDP forwards run collective ops; pad short ranks with no-save duplicates so
+            # every rank executes the same number of forwards. A rank with empty
+            # rank_indices (world_size > num_samples) still pads with sample 0 to stay in
+            # sync with the others.
+            max_per_rank = math.ceil(num_samples / world_size)
+            pad_seed = rank_indices[-1] if rank_indices else 0
+            work_items += [(pad_seed, False)] * (max_per_rank - len(work_items))
+
+        # W&B logging is handled by the trainer (after gathering across ranks),
+        # so we always pass wandb_run=None to the runner.
+        sampled = self._validation_runner.run(
+            transformer=self._transformer,
+            step=self._global_step,
+            output_dir=Path(self._config.output_dir),
+            device=self._accelerator.device,
+            progress=progress,
+            wandb_run=None,
+            work_items=work_items,
+        )
+
+        if world_size > 1:
+            sampled = sorted(gather_object(sampled), key=lambda x: x[0])
+
+        paths = [p for _, p in sampled]
+
+        if (
+            self._accelerator.is_main_process
+            and paths
+            and self._config.wandb.log_validation_videos
+            and self._wandb_run is not None
+        ):
+            self._validation_runner.log_to_wandb(self._wandb_run, paths, self._global_step)
+
+        # Non-main ranks must not reach checkpoint collectives while main is still logging to W&B.
+        self._accelerator.wait_for_everyone()
+
+        return paths
+
+    @staticmethod
+    def _log_training_stats(stats: TrainingStats) -> None:
+        """Log training statistics."""
+        stats_str = (
+            "📊 Training Statistics:\n"
+            f" - Total time: {stats.total_time_seconds / 60:.1f} minutes\n"
+            f" - Training speed: {stats.steps_per_second:.2f} steps/second\n"
+            f" - Samples/second: {stats.samples_per_second:.2f}\n"
+            f" - Peak GPU memory: {stats.peak_gpu_memory_gb:.2f} GB"
+        )
+        if stats.num_processes > 1:
+            stats_str += f"\n - Number of processes: {stats.num_processes}\n"
+            stats_str += f" - Global batch size: {stats.global_batch_size}"
+        logger.info(stats_str)
+
+    def _save_checkpoint(self) -> Path | None:
+        """Save the model weights."""
+        is_lora = self._config.model.training_mode == "lora"
+        is_fsdp = self._accelerator.distributed_type == DistributedType.FSDP
+
+        # Prepare paths
+        save_dir = Path(self._config.output_dir) / "checkpoints"
+        prefix = "lora" if is_lora else "model"
+        filename = f"{prefix}_weights_step_{self._global_step:05d}.safetensors"
+        saved_weights_path = save_dir / filename
+
+        # Get state dict (collective operation - all processes must participate)
+        self._accelerator.wait_for_everyone()
+        full_state_dict = self._accelerator.get_state_dict(self._transformer)
+
+        if not IS_MAIN_PROCESS:
+            return None
+
+        save_dir.mkdir(exist_ok=True, parents=True)
+
+        # Determine save precision
+        save_dtype = torch.bfloat16 if self._config.checkpoints.precision == "bfloat16" else torch.float32
+
+        # For LoRA: extract only adapter weights; for full: use as-is
+        if is_lora:
+            unwrapped = self._accelerator.unwrap_model(self._transformer, keep_torch_compile=False)
+            # For FSDP, pass full_state_dict since model params aren't directly accessible
+            state_dict = get_peft_model_state_dict(unwrapped, state_dict=full_state_dict if is_fsdp else None)
+
+            # Remove "base_model.model." prefix added by PEFT
+            state_dict = {k.replace("base_model.model.", "", 1): v for k, v in state_dict.items()}
+
+            # Convert to ComfyUI-compatible format (add "diffusion_model." prefix)
+            state_dict = {f"diffusion_model.{k}": v for k, v in state_dict.items()}
+
+            # Cast to configured precision
+            state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in state_dict.items()}
+
+            # Build metadata for safetensors file
+            metadata = self._build_checkpoint_metadata()
+
+            # Save to disk with metadata
+            save_file(state_dict, saved_weights_path, metadata=metadata)
+        else:
+            # Cast to configured precision
+            full_state_dict = {k: v.to(save_dtype) if isinstance(v, Tensor) else v for k, v in full_state_dict.items()}
+
+            # Save to disk
+            self._accelerator.save(full_state_dict, saved_weights_path)
+
+        rel_path = saved_weights_path.relative_to(self._config.output_dir)
+        logger.info(f"💾 {prefix.capitalize()} weights for step {self._global_step} saved in {rel_path}")
+
+        self._checkpoint_paths.append(saved_weights_path)
+        self._cleanup_checkpoints()
+
+        self._save_training_state(save_dir)
+
+        return saved_weights_path
+
+    def _cleanup_checkpoints(self) -> None:
+        """Clean up old checkpoints."""
+        if 0 < self._config.checkpoints.keep_last_n < len(self._checkpoint_paths):
+            checkpoints_to_remove = self._checkpoint_paths[: -self._config.checkpoints.keep_last_n]
+            for old_checkpoint in checkpoints_to_remove:
+                if old_checkpoint.exists():
+                    old_checkpoint.unlink()
+                    logger.info(f"Removed old checkpoint: {old_checkpoint}")
+            self._checkpoint_paths = self._checkpoint_paths[-self._config.checkpoints.keep_last_n :]
+
+    def _save_training_state(self, save_dir: Path) -> None:
+        """Save training state alongside checkpoint for resume.
+        Respects checkpoints.save_training_state config:
+        - "full": optimizer + scheduler + RNG + step
+        - "minimal": scheduler + RNG + step only
+        - "off": skip entirely
+        """
+        if not IS_MAIN_PROCESS:
+            return
+
+        mode = self._config.checkpoints.save_training_state
+        if mode == "off":
+            return
+
+        is_fsdp = self._accelerator.distributed_type == DistributedType.FSDP
+
+        optimizer_state = None
+        if mode == "full":
+            if is_fsdp:
+                logger.warning(
+                    "⚠️ save_training_state='full' is not supported with FSDP. "
+                    "Saving 'minimal' state (scheduler + RNG only)."
+                )
+            else:
+                optimizer_state = self._optimizer.state_dict()
+
+        state = TrainingState(
+            global_step=self._global_step,
+            config_fingerprint=ConfigFingerprint(
+                optimizer_type=self._config.optimization.optimizer_type,
+                scheduler_type=self._config.optimization.scheduler_type,
+                training_mode=self._config.model.training_mode,
+                lora_rank=self._config.lora.rank if self._config.lora is not None else None,
+            ),
+            rng_states=RngStates(
+                torch_state=torch.random.get_rng_state(),
+                cuda_state=torch.cuda.get_rng_state() if torch.cuda.is_available() else None,
+            ),
+            lr_scheduler_state_dict=self._lr_scheduler.state_dict() if self._lr_scheduler is not None else None,
+            optimizer_state_dict=optimizer_state,
+            wandb_run_id=self._wandb_run.id if self._wandb_run is not None else None,
+        )
+
+        state_path = save_dir / f"training_state_step_{self._global_step:05d}.pt"
+        tmp_path = state_path.with_suffix(".pt.tmp")
+        try:
+            torch.save(state.to_save_dict(), tmp_path)
+        except Exception:
+            if tmp_path.exists():
+                tmp_path.unlink()
+            raise
+        tmp_path.rename(state_path)
+
+        file_size_gb = state_path.stat().st_size / (1024**3)
+        if file_size_gb > 1.0 and not self._training_state_size_warned:
+            self._training_state_size_warned = True
+            logger.warning(
+                f"⚠️ Training state file is {file_size_gb:.1f} GB (full mode includes optimizer state). "
+                f'Set checkpoints.save_training_state="minimal" to save only scheduler/RNG/step (~few KB), '
+                f'or "off" to disable entirely.'
+            )
+
+        if not self._training_state_paths or self._training_state_paths[-1] != state_path:
+            self._training_state_paths.append(state_path)
+        self._cleanup_training_states()
+
+        rel_path = state_path.relative_to(self._config.output_dir)
+        logger.debug(f"Training state saved to {rel_path}")
+
+    def _cleanup_training_states(self) -> None:
+        """Clean up old training state files, using the same keep_last_n as checkpoints."""
+        keep_n = self._config.checkpoints.keep_last_n
+        if 0 < keep_n < len(self._training_state_paths):
+            to_remove = self._training_state_paths[:-keep_n]
+            for old_state in to_remove:
+                if old_state.exists():
+                    old_state.unlink()
+                    logger.debug(f"Removed old training state: {old_state}")
+            self._training_state_paths = self._training_state_paths[-keep_n:]
+
+    def _build_checkpoint_metadata(self) -> dict[str, str]:
+        """Build metadata dictionary for safetensors checkpoint.
+        Delegates to the training strategy to get strategy-specific metadata
+        that downstream inference pipelines may need.
+        Returns:
+            Dictionary of string key-value pairs for safetensors metadata.
+            Values are converted to strings for safetensors compatibility.
+        """
+        raw_metadata = self._training_strategy.get_checkpoint_metadata()
+        # Convert all values to strings for safetensors compatibility
+        metadata = {k: str(v) for k, v in raw_metadata.items()}
+        if metadata:
+            logger.info(f"Saving checkpoint metadata: {metadata}")
+        return metadata
+
+    def _save_config(self) -> None:
+        """Save the training configuration as a YAML file in the output directory."""
+        if not IS_MAIN_PROCESS:
+            return
+
+        config_path = Path(self._config.output_dir) / "training_config.yaml"
+        with open(config_path, "w") as f:
+            yaml.dump(self._config.model_dump(), f, default_flow_style=False, indent=2)
+
+        logger.info(f"💾 Training configuration saved to: {config_path.relative_to(self._config.output_dir)}")
+
+    def _init_wandb(self, resume_run_id: str | None = None) -> None:
+        """Initialize Weights & Biases run, resuming an existing run if its id is provided."""
+        if not self._config.wandb.enabled or not IS_MAIN_PROCESS:
+            self._wandb_run = None
+            return
+
+        wandb_config = self._config.wandb
+        init_kwargs: dict[str, Any] = {
+            "project": wandb_config.project,
+            "entity": wandb_config.entity,
+            "name": Path(self._config.output_dir).name,
+            "tags": wandb_config.tags,
+            "config": self._config.model_dump(),
+        }
+        if resume_run_id is not None:
+            init_kwargs["id"] = resume_run_id
+            init_kwargs["resume"] = "must"
+        run = wandb.init(**init_kwargs)
+        self._wandb_run = run
+
+    def _log_metrics(self, metrics: dict[str, float]) -> None:
+        """Log metrics to Weights & Biases."""
+        if self._wandb_run is not None:
+            self._wandb_run.log(metrics)
diff --git a/packages/ltx-trainer/src/ltx_trainer/training_state.py b/packages/ltx-trainer/src/ltx_trainer/training_state.py
new file mode 100644
index 0000000000000000000000000000000000000000..45dd2df90407795e0db99a04e89bd38ce0974eb7
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/training_state.py
@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from pydantic import BaseModel, ConfigDict
+
+
+class ConfigFingerprint(BaseModel):
+    optimizer_type: str
+    scheduler_type: str
+    training_mode: str
+    lora_rank: int | None = None
+
+
+class RngStates(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    torch_state: torch.Tensor
+    cuda_state: torch.Tensor | None = None
+
+
+class TrainingState(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    global_step: int
+    config_fingerprint: ConfigFingerprint
+    rng_states: RngStates
+    lr_scheduler_state_dict: dict[str, Any] | None = None
+    optimizer_state_dict: dict[str, Any] | None = None
+    wandb_run_id: str | None = None
+
+    def to_save_dict(self) -> dict[str, Any]:
+        """Build dict suitable for torch.save -- recurses BaseModel sub-models, passes tensors/dicts through."""
+
+        def _convert(value: object) -> object:
+            if isinstance(value, BaseModel):
+                return {k: _convert(v) for k, v in value if v is not None}
+            return value
+
+        return {k: _convert(v) for k, v in self if v is not None}
+
+    @classmethod
+    def from_save_dict(cls, data: dict[str, Any]) -> TrainingState:
+        """Construct from torch.load output with Pydantic validation."""
+        return cls(
+            global_step=data["global_step"],
+            config_fingerprint=ConfigFingerprint(**data["config_fingerprint"]),
+            rng_states=RngStates(**data["rng_states"]),
+            lr_scheduler_state_dict=data.get("lr_scheduler_state_dict"),
+            optimizer_state_dict=data.get("optimizer_state_dict"),
+            wandb_run_id=data.get("wandb_run_id"),
+        )
diff --git a/packages/ltx-trainer/src/ltx_trainer/training_strategies/__init__.py b/packages/ltx-trainer/src/ltx_trainer/training_strategies/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..71857c7d71bf79377c43bd54ef31468e1f517c49
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/training_strategies/__init__.py
@@ -0,0 +1,90 @@
+"""Training strategies for different conditioning modes.
+This package implements the Strategy Pattern to handle different training modes:
+- Text-to-video training (standard generation, optionally with audio) [DEPRECATED]
+- Video-to-video training (IC-LoRA mode with reference videos) [DEPRECATED]
+- Flexible training (unified conditioning framework supporting all scenarios) [RECOMMENDED]
+Each strategy encapsulates the specific logic for preparing model inputs and computing loss.
+"""
+
+import warnings
+
+from ltx_trainer import logger
+from ltx_trainer.training_strategies.base_strategy import (
+    DEFAULT_FPS,
+    VIDEO_SCALE_FACTORS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+from ltx_trainer.training_strategies.flexible import FlexibleStrategy, FlexibleStrategyConfig
+from ltx_trainer.training_strategies.text_to_video import TextToVideoConfig, TextToVideoStrategy
+from ltx_trainer.training_strategies.video_to_video import VideoToVideoConfig, VideoToVideoStrategy
+
+# Type alias for all strategy config types
+TrainingStrategyConfig = TextToVideoConfig | VideoToVideoConfig | FlexibleStrategyConfig
+
+__all__ = [
+    "DEFAULT_FPS",
+    "VIDEO_SCALE_FACTORS",
+    "FlexibleStrategy",
+    "FlexibleStrategyConfig",
+    "ModelInputs",
+    "TextToVideoConfig",
+    "TextToVideoStrategy",
+    "TrainingStrategy",
+    "TrainingStrategyConfig",
+    "TrainingStrategyConfigBase",
+    "VideoToVideoConfig",
+    "VideoToVideoStrategy",
+    "get_training_strategy",
+]
+
+
+def get_training_strategy(config: TrainingStrategyConfig) -> TrainingStrategy:
+    """Factory function to create the appropriate training strategy.
+    The strategy is determined by the `name` field in the configuration.
+    Args:
+        config: Strategy-specific configuration with a `name` field
+    Returns:
+        The appropriate training strategy instance
+    Raises:
+        ValueError: If strategy name is not supported
+    Note:
+        The `text_to_video` and `video_to_video` strategies are deprecated.
+        Please use the `flexible` strategy instead.
+    """
+
+    match config:
+        case TextToVideoConfig():
+            warnings.warn(
+                "The 'text_to_video' training strategy is deprecated and will be removed "
+                "in a future version. Please migrate to the 'flexible' strategy. "
+                "See the migration guide in the documentation.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            strategy = TextToVideoStrategy(config)
+        case VideoToVideoConfig():
+            warnings.warn(
+                "The 'video_to_video' training strategy is deprecated and will be removed "
+                "in a future version. Please migrate to the 'flexible' strategy. "
+                "See the migration guide in the documentation.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            strategy = VideoToVideoStrategy(config)
+        case FlexibleStrategyConfig():
+            strategy = FlexibleStrategy(config)
+        case _:
+            raise ValueError(f"Unknown training strategy config type: {type(config).__name__}")
+
+    # Determine audio mode for logging
+    if hasattr(config, "with_audio"):
+        audio_mode = "(audio enabled 🔈)" if config.with_audio else "(audio disabled 🔇)"
+    elif hasattr(config, "audio") and config.audio is not None:
+        audio_mode = "(audio enabled 🔈)"
+    else:
+        audio_mode = "(audio disabled 🔇)"
+
+    logger.debug(f"🎯 Using {strategy.__class__.__name__} training strategy {audio_mode}")
+    return strategy
diff --git a/packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py b/packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py
new file mode 100644
index 0000000000000000000000000000000000000000..d74083f0797262d95ac55172835d36c485af7fd3
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/training_strategies/base_strategy.py
@@ -0,0 +1,249 @@
+"""Base class for training strategies.
+This module defines the abstract base class that all training strategies must implement,
+along with the base configuration class.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import torch
+from pydantic import BaseModel, ConfigDict, Field
+from torch import Tensor
+
+from ltx_core.components.patchifiers import (
+    AudioPatchifier,
+    VideoLatentPatchifier,
+    get_pixel_coords,
+)
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.types import AudioLatentShape, SpatioTemporalScaleFactors, VideoLatentShape
+from ltx_trainer.timestep_samplers import TimestepSampler
+
+# Default frames per second for video missing in the FPS metadata
+DEFAULT_FPS = 24
+
+# VAE scale factors for LTX-2
+VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
+
+
+class TrainingStrategyConfigBase(BaseModel):
+    """Base configuration class for training strategies.
+    All strategy-specific configuration classes should inherit from this.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    name: Literal["text_to_video", "video_to_video", "flexible"] = Field(
+        description="Unique name identifying the training strategy type"
+    )
+
+    @abstractmethod
+    def get_data_sources(self) -> dict[str, str]:
+        """Get the required data sources for this strategy.
+        Returns a mapping of directory name (relative to ``preprocessed_data_root``)
+        to the dataset output key under which that directory's contents are exposed.
+        This is the single source of truth for which directories the strategy needs:
+        it drives both dataset wiring (in the trainer) and existence validation
+        (in ``LtxTrainerConfig``).
+        """
+
+
+@dataclass
+class ModelInputs:
+    """Container for model inputs using the Modality-based interface."""
+
+    video: Modality | None
+    audio: Modality | None
+
+    # Training targets (for loss computation)
+    video_targets: Tensor | None
+    audio_targets: Tensor | None
+
+    # Masks for loss computation (True = compute loss for this token)
+    video_loss_mask: Tensor | None
+    audio_loss_mask: Tensor | None
+
+
+class TrainingStrategy(ABC):
+    """Abstract base class for training strategies.
+    Each strategy encapsulates the logic for a specific training mode,
+    handling input preparation and loss computation.
+    """
+
+    def __init__(self, config: TrainingStrategyConfigBase):
+        """Initialize strategy with configuration.
+        Args:
+            config: Strategy-specific configuration
+        """
+        self.config = config
+        self._video_patchifier = VideoLatentPatchifier(patch_size=1)
+        self._audio_patchifier = AudioPatchifier(patch_size=1)
+
+    @abstractmethod
+    def prepare_training_inputs(
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare training inputs from a raw data batch.
+        Args:
+            batch: Raw batch data from the dataset. Contains:
+                - "latents": Video latent data
+                - "conditions": Text embeddings with keys:
+                    - "video_prompt_embeds": Already processed by embedding connectors
+                    - "audio_prompt_embeds": Already processed by embedding connectors
+                    - "prompt_attention_mask": Attention mask
+                - Additional keys depending on strategy (e.g., "ref_latents" for IC-LoRA)
+            timestep_sampler: Sampler for generating timesteps and noise
+        Returns:
+            ModelInputs containing Modality objects and training targets
+        """
+
+    @abstractmethod
+    def compute_loss(
+        self,
+        video_pred: Tensor,
+        audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute the training loss.
+        Args:
+            video_pred: Video prediction from the transformer model
+            audio_pred: Audio prediction from the transformer model (None for video-only)
+            inputs: The prepared model inputs containing targets and masks
+        Returns:
+            Per-element loss tensor of shape [B,]. The trainer reduces to a scalar
+            before backward(). Returning unreduced loss enables per-sigma-bucket tracking.
+        """
+
+    def get_checkpoint_metadata(self) -> dict[str, Any]:
+        """Get strategy-specific metadata to include in checkpoint files.
+        Override this method in subclasses to add custom metadata,
+        e.g. any parameters that a downstream inference pipeline may need.
+        Returns:
+            Dictionary of metadata key-value pairs (values must be JSON-serializable)
+        """
+        return {}
+
+    def _get_video_positions(
+        self,
+        num_frames: int,
+        height: int,
+        width: int,
+        batch_size: int,
+        fps: float,
+        device: torch.device,
+    ) -> Tensor:
+        """Generate video position embeddings using ltx_core's native implementation.
+        Args:
+            num_frames: Number of latent frames
+            height: Latent height
+            width: Latent width
+            batch_size: Batch size
+            fps: Frames per second
+            device: Target device
+        Returns:
+            Position tensor of shape [B, 3, seq_len, 2] (float32)
+        """
+        latent_coords = self._video_patchifier.get_patch_grid_bounds(
+            output_shape=VideoLatentShape(
+                frames=num_frames,
+                height=height,
+                width=width,
+                batch=batch_size,
+                channels=128,  # Video latent channels
+            ),
+            device=device,
+        )
+
+        # Convert latent coords to pixel coords with causal fix
+        pixel_coords = get_pixel_coords(
+            latent_coords=latent_coords,
+            scale_factors=VIDEO_SCALE_FACTORS,
+            causal_fix=True,
+        ).float()
+
+        # Scale temporal dimension by 1/fps to get time in seconds
+        pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / fps
+
+        return pixel_coords
+
+    def _get_audio_positions(
+        self,
+        num_time_steps: int,
+        batch_size: int,
+        device: torch.device,
+    ) -> Tensor:
+        """Generate audio position embeddings using ltx_core's native implementation.
+        Args:
+            num_time_steps: Number of audio time steps (T, not T*mel_bins)
+            batch_size: Batch size
+            device: Target device
+        Returns:
+            Position tensor of shape [B, 1, num_time_steps, 2]
+        Note:
+            Audio latents should be in patchified format [B, T, C*F] = [B, T, 128]
+            where T is the number of time steps, C=8 channels, F=16 mel bins.
+            This matches the format produced by AudioPatchifier.patchify().
+        """
+        mel_bins = 16
+
+        return self._audio_patchifier.get_patch_grid_bounds(
+            output_shape=AudioLatentShape(
+                frames=num_time_steps,
+                mel_bins=mel_bins,
+                batch=batch_size,
+                channels=8,  # Audio latent channels
+            ),
+            device=device,
+        )
+
+    @staticmethod
+    def _create_per_token_timesteps(conditioning_mask: Tensor, sampled_sigma: Tensor) -> Tensor:
+        """Create per-token timesteps based on conditioning mask.
+        Args:
+            conditioning_mask: Boolean mask of shape (batch_size, sequence_length),
+                where True = conditioning token (timestep=0), False = target token (use sigma)
+            sampled_sigma: Sampled sigma values of shape (batch_size,) or (batch_size, 1, 1)
+        Returns:
+            Timesteps tensor of shape [batch_size, sequence_length]
+        """
+        # Expand to match conditioning mask shape [B, seq_len]
+        expanded_sigma = sampled_sigma.view(-1, 1).expand_as(conditioning_mask)
+
+        # Conditioning tokens get 0, target tokens get the sampled sigma
+        return torch.where(conditioning_mask, torch.zeros_like(expanded_sigma), expanded_sigma)
+
+    @staticmethod
+    def _create_first_frame_conditioning_mask(
+        batch_size: int,
+        sequence_length: int,
+        height: int,
+        width: int,
+        device: torch.device,
+        first_frame_conditioning_p: float = 0.0,
+    ) -> Tensor:
+        """Create conditioning mask for first frame conditioning.
+        Args:
+            batch_size: Batch size
+            sequence_length: Total sequence length
+            height: Latent height
+            width: Latent width
+            device: Target device
+            first_frame_conditioning_p: Probability of conditioning on the first frame
+        Returns:
+            Boolean mask where True indicates first frame tokens (if conditioning is enabled).
+            The conditioning decision is drawn independently per batch element so the training
+            signal across samples in a batch is i.i.d.
+        """
+        conditioning_mask = torch.zeros(batch_size, sequence_length, dtype=torch.bool, device=device)
+
+        if first_frame_conditioning_p > 0:
+            first_frame_end_idx = height * width
+            if first_frame_end_idx < sequence_length:
+                # Per-sample Bernoulli draw so each batch element is independently conditioned.
+                per_sample_condition = torch.rand(batch_size, device=device) < first_frame_conditioning_p
+                conditioning_mask[per_sample_condition, :first_frame_end_idx] = True
+
+        return conditioning_mask
diff --git a/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py b/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py
new file mode 100644
index 0000000000000000000000000000000000000000..82d0d9ae5c9d69628e58b03035d5efe26db4650f
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/training_strategies/flexible.py
@@ -0,0 +1,718 @@
+"""Flexible training strategy for a unified conditioning framework.
+This strategy implements the Unified Conditioning Framework that supports:
+- Simple fine-tuning with text conditioning (text-to-video/audio)
+- Intrinsic conditioning (first_frame, prefix, suffix, spatial_crop, mask)
+- Extrinsic conditioning (concatenation-based, IC-LoRA style)
+The flexible strategy replaces TextToVideoStrategy and VideoToVideoStrategy by expressing
+all conditioning scenarios through configuration rather than code.
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Annotated, Any, Literal, Union
+
+import torch
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+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 (
+    DEFAULT_FPS,
+    VIDEO_SCALE_FACTORS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+
+# =============================================================================
+# Configuration Classes
+# =============================================================================
+
+
+class IntrinsicConditionBase(BaseModel):
+    """Base for intrinsic conditioning — tokens get clean latents, timestep=0, no loss."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    probability: float = Field(
+        default=1.0,
+        ge=0.0,
+        le=1.0,
+        description="Probability of applying this condition",
+    )
+
+
+class FirstFrameConditionConfig(IntrinsicConditionBase):
+    """First frame conditioning — frame 0 is clean, excluded from loss."""
+
+    type: Literal["first_frame"] = "first_frame"
+
+
+class PrefixConditionConfig(IntrinsicConditionBase):
+    """Prefix conditioning — first N temporal units are clean, excluded from loss."""
+
+    type: Literal["prefix"] = "prefix"
+    temporal_boundary: int = Field(
+        ...,
+        ge=1,
+        description="Number of temporal units for prefix region. "
+        "For video: number of latent frames. For audio: number of audio latent timesteps.",
+    )
+
+
+class SuffixConditionConfig(IntrinsicConditionBase):
+    """Suffix conditioning — last N temporal units are clean, excluded from loss."""
+
+    type: Literal["suffix"] = "suffix"
+    temporal_boundary: int = Field(
+        ...,
+        ge=1,
+        description="Number of temporal units for suffix region. "
+        "For video: number of latent frames. For audio: number of audio latent timesteps.",
+    )
+
+
+class SpatialCropConditionConfig(IntrinsicConditionBase):
+    """Spatial crop conditioning — rectangular pixel region is clean, excluded from loss."""
+
+    type: Literal["spatial_crop"] = "spatial_crop"
+    spatial_region: tuple[int, int, int, int] = Field(
+        ...,
+        description="Spatial crop region as (y1, x1, y2, x2) in pixel coordinates",
+    )
+
+
+class MaskConditionConfig(IntrinsicConditionBase):
+    """Mask conditioning — per-sample binary mask determines conditioning tokens."""
+
+    type: Literal["mask"] = "mask"
+    mask_dir: str = Field(
+        ...,
+        description="Directory containing per-sample masks",
+    )
+
+
+class ReferenceConditionConfig(BaseModel):
+    """Reference conditioning (IC-LoRA style concatenation).
+    External reference latents are concatenated to the target sequence.
+    Reference tokens are clean (timestep=0), excluded from loss, and
+    participate in bidirectional self-attention.
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    type: Literal["reference"] = "reference"
+    latents_dir: str = Field(..., description="Directory for reference latents")
+    probability: float = Field(default=1.0, ge=0.0, le=1.0, description="Probability of applying this condition")
+
+
+# Discriminated union for condition configs
+ConditionConfig = Annotated[
+    Union[
+        FirstFrameConditionConfig,
+        PrefixConditionConfig,
+        SuffixConditionConfig,
+        SpatialCropConditionConfig,
+        MaskConditionConfig,
+        ReferenceConditionConfig,
+    ],
+    Field(discriminator="type"),
+]
+
+
+class ModalityConfig(BaseModel):
+    """Configuration for a single modality (video or audio)."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    is_generated: bool = Field(
+        ...,
+        description="True = generated modality (denoised, contributes to loss), False = conditioning-only modality",
+    )
+
+    latents_dir: str = Field(
+        ...,
+        description="Directory for latents",
+    )
+
+    conditions: list[ConditionConfig] = Field(
+        default_factory=list,
+        description="List of conditions (e.g. first_frame, prefix, reference). Text conditioning is always applied.",
+    )
+
+
+class FlexibleStrategyConfig(TrainingStrategyConfigBase):
+    """Configuration for the flexible training strategy.
+    This strategy supports all conditioning scenarios through configuration:
+    - Text-to-video/audio with simple fine-tuning
+    - Intrinsic conditioning like first-frame, extension, outpainting
+    - Reference conditioning like IC-LoRA (concatenation-based reference)
+    """
+
+    name: Literal["flexible"] = "flexible"
+
+    video: ModalityConfig | None = Field(
+        default=None,
+        description="Video modality configuration",
+    )
+
+    audio: ModalityConfig | None = Field(
+        default=None,
+        description="Audio modality configuration",
+    )
+
+    @model_validator(mode="after")
+    def validate_at_least_one_generated(self) -> "FlexibleStrategyConfig":
+        """Ensure at least one modality has is_generated=true."""
+        has_video_target = self.video is not None and self.video.is_generated
+        has_audio_target = self.audio is not None and self.audio.is_generated
+        if not has_video_target and not has_audio_target:
+            raise ValueError("At least one modality must have is_generated=true")
+        return self
+
+    @model_validator(mode="after")
+    def validate_audio_intrinsic_regions(self) -> "FlexibleStrategyConfig":
+        """Reject video-only intrinsic regions on the audio modality."""
+        if self.audio is None:
+            return self
+        for cond in self.audio.conditions:
+            if isinstance(cond, (FirstFrameConditionConfig, SpatialCropConditionConfig)):
+                raise ValueError(
+                    f"Intrinsic condition '{cond.type}' is not supported for audio. "
+                    f"Audio supports: prefix, suffix, mask."
+                )
+        return self
+
+    def get_data_sources(self) -> dict[str, str]:
+        """Dynamically determine required data sources from config.
+        Returns a mapping of directory name (under ``preprocessed_data_root``) to
+        the dataset output key.
+        """
+        sources: dict[str, str] = {"conditions": "conditions"}
+
+        if self.video is not None:
+            sources[self.video.latents_dir] = "video_latents"
+        if self.audio is not None:
+            sources[self.audio.latents_dir] = "audio_latents"
+
+        for modality_config in (self.video, self.audio):
+            if modality_config is None:
+                continue
+            for cond in modality_config.conditions:
+                if isinstance(cond, ReferenceConditionConfig):
+                    sources[cond.latents_dir] = cond.latents_dir
+                elif isinstance(cond, MaskConditionConfig):
+                    sources[cond.mask_dir] = cond.mask_dir
+
+        return sources
+
+
+# =============================================================================
+# Helper Data Structures
+# =============================================================================
+
+
+@dataclass
+class ModalityProcessingResult:
+    """Result of processing a single modality."""
+
+    modality: Modality
+    targets: Tensor | None
+    loss_mask: Tensor | None
+
+
+@dataclass
+class LatentData:
+    """Loaded and patchified latents with metadata."""
+
+    latents: Tensor  # [B, seq_len, C]
+    num_frames: int
+    height: int
+    width: int
+    fps: float
+
+
+# =============================================================================
+# FlexibleStrategy Implementation
+# =============================================================================
+
+
+class FlexibleStrategy(TrainingStrategy):
+    """Unified training strategy supporting all conditioning scenarios.
+    This strategy implements the Unified Conditioning Framework, allowing
+    any training scenario to be expressed through configuration.
+    """
+
+    config: FlexibleStrategyConfig
+
+    def __init__(self, config: FlexibleStrategyConfig):
+        """Initialize strategy with configuration.
+        Args:
+            config: Flexible strategy configuration
+        """
+        super().__init__(config)
+        self.config = config
+        self.reference_spatial_scale_factor, self.reference_temporal_scale_factor = (
+            self._infer_reference_scale_factors_from_config()
+        )
+
+    def prepare_training_inputs(
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare training inputs by processing video and audio modalities."""
+        video_result = self._process_modality(self.config.video, batch, "video", timestep_sampler)
+        audio_result = self._process_modality(self.config.audio, batch, "audio", timestep_sampler)
+
+        return ModelInputs(
+            video=video_result.modality if video_result else None,
+            audio=audio_result.modality if audio_result else None,
+            video_targets=video_result.targets if video_result else None,
+            audio_targets=audio_result.targets if audio_result else None,
+            video_loss_mask=video_result.loss_mask if video_result else None,
+            audio_loss_mask=audio_result.loss_mask if audio_result else None,
+        )
+
+    def compute_loss(
+        self,
+        video_pred: Tensor | None,
+        audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute masked MSE loss for video and audio predictions. Returns [B,]."""
+        total_loss = None
+
+        if video_pred is not None and inputs.video_targets is not None:
+            video_loss = self._compute_modality_loss(
+                pred=video_pred,
+                targets=inputs.video_targets,
+                loss_mask=inputs.video_loss_mask,
+            )
+            total_loss = video_loss
+
+        if audio_pred is not None and inputs.audio_targets is not None:
+            audio_loss = self._compute_modality_loss(
+                pred=audio_pred,
+                targets=inputs.audio_targets,
+                loss_mask=inputs.audio_loss_mask,
+            )
+            total_loss = audio_loss if total_loss is None else total_loss + audio_loss
+
+        if total_loss is None:
+            raise ValueError("No valid predictions and targets provided for loss computation")
+
+        return total_loss
+
+    def get_checkpoint_metadata(self) -> dict[str, Any]:
+        """Include reference scale factors in checkpoint metadata for inference pipelines."""
+        metadata: dict[str, Any] = {}
+        spatial = self.reference_spatial_scale_factor
+        temporal = self.reference_temporal_scale_factor
+        if spatial is not None and spatial != 1:
+            metadata["reference_spatial_scale_factor"] = spatial
+            metadata["reference_downscale_factor"] = spatial  # backward compat
+        if temporal is not None and temporal != 1:
+            metadata["reference_temporal_scale_factor"] = temporal
+        return metadata
+
+    def _infer_reference_scale_factors_from_config(self) -> tuple[int | None, int | None]:
+        """Infer spatial and temporal scale factors by peeking at one sample pair."""
+        if self.config.video is None:
+            return None, None
+        for cond in self.config.video.conditions:
+            if not isinstance(cond, ReferenceConditionConfig):
+                continue
+            target_dir = Path(self.config.video.latents_dir)
+            ref_dir = Path(cond.latents_dir)
+            for sample_file in target_dir.rglob("*.pt"):
+                ref_file = ref_dir / sample_file.relative_to(target_dir)
+                if not ref_file.exists():
+                    continue
+                target_data = torch.load(sample_file, map_location="cpu", weights_only=True)
+                ref_data = torch.load(ref_file, map_location="cpu", weights_only=True)
+                if "height" not in ref_data or "height" not in target_data:
+                    continue
+                spatial = self._infer_scale_factor(
+                    ref_data["height"],
+                    ref_data["width"],
+                    target_data["height"],
+                    target_data["width"],
+                )
+                temporal = self._infer_temporal_scale_factor(
+                    ref_data["num_frames"],
+                    target_data["num_frames"],
+                )
+                return spatial, temporal
+        return None, None
+
+    def _process_modality(
+        self,
+        modality_config: ModalityConfig | None,
+        batch: dict[str, Any],
+        modality_key: str,
+        timestep_sampler: TimestepSampler,
+    ) -> ModalityProcessingResult | None:
+        """Process a single modality: load latents, add noise, apply conditions, build Modality."""
+        if modality_config is None:
+            return None
+
+        # Step 1: Load and patchify latents
+        data = self._patchify_latent_data(batch[f"{modality_key}_latents"], modality_key)
+        latents = data.latents
+
+        batch_size, seq_len, _ = latents.shape
+        device = latents.device
+        dtype = latents.dtype
+
+        # Step 2: Get text embeddings
+        conditions = batch["conditions"]
+        prompt_embeds = conditions[f"{modality_key}_prompt_embeds"]
+        prompt_attention_mask = conditions["prompt_attention_mask"]
+
+        # Step 3: Initialize noise, timesteps, and loss mask based on is_generated flag
+        if modality_config.is_generated:
+            noisy_latents, targets, timesteps, loss_mask, sigmas = self._initialize_noisy_target(
+                latents, timestep_sampler
+            )
+        else:
+            # Conditioning modality: keep clean (sigma=0), no loss
+            noisy_latents = latents
+            targets = None
+            timesteps = torch.zeros(batch_size, seq_len, device=device, dtype=dtype)
+            loss_mask = None
+            sigmas = torch.zeros(batch_size, device=device, dtype=dtype)
+
+        # Step 4: Generate positions
+        if modality_key == "video":
+            positions = self._get_video_positions(
+                num_frames=data.num_frames,
+                height=data.height,
+                width=data.width,
+                batch_size=batch_size,
+                fps=data.fps,
+                device=device,
+            )
+        else:
+            positions = self._get_audio_positions(
+                num_time_steps=seq_len,
+                batch_size=batch_size,
+                device=device,
+            )
+
+        # Step 5: Apply conditions (intrinsic first, then extrinsic)
+        for cond in modality_config.conditions:
+            if isinstance(cond, IntrinsicConditionBase) and modality_config.is_generated:
+                noisy_latents, timesteps, loss_mask = self._apply_intrinsic_condition(
+                    noisy_latents=noisy_latents,
+                    clean_latents=latents,
+                    timesteps=timesteps,
+                    loss_mask=loss_mask,
+                    config=cond,
+                    height=data.height,
+                    width=data.width,
+                    batch=batch,
+                )
+
+        for cond in modality_config.conditions:
+            if isinstance(cond, ReferenceConditionConfig):
+                noisy_latents, positions, timesteps, loss_mask, targets = self._apply_reference_condition(
+                    noisy_latents=noisy_latents,
+                    positions=positions,
+                    timesteps=timesteps,
+                    loss_mask=loss_mask,
+                    targets=targets,
+                    batch=batch,
+                    config=cond,
+                    modality_key=modality_key,
+                )
+
+        # Step 6: Build Modality
+        modality = Modality(
+            enabled=True,
+            latent=noisy_latents,
+            sigma=sigmas,
+            timesteps=timesteps,
+            positions=positions,
+            context=prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        return ModalityProcessingResult(
+            modality=modality,
+            targets=targets,
+            loss_mask=loss_mask,
+        )
+
+    @staticmethod
+    def _initialize_noisy_target(
+        latents: Tensor,
+        timestep_sampler: TimestepSampler,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+        """Add noise to latents and create training targets. Returns (noisy, targets, timesteps, mask, sigmas)."""
+        batch_size, seq_len, _ = latents.shape
+        sigmas = timestep_sampler.sample_for(latents)
+        noise = torch.randn_like(latents)
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+        noisy_latents = (1 - sigmas_expanded) * latents + sigmas_expanded * noise
+        targets = noise - latents  # velocity prediction
+        timesteps = sigmas.view(-1, 1).expand(batch_size, seq_len).clone()
+        loss_mask = torch.ones(batch_size, seq_len, dtype=torch.bool, device=latents.device)
+        return noisy_latents, targets, timesteps, loss_mask, sigmas
+
+    def _apply_intrinsic_condition(
+        self,
+        noisy_latents: Tensor,
+        clean_latents: Tensor,
+        timesteps: Tensor,
+        loss_mask: Tensor,
+        config: IntrinsicConditionBase,
+        height: int,
+        width: int,
+        batch: dict[str, Any],
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Apply intrinsic conditioning using a binary mask.
+        For each token, the mask value determines conditioning strength:
+        - mask=1: conditioned (clean latent, timestep=0, excluded from loss)
+        - mask=0: generated (noisy latent, original timestep, contributes to loss)
+        The conditioning decision is drawn independently per batch element so the training
+        signal across samples in a batch is i.i.d. -- a single batch-wide draw would
+        correlate gradient updates across the batch.
+        """
+        batch_size, seq_len, _ = noisy_latents.shape
+        device = noisy_latents.device
+
+        # Per-sample Bernoulli draw -- each element is independently conditioned.
+        apply_per_sample = torch.rand(batch_size, device=device) < config.probability
+        if not apply_per_sample.any():
+            return noisy_latents, timesteps, loss_mask
+
+        if isinstance(config, FirstFrameConditionConfig):
+            mask = self._compute_temporal_mask(batch_size, seq_len, height, width, 1, False, device)
+        elif isinstance(config, PrefixConditionConfig):
+            mask = self._compute_temporal_mask(
+                batch_size, seq_len, height, width, config.temporal_boundary, False, device
+            )
+        elif isinstance(config, SuffixConditionConfig):
+            mask = self._compute_temporal_mask(
+                batch_size, seq_len, height, width, config.temporal_boundary, True, device
+            )
+        elif isinstance(config, SpatialCropConditionConfig):
+            mask = self._compute_spatial_crop_mask(batch_size, seq_len, height, width, config.spatial_region, device)
+        elif isinstance(config, MaskConditionConfig):
+            # Binarize to match inference, which thresholds masks at load time
+            # (validation_runner._load_and_downsample_mask / _load_audio_mask).
+            mask = (batch[config.mask_dir]["mask"].reshape(batch_size, seq_len) > 0.5).float()
+        else:
+            raise ValueError(f"Unknown intrinsic condition type: {type(config).__name__}")
+
+        # Zero the mask for samples the per-sample draw did not select.
+        mask = mask * apply_per_sample.view(-1, 1).to(mask.dtype)
+
+        # Apply binary mask: clean conditioned tokens, noisy generated tokens.
+        m = mask.unsqueeze(-1)
+        noisy_latents = m * clean_latents + (1 - m) * noisy_latents
+        timesteps = (1 - mask) * timesteps
+        loss_mask = loss_mask & (mask == 0)
+
+        return noisy_latents, timesteps, loss_mask
+
+    @staticmethod
+    def _compute_temporal_mask(
+        batch_size: int,
+        seq_len: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        from_end: bool,
+        device: torch.device,
+    ) -> Tensor:
+        """Compute float mask for temporal region (prefix or suffix). Returns [B, seq_len] in {0, 1}."""
+        tokens_per_frame = height * width
+        num_tokens = num_frames * tokens_per_frame
+        mask = torch.zeros(batch_size, seq_len, device=device)
+        if from_end:
+            mask[:, -num_tokens:] = 1.0
+        else:
+            mask[:, :num_tokens] = 1.0
+        return mask
+
+    @staticmethod
+    def _compute_spatial_crop_mask(
+        batch_size: int,
+        seq_len: int,
+        height: int,
+        width: int,
+        region: tuple[int, int, int, int],
+        device: torch.device,
+    ) -> Tensor:
+        """Compute float mask for spatial crop region (y1, x1, y2, x2) in pixel coords.
+        Returns [B, seq_len] in {0, 1}.
+        """
+        y1, x1, y2, x2 = region
+        num_frames = seq_len // (height * width)
+
+        # Convert pixel to latent coordinates and clamp (per-axis VAE scale factor).
+        def to_latent(v: int, scale: int, max_v: int) -> int:
+            return max(0, min(v // scale, max_v))
+
+        ly1 = to_latent(y1, VIDEO_SCALE_FACTORS.height, height)
+        ly2 = to_latent(y2, VIDEO_SCALE_FACTORS.height, height)
+        lx1 = to_latent(x1, VIDEO_SCALE_FACTORS.width, width)
+        lx2 = to_latent(x2, VIDEO_SCALE_FACTORS.width, width)
+
+        # Create spatial mask and tile across frames
+        spatial_mask = torch.zeros(height, width, device=device)
+        spatial_mask[ly1:ly2, lx1:lx2] = 1.0
+        full_mask = spatial_mask.flatten().repeat(num_frames)
+
+        return full_mask.unsqueeze(0).expand(batch_size, -1)
+
+    def _patchify_latent_data(self, latent_data: dict[str, Any], modality_key: str) -> LatentData:
+        """Patchify latent data and extract metadata."""
+        latents = latent_data["latents"]
+
+        if modality_key == "video":
+            num_frames = latent_data["num_frames"][0].item()
+            height = latent_data["height"][0].item()
+            width = latent_data["width"][0].item()
+            fps = latent_data.get("fps")
+            fps = fps[0].item() if fps is not None else DEFAULT_FPS
+            latents = self._video_patchifier.patchify(latents)
+        else:
+            num_frames = latent_data.get("num_frames", [latents.shape[2]])[0]
+            if isinstance(num_frames, Tensor):
+                num_frames = num_frames.item()
+            height = 1
+            width = 1
+            fps = 1.0
+            latents = self._audio_patchifier.patchify(latents)
+
+        return LatentData(latents=latents, num_frames=num_frames, height=height, width=width, fps=fps)
+
+    def _apply_reference_condition(
+        self,
+        noisy_latents: Tensor,
+        positions: Tensor,
+        timesteps: Tensor,
+        loss_mask: Tensor | None,
+        targets: Tensor | None,
+        batch: dict[str, Any],
+        config: ReferenceConditionConfig,
+        modality_key: str,
+    ) -> tuple[Tensor, Tensor, Tensor, Tensor | None, Tensor | None]:
+        """Concatenate reference latents to target sequence for reference conditioning (IC-LoRA style).
+        The apply/skip decision is batch-wide (reference conditioning changes the sequence
+        length, so it cannot be applied to only part of a batch) but is drawn from the torch
+        RNG so runs are reproducible under ``torch.manual_seed`` — mirroring the intrinsic
+        per-sample draw rather than Python's unseeded ``random``.
+        """
+        if torch.rand((), device=noisy_latents.device).item() >= config.probability:
+            return noisy_latents, positions, timesteps, loss_mask, targets
+
+        # Load and patchify condition latents
+        cond = self._patchify_latent_data(batch[config.latents_dir], modality_key)
+        cond_latents = cond.latents
+
+        batch_size, cond_seq_len, _ = cond_latents.shape
+        device = cond_latents.device
+        dtype = cond_latents.dtype
+
+        # Generate condition positions
+        if modality_key == "video":
+            cond_positions = self._get_video_positions(
+                num_frames=cond.num_frames,
+                height=cond.height,
+                width=cond.width,
+                batch_size=batch_size,
+                fps=cond.fps,
+                device=device,
+            )
+        else:
+            cond_positions = self._get_audio_positions(
+                num_time_steps=cond_seq_len,
+                batch_size=batch_size,
+                device=device,
+            )
+
+        # Translate / rescale ref positions into the target's frame (video only).
+        if modality_key == "video":
+            spatial_sf = self.reference_spatial_scale_factor or 1
+            temporal_sf = self.reference_temporal_scale_factor or 1
+            if spatial_sf != 1 or temporal_sf != 1:
+                cond_positions = cond_positions.clone()
+                if temporal_sf != 1:
+                    # Ref positions are already at the ref's effective fps (source_fps / S,
+                    # stored by process_videos.py). Shift by (S - 1) / target_fps so ref's
+                    # last patch aligns with target's last; clamp the causal patch at 0.
+                    t_target = positions[:, 0, 0:1, 1:2]  # = 1 / target_fps
+                    cond_positions[:, 0, ...] = torch.clamp(
+                        cond_positions[:, 0, ...] - (temporal_sf - 1) * t_target, min=0
+                    )
+                if spatial_sf != 1:
+                    cond_positions[:, 1, ...] *= spatial_sf
+                    cond_positions[:, 2, ...] *= spatial_sf
+
+        # Condition tokens: clean, timestep=0, no loss
+        cond_timesteps = torch.zeros(batch_size, cond_seq_len, device=device, dtype=dtype)
+        cond_loss_mask = torch.zeros(batch_size, cond_seq_len, dtype=torch.bool, device=device)
+
+        # Concatenate condition and target sequences (condition first, then target)
+        combined_latents = torch.cat([cond_latents, noisy_latents], dim=1)
+        combined_positions = torch.cat([cond_positions, positions], dim=2)
+        combined_timesteps = torch.cat([cond_timesteps, timesteps], dim=1)
+
+        combined_loss_mask = torch.cat([cond_loss_mask, loss_mask], dim=1) if loss_mask is not None else None
+
+        # Targets remain unchanged (only for target portion, not condition portion)
+
+        return combined_latents, combined_positions, combined_timesteps, combined_loss_mask, targets
+
+    @staticmethod
+    def _compute_modality_loss(pred: Tensor, targets: Tensor, loss_mask: Tensor) -> Tensor:
+        """Compute per-element MSE loss for a single modality. Returns [B,]."""
+        # Slice prediction to match targets length (removes any prepended condition tokens)
+        target_len = targets.shape[1]
+        pred = pred[:, -target_len:, :]
+        mask = loss_mask[:, -target_len:]
+
+        # Compute masked MSE loss, reduce per-element [B,] over (seq, channels)
+        mask_expanded = mask.unsqueeze(-1).float()
+        squared_error = (pred - targets).pow(2)
+        masked_loss = squared_error * mask_expanded
+        return masked_loss.mean(dim=[-2, -1]) / mask_expanded.mean(dim=[-2, -1]).clamp(min=1e-8)
+
+    @staticmethod
+    def _infer_scale_factor(cond_height: int, cond_width: int, target_height: int, target_width: int) -> int:
+        """Infer spatial scale factor between condition and target resolutions."""
+        if target_height == cond_height and target_width == cond_width:
+            return 1
+        scale_h = target_height // cond_height if cond_height > 0 else 1
+        scale_w = target_width // cond_width if cond_width > 0 else 1
+        if scale_h != scale_w:
+            raise ValueError(
+                f"Non-uniform scale factors between condition and target: height={scale_h}, width={scale_w}. "
+                "Condition and target resolutions must scale uniformly."
+            )
+        return scale_h
+
+    @staticmethod
+    def _infer_temporal_scale_factor(cond_num_frames: int, target_num_frames: int) -> int:
+        """Infer temporal scale factor between condition and target latent frame counts.
+        The first latent frame encodes a single pixel frame (the VAE's causal structure),
+        so the temporal groups count is (num_frames - 1).  The scale factor is the ratio
+        of target groups to condition groups.
+        """
+        if target_num_frames == cond_num_frames:
+            return 1
+        target_groups = target_num_frames - 1
+        cond_groups = cond_num_frames - 1
+        if cond_groups <= 0 or target_groups <= 0:
+            return 1
+        if target_groups % cond_groups != 0:
+            raise ValueError(
+                f"Target temporal groups ({target_groups}) is not evenly divisible by "
+                f"condition temporal groups ({cond_groups})."
+            )
+        return target_groups // cond_groups
diff --git a/packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py b/packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py
new file mode 100644
index 0000000000000000000000000000000000000000..66cf2297caf597947fadb4d34e071850477f86f8
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/training_strategies/text_to_video.py
@@ -0,0 +1,277 @@
+"""Text-to-video training strategy.
+This strategy implements standard text-to-video generation training where:
+- Only target latents are used (no reference videos)
+- Standard noise application and loss computation
+- Supports first frame conditioning
+- Optionally supports joint audio-video training
+"""
+
+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 import logger
+from ltx_trainer.timestep_samplers import TimestepSampler
+from ltx_trainer.training_strategies.base_strategy import (
+    DEFAULT_FPS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+
+
+class TextToVideoConfig(TrainingStrategyConfigBase):
+    """Configuration for text-to-video training strategy."""
+
+    name: Literal["text_to_video"] = "text_to_video"
+
+    first_frame_conditioning_p: float = Field(
+        default=0.1,
+        description="Probability of conditioning on the first frame during training",
+        ge=0.0,
+        le=1.0,
+    )
+
+    with_audio: bool = Field(
+        default=False,
+        description="Whether to include audio in training (joint audio-video generation)",
+    )
+
+    audio_latents_dir: str = Field(
+        default="audio_latents",
+        description="Directory name for audio latents when with_audio is True",
+    )
+
+    def get_data_sources(self) -> dict[str, str]:
+        """Text-to-video training requires latents and text conditions.
+        When ``with_audio`` is True, also requires audio latents.
+        """
+        sources = {
+            "latents": "latents",
+            "conditions": "conditions",
+        }
+        if self.with_audio:
+            sources[self.audio_latents_dir] = "audio_latents"
+        return sources
+
+
+class TextToVideoStrategy(TrainingStrategy):
+    """Text-to-video training strategy.
+    This strategy implements regular video generation training where:
+    - Only target latents are used (no reference videos)
+    - Standard noise application and loss computation
+    - Supports first frame conditioning
+    - Optionally supports joint audio-video training when with_audio=True
+    """
+
+    config: TextToVideoConfig
+
+    def __init__(self, config: TextToVideoConfig):
+        """Initialize strategy with configuration.
+        Args:
+            config: Text-to-video configuration
+        """
+        super().__init__(config)
+
+    def prepare_training_inputs(
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare inputs for text-to-video training."""
+        # Get pre-encoded latents - dataset provides uniform non-patchified format [B, C, F, H, W]
+        latents = batch["latents"]
+        video_latents = latents["latents"]
+
+        # Get video dimensions (assume same for all batch elements)
+        num_frames = latents["num_frames"][0].item()
+        height = latents["height"][0].item()
+        width = latents["width"][0].item()
+
+        # Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
+        video_latents = self._video_patchifier.patchify(video_latents)
+
+        # Handle FPS with backward compatibility
+        fps = latents.get("fps", None)
+        if fps is not None and not torch.all(fps == fps[0]):
+            logger.warning(
+                f"Different FPS values found in the batch. Found: {fps.tolist()}, using the first one: {fps[0].item()}"
+            )
+        fps = fps[0].item() if fps is not None else DEFAULT_FPS
+
+        # Get text embeddings (already processed by embedding connectors in trainer)
+        conditions = batch["conditions"]
+        video_prompt_embeds = conditions["video_prompt_embeds"]
+        audio_prompt_embeds = conditions["audio_prompt_embeds"]
+        prompt_attention_mask = conditions["prompt_attention_mask"]
+
+        batch_size = video_latents.shape[0]
+        video_seq_len = video_latents.shape[1]
+        device = video_latents.device
+
+        # Create conditioning mask (first frame conditioning)
+        video_conditioning_mask = self._create_first_frame_conditioning_mask(
+            batch_size=batch_size,
+            sequence_length=video_seq_len,
+            height=height,
+            width=width,
+            device=device,
+            first_frame_conditioning_p=self.config.first_frame_conditioning_p,
+        )
+
+        # Sample noise and sigmas
+        sigmas = timestep_sampler.sample_for(video_latents)
+        video_noise = torch.randn_like(video_latents)
+
+        # Apply noise: noisy = (1 - sigma) * clean + sigma * noise
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+        noisy_video = (1 - sigmas_expanded) * video_latents + sigmas_expanded * video_noise
+
+        # For conditioning tokens, use clean latents
+        conditioning_mask_expanded = video_conditioning_mask.unsqueeze(-1)
+        noisy_video = torch.where(conditioning_mask_expanded, video_latents, noisy_video)
+
+        # Compute video targets (velocity prediction)
+        video_targets = video_noise - video_latents
+
+        # Create per-token timesteps
+        video_timesteps = self._create_per_token_timesteps(video_conditioning_mask, sigmas.squeeze())
+
+        # Generate video positions using ltx_core's native implementation
+        video_positions = self._get_video_positions(
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            batch_size=batch_size,
+            fps=fps,
+            device=device,
+        )
+
+        # Create video Modality
+        video_modality = Modality(
+            enabled=True,
+            sigma=sigmas,
+            latent=noisy_video,
+            timesteps=video_timesteps,
+            positions=video_positions,
+            context=video_prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        # Video loss mask: True for tokens we want to compute loss on (non-conditioning tokens)
+        video_loss_mask = ~video_conditioning_mask
+
+        # Handle audio if enabled
+        audio_modality = None
+        audio_targets = None
+        audio_loss_mask = None
+
+        if self.config.with_audio:
+            audio_modality, audio_targets, audio_loss_mask = self._prepare_audio_inputs(
+                batch=batch,
+                sigmas=sigmas,
+                audio_prompt_embeds=audio_prompt_embeds,
+                prompt_attention_mask=prompt_attention_mask,
+                batch_size=batch_size,
+                device=device,
+            )
+
+        return ModelInputs(
+            video=video_modality,
+            audio=audio_modality,
+            video_targets=video_targets,
+            audio_targets=audio_targets,
+            video_loss_mask=video_loss_mask,
+            audio_loss_mask=audio_loss_mask,
+        )
+
+    def _prepare_audio_inputs(
+        self,
+        batch: dict[str, Any],
+        sigmas: Tensor,
+        audio_prompt_embeds: Tensor,
+        prompt_attention_mask: Tensor,
+        batch_size: int,
+        device: torch.device,
+    ) -> tuple[Modality, Tensor, Tensor]:
+        """Prepare audio inputs for joint audio-video training.
+        Args:
+            batch: Raw batch data containing audio_latents
+            sigmas: Sampled sigma values (same as video)
+            audio_prompt_embeds: Audio context embeddings
+            prompt_attention_mask: Attention mask for context
+            batch_size: Batch size
+            device: Target device
+        Returns:
+            Tuple of (audio_modality, audio_targets, audio_loss_mask)
+        """
+        # Get audio latents - dataset provides uniform non-patchified format [B, C, T, F]
+        audio_data = batch["audio_latents"]
+        audio_latents = audio_data["latents"]
+
+        # Patchify audio latents: [B, C, T, F] -> [B, T, C*F]
+        audio_latents = self._audio_patchifier.patchify(audio_latents)
+
+        audio_seq_len = audio_latents.shape[1]
+
+        # Sample audio noise
+        audio_noise = torch.randn_like(audio_latents)
+
+        # Apply noise to audio (same sigma as video)
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+        noisy_audio = (1 - sigmas_expanded) * audio_latents + sigmas_expanded * audio_noise
+
+        # Compute audio targets
+        audio_targets = audio_noise - audio_latents
+
+        # Audio timesteps: all tokens use the sampled sigma (no conditioning mask)
+        audio_timesteps = sigmas.view(-1, 1).expand(-1, audio_seq_len)
+
+        # Generate audio positions
+        audio_positions = self._get_audio_positions(
+            num_time_steps=audio_seq_len,
+            batch_size=batch_size,
+            device=device,
+        )
+
+        # Create audio Modality
+        audio_modality = Modality(
+            enabled=True,
+            latent=noisy_audio,
+            sigma=sigmas,
+            timesteps=audio_timesteps,
+            positions=audio_positions,
+            context=audio_prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        # Audio loss mask: all tokens contribute to loss (no conditioning)
+        audio_loss_mask = torch.ones(batch_size, audio_seq_len, dtype=torch.bool, device=device)
+
+        return audio_modality, audio_targets, audio_loss_mask
+
+    def compute_loss(
+        self,
+        video_pred: Tensor,
+        audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute masked MSE loss for video and optionally audio. Returns [B,]."""
+        # Video loss: per-element mean over (seq, channels), [B,]
+        video_loss = (video_pred - inputs.video_targets).pow(2)
+        video_loss_mask = inputs.video_loss_mask.unsqueeze(-1).float()
+        masked = video_loss.mul(video_loss_mask)
+        video_loss = masked.mean(dim=[-2, -1]) / video_loss_mask.mean(dim=[-2, -1]).clamp(min=1e-8)
+
+        # If no audio, return video loss only
+        if not self.config.with_audio or audio_pred is None or inputs.audio_targets is None:
+            return video_loss
+
+        # Audio loss: per-element mean over (seq, channels), [B,]
+        audio_loss = (audio_pred - inputs.audio_targets).pow(2).mean(dim=[-2, -1])
+
+        # Combined loss [B,]
+        return video_loss + audio_loss
diff --git a/packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py b/packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py
new file mode 100644
index 0000000000000000000000000000000000000000..2d133b3af3baa35f772c373bccea023aa3df5b33
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/training_strategies/video_to_video.py
@@ -0,0 +1,294 @@
+"""Video-to-video training strategy for IC-LoRA.
+This strategy implements training with reference video conditioning where:
+- Reference latents (clean) are concatenated with target latents (noised)
+- Video coordinates handle both reference and target sequences
+- Loss is computed only on the target portion
+"""
+
+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 import logger
+from ltx_trainer.timestep_samplers import TimestepSampler
+from ltx_trainer.training_strategies.base_strategy import (
+    DEFAULT_FPS,
+    ModelInputs,
+    TrainingStrategy,
+    TrainingStrategyConfigBase,
+)
+
+
+class VideoToVideoConfig(TrainingStrategyConfigBase):
+    """Configuration for video-to-video (IC-LoRA) training strategy."""
+
+    name: Literal["video_to_video"] = "video_to_video"
+
+    first_frame_conditioning_p: float = Field(
+        default=0.1,
+        description="Probability of conditioning on the first frame during training",
+        ge=0.0,
+        le=1.0,
+    )
+
+    reference_latents_dir: str = Field(
+        default="reference_latents",
+        description="Directory name for latents of reference videos",
+    )
+
+    def get_data_sources(self) -> dict[str, str]:
+        """IC-LoRA training requires latents, conditions, and reference latents."""
+        return {
+            "latents": "latents",
+            "conditions": "conditions",
+            self.reference_latents_dir: "ref_latents",
+        }
+
+
+class VideoToVideoStrategy(TrainingStrategy):
+    """Video-to-video training strategy for IC-LoRA.
+    This strategy implements training with reference video conditioning where:
+    - Reference latents (clean) are concatenated with target latents (noised)
+    - Video coordinates handle both reference and target sequences
+    - Loss is computed only on the target portion
+    Attributes:
+        reference_downscale_factor: The inferred downscale factor of reference videos.
+            This is computed from the first batch and cached for metadata export.
+    """
+
+    config: VideoToVideoConfig
+    reference_downscale_factor: int | None
+
+    def __init__(self, config: VideoToVideoConfig):
+        """Initialize strategy with configuration.
+        Args:
+            config: Video-to-video configuration
+        """
+        super().__init__(config)
+        self.reference_downscale_factor = None  # Will be inferred from first batch
+
+    def prepare_training_inputs(  # noqa: PLR0915
+        self,
+        batch: dict[str, Any],
+        timestep_sampler: TimestepSampler,
+    ) -> ModelInputs:
+        """Prepare inputs for IC-LoRA training with reference videos."""
+        # Get pre-encoded latents - dataset provides uniform non-patchified format [B, C, F, H, W]
+        latents = batch["latents"]
+        target_latents = latents["latents"]
+        ref_latents = batch["ref_latents"]["latents"]
+
+        # Get dimensions
+        num_frames = latents["num_frames"][0].item()
+        height = latents["height"][0].item()
+        width = latents["width"][0].item()
+
+        ref_latents_info = batch["ref_latents"]
+        ref_frames = ref_latents_info["num_frames"][0].item()
+        ref_height = ref_latents_info["height"][0].item()
+        ref_width = ref_latents_info["width"][0].item()
+
+        # Infer reference downscale factor from dimension ratios
+        # This allows training with downscaled reference videos for efficiency
+        reference_downscale_factor = self._infer_reference_downscale_factor(
+            target_height=height,
+            target_width=width,
+            ref_height=ref_height,
+            ref_width=ref_width,
+        )
+
+        # Cache the scale factor for metadata export (only on first batch)
+        if self.reference_downscale_factor is None:
+            self.reference_downscale_factor = reference_downscale_factor
+        elif self.reference_downscale_factor != reference_downscale_factor:
+            raise ValueError(
+                f"Inconsistent reference downscale factor across batches. "
+                f"First batch had factor={self.reference_downscale_factor}, "
+                f"but current batch has factor={reference_downscale_factor}. "
+                f"All training samples must use the same reference/target resolution ratio."
+            )
+
+        # Patchify latents: [B, C, F, H, W] -> [B, seq_len, C]
+        target_latents = self._video_patchifier.patchify(target_latents)
+        ref_latents = self._video_patchifier.patchify(ref_latents)
+
+        # Handle FPS
+        fps = latents.get("fps", None)
+        if fps is not None and not torch.all(fps == fps[0]):
+            logger.warning(
+                f"Different FPS values found in the batch. Found: {fps.tolist()}, using the first one: {fps[0].item()}"
+            )
+        fps = fps[0].item() if fps is not None else DEFAULT_FPS
+
+        # Get text embeddings (already processed by embedding connectors in trainer)
+        # Video-to-video uses only video embeddings
+        conditions = batch["conditions"]
+        prompt_embeds = conditions["video_prompt_embeds"]
+        prompt_attention_mask = conditions["prompt_attention_mask"]
+
+        batch_size = target_latents.shape[0]
+        ref_seq_len = ref_latents.shape[1]
+        target_seq_len = target_latents.shape[1]
+        device = target_latents.device
+
+        # Create conditioning mask
+        # Reference tokens are always conditioning (timestep=0)
+        ref_conditioning_mask = torch.ones(batch_size, ref_seq_len, dtype=torch.bool, device=device)
+
+        # Target tokens: check for first frame conditioning
+        target_conditioning_mask = self._create_first_frame_conditioning_mask(
+            batch_size=batch_size,
+            sequence_length=target_seq_len,
+            height=height,
+            width=width,
+            device=device,
+            first_frame_conditioning_p=self.config.first_frame_conditioning_p,
+        )
+
+        # Combined conditioning mask
+        conditioning_mask = torch.cat([ref_conditioning_mask, target_conditioning_mask], dim=1)
+
+        # Sample noise and sigmas for target
+        sigmas = timestep_sampler.sample_for(target_latents)
+        noise = torch.randn_like(target_latents)
+        sigmas_expanded = sigmas.view(-1, 1, 1)
+
+        # Apply noise to target
+        noisy_target = (1 - sigmas_expanded) * target_latents + sigmas_expanded * noise
+
+        # For first frame conditioning in target, use clean latents
+        target_conditioning_mask_expanded = target_conditioning_mask.unsqueeze(-1)
+        noisy_target = torch.where(target_conditioning_mask_expanded, target_latents, noisy_target)
+
+        # Targets for loss computation (velocity prediction) - only for target portion
+        targets = noise - target_latents
+
+        # Concatenate reference (clean) and target (noisy)
+        combined_latents = torch.cat([ref_latents, noisy_target], dim=1)
+
+        # Create per-token timesteps
+        timesteps = self._create_per_token_timesteps(conditioning_mask, sigmas.squeeze())
+
+        # Generate positions for reference and target separately, then concatenate
+        ref_positions = self._get_video_positions(
+            num_frames=ref_frames,
+            height=ref_height,
+            width=ref_width,
+            batch_size=batch_size,
+            fps=fps,
+            device=device,
+        )
+
+        # Scale reference positions to match target coordinate space
+        # This maps ref positions from (0, ref_H, ref_W) to (0, target_H, target_W)
+        # Position tensor shape: [B, 3, seq_len, 2] where dim 1 is (time, height, width)
+        if reference_downscale_factor != 1:
+            ref_positions = ref_positions.clone()
+            ref_positions[:, 1, ...] *= reference_downscale_factor  # height axis
+            ref_positions[:, 2, ...] *= reference_downscale_factor  # width axis
+            # Time axis (index 0) remains unchanged
+
+        target_positions = self._get_video_positions(
+            num_frames=num_frames,
+            height=height,
+            width=width,
+            batch_size=batch_size,
+            fps=fps,
+            device=device,
+        )
+
+        # Concatenate positions along sequence dimension
+        positions = torch.cat([ref_positions, target_positions], dim=2)
+
+        # Create video Modality
+        video_modality = Modality(
+            enabled=True,
+            latent=combined_latents,
+            sigma=sigmas,
+            timesteps=timesteps,
+            positions=positions,
+            context=prompt_embeds,
+            context_mask=prompt_attention_mask,
+        )
+
+        # Loss mask: only compute loss on non-conditioning target tokens
+        # Reference tokens: all False (no loss)
+        # Target tokens: True where not conditioning
+        ref_loss_mask = torch.zeros(batch_size, ref_seq_len, dtype=torch.bool, device=device)
+        target_loss_mask = ~target_conditioning_mask
+        video_loss_mask = torch.cat([ref_loss_mask, target_loss_mask], dim=1)
+
+        return ModelInputs(
+            video=video_modality,
+            audio=None,
+            video_targets=targets,
+            audio_targets=None,
+            video_loss_mask=video_loss_mask,
+            audio_loss_mask=None,
+        )
+
+    def compute_loss(
+        self,
+        video_pred: Tensor,
+        _audio_pred: Tensor | None,
+        inputs: ModelInputs,
+    ) -> Tensor:
+        """Compute masked loss on target portion only. Returns [B,]."""
+        # Slice prediction to match targets length (removes prepended reference tokens)
+        target_len = inputs.video_targets.shape[1]
+        target_pred = video_pred[:, -target_len:, :]
+        target_loss_mask = inputs.video_loss_mask[:, -target_len:]
+
+        # Compute per-element loss [B,]
+        loss = (target_pred - inputs.video_targets).pow(2)
+        loss_mask = target_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)
+
+    def get_checkpoint_metadata(self) -> dict[str, Any]:
+        """Get metadata for checkpoint files."""
+        metadata: dict[str, Any] = {}
+        # Always include reference_downscale_factor for IC-LoRAs so inference
+        # pipelines know the expected scale factor for reference videos.
+        if self.reference_downscale_factor is not None:
+            metadata["reference_downscale_factor"] = self.reference_downscale_factor
+        return metadata
+
+    @staticmethod
+    def _infer_reference_downscale_factor(
+        target_height: int,
+        target_width: int,
+        ref_height: int,
+        ref_width: int,
+    ) -> int:
+        """Infer the reference downscale factor from target and reference dimensions."""
+        # If dimensions match, no scaling needed
+        if target_height == ref_height and target_width == ref_width:
+            return 1
+
+        # Calculate scale factors for each dimension
+        if target_height % ref_height != 0 or target_width % ref_width != 0:
+            raise ValueError(
+                f"Target dimensions ({target_height}x{target_width}) must be exact multiples "
+                f"of reference dimensions ({ref_height}x{ref_width})"
+            )
+
+        scale_h = target_height // ref_height
+        scale_w = target_width // ref_width
+
+        if scale_h != scale_w:
+            raise ValueError(
+                f"Reference scale must be uniform. Got height scale {scale_h} and width scale {scale_w}. "
+                f"Target: {target_height}x{target_width}, Reference: {ref_height}x{ref_width}"
+            )
+
+        if scale_h < 1:
+            raise ValueError(
+                f"Reference dimensions ({ref_height}x{ref_width}) cannot be larger than "
+                f"target dimensions ({target_height}x{target_width})"
+            )
+
+        return scale_h
diff --git a/packages/ltx-trainer/src/ltx_trainer/utils.py b/packages/ltx-trainer/src/ltx_trainer/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..85bf1ddf10f1cf370beb102bad19bf2b7f791985
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/utils.py
@@ -0,0 +1,88 @@
+import io
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import ExifTags, Image, ImageCms, ImageOps
+from PIL.Image import Image as PilImage
+
+
+def open_image_as_srgb(image_path: str | Path | io.BytesIO) -> PilImage:
+    """
+    Opens an image file, applies rotation (if it's set in metadata) and converts it
+    to the sRGB color space respecting the original image color space .
+    Args:
+        image_path: Path to the image file
+    Returns:
+        PIL Image in sRGB color space
+    """
+    exif_colorspace_srgb = 1
+
+    with Image.open(image_path) as img_raw:
+        img = ImageOps.exif_transpose(img_raw)
+
+    input_icc_profile = img.info.get("icc_profile")
+
+    # Try to convert to sRGB if the image has ICC profile metadata
+    srgb_profile = ImageCms.createProfile(colorSpace="sRGB")
+    if input_icc_profile is not None:
+        input_profile = ImageCms.ImageCmsProfile(io.BytesIO(input_icc_profile))
+        srgb_img = ImageCms.profileToProfile(img, input_profile, srgb_profile, outputMode="RGB")
+    else:
+        # Try fall back to checking EXIF
+        exif_data = img.getexif()
+        if exif_data is not None:
+            # Assume sRGB if no ICC profile and EXIF has no ColorSpace tag
+            color_space_value = exif_data.get(ExifTags.Base.ColorSpace.value)
+            if color_space_value is not None and color_space_value != exif_colorspace_srgb:
+                raise ValueError(
+                    "Image has colorspace tag in EXIF but it isn't set to sRGB,"
+                    " conversion is not supported."
+                    f" EXIF ColorSpace tag value is {color_space_value}",
+                )
+
+        srgb_img = img.convert("RGB")
+
+        # Set sRGB profile in metadata since now the image is assumed to be in sRGB.
+        srgb_profile_data = ImageCms.ImageCmsProfile(srgb_profile).tobytes()
+        srgb_img.info["icc_profile"] = srgb_profile_data
+
+    return srgb_img
+
+
+def save_image(image_tensor: torch.Tensor, output_path: Path | str) -> None:
+    """Save an image tensor to a file.
+    Args:
+        image_tensor: Image tensor of shape [C, H, W] or [C, 1, H, W] in range [0, 1] or [0, 255].
+            C must be 3 (RGB).
+        output_path: Path to save the image (any PIL-supported format, e.g., .png or .jpg)
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Handle [C, 1, H, W] format (single frame from video tensor)
+    if image_tensor.ndim == 4:
+        # Squeeze frame dimension: [C, 1, H, W] -> [C, H, W]
+        if image_tensor.shape[1] == 1:
+            image_tensor = image_tensor.squeeze(1)
+        else:
+            raise ValueError(f"Expected single-frame tensor with shape [C, 1, H, W], got shape {image_tensor.shape}")
+
+    if image_tensor.ndim != 3:
+        raise ValueError(f"Expected 3D tensor [C, H, W], got {image_tensor.ndim}D tensor")
+
+    if image_tensor.shape[0] != 3:
+        raise ValueError(f"Expected 3 channels (RGB), got {image_tensor.shape[0]} channels")
+
+    # Normalize to [0, 255] uint8
+    if torch.is_floating_point(image_tensor) and image_tensor.max() <= 1.0:
+        image_tensor = image_tensor * 255
+
+    # Clamp to valid uint8 range to prevent overflow
+    image_tensor = image_tensor.clamp(0, 255)
+
+    # [C, H, W] -> [H, W, C]
+    image_np: np.ndarray = image_tensor.permute(1, 2, 0).to(torch.uint8).cpu().numpy()
+
+    # Save using PIL
+    Image.fromarray(image_np).save(output_path)
diff --git a/packages/ltx-trainer/src/ltx_trainer/validation_runner.py b/packages/ltx-trainer/src/ltx_trainer/validation_runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f54a758221377971e340c7b05da0951de9649f6
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/validation_runner.py
@@ -0,0 +1,1216 @@
+"""Self-contained validation runner for LTX-2 training.
+Encapsulates all validation concerns: model loading, embedding caching,
+conditioning media encoding, sample generation (denoising + decoding),
+output saving, and W&B logging.
+The trainer only needs to call ``run()`` with the current transformer state.
+"""
+
+import os
+from collections.abc import Callable
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import torch
+from einops import rearrange
+from torch import Tensor
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms import functional as TF  # noqa: N812
+from torchvision.transforms.functional import to_tensor
+
+from ltx_core.components.diffusion_steps import EulerDiffusionStep
+from ltx_core.components.guiders import CFGGuider, STGGuider
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.conditioning.types.latent_cond import VideoConditionByLatentIndex
+from ltx_core.conditioning.types.mask_cond import VideoConditionByMask
+from ltx_core.conditioning.types.reference_video_cond import VideoConditionByReferenceLatent
+from ltx_core.guidance.perturbations import (
+    BatchedPerturbationConfig,
+    Perturbation,
+    PerturbationConfig,
+    PerturbationType,
+)
+from ltx_core.model.audio_vae.audio_vae import encode_audio as ltx_encode_audio
+from ltx_core.model.transformer.modality import Modality
+from ltx_core.model.transformer.model import X0Model
+from ltx_core.model.video_vae import SpatialTilingConfig, TemporalTilingConfig, TilingConfig
+from ltx_core.tools import AudioLatentTools, VideoLatentTools
+from ltx_core.types import (
+    Audio,
+    AudioLatentShape,
+    LatentState,
+    SpatioTemporalScaleFactors,
+    VideoLatentShape,
+    VideoPixelShape,
+)
+from ltx_trainer import logger
+from ltx_trainer.config import ValidationConfig, ValidationSample
+from ltx_trainer.gpu_utils import free_gpu_memory_context
+from ltx_trainer.model_loader import (
+    load_audio_vae_decoder,
+    load_audio_vae_encoder,
+    load_embeddings_processor,
+    load_text_encoder,
+    load_video_vae_decoder,
+    load_video_vae_encoder,
+    load_vocoder,
+)
+from ltx_trainer.progress import SamplingContext, TrainingProgress
+from ltx_trainer.utils import open_image_as_srgb, save_image
+from ltx_trainer.video_utils import read_video, save_video
+
+if TYPE_CHECKING:
+    from ltx_core.model.transformer import LTXModel
+
+VIDEO_SCALE_FACTORS = SpatioTemporalScaleFactors.default()
+_DEFAULT_TILING = TilingConfig(
+    spatial_config=SpatialTilingConfig(tile_size_in_pixels=192, tile_overlap_in_pixels=64),
+    temporal_config=TemporalTilingConfig(tile_size_in_frames=48, tile_overlap_in_frames=24),
+)
+
+
+def _local_rank_device() -> torch.device:
+    """Per-rank CUDA device for early init (before Accelerator exists).
+    DDP-safe: ``LOCAL_RANK`` is set by accelerate before trainer init; loading on bare
+    ``"cuda"`` would resolve to ``cuda:0`` on every rank and crash with a device mismatch.
+    """
+    if not torch.cuda.is_available():
+        return torch.device("cpu")
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    return torch.device(f"cuda:{local_rank}")
+
+
+# =============================================================================
+# Data classes
+# =============================================================================
+
+
+@dataclass
+class CachedPromptEmbeddings:
+    """Pre-computed text embeddings for a validation prompt."""
+
+    video_context_positive: Tensor
+    audio_context_positive: Tensor
+    video_context_negative: Tensor | None = None
+    audio_context_negative: Tensor | None = None
+
+
+@dataclass
+class CachedConditionMedia:
+    """Pre-encoded media for a single validation condition."""
+
+    latent: Tensor
+    pixels: Tensor | None = None
+    mask: Tensor | None = None
+
+
+@dataclass
+class CachedSampleMedia:
+    """Pre-encoded conditioning media for one validation sample.
+    Keyed by condition index (position in the sample's conditions list).
+    """
+
+    conditions: dict[int, CachedConditionMedia] = field(default_factory=dict)
+
+
+# =============================================================================
+# ValidationRunner
+# =============================================================================
+
+
+class ValidationRunner:
+    """Self-contained validation: loads models, caches media, generates samples.
+    Lifecycle:
+        1. ``__init__``: loads text encoder + embeddings processor, caches prompt
+           embeddings, unloads both.  Loads VAE encoder, encodes conditioning media,
+           unloads.  Loads VAE decoder / audio decoder / vocoder and keeps them on CPU.
+        2. ``run()``: called at each validation step with the current transformer.
+           Generates all samples, saves outputs, returns file paths.
+    """
+
+    def __init__(
+        self,
+        config: ValidationConfig,
+        model_path: str | Path,
+        text_encoder_path: str | Path | None,
+        load_text_encoder_in_8bit: bool = False,
+    ):
+        self._config = config
+        self._model_path = Path(model_path)
+
+        self._video_patchifier = VideoLatentPatchifier(patch_size=1)
+        self._audio_patchifier = AudioPatchifier(patch_size=1)
+
+        self._cached_embeddings = self._cache_prompt_embeddings(text_encoder_path, load_text_encoder_in_8bit)
+        self._cached_media = self._encode_conditioning_media()
+        self._load_decoder_components()
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    @torch.no_grad()
+    @free_gpu_memory_context(after=True)
+    def run(
+        self,
+        transformer: "LTXModel",
+        step: int,
+        output_dir: Path,
+        device: torch.device,
+        progress: TrainingProgress,
+        wandb_run: object | None = None,
+        work_items: list[tuple[int, bool]] | None = None,
+    ) -> list[tuple[int, Path]]:
+        """Generate validation samples, save outputs, and optionally log to W&B.
+        Args:
+            work_items: Optional list of ``(sample_index, save_output)`` tuples for distributed
+                validation. Each rank passes its assigned indices so prompts are split across
+                ranks. ``save_output=False`` runs the generation but skips writing to disk; this
+                is used for FSDP padding so every rank executes the same number of forwards and
+                collective ops stay aligned. When ``None``, all samples are generated and saved.
+        Returns:
+            List of ``(sample_index, path)`` tuples for samples that were saved. Callers can
+            ``gather_object`` and sort by index to reconstruct the global ordering.
+        """
+        samples = self._config.samples
+        if not samples:
+            return []
+
+        if work_items is None:
+            work_items = [(i, True) for i in range(len(samples))]
+        if not work_items:
+            return []
+
+        inference_steps = self._config.inference_steps
+        sampling_ctx = progress.start_sampling(num_prompts=len(work_items), num_steps=inference_steps)
+
+        samples_dir = output_dir / "samples"
+        samples_dir.mkdir(exist_ok=True, parents=True)
+
+        results: list[tuple[int, Path]] = []
+
+        for local_i, (sample_idx, save_output) in enumerate(work_items):
+            sample = samples[sample_idx]
+            sampling_ctx.start_video(local_i)
+
+            cached_embeddings = self._cached_embeddings[sample_idx] if self._cached_embeddings else None
+            cached_media = self._cached_media[sample_idx] if self._cached_media else CachedSampleMedia()
+
+            video, audio = self._generate_sample(
+                sample=sample,
+                cached_embeddings=cached_embeddings,
+                cached_media=cached_media,
+                transformer=transformer,
+                device=device,
+                sampling_ctx=sampling_ctx,
+            )
+
+            if not save_output:
+                continue
+
+            dims = sample.video_dims or self._config.video_dims
+            num_frames = dims[2]
+            ext = "png" if num_frames == 1 else "mp4"
+            out_path = samples_dir / f"step_{step:06d}_{sample_idx + 1}.{ext}"
+
+            if video is not None:
+                if num_frames == 1:
+                    save_image(video, out_path)
+                else:
+                    save_video(
+                        video_tensor=video,
+                        output_path=out_path,
+                        fps=self._config.frame_rate,
+                        audio=audio,
+                        audio_sample_rate=self._vocoder.output_sampling_rate if audio is not None else None,
+                        video_format="CFHW",
+                    )
+                results.append((sample_idx, out_path))
+            elif audio is not None:
+                out_path = out_path.with_suffix(".wav")
+                self._save_audio(audio, out_path)
+                results.append((sample_idx, out_path))
+
+        sampling_ctx.cleanup()
+
+        rel_path = samples_dir.relative_to(output_dir)
+        logger.info(f"🎥 Validation samples for step {step} saved in {rel_path}")
+
+        if wandb_run is not None and results:
+            self.log_to_wandb(wandb_run, [p for _, p in results], step)
+
+        return results
+
+    # --- Initialization: caching ---
+
+    @torch.no_grad()
+    @free_gpu_memory_context(after=True)
+    def _cache_prompt_embeddings(
+        self,
+        text_encoder_path: str | Path | None,
+        load_in_8bit: bool,
+    ) -> list[CachedPromptEmbeddings]:
+        """Load text encoder, encode all validation prompts, cache on CPU, then unload."""
+        prompts = [s.prompt for s in self._config.samples]
+        if not prompts:
+            return []
+
+        init_device = _local_rank_device()
+
+        logger.debug("Loading text encoder for validation embedding caching...")
+        text_encoder = load_text_encoder(
+            gemma_model_path=text_encoder_path, device=init_device, dtype=torch.bfloat16, load_in_8bit=load_in_8bit
+        )
+
+        logger.debug("Loading embeddings processor for validation embedding caching...")
+        embeddings_processor = load_embeddings_processor(
+            checkpoint_path=self._model_path, device=init_device, dtype=torch.bfloat16
+        )
+
+        logger.info(f"Pre-computing embeddings for {len(prompts)} validation prompts...")
+        cached: list[CachedPromptEmbeddings] = []
+        neg_hs, neg_mask = text_encoder.encode([self._config.negative_prompt])[0]
+        neg_out = embeddings_processor.process_hidden_states(neg_hs, neg_mask)
+
+        for prompt in prompts:
+            pos_hs, pos_mask = text_encoder.encode([prompt])[0]
+            pos_out = embeddings_processor.process_hidden_states(pos_hs, pos_mask)
+
+            cached.append(
+                CachedPromptEmbeddings(
+                    video_context_positive=pos_out.video_encoding.cpu(),
+                    audio_context_positive=pos_out.audio_encoding.cpu(),
+                    video_context_negative=neg_out.video_encoding.cpu(),
+                    audio_context_negative=(
+                        neg_out.audio_encoding.cpu() if neg_out.audio_encoding is not None else None
+                    ),
+                )
+            )
+
+        del text_encoder, embeddings_processor
+        logger.debug("Validation prompt embeddings cached. Text encoder unloaded.")
+        return cached
+
+    @torch.no_grad()
+    @free_gpu_memory_context(after=True)
+    def _encode_conditioning_media(self) -> list[CachedSampleMedia]:
+        """Load VAE encoders, encode all conditioning media, cache on CPU, then unload."""
+        samples = self._config.samples
+        if not samples:
+            return []
+
+        def _cond_needs_video_encoder(cond: object) -> bool:
+            if getattr(cond, "type", None) in ("first_frame", "video_to_audio", "spatial_crop"):
+                return True
+            return getattr(cond, "video", None) is not None
+
+        def _cond_needs_audio_encoder(cond: object) -> bool:
+            if getattr(cond, "type", None) == "audio_to_video":
+                return True
+            return getattr(cond, "type", None) in ("prefix", "suffix", "mask", "reference") and (
+                getattr(cond, "audio", None) is not None
+            )
+
+        needs_video_encoder = any(_cond_needs_video_encoder(c) for s in samples for c in s.conditions)
+        needs_audio_encoder = any(_cond_needs_audio_encoder(c) for s in samples for c in s.conditions)
+
+        if not needs_video_encoder and not needs_audio_encoder:
+            return [CachedSampleMedia() for _ in samples]
+
+        device = _local_rank_device()
+        vae_encoder = None
+        audio_encoder = None
+
+        if needs_video_encoder:
+            logger.debug("Loading VAE encoder for validation media encoding...")
+            vae_encoder = load_video_vae_encoder(self._model_path, device="cpu", dtype=torch.bfloat16)
+        if needs_audio_encoder:
+            logger.debug("Loading audio VAE encoder for validation media encoding...")
+            audio_encoder = load_audio_vae_encoder(self._model_path, device="cpu", dtype=torch.bfloat16)
+
+        logger.info(f"Pre-encoding conditioning media for {len(samples)} validation samples...")
+        cached: list[CachedSampleMedia] = []
+
+        for sample in samples:
+            sample_dims = sample.video_dims or self._config.video_dims
+            cached.append(self._encode_sample_conditions(sample, sample_dims, vae_encoder, audio_encoder, device))
+
+        del vae_encoder, audio_encoder
+        logger.info("Validation conditioning media cached on CPU.")
+        return cached
+
+    def _encode_sample_conditions(
+        self,
+        sample: ValidationSample,
+        dims: tuple[int, int, int],
+        vae_encoder: torch.nn.Module | None,
+        audio_encoder: torch.nn.Module | None,
+        device: torch.device,
+    ) -> CachedSampleMedia:
+        """Encode all conditioning media for a single validation sample."""
+        s_width, s_height, s_num_frames = dims
+        sample_media = CachedSampleMedia()
+
+        for cond_idx, cond in enumerate(sample.conditions):
+            if cond.type == "first_frame":
+                image = self._load_first_frame(Path(cond.image_or_video))
+                latent = self._encode_image(image, s_height, s_width, vae_encoder, device)
+                sample_media.conditions[cond_idx] = CachedConditionMedia(latent=latent)
+
+            elif cond.type == "reference" and cond.video is not None:
+                ref_video, _ = read_video(cond.video, max_frames=s_num_frames)
+                preprocessed, pixels = self._preprocess_reference(
+                    ref_video, s_height, s_width, cond.downscale_factor, cond.temporal_scale_factor
+                )
+                latent = self._encode_video(preprocessed, vae_encoder, device)
+                sample_media.conditions[cond_idx] = CachedConditionMedia(
+                    latent=latent, pixels=pixels if cond.include_in_output else None
+                )
+
+            elif cond.type == "reference" and getattr(cond, "audio", None) is not None:
+                video_duration = s_num_frames / self._config.frame_rate
+                latent = self._encode_audio(cond.audio, audio_encoder, device, max_duration=video_duration)
+                sample_media.conditions[cond_idx] = CachedConditionMedia(latent=latent)
+
+            elif cond.type in ("prefix", "suffix", "mask", "spatial_crop", "video_to_audio"):
+                latent = self._encode_temporal_condition_media(
+                    cond=cond,
+                    target_width=s_width,
+                    target_height=s_height,
+                    max_frames=s_num_frames,
+                    vae_encoder=vae_encoder,
+                    audio_encoder=audio_encoder,
+                    device=device,
+                )
+                mask_tensor = None
+                if latent is not None and cond.type == "mask" and hasattr(cond, "mask") and cond.mask is not None:
+                    if getattr(cond, "audio", None) is not None:
+                        mask_tensor = self._load_audio_mask(cond.mask, s_num_frames, self._config.frame_rate)
+                    else:
+                        mask_tensor = self._load_and_downsample_mask(cond.mask, s_width, s_height, s_num_frames)
+                if latent is not None:
+                    sample_media.conditions[cond_idx] = CachedConditionMedia(latent=latent, mask=mask_tensor)
+
+            elif cond.type == "audio_to_video":
+                video_duration = s_num_frames / self._config.frame_rate
+                latent = self._encode_audio(cond.audio, audio_encoder, device, max_duration=video_duration)
+                sample_media.conditions[cond_idx] = CachedConditionMedia(latent=latent)
+
+        return sample_media
+
+    def _load_decoder_components(self) -> None:
+        """Load VAE decoder, audio decoder, and vocoder. Kept on CPU until generation."""
+        self._vae_decoder = None
+        needs_video_decoder = self._config.generate_video or any(
+            c.type == "video_to_audio" for s in self._config.samples for c in s.conditions
+        )
+        if needs_video_decoder:
+            logger.debug("Loading video VAE decoder for validation...")
+            self._vae_decoder = load_video_vae_decoder(self._model_path, device="cpu", dtype=torch.bfloat16)
+            if self._vae_decoder is not None:
+                self._vae_decoder.requires_grad_(False)
+
+        self._audio_decoder = None
+        self._vocoder = None
+        needs_audio_decoder = self._config.generate_audio or any(
+            c.type == "audio_to_video" for s in self._config.samples for c in s.conditions
+        )
+        if needs_audio_decoder:
+            logger.debug("Loading audio decoder and vocoder for validation...")
+            self._audio_decoder = load_audio_vae_decoder(self._model_path, device="cpu", dtype=torch.bfloat16)
+            if self._audio_decoder is not None:
+                self._audio_decoder.requires_grad_(False)
+            self._vocoder = load_vocoder(self._model_path, device="cpu", dtype=torch.bfloat16)
+            if self._vocoder is not None:
+                self._vocoder.requires_grad_(False)
+
+    def _encode_temporal_condition_media(
+        self,
+        cond: object,
+        *,
+        target_width: int,
+        target_height: int,
+        max_frames: int,
+        vae_encoder: torch.nn.Module | None,
+        audio_encoder: torch.nn.Module | None,
+        device: torch.device,
+    ) -> Tensor | None:
+        """Encode media for prefix/suffix/mask/spatial_crop/frozen-video conditions."""
+        video_path = getattr(cond, "video", None)
+        audio_path = getattr(cond, "audio", None)
+        if video_path is None and audio_path is not None:
+            max_dur = getattr(cond, "duration", None)
+            if max_dur is None and cond.type == "mask":
+                max_dur = max_frames / self._config.frame_rate
+            return self._encode_audio(
+                audio_path,
+                audio_encoder,
+                device,
+                max_duration=max_dur,
+                from_end=cond.type == "suffix",
+            )
+        if video_path is None:
+            return None
+
+        video, _ = read_video(video_path, max_frames=max_frames)
+        video = self._resize_and_center_crop(video, target_height, target_width)
+        requested_frames = getattr(cond, "num_frames", None)
+
+        if cond.type == "suffix" and requested_frames is not None:
+            # The VAE encoder has temporal receptive fields, so encoding a short
+            # suffix clip independently produces different latents than encoding the
+            # full video. Encode the full clip and extract the last N latent frames.
+            video = rearrange(video, "f c h w -> 1 c f h w")
+            valid_frames = (video.shape[2] - 1) // 8 * 8 + 1
+            video = video[:, :, :valid_frames]
+            latent = self._encode_video(video * 2.0 - 1.0, vae_encoder, device)
+            num_suffix_latent_frames = requested_frames // 8
+            return latent[:, :, -num_suffix_latent_frames:]
+
+        if requested_frames is not None:
+            requested_frames = min(requested_frames, video.shape[0])
+            video = video[:requested_frames]
+        video = rearrange(video, "f c h w -> 1 c f h w")
+        valid_frames = (video.shape[2] - 1) // 8 * 8 + 1
+        video = video[:, :, :valid_frames]
+        return self._encode_video(video * 2.0 - 1.0, vae_encoder, device)
+
+    # ------------------------------------------------------------------
+    # Unified generation pipeline
+    # ------------------------------------------------------------------
+
+    def _generate_sample(
+        self,
+        sample: ValidationSample,
+        cached_embeddings: CachedPromptEmbeddings | None,
+        cached_media: CachedSampleMedia,
+        transformer: "LTXModel",
+        device: torch.device,
+        sampling_ctx: SamplingContext,
+    ) -> tuple[Tensor | None, Tensor | None]:
+        """Generate one sample: build conditioned states, denoise, decode."""
+        dims = sample.video_dims or self._config.video_dims
+        width, height, num_frames = dims
+        seed = sample.seed or self._config.seed
+        generate_audio = self._config.generate_audio
+        generate_video = self._config.generate_video
+
+        # Determine frozen modalities from conditions
+        video_frozen = any(c.type == "video_to_audio" for c in sample.conditions)
+        audio_frozen = any(c.type == "audio_to_video" for c in sample.conditions)
+
+        # 1. Prompt embeddings
+        v_ctx_pos, a_ctx_pos, v_ctx_neg, a_ctx_neg = self._get_prompt_embeddings(cached_embeddings, device)
+
+        # 2. Build video state
+        generator = torch.Generator(device=device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+
+        video_state: LatentState | None = None
+        video_clean: LatentState | None = None
+        video_tools: VideoLatentTools | None = None
+
+        if generate_video or video_frozen:
+            video_tools = self._create_video_tools(width, height, num_frames)
+
+            if video_frozen:
+                frozen_media = self._find_condition_media(sample, cached_media, "video_to_audio")
+                video_state = video_tools.create_initial_state(
+                    device=device, dtype=torch.bfloat16, initial_latent=frozen_media.latent.to(device)
+                )
+                video_state = replace(video_state, denoise_mask=torch.zeros_like(video_state.denoise_mask))
+                video_clean = video_state
+            else:
+                video_state = video_tools.create_initial_state(device=device, dtype=torch.bfloat16)
+                video_state = self._apply_video_conditionings(video_state, video_tools, sample, cached_media, device)
+                video_clean = video_state
+                video_state = noiser(video_state, noise_scale=1.0)
+
+        # 3. Build audio state
+        audio_state: LatentState | None = None
+        audio_clean: LatentState | None = None
+        audio_tools: AudioLatentTools | None = None
+
+        if generate_audio or audio_frozen:
+            if audio_frozen:
+                frozen_media = self._find_condition_media(sample, cached_media, "audio_to_video")
+                frozen_latent = frozen_media.latent.to(device)
+                audio_tools = AudioLatentTools(
+                    patchifier=self._audio_patchifier,
+                    target_shape=AudioLatentShape.from_torch_shape(frozen_latent.shape),
+                )
+                audio_state = audio_tools.create_initial_state(
+                    device=device, dtype=torch.bfloat16, initial_latent=frozen_latent
+                )
+                audio_state = replace(audio_state, denoise_mask=torch.zeros_like(audio_state.denoise_mask))
+                audio_clean = audio_state
+            else:
+                audio_tools = self._create_audio_tools(num_frames, self._config.frame_rate)
+                audio_state = audio_tools.create_initial_state(device=device, dtype=torch.bfloat16)
+                audio_state = self._apply_audio_conditionings(audio_state, audio_tools, sample, cached_media, device)
+                audio_clean = audio_state
+                audio_state = noiser(audio_state, noise_scale=1.0)
+
+        # 4. Denoising loop
+        video_state, audio_state = self._run_denoising(
+            transformer=transformer,
+            video_state=video_state,
+            audio_state=audio_state,
+            video_clean=video_clean,
+            audio_clean=audio_clean,
+            video_frozen=video_frozen,
+            audio_frozen=audio_frozen,
+            v_ctx_pos=v_ctx_pos,
+            a_ctx_pos=a_ctx_pos,
+            v_ctx_neg=v_ctx_neg,
+            a_ctx_neg=a_ctx_neg,
+            device=device,
+            sampling_ctx=sampling_ctx,
+        )
+
+        # 5. Decode modalities (both generated and frozen — frozen audio/video is included in output)
+        video_output = self._finalize_modality(video_state, video_tools, self._decode_video, device)
+        audio_output = self._finalize_modality(audio_state, audio_tools, self._decode_audio, device)
+
+        # 6. Side-by-side reference output
+        if video_output is not None:
+            video_output = self._apply_reference_side_by_side(video_output, sample, cached_media)
+
+        return video_output, audio_output
+
+    # ------------------------------------------------------------------
+    # Conditioning application
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _apply_video_conditionings(
+        state: LatentState,
+        tools: VideoLatentTools,
+        sample: ValidationSample,
+        cached_media: CachedSampleMedia,
+        device: torch.device,
+    ) -> LatentState:
+        """Apply all video-targeting conditionings from the sample's conditions list."""
+        for cond_idx, cond in enumerate(sample.conditions):
+            media = cached_media.conditions.get(cond_idx)
+            if media is None:
+                continue
+
+            latent = media.latent.to(device=device, dtype=torch.bfloat16)
+
+            if cond.type == "first_frame" or (cond.type == "prefix" and getattr(cond, "video", None) is not None):
+                state = VideoConditionByLatentIndex(latent=latent, strength=1.0, latent_idx=0).apply_to(state, tools)
+
+            elif cond.type == "suffix" and getattr(cond, "video", None) is not None:
+                _, _, tgt_frames, _, _ = tools.target_shape.to_torch_shape()
+                _, _, cond_frames, _, _ = latent.shape
+                suffix_start = tgt_frames - cond_frames
+                state = VideoConditionByLatentIndex(latent=latent, strength=1.0, latent_idx=suffix_start).apply_to(
+                    state, tools
+                )
+
+            elif cond.type == "reference" and cond.video is not None:
+                state = VideoConditionByReferenceLatent(
+                    latent=latent,
+                    downscale_factor=cond.downscale_factor,
+                    temporal_scale_factor=cond.temporal_scale_factor,
+                    strength=1.0,
+                ).apply_to(state, tools)
+
+            elif cond.type == "mask":
+                mask = media.mask.to(device=device)
+                state = VideoConditionByMask(latent=latent, mask=mask, strength=1.0).apply_to(state, tools)
+
+            elif cond.type == "spatial_crop":
+                mask = _build_spatial_crop_mask(cond.spatial_region, tools.target_shape, device)
+                state = VideoConditionByMask(latent=latent, mask=mask, strength=1.0).apply_to(state, tools)
+
+        return state
+
+    @staticmethod
+    def _apply_audio_conditionings(
+        state: LatentState,
+        tools: AudioLatentTools,
+        sample: ValidationSample,
+        cached_media: CachedSampleMedia,
+        device: torch.device,
+    ) -> LatentState:
+        """Apply all audio-targeting conditionings from the sample's conditions list."""
+        for cond_idx, cond in enumerate(sample.conditions):
+            media = cached_media.conditions.get(cond_idx)
+            if media is None:
+                continue
+
+            latent = media.latent.to(device=device, dtype=torch.bfloat16)
+
+            if cond.type in ("prefix", "suffix") and getattr(cond, "audio", None) is not None:
+                tokens = tools.patchifier.patchify(latent)
+                num_cond_tokens = tokens.shape[1]
+
+                state = state.clone()
+                start = 0 if cond.type == "prefix" else tools.target_shape.token_count() - num_cond_tokens
+                stop = start + num_cond_tokens
+
+                state.clean_latent[:, start:stop] = tokens
+                state.denoise_mask[:, start:stop] = 0.0
+
+            elif cond.type == "reference" and getattr(cond, "audio", None) is not None:
+                tokens = tools.patchifier.patchify(latent)
+                ref_shape = AudioLatentShape.from_torch_shape(latent.shape)
+
+                positions = tools.patchifier.get_patch_grid_bounds(
+                    output_shape=ref_shape,
+                    device=device,
+                )
+
+                denoise_mask = torch.zeros(
+                    *tokens.shape[:2],
+                    1,
+                    device=device,
+                    dtype=torch.float32,
+                )
+
+                state = LatentState(
+                    latent=torch.cat([state.latent, torch.zeros_like(tokens, dtype=state.latent.dtype)], dim=1),
+                    denoise_mask=torch.cat([state.denoise_mask, denoise_mask], dim=1),
+                    positions=torch.cat([state.positions, positions], dim=2),
+                    clean_latent=torch.cat([state.clean_latent, tokens], dim=1),
+                    attention_mask=None,
+                )
+
+            elif cond.type == "mask" and getattr(cond, "audio", None) is not None:
+                tokens = tools.patchifier.patchify(latent)
+                target_len = state.latent.shape[1]
+                tokens = tokens[:, :target_len]
+
+                mask = media.mask.to(device=device)
+                if mask.shape[-1] != target_len:
+                    mask = torch.nn.functional.interpolate(
+                        mask.unsqueeze(0).float(), size=target_len, mode="nearest"
+                    ).squeeze(0)
+                    mask = (mask > 0.5).float()
+                if mask.dim() == 2:
+                    mask = mask.unsqueeze(-1)
+
+                m = mask.to(dtype=state.latent.dtype)
+                inv = 1 - m
+
+                state = LatentState(
+                    latent=state.latent,
+                    denoise_mask=state.denoise_mask * inv,
+                    positions=state.positions,
+                    clean_latent=state.clean_latent * inv + tokens * m,
+                    attention_mask=state.attention_mask,
+                )
+
+        return state
+
+    @staticmethod
+    def _find_condition_media(
+        sample: ValidationSample, cached_media: CachedSampleMedia, condition_type: str
+    ) -> CachedConditionMedia:
+        """Look up the cached media for the first condition matching the given type."""
+        for cond_idx, cond in enumerate(sample.conditions):
+            if cond.type == condition_type and cond_idx in cached_media.conditions:
+                return cached_media.conditions[cond_idx]
+        raise ValueError(f"No cached media found for condition type '{condition_type}'")
+
+    @staticmethod
+    def _apply_reference_side_by_side(
+        video_output: Tensor, sample: ValidationSample, cached_media: CachedSampleMedia
+    ) -> Tensor:
+        """Concatenate reference video pixels side-by-side with generated output if requested."""
+        for cond_idx, cond in enumerate(sample.conditions):
+            if cond.type == "reference" and cond.include_in_output:
+                media = cached_media.conditions.get(cond_idx)
+                if media is not None and media.pixels is not None:
+                    ref_pixels = media.pixels
+                    output_frames = video_output.shape[1]
+                    ref_frames = ref_pixels.shape[1]
+                    if ref_frames < output_frames:
+                        temporal_sf = cond.temporal_scale_factor
+                        if temporal_sf > 1:
+                            # VAE-aligned stretch mirroring training-time subsampling:
+                            # first frame preserved once (aligns with VAE's causal
+                            # first latent), each subsequent ref frame held for
+                            # `temporal_sf` output frames.
+                            indices = torch.tensor([0] + [1 + (i - 1) // temporal_sf for i in range(1, output_frames)])
+                            indices = indices.clamp(max=ref_frames - 1)
+                        else:
+                            indices = torch.linspace(0, ref_frames - 1, output_frames).round().long()
+                        ref_pixels = ref_pixels[:, indices]
+                    video_output = _concatenate_videos_side_by_side(ref_pixels, video_output)
+        return video_output
+
+    # ------------------------------------------------------------------
+    # Denoising loop
+    # ------------------------------------------------------------------
+
+    def _run_denoising(  # noqa: PLR0913
+        self,
+        transformer: "LTXModel",
+        video_state: LatentState | None,
+        audio_state: LatentState | None,
+        video_clean: LatentState | None,
+        audio_clean: LatentState | None,
+        *,
+        video_frozen: bool,
+        audio_frozen: bool,
+        v_ctx_pos: Tensor,
+        a_ctx_pos: Tensor,
+        v_ctx_neg: Tensor | None,
+        a_ctx_neg: Tensor | None,
+        device: torch.device,
+        sampling_ctx: SamplingContext,
+    ) -> tuple[LatentState | None, LatentState | None]:
+        """Run the Euler denoising loop with CFG/STG, handling frozen modalities."""
+        cfg = self._config
+        scheduler = LTX2Scheduler()
+        sigmas = scheduler.execute(steps=cfg.inference_steps).to(device).float()
+        stepper = EulerDiffusionStep()
+        cfg_guider = CFGGuider(cfg.guidance_scale)
+        stg_guider = STGGuider(cfg.stg_scale)
+
+        stg_perturbation_config = (
+            self._build_stg_perturbation_config(
+                cfg.stg_blocks, cfg.stg_mode, transformer.num_blocks, device, next(transformer.parameters()).dtype
+            )
+            if stg_guider.enabled()
+            else None
+        )
+
+        x0_model = X0Model(transformer)
+
+        for step_idx, sigma in enumerate(sigmas[:-1]):
+            v_sigma = torch.zeros_like(sigma) if video_frozen else sigma
+            a_sigma = torch.zeros_like(sigma) if audio_frozen else sigma
+
+            video = (
+                self._modality_from_latent_state(video_state, v_ctx_pos, v_sigma.unsqueeze(0))
+                if video_state is not None
+                else None
+            )
+            audio = (
+                self._modality_from_latent_state(audio_state, a_ctx_pos, a_sigma.unsqueeze(0))
+                if audio_state is not None
+                else None
+            )
+
+            pos_video, pos_audio = x0_model(video=video, audio=audio, perturbations=None)
+            denoised_video, denoised_audio = pos_video, pos_audio
+
+            # CFG
+            if cfg_guider.enabled() and v_ctx_neg is not None:
+                video_neg = replace(video, context=v_ctx_neg) if video is not None else None
+                audio_neg = replace(audio, context=a_ctx_neg) if audio is not None else None
+                neg_video, neg_audio = x0_model(video=video_neg, audio=audio_neg, perturbations=None)
+
+                if not video_frozen and denoised_video is not None:
+                    denoised_video = denoised_video + cfg_guider.delta(pos_video, neg_video)
+                if not audio_frozen and denoised_audio is not None:
+                    denoised_audio = denoised_audio + cfg_guider.delta(pos_audio, neg_audio)
+
+            # STG
+            if stg_perturbation_config is not None:
+                ptb_video, ptb_audio = x0_model(video=video, audio=audio, perturbations=stg_perturbation_config)
+                if not video_frozen and denoised_video is not None:
+                    denoised_video = denoised_video + stg_guider.delta(pos_video, ptb_video)
+                if not audio_frozen and denoised_audio is not None and ptb_audio is not None:
+                    denoised_audio = denoised_audio + stg_guider.delta(pos_audio, ptb_audio)
+
+            # Re-apply conditioning mask
+            if denoised_video is not None and video_clean is not None:
+                denoised_video = self._post_process_latent(
+                    denoised_video,
+                    video_state.denoise_mask,
+                    video_clean.clean_latent,
+                )
+            if denoised_audio is not None and audio_clean is not None:
+                denoised_audio = self._post_process_latent(
+                    denoised_audio,
+                    audio_state.denoise_mask,
+                    audio_clean.clean_latent,
+                )
+
+            # Euler step (skip for frozen modalities)
+            if video_state is not None and not video_frozen:
+                video_state = replace(
+                    video_state,
+                    latent=stepper.step(video.latent, denoised_video, sigmas, step_idx),
+                )
+            if audio_state is not None and not audio_frozen:
+                audio_state = replace(
+                    audio_state,
+                    latent=stepper.step(audio.latent, denoised_audio, sigmas, step_idx),
+                )
+
+            sampling_ctx.advance_step()
+
+        return video_state, audio_state
+
+    # --- Finalization (decode) ---
+
+    def _finalize_modality(
+        self,
+        state: LatentState | None,
+        tools: VideoLatentTools | AudioLatentTools | None,
+        decode_fn: Callable[[LatentState, torch.device], Tensor],
+        device: torch.device,
+    ) -> Tensor | None:
+        """Clear conditioning tokens, unpatchify, and decode a modality (generated or frozen)."""
+        if state is None or tools is None:
+            return None
+        state = tools.clear_conditioning(state)
+        state = tools.unpatchify(state)
+        return decode_fn(state, device)
+
+    def _decode_video(self, video_state: LatentState, device: torch.device) -> Tensor:
+        """Decode video latents to pixels using tiled VAE decoding."""
+        self._vae_decoder.to(device)
+        latent = video_state.latent.to(dtype=torch.bfloat16)
+
+        chunks = list(self._vae_decoder.tiled_decode(latent, tiling_config=_DEFAULT_TILING))
+        decoded_video = torch.cat(chunks, dim=2)
+
+        decoded_video = ((decoded_video + 1.0) / 2.0).clamp(0.0, 1.0)
+        self._vae_decoder.to("cpu")
+        return decoded_video[0].float().cpu()
+
+    def _decode_audio(self, audio_state: LatentState, device: torch.device) -> Tensor:
+        """Decode audio latents to waveform via audio VAE + vocoder."""
+        self._audio_decoder.to(device)
+        first_param = next(self._audio_decoder.parameters(), None)
+        decoder_dtype = first_param.dtype if first_param is not None else audio_state.latent.dtype
+        latent = audio_state.latent.to(dtype=decoder_dtype, device=device)
+        decoded_audio = self._audio_decoder(latent)
+        self._audio_decoder.to("cpu")
+
+        self._vocoder.to(device)
+        audio_waveform = self._vocoder(decoded_audio)
+        self._vocoder.to("cpu")
+
+        return audio_waveform.squeeze(0).float().cpu()
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _save_audio(self, audio: Tensor, path: Path) -> None:
+        """Save an audio waveform tensor to a .wav file using stdlib wave."""
+        import wave  # noqa: PLC0415
+
+        if audio.dim() == 1:
+            audio = audio.unsqueeze(0)
+        pcm = (audio.clamp(-1.0, 1.0).cpu() * 32767).to(torch.int16)
+        with wave.open(str(path), "wb") as wf:
+            wf.setnchannels(pcm.shape[0])
+            wf.setsampwidth(2)
+            wf.setframerate(self._vocoder.output_sampling_rate)
+            wf.writeframes(pcm.T.contiguous().numpy().tobytes())
+
+    @staticmethod
+    def _encode_audio(
+        audio_path: str,
+        audio_encoder: torch.nn.Module,
+        device: torch.device,
+        max_duration: float | None = None,
+        *,
+        from_end: bool = False,
+    ) -> Tensor:
+        """Encode an audio file through the audio VAE encoder. Returns latent on CPU."""
+        import torchaudio  # noqa: PLC0415
+
+        waveform, sr = torchaudio.load(audio_path)
+        if max_duration is not None:
+            num_samples = min(round(max_duration * sr), waveform.shape[-1])
+            waveform = waveform[:, -num_samples:] if from_end else waveform[:, :num_samples]
+        audio = Audio(waveform=waveform.unsqueeze(0), sampling_rate=sr)
+        audio_encoder.to(device)
+        latent = ltx_encode_audio(audio, audio_encoder)
+        audio_encoder.to("cpu")
+        return latent.cpu()
+
+    @staticmethod
+    def _load_and_downsample_mask(
+        mask_path: str | Path,
+        target_width: int,
+        target_height: int,
+        target_num_frames: int,
+    ) -> Tensor:
+        """Load a mask image/video and downsample to latent-space dimensions.
+        Returns a binary float tensor of shape [1, F', H', W'] where F', H', W' are
+        the latent-space dimensions corresponding to the target video dims.
+        """
+        from ltx_core.types import SpatioTemporalScaleFactors  # noqa: PLC0415
+
+        sf = SpatioTemporalScaleFactors.default()
+        latent_f = (target_num_frames - 1) // sf.time + 1
+
+        mask_path = Path(mask_path)
+        if mask_path.suffix.lower() in (".png", ".jpg", ".jpeg", ".bmp", ".webp"):
+            img = to_tensor(open_image_as_srgb(mask_path)).mean(dim=0, keepdim=True)  # [1, H, W]
+            img = torch.nn.functional.interpolate(
+                img.unsqueeze(0), size=(target_height, target_width), mode="nearest"
+            ).squeeze(0)  # [1, H, W]
+            mask_pixels = img.expand(target_num_frames, -1, -1)  # [F, H, W]
+        else:
+            frames, _ = read_video(str(mask_path), max_frames=target_num_frames)  # [F, C, H, W]
+            frames = frames[:target_num_frames].mean(dim=1)  # grayscale [F, H, W]
+            mask_pixels = torch.nn.functional.interpolate(
+                frames.unsqueeze(1), size=(target_height, target_width), mode="nearest"
+            ).squeeze(1)  # [F, H, W]
+
+        mask_latent = torch.nn.functional.avg_pool2d(
+            mask_pixels.unsqueeze(1), kernel_size=(sf.height, sf.width)
+        ).squeeze(1)  # [F, H', W']
+
+        # Temporal: max-pool over groups of sf.time frames (conservative — any masked frame masks the group)
+        f_spatial = mask_latent.shape[0]
+        pad_f = (sf.time - f_spatial % sf.time) % sf.time
+        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, sf.time, h_prime, w_prime).amax(dim=1)[:latent_f]
+
+        mask_latent = (mask_latent > 0.5).float()
+        return mask_latent.unsqueeze(0)  # [1, F', H', W'] — unpatchified latent space
+
+    @staticmethod
+    def _load_audio_mask(
+        mask_path: str | Path,
+        target_num_frames: int,
+        frame_rate: float,
+    ) -> Tensor:
+        """Load an audio mask (.pt or .wav) and resample to target audio latent length.
+        Returns a float tensor of shape [1, num_tokens] in patchified token space.
+        """
+        mask_path = Path(mask_path)
+        if mask_path.suffix == ".pt":
+            raw = torch.load(mask_path, map_location="cpu", weights_only=True)
+            if isinstance(raw, dict):
+                raw = raw.get("mask", next(iter(raw.values())))
+            raw = raw.float().flatten()
+        else:
+            import torchaudio  # noqa: PLC0415
+
+            waveform, _sr = torchaudio.load(mask_path)
+            raw = waveform.abs().mean(dim=0)
+
+        target_duration = target_num_frames / frame_rate
+        target_shape = AudioLatentShape.from_duration(batch=1, duration=target_duration)
+        target_len = target_shape.frames
+
+        resampled = torch.nn.functional.interpolate(
+            raw.unsqueeze(0).unsqueeze(0), size=target_len, mode="nearest"
+        ).squeeze()
+        mask = (resampled > 0.5).float()
+        return mask.unsqueeze(0)  # [1, num_tokens]
+
+    def _create_video_tools(self, width: int, height: int, num_frames: int) -> VideoLatentTools:
+        """Create VideoLatentTools for the given output dimensions."""
+        pixel_shape = VideoPixelShape(
+            batch=1, frames=num_frames, height=height, width=width, fps=self._config.frame_rate
+        )
+        return VideoLatentTools(
+            patchifier=self._video_patchifier,
+            target_shape=VideoLatentShape.from_pixel_shape(shape=pixel_shape),
+            fps=self._config.frame_rate,
+            scale_factors=VIDEO_SCALE_FACTORS,
+            causal_fix=True,
+        )
+
+    def _create_audio_tools(self, num_frames: int, frame_rate: float) -> AudioLatentTools:
+        """Create AudioLatentTools for the given video duration."""
+        return AudioLatentTools(
+            patchifier=self._audio_patchifier,
+            target_shape=AudioLatentShape.from_duration(batch=1, duration=num_frames / frame_rate),
+        )
+
+    @staticmethod
+    def _get_prompt_embeddings(
+        cached: CachedPromptEmbeddings | None, device: torch.device
+    ) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]:
+        """Move cached prompt embeddings to device, returning (v_pos, a_pos, v_neg, a_neg)."""
+        if cached is None:
+            raise ValueError("Cached prompt embeddings are required for validation generation")
+        return (
+            cached.video_context_positive.to(device),
+            cached.audio_context_positive.to(device),
+            cached.video_context_negative.to(device) if cached.video_context_negative is not None else None,
+            cached.audio_context_negative.to(device) if cached.audio_context_negative is not None else None,
+        )
+
+    @staticmethod
+    def log_to_wandb(wandb_run: object, sample_paths: list[Path], step: int) -> None:
+        """Log validation outputs (images, videos, or audio) to Weights & Biases."""
+        import wandb  # noqa: PLC0415
+
+        suffix = sample_paths[0].suffix.lower()
+        if suffix in (".png", ".jpg", ".jpeg", ".heic", ".webp"):
+            media = [wandb.Image(str(path)) for path in sample_paths]
+        elif suffix in (".wav", ".mp3", ".flac", ".ogg"):
+            media = [wandb.Audio(str(path)) for path in sample_paths]
+        else:
+            media = [wandb.Video(str(path), format=suffix.lstrip(".")) for path in sample_paths]
+        wandb_run.log({"validation_samples": media}, step=step)
+
+    @staticmethod
+    def _post_process_latent(denoised: Tensor, denoise_mask: Tensor, clean_latent: Tensor) -> Tensor:
+        """Blend denoised output with clean latent according to the denoise mask."""
+        return (denoised * denoise_mask + clean_latent.float() * (1 - denoise_mask)).to(denoised.dtype)
+
+    @staticmethod
+    def _modality_from_latent_state(state: LatentState, context: Tensor, sigma: Tensor) -> Modality:
+        """Build a Modality object from a LatentState, text context, and sigma."""
+        return Modality(
+            enabled=True,
+            latent=state.latent,
+            sigma=sigma,
+            timesteps=state.denoise_mask * sigma,
+            positions=state.positions,
+            context=context,
+            context_mask=None,
+        )
+
+    @staticmethod
+    def _build_stg_perturbation_config(
+        stg_blocks: list[int] | None,
+        stg_mode: Literal["stg_av", "stg_v"],
+        num_blocks: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> BatchedPerturbationConfig:
+        """Build STG perturbation config that skips self-attention in the specified blocks."""
+        perturbations: list[Perturbation] = [
+            Perturbation(type=PerturbationType.SKIP_VIDEO_SELF_ATTN, blocks=stg_blocks)
+        ]
+        if stg_mode == "stg_av":
+            perturbations.append(Perturbation(type=PerturbationType.SKIP_AUDIO_SELF_ATTN, blocks=stg_blocks))
+        return BatchedPerturbationConfig([PerturbationConfig(perturbations=perturbations)], num_blocks, device, dtype)
+
+    @staticmethod
+    def _load_first_frame(media_path: Path) -> Tensor:
+        """Load the first frame from an image or video file as [C, H, W] in [0, 1]."""
+        video_extensions = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
+        if media_path.suffix.lower() in video_extensions:
+            frames, _ = read_video(str(media_path), max_frames=1)
+            return frames[0]
+        image = open_image_as_srgb(str(media_path))
+        return TF.to_tensor(image)
+
+    @staticmethod
+    def _resize_and_center_crop(tensor: Tensor, target_height: int, target_width: int) -> Tensor:
+        """Resize [N, C, H, W] tensor to cover target dims (preserving aspect ratio) and center-crop."""
+        current_height, current_width = tensor.shape[2:]
+        if current_height == target_height and current_width == target_width:
+            return tensor
+
+        aspect_ratio = current_width / current_height
+        target_aspect_ratio = target_width / target_height
+        if aspect_ratio > target_aspect_ratio:
+            resize_height, resize_width = target_height, int(target_height * aspect_ratio)
+        else:
+            resize_height, resize_width = int(target_width / aspect_ratio), target_width
+
+        tensor = TF.resize(tensor, size=[resize_height, resize_width], interpolation=InterpolationMode.BICUBIC)
+        tensor = tensor.clamp(0, 1)
+        h_start = (resize_height - target_height) // 2
+        w_start = (resize_width - target_width) // 2
+        return tensor[:, :, h_start : h_start + target_height, w_start : w_start + target_width]
+
+    @staticmethod
+    def _encode_image(
+        image: Tensor, target_height: int, target_width: int, vae_encoder: torch.nn.Module, device: torch.device
+    ) -> Tensor:
+        """Encode a conditioning image to latent space. Returns [B, C, 1, H', W'] on CPU."""
+        image = ValidationRunner._resize_and_center_crop(image.unsqueeze(0), target_height, target_width)
+        image = image.unsqueeze(2)  # [B, C, 1, H, W]
+        image = (image * 2.0 - 1.0).to(device=device, dtype=torch.bfloat16)
+        vae_encoder.to(device)
+        encoded = vae_encoder(image)
+        vae_encoder.to("cpu")
+        return encoded.cpu()
+
+    @staticmethod
+    def _preprocess_reference(
+        video: Tensor,
+        target_height: int,
+        target_width: int,
+        downscale_factor: int = 1,
+        temporal_scale_factor: int = 1,
+    ) -> tuple[Tensor, Tensor]:
+        """Preprocess reference video. Returns (preprocessed [-1,1], pixels [0,1])."""
+        ref_height = target_height // downscale_factor
+        ref_width = target_width // downscale_factor
+        if ref_height % 32 != 0 or ref_width % 32 != 0:
+            raise ValueError(
+                f"Scaled reference dimensions ({ref_height}x{ref_width}) must be divisible by 32. "
+                f"Original: {target_height}x{target_width}, downscale_factor: {downscale_factor}"
+            )
+
+        video = ValidationRunner._resize_and_center_crop(video, ref_height, ref_width)
+        video = rearrange(video, "f c h w -> 1 c f h w")
+        valid_frames = (video.shape[2] - 1) // 8 * 8 + 1
+        video = video[:, :, :valid_frames]
+
+        # VAE-aligned temporal subsampling: keep frame 0, then every Nth frame
+        if temporal_scale_factor > 1:
+            indices = [0, *list(range(1, video.shape[2], temporal_scale_factor))]
+            video = video[:, :, indices]
+
+        pixels = video[0].clone()
+        preprocessed = video * 2.0 - 1.0
+        return preprocessed, pixels
+
+    @staticmethod
+    def _encode_video(video: Tensor, vae_encoder: torch.nn.Module, device: torch.device) -> Tensor:
+        """Encode a [B, C, F, H, W] video tensor through the VAE. Returns latent on CPU."""
+        vae_encoder.to(device)
+        latent = vae_encoder.tiled_encode(video.to(dtype=torch.bfloat16), TilingConfig.default())
+        vae_encoder.to("cpu")
+        return latent.cpu()
+
+
+def _build_spatial_crop_mask(
+    region: tuple[int, int, int, int], target_shape: VideoLatentShape, device: torch.device
+) -> Tensor:
+    """Build a binary mask from a pixel-space spatial region (y1, x1, y2, x2)."""
+    y1, x1, y2, x2 = region
+    _, _, frames, height, width = target_shape.to_torch_shape()
+
+    def to_latent(v: int, scale: int, max_v: int) -> int:
+        return max(0, min(v // scale, max_v))
+
+    ly1 = to_latent(y1, VIDEO_SCALE_FACTORS.height, height)
+    ly2 = to_latent(y2, VIDEO_SCALE_FACTORS.height, height)
+    lx1 = to_latent(x1, VIDEO_SCALE_FACTORS.width, width)
+    lx2 = to_latent(x2, VIDEO_SCALE_FACTORS.width, width)
+
+    spatial_mask = torch.zeros(height, width, dtype=torch.float32, device=device)
+    spatial_mask[ly1:ly2, lx1:lx2] = 1.0
+
+    return spatial_mask.unsqueeze(0).unsqueeze(0).expand(1, frames, -1, -1)  # [1, F, H, W]
+
+
+def _concatenate_videos_side_by_side(left_video: Tensor, right_video: Tensor) -> Tensor:
+    """Concatenate two [C, F, H, W] videos horizontally, matching height and padding frames."""
+    left_height, left_width = left_video.shape[2], left_video.shape[3]
+    right_height = right_video.shape[2]
+    if left_height != right_height:
+        scale = right_height / left_height
+        new_width = int(left_width * scale)
+        c, f, h, w = left_video.shape
+        left_video = left_video.reshape(c * f, 1, h, w)
+        left_video = TF.resize(left_video, size=[right_height, new_width], interpolation=InterpolationMode.BICUBIC)
+        left_video = left_video.clamp(0, 1)
+        left_video = left_video.reshape(c, f, right_height, new_width)
+    left_frames, right_frames = left_video.shape[1], right_video.shape[1]
+    if left_frames < right_frames:
+        padding = left_video[:, -1:].expand(-1, right_frames - left_frames, -1, -1)
+        left_video = torch.cat([left_video, padding], dim=1)
+    elif right_frames < left_frames:
+        padding = right_video[:, -1:].expand(-1, left_frames - right_frames, -1, -1)
+        right_video = torch.cat([right_video, padding], dim=1)
+    return torch.cat([left_video, right_video], dim=3)
diff --git a/packages/ltx-trainer/src/ltx_trainer/video_utils.py b/packages/ltx-trainer/src/ltx_trainer/video_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e5df1bff6b127d016ebd84db2f540e43e9d53f0
--- /dev/null
+++ b/packages/ltx-trainer/src/ltx_trainer/video_utils.py
@@ -0,0 +1,188 @@
+"""Video I/O utilities using PyAV.
+This module provides functions for reading and writing video files using PyAV,
+with optional audio support.
+"""
+
+from fractions import Fraction
+from pathlib import Path
+from typing import Literal
+
+import av
+import numpy as np
+import torch
+from torch import Tensor
+
+VideoFormat = Literal["CFHW", "FCHW"]
+
+
+def get_video_frame_count(video_path: str | Path) -> int:
+    """Get the number of frames in a video file.
+    Tries three approaches in order: stream metadata, duration*fps estimate,
+    full decode. The estimate may be off by a few frames for VFR videos or
+    containers with edit lists — exact for the min_frames filtering use case.
+    Args:
+        video_path: Path to the video file
+    Returns:
+        Number of frames in the video
+    """
+    with av.open(str(video_path)) as container:
+        video_stream = container.streams.video[0]
+
+        if video_stream.frames > 0:
+            return video_stream.frames
+
+        # Fast estimate from container metadata (avoids full decode).
+        # Uses Fraction arithmetic to prevent float precision loss.
+        rate = video_stream.average_rate or video_stream.base_rate
+        if video_stream.duration and video_stream.time_base and rate:
+            duration = Fraction(video_stream.duration) * Fraction(video_stream.time_base)
+            return round(duration * Fraction(rate))
+
+        # Last resort: full decode (very slow for 4K)
+        return sum(1 for _ in container.decode(video=0))
+
+
+def read_video(video_path: str | Path, max_frames: int | None = None) -> tuple[Tensor, float]:
+    """Load frames from a video file using PyAV.
+    Args:
+        video_path: Path to the video file
+        max_frames: Maximum number of frames to read. If None, reads all frames.
+    Returns:
+        Video tensor with shape [F, C, H, W] in range [0, 1] and frames per second (fps).
+    """
+    with av.open(str(video_path)) as container:
+        video_stream = container.streams.video[0]
+        fps = float(video_stream.average_rate or video_stream.base_rate or 24)
+
+        frames = []
+        for frame in container.decode(video=0):
+            if max_frames is not None and len(frames) >= max_frames:
+                break
+            frames.append(frame.to_ndarray(format="rgb24"))
+
+    frames_np = np.stack(frames, axis=0)  # [F, H, W, C]
+    video = torch.from_numpy(frames_np).float().div(255.0)  # [F, H, W, C] in [0, 1]
+    return video.permute(0, 3, 1, 2), fps  # [F, C, H, W]
+
+
+def save_video(
+    video_tensor: torch.Tensor,
+    output_path: Path | str,
+    fps: float = 24.0,
+    audio: torch.Tensor | None = None,
+    audio_sample_rate: int | None = None,
+    video_format: VideoFormat | None = None,
+) -> None:
+    """Save a video tensor to a file using PyAV, optionally with audio.
+    Args:
+        video_tensor: Video tensor of shape [C, F, H, W] or [F, C, H, W] in range [0, 1] or [0, 255]
+        output_path: Path to save the video
+        fps: Frames per second for the output video
+        audio: Optional audio tensor of shape [C, samples] or [samples, C] in range [-1, 1]
+        audio_sample_rate: Sample rate for the audio (required if audio is provided)
+        video_format: Explicit layout of ``video_tensor``, either ``"CFHW"`` or ``"FCHW"``.
+            When ``None`` (default), the layout is auto-detected using a heuristic that only
+            works when ``shape[1] > 3`` — the ambiguous ``[C=3, F=3, H, W]`` / ``[F=3, C=3, H, W]``
+            case requires passing this argument explicitly.
+    """
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Normalize to [F, H, W, C] uint8 numpy array
+    video_np = _prepare_video_array(video_tensor, video_format=video_format)
+    _, height, width, _ = video_np.shape
+
+    with av.open(str(output_path), mode="w") as container:
+        # Setup video stream
+        video_stream = container.add_stream("libx264", rate=int(fps))
+        video_stream.width = width
+        video_stream.height = height
+        video_stream.pix_fmt = "yuv420p"
+        video_stream.options = {"crf": "18"}
+
+        # Setup audio stream if needed
+        if audio is not None:
+            if audio_sample_rate is None:
+                raise ValueError("audio_sample_rate must be provided when audio is given")
+            audio_stream = container.add_stream("aac", rate=audio_sample_rate)
+            audio_stream.layout = "stereo"
+            audio_stream.time_base = Fraction(1, audio_sample_rate)
+
+        # Write video frames
+        for frame_array in video_np:
+            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
+            for packet in video_stream.encode(frame):
+                container.mux(packet)
+        for packet in video_stream.encode():
+            container.mux(packet)
+
+        # Write audio if provided
+        if audio is not None:
+            _write_audio(container, audio_stream, audio, audio_sample_rate)
+
+
+def _prepare_video_array(
+    video_tensor: torch.Tensor,
+    video_format: VideoFormat | None = None,
+) -> np.ndarray:
+    """Convert video tensor to [F, H, W, C] uint8 numpy array.
+    If ``video_format`` is provided, it is trusted. Otherwise, the layout is auto-detected
+    using a heuristic that only fires when ``shape[0] == 3 and shape[1] > 3`` (CFHW). The
+    ambiguous ``[C=3, F=3, H, W]`` / ``[F=3, C=3, H, W]`` case cannot be disambiguated and
+    defaults to the FCHW interpretation — callers must pass ``video_format`` explicitly for
+    3-frame CFHW tensors.
+    """
+    if video_format == "CFHW":
+        video_tensor = video_tensor.permute(1, 0, 2, 3)  # [C, F, H, W] -> [F, C, H, W]
+    elif video_format is None and video_tensor.shape[0] == 3 and video_tensor.shape[1] > 3:
+        video_tensor = video_tensor.permute(1, 0, 2, 3)
+
+    # Normalize to [0, 255] uint8
+    if video_tensor.max() <= 1.0:
+        video_tensor = video_tensor * 255
+
+    # [F, C, H, W] -> [F, H, W, C]
+    return video_tensor.permute(0, 2, 3, 1).to(torch.uint8).cpu().numpy()
+
+
+def _write_audio(
+    container: av.container.Container,
+    audio_stream: av.audio.AudioStream,
+    audio: torch.Tensor,
+    sample_rate: int,
+) -> None:
+    """Write audio tensor to container as stereo AAC."""
+    audio = audio.cpu().float()
+
+    # Normalize to [samples, 2] stereo format
+    if audio.ndim == 1:
+        audio = audio.unsqueeze(1).repeat(1, 2)  # Mono -> stereo
+    elif audio.shape[0] == 2 and audio.shape[1] != 2:
+        audio = audio.T  # [2, samples] -> [samples, 2]
+    if audio.shape[1] == 1:
+        audio = audio.repeat(1, 2)  # Mono -> stereo
+
+    # Convert to int16 interleaved: [samples, 2] -> [1, samples*2]
+    audio_int16 = (audio.clamp(-1, 1) * 32767).to(torch.int16)
+    audio_interleaved = audio_int16.contiguous().view(1, -1).numpy()
+
+    # Create audio frame
+    frame = av.AudioFrame.from_ndarray(audio_interleaved, format="s16", layout="stereo")
+    frame.sample_rate = sample_rate
+
+    # Resample to encoder format and write
+    resampler = av.audio.resampler.AudioResampler(
+        format=audio_stream.codec_context.format,
+        layout=audio_stream.codec_context.layout,
+        rate=sample_rate,
+    )
+
+    pts = 0
+    for resampled_frame in resampler.resample(frame):
+        resampled_frame.pts = pts
+        pts += resampled_frame.samples
+        for packet in audio_stream.encode(resampled_frame):
+            container.mux(packet)
+
+    for packet in audio_stream.encode():
+        container.mux(packet)
diff --git a/packages/ltx-trainer/templates/model_card.md b/packages/ltx-trainer/templates/model_card.md
new file mode 100644
index 0000000000000000000000000000000000000000..accb8c448062489d2a21a55f1a743c86e14d0420
--- /dev/null
+++ b/packages/ltx-trainer/templates/model_card.md
@@ -0,0 +1,59 @@
+---
+tags:
+  - ltx-2
+  - ltx-video
+  - text-to-video
+  - audio-video
+pinned: true
+language:
+  - en
+license: other
+pipeline_tag: text-to-video
+library_name: diffusers
+---
+
+# {model_name}
+
+This is a fine-tuned version of [`{base_model}`]({base_model_link}) trained on custom data.
+
+## Model Details
+
+- **Base Model:** [`{base_model}`]({base_model_link})
+- **Training Type:** {training_type}
+- **Training Steps:** {training_steps}
+- **Learning Rate:** {learning_rate}
+- **Batch Size:** {batch_size}
+
+## Sample Outputs
+
+| | | | |
+|:---:|:---:|:---:|:---:|
+{sample_grid}
+
+## Usage
+
+This model is designed to be used with the LTX-2 (Lightricks Audio-Video) pipeline.
+
+### 🔌 Using Trained LoRAs in ComfyUI
+
+In order to use the trained LoRA in ComfyUI, follow these steps:
+
+1. Copy your trained LoRA checkpoint (`.safetensors` file) to the `models/loras` folder in your ComfyUI installation.
+2. In your ComfyUI workflow:
+    - Add the "Load LoRA" node to choose your LoRA file
+    - Connect it to the "Load Checkpoint" node to apply the LoRA to the base model
+
+You can find reference Text-to-Video (T2V) and Image-to-Video (I2V) workflows in the
+official [LTX-2 repository](https://github.com/Lightricks/LTX-2).
+
+### Example Prompts
+
+{validation_prompts}
+
+
+This model inherits the license of the base model ([`{base_model}`]({base_model_link})).
+
+## Acknowledgments
+
+- Base model: [Lightricks](https://huggingface.co/Lightricks/LTX-2)
+- Trainer: [LTX-2](https://github.com/Lightricks/LTX-2)