# ==============================================================================
# LLM-TTS-Factory Configuration
# ==============================================================================

# ------------------------------------------------------------------------------
# Global Settings
# ------------------------------------------------------------------------------
global:
  seed: 1337
  device: "cuda:0"
  num_workers: 4
  use_wandb: true
  wandb_project: "soprano-tts"
  tokenizer_name: "ekwek/Soprano-80M"

# ------------------------------------------------------------------------------
# File Paths
# Use forward slashes (/) for both Windows and Linux, or standard OS paths.
# Relative paths are evaluated from the directory where the script is run.
#
# Examples:
#   Linux:   "/home/ubuntu/data/lj_speech/LJSpeech-1.1"
#   Windows: "C:/Users/Name/Documents/datasets/LJSpeech-1.1"
#   Relative: "./data/LJSpeech-1.1"
# ------------------------------------------------------------------------------
paths:
  # Path to LJSpeech-1.1 directory
  dataset_root: "./data/LJSpeech-1.1"
  
  # Base directory to save all checkpoints, logs, and generated datasets
  save_dir: "./checkpoints"
  
  # Pretrained model paths (set to null if training from scratch or not applicable)
  pretrained_codec_path: null
  pretrained_llm_path: null
  pretrained_decoder_path: null
  pretrained_discriminator_path: null

# ------------------------------------------------------------------------------
# Data Generation Configuration (generate_dataset*.py)
# ------------------------------------------------------------------------------
data_generation:
  val_prop: 0.1
  val_max: 512

# ------------------------------------------------------------------------------
# Codec Training Configuration (codec_train.py)
# ------------------------------------------------------------------------------
codec:
  sample_rate: 32000
  batch_size: 16
  num_epochs: 100
  learning_rate: 1.0e-4
  freeze_encoder: false
  # Encoder/Decoder architecture params
  encoder_dim: 768
  encoder_num_layers: 8
  bottleneck_channels: 5
  decoder_num_layers: 8

# ------------------------------------------------------------------------------
# LLM Training Configuration (train_llm.py)
# ------------------------------------------------------------------------------
llm:
  from_scratch: false
  batch_size: 64
  max_steps: 150000
  max_lr: 2.0e-5
  min_lr_ratio: 0.3
  warmup_ratio: 0.3
  cooldown_ratio: 0.1
  grad_accum_steps: 1
  seq_len: 1024
  val_freq: 250
  save_freq: 5000
  text_factor: 0.5
  betas: [0.9, 0.95]
  weight_decay: 0.1

# ------------------------------------------------------------------------------
# Decoder (Vocos) Training Configuration (train_decoder.py)
# ------------------------------------------------------------------------------
decoder:
  use_discriminator: true
  batch_size: 8  # Reduced to avoid OOM, adjust based on your GPU
  max_steps: 200000
  max_lr: 2.0e-4
  min_lr_ratio: 0.1
  warmup_ratio: 0.2
  cooldown_ratio: 0.1
  grad_accum_steps: 1
  seq_len: 1024
  segment_size_samples: 32768  # ~1 sec (16 tokens)
  val_freq: 250
  val_steps: 10
  save_freq: 3000
  text_factor: 0.0
  betas: [0.8, 0.99]
  weight_decay: 0.1
  start_step: 0
  
  # Loss Weights
  lambda_mel: 45.0
  lambda_fm: 2.0
  lambda_gen: 1.0
  lambda_stft: 1.0

# ------------------------------------------------------------------------------
# Inference Configuration (simple_inference.py)
# ------------------------------------------------------------------------------
inference:
  temperature: 0.8
  top_k: 50
  top_p: 0.95
  repetition_penalty: 1.2
  max_new_tokens: 512