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