# config_codec_v2.yml
# ==============================================================================
# Configuration for TTS Codec V2 Training
# 
# Key features:
#   - Style is tokenized from noised mel spectrogram
#   - Style tokens are very coarse temporal (~2-8 Hz)
#   - Content tokens from prosody + text (separate from style)
# ==============================================================================

log_dir: "codec_v2_training"
save_freq: 2
log_interval: 10
device: "cuda"
epochs_codec: 100
batch_size: 16
max_len: 8000

# Pretrained models
pretrained_model: "again_wavlm/epoch_flow_00108.pth"  # For text encoder
pretrained_codec: ""  # Path to resume codec training

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"
  min_length: 50

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

model_params:
  multispeaker: false

  # Base model dimensions
  dim_in: 64 
  hidden_dim: 512
  max_conv_dim: 512
  n_layer: 3
  n_mels: 80  # 80 mel bands for style encoder input

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

  # ==========================================================================
  # CODEC V2 SPECIFIC PARAMETERS
  # ==========================================================================
  
  # Style encoder settings
  # style_compression=16 at hop_length=600 (44.1kHz) = ~4.6 Hz style tokens
  # style_compression=32 = ~2.3 Hz style tokens (more coarse)
  style_compression: 16
  style_dim_codec: 64
  style_fsq_levels: [8, 8, 8, 8]  # 4096 style codes
  
  # Content encoder settings  
  # content_strides=[2] = 2x compression
  content_strides: [2]
  content_fsq_levels: [8, 8, 8, 8, 8, 8]  # 262144 content codes
  
  # Codec latent dimensions
  codec_latent_dim: 256
  codec_hidden_dim: 512

  # Decoder config (same as ringformer)
  decoder: 
    type: 'ringformer'
    resblock_kernel_sizes: [3, 7, 11]
    upsample_rates: [4, 6, 10]  # Will be auto-computed based on compression
    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 WavLM loss
  slm:
    model: 'microsoft/wavlm-base-plus'
    sr: 16000
    hidden: 768
    nlayers: 13
    initial_channel: 64
  
  # Diffusion config (not used in codec 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_params:
  lr: 1.0e-4
  pct_start: 0.0
  
loss_params:
  lambda_mel: 45.0           # Mel reconstruction loss
  lambda_gen: 1.0            # Generator adversarial loss
  lambda_slm: 1.0            # SLM (WavLM) feature matching loss
  lambda_mag: 1.0            # Magnitude/phase loss
  lambda_style_entropy: 1.0  # Style FSQ entropy regularization
  lambda_content_entropy: 1.0 # Content FSQ entropy regularization
  TMA_epoch: 5               # Epoch to start adversarial training
  
  # Style noise schedule (for robustness)
  style_noise_start: 0.5     # Initial noise level for style encoder
  style_noise_end: 0.1       # Final noise level (decays over epochs)