File size: 3,384 Bytes
fb0011a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
# 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