add Speaker Conditional GST support

TTS/bin/train_tts.py

@@ -10,6 +10,8 @@ import traceback
 
 import numpy as np
 import torch
+
+from random import randrange
 from torch.utils.data import DataLoader
 from TTS.tts.datasets.preprocess import load_meta_data
 from TTS.tts.datasets.TTSDataset import MyDataset
@@ -39,7 +41,6 @@ from TTS.utils.training import (NoamLR, adam_weight_decay, check_update,
 
 use_cuda, num_gpus = setup_torch_training_env(True, False)
 
-
 def setup_loader(ap, r, is_val=False, verbose=False, speaker_mapping=None):
     if is_val and not c.run_eval:
         loader = None
@@ -432,6 +433,14 @@ def evaluate(model, criterion, ap, global_step, epoch, speaker_mapping=None):
         test_figures = {}
         print(" | > Synthesizing test sentences")
         speaker_id = 0 if c.use_speaker_embedding else None
+        speaker_embedding = speaker_mapping[list(speaker_mapping.keys())[randrange(len(speaker_mapping)-1)]]['embedding'] if c.use_external_speaker_embedding_file and c.use_speaker_embedding else None
+        style_wav = c.get("gst_style_input")
+        if style_wav is None and c.use_gst:
+            # inicialize GST with zero dict.
+            style_wav = {}
+            print("WARNING: You don't provided a gst style wav, for this reason we use a zero tensor!")
+            for i in range(c.gst['gst_style_tokens']):
+                style_wav[str(i)] = 0
         style_wav = c.get("gst_style_input")
         for idx, test_sentence in enumerate(test_sentences):
             try:
@@ -442,6 +451,7 @@ def evaluate(model, criterion, ap, global_step, epoch, speaker_mapping=None):
                     use_cuda,
                     ap,
                     speaker_id=speaker_id,
+                    speaker_embedding=speaker_embedding,
                     style_wav=style_wav,
                     truncated=False,
                     enable_eos_bos_chars=c.enable_eos_bos_chars, #pylint: disable=unused-argument

TTS/tts/layers/gst_layers.py

@@ -8,14 +8,17 @@ class GST(nn.Module):
 
     See https://arxiv.org/pdf/1803.09017"""
 
-    def __init__(self, num_mel, num_heads, num_style_tokens, embedding_dim):
+    def __init__(self, num_mel, num_heads, num_style_tokens, gst_embedding_dim, speaker_embedding_dim=None):
         super().__init__()
-        self.encoder = ReferenceEncoder(num_mel, embedding_dim)
+        self.encoder = ReferenceEncoder(num_mel, gst_embedding_dim)
         self.style_token_layer = StyleTokenLayer(num_heads, num_style_tokens,
-                                                 embedding_dim)
+                                                 gst_embedding_dim, speaker_embedding_dim)
 
-    def forward(self, inputs):
+    def forward(self, inputs, speaker_embedding=None):
         enc_out = self.encoder(inputs)
+        # concat speaker_embedding
+        if speaker_embedding is not None:
+            enc_out = torch.cat([enc_out, speaker_embedding], dim=-1)
         style_embed = self.style_token_layer(enc_out)
 
         return style_embed
@@ -72,7 +75,7 @@ class ReferenceEncoder(nn.Module):
         # x: 3D tensor [batch_size, post_conv_width,
         #               num_channels*post_conv_height]
         self.recurrence.flatten_parameters()
-        _, out = self.recurrence(x)
+        memory, out = self.recurrence(x)
         # out: 3D tensor [seq_len==1, batch_size, encoding_size=128]
 
         return out.squeeze(0)
@@ -90,9 +93,14 @@ class StyleTokenLayer(nn.Module):
     """NN Module attending to style tokens based on prosody encodings."""
 
     def __init__(self, num_heads, num_style_tokens,
-                 embedding_dim):
+                 embedding_dim, speaker_embedding_dim=None):
         super().__init__()
+
         self.query_dim = embedding_dim // 2
+
+        if speaker_embedding_dim:
+            self.query_dim += speaker_embedding_dim
+
         self.key_dim = embedding_dim // num_heads
         self.style_tokens = nn.Parameter(
             torch.FloatTensor(num_style_tokens, self.key_dim))
@@ -115,7 +123,6 @@ class StyleTokenLayer(nn.Module):
 
         return style_embed
 
