Medix-AI/coqui-tts
Medix-AI
/
coqui-tts
mirror of https://github.com/coqui-ai/TTS.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add Speaker Conditional GST support

This commit is contained in:
Edresson 2020-09-29 16:09:27 -03:00
parent cf02ace5b7
commit 99d5a0ac07
6 changed files with 62 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
TTS/bin/train_tts.py
View File

@ -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

23
TTS/tts/layers/gst_layers.py
View File

@ -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:
@ -166,4 +173,4 @@ class MultiHeadAttention(nn.Module):
torch.split(out, 1, dim=0),
dim=3).squeeze(0) # [N, T_q, num_units]
return out
return out

16
TTS/tts/models/tacotron.py
View File

@ -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

23
TTS/tts/models/tacotron2.py
View File

@ -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:

11
TTS/tts/models/tacotron_abstract.py
View File

@ -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

3
TTS/tts/utils/generic_utils.py
View File

@ -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)