Update model.py

model.py

@@ -149,7 +149,6 @@ class Postnet(nn.Module):
 class Encoder(nn.Module):
     def __init__(self, hparams):
         super(Encoder, self).__init__()
-
         convolutions = []
         for _ in range(hparams.encoder_n_convolutions):
             conv_layer = nn.Sequential(
@@ -162,28 +161,19 @@ class Encoder(nn.Module):
             convolutions.append(conv_layer)
         self.convolutions = nn.ModuleList(convolutions)
 
-        # Modify the input dimensionality for LSTM
-        lstm_input_dim = hparams.encoder_embedding_dim
-        self.lstm = nn.LSTM(lstm_input_dim,
+        self.lstm = nn.LSTM(hparams.encoder_embedding_dim,
                             int(hparams.encoder_embedding_dim / 2), 1,
                             batch_first=True, bidirectional=True)
 
     def forward(self, x, input_lengths, speaker_embedding):
+        # Modify the input x to concatenate the speaker embedding
+        x = torch.cat((x, speaker_embedding.unsqueeze(1).expand(-1, x.size(1), -1)), dim=-1)
+        
         for conv in self.convolutions:
             x = F.dropout(F.relu(conv(x)), 0.5, self.training)
 
         x = x.transpose(1, 2)
 
-        # No changes in the concatenation part
-        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)
-
         input_lengths = input_lengths.cpu().numpy()
         x = nn.utils.rnn.pack_padded_sequence(
             x, input_lengths, batch_first=True)
@@ -193,31 +183,21 @@ class Encoder(nn.Module):
 
         outputs, _ = nn.utils.rnn.pad_packed_sequence(
             outputs, batch_first=True)
+
         return outputs
 
-    def inference(self, x, speaker_embedding=None):
+    def inference(self, x):
         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)
 
         return outputs
 
 
-
 class Decoder(nn.Module):
     def __init__(self, hparams):
         super(Decoder, self).__init__()
@@ -409,6 +389,7 @@ 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)
@@ -488,17 +469,15 @@ class Tacotron2(nn.Module):
 
     def parse_batch(self, batch):
         text_padded, input_lengths, mel_padded, gate_padded, \
-            output_lengths, mel_speaker = batch
+            output_lengths = 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))
 
@@ -520,13 +499,13 @@ class Tacotron2(nn.Module):
 
         embedded_inputs = self.embedding(text_inputs).transpose(1, 2)
 
+        # Pass the speaker embedding to the Encoder
         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(
@@ -535,7 +514,9 @@ class Tacotron2(nn.Module):
 
     def inference(self, inputs, speaker_embedding):
         embedded_inputs = self.embedding(inputs).transpose(1, 2)
-        encoder_outputs = self.encoder.inference(embedded_inputs,speaker_embedding)
+        # Pass the speaker embedding to the Encoder
+        encoder_outputs = self.encoder.inference(embedded_inputs, speaker_embedding)
+
         mel_outputs, gate_outputs, alignments = self.decoder.inference(
             encoder_outputs)
 
@@ -545,4 +526,5 @@ class Tacotron2(nn.Module):
         outputs = self.parse_output(
             [mel_outputs, mel_outputs_postnet, gate_outputs, alignments])
 
-        return outputs
+        return outputs
+