Speaker Embeddings integration

model.py

@@ -166,16 +166,28 @@ class Encoder(nn.Module):
             convolutions.append(conv_layer)
         self.convolutions = nn.ModuleList(convolutions)
 
-        self.lstm = nn.LSTM(hparams.encoder_embedding_dim,
+        self.lstm = nn.LSTM(hparams.encoder_embedding_dim + hparams.speaker_embedding_dim,
                             int(hparams.encoder_embedding_dim / 2), 1,
                             batch_first=True, bidirectional=True)
 
-    def forward(self, x, input_lengths):
+    def forward(self, x, input_lengths, speaker_embedding):
         for conv in self.convolutions:
             x = F.dropout(F.relu(conv(x)), 0.5, self.training)
 
         x = x.transpose(1, 2)
 
+        # this concatenation part is largely from https://github.com/CorentinJ/Real-Time-Voice-Cloning
+        batch_size = x.size()[0]
+        num_chars = x.size()[1]
+        idx = 0 if speaker_embedding.dim() == 1 else 1
+        speaker_embedding_size = speaker_embedding.size()[idx] 
+        e = speaker_embedding.repeat_interleave(num_chars, dim=idx)
+        # Reshape & transpose
+        e = e.reshape(batch_size, speaker_embedding_size, num_chars)
+        e = e.transpose(1, 2)
+        # Concatenate the tiled speaker embedding with the encoder output
+        x = torch.cat((x, e), 2)
+
         # pytorch tensor are not reversible, hence the conversion
         input_lengths = input_lengths.cpu().numpy()
         x = nn.utils.rnn.pack_padded_sequence(
@@ -186,15 +198,24 @@ class Encoder(nn.Module):
 
         outputs, _ = nn.utils.rnn.pad_packed_sequence(
             outputs, batch_first=True)
-
         return outputs
 
-    def inference(self, x):
+    def inference(self, x, speaker_embedding=None):
         for conv in self.convolutions:
             x = F.dropout(F.relu(conv(x)), 0.5, self.training)
 
         x = x.transpose(1, 2)
 
+        if speaker_embedding is not None:
+            batch_size = x.size()[0]
+            num_chars = x.size()[1]
+            idx = 0 if speaker_embedding.dim() == 1 else 1
+            speaker_embedding_size = speaker_embedding.size()[idx] 
+            e = speaker_embedding.repeat_interleave(num_chars, dim=idx)
+            e = e.reshape(batch_size, speaker_embedding_size, num_chars)
+            e = e.transpose(1, 2)
+            x = torch.cat((x, e), 2)
+
         self.lstm.flatten_parameters()
         outputs, _ = self.lstm(x)
 
@@ -392,7 +413,6 @@ class Decoder(nn.Module):
         gate_outputs: gate outputs from the decoder
         alignments: sequence of attention weights from the decoder
         """
-
         decoder_input = self.get_go_frame(memory).unsqueeze(0)
         decoder_inputs = self.parse_decoder_inputs(decoder_inputs)
         decoder_inputs = torch.cat((decoder_input, decoder_inputs), dim=0)
@@ -472,15 +492,17 @@ class Tacotron2(nn.Module):
 
     def parse_batch(self, batch):
         text_padded, input_lengths, mel_padded, gate_padded, \
-            output_lengths = batch
+            output_lengths, mel_speaker = batch
         text_padded = to_gpu(text_padded).long()
         input_lengths = to_gpu(input_lengths).long()
         max_len = torch.max(input_lengths.data).item()
         mel_padded = to_gpu(mel_padded).float()
         gate_padded = to_gpu(gate_padded).float()
         output_lengths = to_gpu(output_lengths).long()
+        mel_speaker = to_gpu(mel_speaker).float()
 
         return (
+            mel_speaker,
             (text_padded, input_lengths, mel_padded, max_len, output_lengths),
             (mel_padded, gate_padded))
 
@@ -496,27 +518,28 @@ class Tacotron2(nn.Module):
 
         return outputs
 
-    def forward(self, inputs):
+    def forward(self, inputs, speaker_embedding):
         text_inputs, text_lengths, mels, max_len, output_lengths = inputs
         text_lengths, output_lengths = text_lengths.data, output_lengths.data
 
         embedded_inputs = self.embedding(text_inputs).transpose(1, 2)
 
-        encoder_outputs = self.encoder(embedded_inputs, text_lengths)
+        encoder_outputs = self.encoder(embedded_inputs, text_lengths, speaker_embedding)
 
         mel_outputs, gate_outputs, alignments = self.decoder(
             encoder_outputs, mels, memory_lengths=text_lengths)
 
         mel_outputs_postnet = self.postnet(mel_outputs)
+
         mel_outputs_postnet = mel_outputs + mel_outputs_postnet
 
         return self.parse_output(
             [mel_outputs, mel_outputs_postnet, gate_outputs, alignments],
             output_lengths)
 
-    def inference(self, inputs):
+    def inference(self, inputs, speaker_embedding):
         embedded_inputs = self.embedding(inputs).transpose(1, 2)
-        encoder_outputs = self.encoder.inference(embedded_inputs)
+        encoder_outputs = self.encoder.inference(embedded_inputs,speaker_embedding)
         mel_outputs, gate_outputs, alignments = self.decoder.inference(
             encoder_outputs)
 
@@ -526,4 +549,4 @@ class Tacotron2(nn.Module):
         outputs = self.parse_output(
             [mel_outputs, mel_outputs_postnet, gate_outputs, alignments])
 
-        return outputs
+        return outputs