-
 class MultiHeadAttention(nn.Module):
     '''
     input:

TTS/tts/models/tacotron.py

@@ -35,7 +35,8 @@ class Tacotron(TacotronAbstract):
                  gst_embedding_dim=256,
                  gst_num_heads=4,
                  gst_style_tokens=10,
-                 memory_size=5):
+                 memory_size=5,
+                 gst_use_speaker_embedding=False):
         super(Tacotron,
               self).__init__(num_chars, num_speakers, r, postnet_output_dim,
                              decoder_output_dim, attn_type, attn_win,
@@ -45,7 +46,7 @@ class Tacotron(TacotronAbstract):
                              bidirectional_decoder, double_decoder_consistency,
                              ddc_r, encoder_in_features, decoder_in_features,
                              speaker_embedding_dim, gst, gst_embedding_dim,
-                             gst_num_heads, gst_style_tokens)
+                             gst_num_heads, gst_style_tokens, gst_use_speaker_embedding)
 
         # speaker embedding layers
         if self.num_speakers > 1:
@@ -78,7 +79,8 @@ class Tacotron(TacotronAbstract):
             self.gst_layer = GST(num_mel=80,
                                  num_heads=gst_num_heads,
                                  num_style_tokens=gst_style_tokens,
-                                 embedding_dim=gst_embedding_dim)
+                                 gst_embedding_dim=self.gst_embedding_dim,
+                                 speaker_embedding_dim=speaker_embedding_dim if self.embeddings_per_sample and self.gst_use_speaker_embedding else None)
         # backward pass decoder
         if self.bidirectional_decoder:
             self._init_backward_decoder()
@@ -108,7 +110,9 @@ class Tacotron(TacotronAbstract):
         # global style token
         if self.gst:
             # B x gst_dim
-            encoder_outputs = self.compute_gst(encoder_outputs, mel_specs)
+            encoder_outputs = self.compute_gst(encoder_outputs,
+                                               mel_specs,
+                                               speaker_embeddings if self.gst_use_speaker_embedding else None)
         # speaker embedding
         if self.num_speakers > 1:
             if not self.embeddings_per_sample:
@@ -149,7 +153,9 @@ class Tacotron(TacotronAbstract):
         encoder_outputs = self.encoder(inputs)
         if self.gst:
             # B x gst_dim
-            encoder_outputs = self.compute_gst(encoder_outputs, style_mel)
+            encoder_outputs = self.compute_gst(encoder_outputs,
+                                               style_mel,
+                                               speaker_embeddings if self.gst_use_speaker_embedding else None)
         if self.num_speakers > 1:
             if not self.embeddings_per_sample:
                 # B x 1 x speaker_embed_dim

TTS/tts/models/tacotron2.py

@@ -33,7 +33,8 @@ class Tacotron2(TacotronAbstract):
                  gst=False,
                  gst_embedding_dim=512,
                  gst_num_heads=4,
-                 gst_style_tokens=10):
+                 gst_style_tokens=10,
+                 gst_use_speaker_embedding=False):
         super(Tacotron2,
               self).__init__(num_chars, num_speakers, r, postnet_output_dim,
                              decoder_output_dim, attn_type, attn_win,
@@ -43,7 +44,7 @@ class Tacotron2(TacotronAbstract):
                              bidirectional_decoder, double_decoder_consistency,
                              ddc_r, encoder_in_features, decoder_in_features,
                              speaker_embedding_dim, gst, gst_embedding_dim,
-                             gst_num_heads, gst_style_tokens)
+                             gst_num_heads, gst_style_tokens, gst_use_speaker_embedding)
 
         # speaker embedding layer
         if self.num_speakers > 1:
@@ -72,7 +73,8 @@ class Tacotron2(TacotronAbstract):
             self.gst_layer = GST(num_mel=80,
                                  num_heads=self.gst_num_heads,
                                  num_style_tokens=self.gst_style_tokens,
-                                 embedding_dim=self.gst_embedding_dim)
+                                 gst_embedding_dim=self.gst_embedding_dim,
+                                 speaker_embedding_dim=speaker_embedding_dim if self.embeddings_per_sample and self.gst_use_speaker_embedding else None)
         # backward pass decoder
         if self.bidirectional_decoder:
             self._init_backward_decoder()
@@ -98,11 +100,11 @@ class Tacotron2(TacotronAbstract):
         embedded_inputs = self.embedding(text).transpose(1, 2)
         # B x T_in_max x D_en
         encoder_outputs = self.encoder(embedded_inputs, text_lengths)
-
         if self.gst:
             # B x gst_dim
-            encoder_outputs = self.compute_gst(encoder_outputs, mel_specs)
-
+            encoder_outputs = self.compute_gst(encoder_outputs,
+                                               mel_specs,
+                                               speaker_embeddings if self.gst_use_speaker_embedding else None)
         if self.num_speakers > 1:
             if not self.embeddings_per_sample:
                 # B x 1 x speaker_embed_dim
@@ -144,8 +146,9 @@ class Tacotron2(TacotronAbstract):
 
         if self.gst:
             # B x gst_dim
-            encoder_outputs = self.compute_gst(encoder_outputs, style_mel)
-
+            encoder_outputs = self.compute_gst(encoder_outputs,
+                                               style_mel,
+                                               speaker_embeddings if self.gst_use_speaker_embedding else None)
         if self.num_speakers > 1:
             if not self.embeddings_per_sample:
                 speaker_embeddings = self.speaker_embedding(speaker_ids)[:, None]
@@ -168,7 +171,9 @@ class Tacotron2(TacotronAbstract):
 
         if self.gst:
             # B x gst_dim
-            encoder_outputs = self.compute_gst(encoder_outputs, style_mel)
+            encoder_outputs = self.compute_gst(encoder_outputs,
+                                               style_mel,
+                                               speaker_embeddings if self.gst_use_speaker_embedding else None)
 
         if self.num_speakers > 1:
             if not self.embeddings_per_sample:

TTS/tts/models/tacotron_abstract.py

@@ -34,7 +34,8 @@ class TacotronAbstract(ABC, nn.Module):
                  gst=False,
                  gst_embedding_dim=512,
                  gst_num_heads=4,
-                 gst_style_tokens=10):
+                 gst_style_tokens=10,
+                 gst_use_speaker_embedding=False):
         """ Abstract Tacotron class """
         super().__init__()
         self.num_chars = num_chars
@@ -45,6 +46,7 @@ class TacotronAbstract(ABC, nn.Module):
         self.gst_embedding_dim = gst_embedding_dim
         self.gst_num_heads = gst_num_heads
         self.gst_style_tokens = gst_style_tokens
+        self.gst_use_speaker_embedding = gst_use_speaker_embedding
         self.num_speakers = num_speakers
         self.bidirectional_decoder = bidirectional_decoder
         self.double_decoder_consistency = double_decoder_consistency
@@ -179,11 +181,14 @@ class TacotronAbstract(ABC, nn.Module):
             self.speaker_embeddings_projected = self.speaker_project_mel(
                 self.speaker_embeddings).squeeze(1)
 
-    def compute_gst(self, inputs, style_input):
+    def compute_gst(self, inputs, style_input, speaker_embedding=None):
         """ Compute global style token """
         device = inputs.device
         if isinstance(style_input, dict):
             query = torch.zeros(1, 1, self.gst_embedding_dim//2).to(device)
+            if speaker_embedding is not None:
+                query = torch.cat([query, speaker_embedding.reshape(1, 1, -1)], dim=-1)
+
             _GST = torch.tanh(self.gst_layer.style_token_layer.style_tokens)
             gst_outputs = torch.zeros(1, 1, self.gst_embedding_dim).to(device)
             for k_token, v_amplifier in style_input.items():
@@ -193,7 +198,7 @@ class TacotronAbstract(ABC, nn.Module):
         elif style_input is None:
             gst_outputs = torch.zeros(1, 1, self.gst_embedding_dim).to(device)
         else:
-            gst_outputs = self.gst_layer(style_input) # pylint: disable=not-callable
+            gst_outputs = self.gst_layer(style_input, speaker_embedding) # pylint: disable=not-callable
         inputs = self._concat_speaker_embedding(inputs, gst_outputs)
         return inputs
 

TTS/tts/utils/generic_utils.py

@@ -59,6 +59,7 @@ def setup_model(num_chars, num_speakers, c, speaker_embedding_dim=None):
                         gst_embedding_dim=c.gst['gst_embedding_dim'],
                         gst_num_heads=c.gst['gst_num_heads'],
                         gst_style_tokens=c.gst['gst_style_tokens'],
+                        gst_use_speaker_embedding=c.gst['gst_use_speaker_embedding'],
                         memory_size=c.memory_size,
                         attn_type=c.attention_type,
                         attn_win=c.windowing,
@@ -85,6 +86,7 @@ def setup_model(num_chars, num_speakers, c, speaker_embedding_dim=None):
                         gst_embedding_dim=c.gst['gst_embedding_dim'],
                         gst_num_heads=c.gst['gst_num_heads'],
                         gst_style_tokens=c.gst['gst_style_tokens'],
+                        gst_use_speaker_embedding=c.gst['gst_use_speaker_embedding'],
                         attn_type=c.attention_type,
                         attn_win=c.windowing,
                         attn_norm=c.attention_norm,
@@ -244,6 +246,7 @@ def check_config_tts(c):
     check_argument('gst', c, restricted=True, val_type=dict)
     check_argument('gst_style_input', c['gst'], restricted=True, val_type=[str, dict])
     check_argument('gst_embedding_dim', c['gst'], restricted=True, val_type=int, min_val=0, max_val=1000)
+    check_argument('gst_use_speaker_embedding', c['gst'], restricted=True, val_type=bool)
     check_argument('gst_num_heads', c['gst'], restricted=True, val_type=int, min_val=2, max_val=10)
     check_argument('gst_style_tokens', c['gst'], restricted=True, val_type=int, min_val=1, max_val=1000)