# =============================================================================
# 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"