# Codec Training Configuration
# ==============================================================================
# This config is for training the RVQ-VAE codec vocoder
# Based on config.yml but with codec-specific settings
# ==============================================================================

log_dir: "codec_training"
save_freq: 2
log_interval: 10
device: "cuda"

# Training epochs
epochs_codec: 100

batch_size: 5  # Reduced due to additional codec memory
max_len: 8000

# Pretrained models (optional - for warm starting)
pretrained_model: "codec_training/epoch_codec_00018.pth"  # Base model for text/style encoders
pretrained_codec: "codec_training/epoch_codec_00018.pth"  # Set to load a pretrained codec

# Quantizer warmup (gradually increase number of quantizers)
quantizer_warmup_steps: 10000
min_quantizers: 1

# Data
data_params:
  train_data: "Data/train_3_filtered_dur_fixed_speakers_test.txt"
  val_data: "Data/val_3_filtered_fixed_speakers.txt"
  root_path: "../tts_model"
  OOD_data: "Data/OOD_texts.txt"
  min_length: 50

preprocess_params:
  sr: 44100
  spect_params:
    n_fft: 4096
    win_length: 2400
    hop_length: 600

# Model parameters
model_params:
  multispeaker: false

  dim_in: 64 
  hidden_dim: 512
  max_conv_dim: 512
  n_layer: 3
  n_mels: 128

  n_token: 178
  max_dur: 75
  style_dim: 64
  style_dim_ac: 128
  dropout: 0.1

  # === CODEC SETTINGS ===
  # Latent dimension for RVQ
  codec_latent_dim: 512
  
  # Number of RVQ codebooks (more = higher quality, more tokens)
  num_quantizers: 8
  
  # Codebook size (vocabulary size per quantizer)
  codebook_size: 4096
  
  # Compression strides (total compression = product)
  # [2, 2, 2] = 8x compression
  codec_strides: [1]

  # Decoder/Vocoder config
  decoder: 
    type: 'ringformer'
    resblock_kernel_sizes: [3, 7, 11]
    upsample_rates: [12, 10]
    upsample_initial_channel: 512
    resblock_dilation_sizes: [[1, 3, 5], [1, 3, 5], [1, 3, 5]]
    upsample_kernel_sizes: [24, 20]
    gen_istft_n_fft: 30
    gen_istft_hop_size: 5
      
  # SLM for perceptual loss
  slm:
    model: 'microsoft/wavlm-base-plus'
    sr: 16000
    hidden: 768
    nlayers: 13
    initial_channel: 64
  
  # Style diffusion (not used in codec training but needed for model building)
  diffusion:
    embedding_mask_proba: 0.2
    transformer:
      num_layers: 4
      num_heads: 4
      head_features: 32
      multiplier: 4
    dist:
      sigma_data: 0.5
      mean_data: 0.1
      acoustic: 0.46
      prosodic: 0.46
      pitch: 0.46
      acoustic_mean: 0.4
      prosodic_mean: 0.4
      pitch_mean: 0.4
      estimate_sigma_data: true
      mean: -0.0
      std: 0.3

# Optimizer
optimizer_params:
  lr: 1.0e-4
  pct_start: 0.0

# Loss weights
loss_params:
  # Waveform losses
  lambda_mel: 10.0        # Multi-resolution STFT
  lambda_gen: 1.0         # GAN generator loss
  lambda_slm: 1.0         # Speech language model loss
  lambda_mag: 0.5         # Magnitude/phase loss
  
  # Codec-specific losses
  lambda_codec: 1.0       # Overall codec loss weight
  lambda_pitch: 1.0       # Pitch reconstruction
  lambda_energy: 1.0      # Energy reconstruction  
  lambda_text: 2.0        # Text embedding reconstruction
  lambda_commitment: 0.25 # VQ commitment loss
  lambda_diversity: 0.02  # Codebook diversity regularization
  
  # GAN training starts after this epoch
  TMA_epoch: 15