variance scaling

Author: yxlllc

configs/acoustic.yaml

@@ -64,6 +64,7 @@ timesteps: 1000
 max_beta: 0.02
 enc_ffn_kernel_size: 3
 use_rope: true
+use_variance_scaling: true
 rel_pos: true
 sampling_algorithm: euler
 sampling_steps: 20

configs/templates/config_acoustic.yaml

@@ -71,6 +71,7 @@ augmentation_args:
 diffusion_type: reflow
 enc_ffn_kernel_size: 3
 use_rope: true
+use_variance_scaling: true
 use_shallow_diffusion: true
 T_start: 0.4
 T_start_infer: 0.4

configs/templates/config_variance.yaml

@@ -65,6 +65,7 @@ tension_logit_max: 10.0
 
 enc_ffn_kernel_size: 3
 use_rope: true
+use_variance_scaling: true
 hidden_size: 256
 dur_prediction_args:
   arch: resnet
@@ -78,7 +79,7 @@ dur_prediction_args:
   lambda_wdur_loss: 1.0
   lambda_sdur_loss: 3.0
 
-use_melody_encoder: false
+use_melody_encoder: true
 melody_encoder_args:
   hidden_size: 128
   enc_layers: 4

configs/variance.yaml

@@ -36,6 +36,7 @@ predict_tension: false
 
 enc_ffn_kernel_size: 3
 use_rope: true
+use_variance_scaling: true
 rel_pos: true
 hidden_size: 256
 
@@ -51,7 +52,7 @@ dur_prediction_args:
   lambda_wdur_loss: 1.0
   lambda_sdur_loss: 3.0
 
-use_melody_encoder: false
+use_melody_encoder: true
 melody_encoder_args:
   hidden_size: 128
   enc_layers: 4

modules/fastspeech/acoustic_encoder.py

@@ -49,6 +49,26 @@ def __init__(self, vocab_size):
                 for v_name in self.variance_embed_list
             })
 
+        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
+        if self.use_variance_scaling:
+            self.variance_scaling_factor = {
+                'energy': 1. / 96,
+                'breathiness': 1. / 96,
+                'voicing': 1. / 96,
+                'tension': 0.1,
+                'key_shift': 1. / 12,
+                'speed': 1.
+            }
+        else:
+            self.variance_scaling_factor = {
+                'energy': 1.,
+                'breathiness': 1.,
+                'voicing': 1.,
+                'tension': 1.,
+                'key_shift': 1.,
+                'speed': 1.
+            }
+
         self.use_key_shift_embed = hparams.get('use_key_shift_embed', False)
         if self.use_key_shift_embed:
             self.key_shift_embed = Linear(1, hparams['hidden_size'])
@@ -64,17 +84,20 @@ def __init__(self, vocab_size):
     def forward_variance_embedding(self, condition, key_shift=None, speed=None, **variances):
         if self.use_variance_embeds:
             variance_embeds = torch.stack([
-                self.variance_embeds[v_name](variances[v_name][:, :, None])
+                self.variance_embeds[v_name](variances[v_name][:, :, None]) 
+                * self.variance_scaling_factor[v_name]
                 for v_name in self.variance_embed_list
             ], dim=-1).sum(-1)
             condition += variance_embeds
 
         if self.use_key_shift_embed:
             key_shift_embed = self.key_shift_embed(key_shift[:, :, None])
+            key_shift_embed *= self.variance_scaling_factor['key_shift']
             condition += key_shift_embed
 
         if self.use_speed_embed:
             speed_embed = self.speed_embed(speed[:, :, None])
+            speed_embed *= self.variance_scaling_factor['speed']
             condition += speed_embed
 
         return condition
@@ -87,7 +110,10 @@ def forward(
     ):
         txt_embed = self.txt_embed(txt_tokens)
         dur = mel2ph_to_dur(mel2ph, txt_tokens.shape[1]).float()
-        dur_embed = self.dur_embed(dur[:, :, None])
+        if self.use_variance_scaling:
+            dur_embed = self.dur_embed(torch.log(1 + dur[:, :, None]))
+        else:
+            dur_embed = self.dur_embed(dur[:, :, None])
         if self.use_lang_id:
             lang_embed = self.lang_embed(languages)
             extra_embed = dur_embed + lang_embed

modules/fastspeech/variance_encoder.py

@@ -17,7 +17,7 @@ def __init__(self, vocab_size):
         self.predict_dur = hparams['predict_dur']
         self.linguistic_mode = 'word' if hparams['predict_dur'] else 'phoneme'
         self.use_lang_id = hparams['use_lang_id']
-
+        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
         self.txt_embed = Embedding(vocab_size, hparams['hidden_size'], PAD_INDEX)
         if self.use_lang_id:
             self.lang_embed = Embedding(hparams['num_lang'] + 1, hparams['hidden_size'], padding_idx=0)
@@ -80,9 +80,11 @@ def forward(
             word_dur = torch.gather(F.pad(word_dur, [1, 0], value=0), 1, ph2word)  # [B, T_w] => [B, T_ph]
             word_dur_embed = self.word_dur_embed(word_dur.float()[:, :, None])
             extra_embed = onset_embed + word_dur_embed
+        elif self.use_variance_scaling:
+            extra_embed = self.ph_dur_embed(torch.log(1 + ph_dur.float())[:, :, None])
         else:
-            ph_dur_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
-            extra_embed = ph_dur_embed
+            extra_embed = self.ph_dur_embed(ph_dur.float()[:, :, None])
+            
         if self.use_lang_id:
             lang_embed = self.lang_embed(languages)
             extra_embed += lang_embed
@@ -109,6 +111,7 @@ def get_hparam(key):
 
         # MIDI inputs
         hidden_size = get_hparam('hidden_size')
+        self.use_variance_scaling = hparams.get('use_variance_scaling', False)
         self.note_midi_embed = Linear(1, hidden_size)
         self.note_dur_embed = Linear(1, hidden_size)
 
@@ -136,8 +139,13 @@ def forward(self, note_midi, note_rest, note_dur, glide=None):
         :param glide: int64 [B, T_n]
         :return: [B, T_n, H]
         """
-        midi_embed = self.note_midi_embed(note_midi[:, :, None]) * ~note_rest[:, :, None]
-        dur_embed = self.note_dur_embed(note_dur.float()[:, :, None])
+        if self.use_variance_scaling:
+            midi_embed = self.note_midi_embed(note_midi[:, :, None] / 128)
+            dur_embed = self.note_dur_embed(torch.log(1 + note_dur.float())[:, :, None])
+        else:
+            midi_embed = self.note_midi_embed(note_midi[:, :, None])
+            dur_embed = self.note_dur_embed(note_dur.float()[:, :, None])
+        midi_embed *= ~note_rest[:, :, None]
         ornament_embed = 0
         if self.use_glide_embed:
             ornament_embed += self.note_glide_embed(glide) * self.glide_embed_scale