Add multilingual support

TTS/bin/resample.py

@@ -11,8 +11,10 @@
 
 def resample_file(func_args):
     filename, output_sr = func_args
-    y, sr = librosa.load(filename, sr=output_sr)
-    librosa.output.write_wav(filename, y, sr)
+    y, sr = librosa.load(filename, sr=None)
+    if output_sr != sr:
+        y = librosa.resample(y, sr, output_sr)
+        librosa.output.write_wav(filename, y, output_sr)
 
 
 if __name__ == "__main__":

TTS/bin/train_tacotron.py

@@ -5,7 +5,7 @@
 import sys
 import time
 import traceback
-from random import randrange
+import random
 
 import numpy as np
 import torch
@@ -17,7 +17,7 @@
 from TTS.tts.utils.generic_utils import setup_model
 from TTS.tts.utils.io import save_best_model, save_checkpoint
 from TTS.tts.utils.measures import alignment_diagonal_score
-from TTS.tts.utils.speakers import parse_speakers
+from TTS.tts.utils.speakers import parse_speakers, parse_languages
 from TTS.tts.utils.synthesis import synthesis
 from TTS.tts.utils.text.symbols import make_symbols, phonemes, symbols
 from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
@@ -70,6 +70,31 @@ def setup_loader(ap, r, is_val=False, verbose=False, dataset=None):
             dataset.sort_items()
 
         sampler = DistributedSampler(dataset) if num_gpus > 1 else None
+        if getattr(c, "weighted_sampler", False) and sampler is None and not is_val:
+            print("Using weighted sampler")
+            # get speaker/language names
+            speaker_names = np.array([item[2] for item in dataset.items])
+            language_names = np.array([item[3] for item in dataset.items])
+
+            unique_speaker_names = np.unique(speaker_names).tolist()
+            unique_language_names = np.unique(language_names).tolist()
+
+            speaker_ids = [unique_speaker_names.index(s) for s in speaker_names]
+            language_ids = [unique_language_names.index(l) for l in language_names]
+
+            # count number samples by speaker/language
+            speaker_count = np.array([len(np.where(speaker_names == s)[0]) for s in unique_speaker_names])
+            language_count = np.array([len(np.where(language_names == l)[0]) for l in unique_language_names])
+
+            # create weight
+            weight_speaker = 1. / speaker_count
+            weight_language = 1. / language_count
+            samples_weight = np.array([weight_speaker[s] for s in speaker_ids]) + np.array([weight_language[l] for l in language_ids]) 
+            dataset_samples_weight = torch.from_numpy(samples_weight).double()
+
+            # create sampler
+            sampler = torch.utils.data.sampler.WeightedRandomSampler(dataset_samples_weight, len(dataset_samples_weight))
+
         loader = DataLoader(
             dataset,
             batch_size=c.eval_batch_size if is_val else c.batch_size,
@@ -92,12 +117,13 @@ def format_data(data):
     mel_input = data[4]
     mel_lengths = data[5]
     stop_targets = data[6]
+    language_names = data[7]
     max_text_length = torch.max(text_lengths.float())
     max_spec_length = torch.max(mel_lengths.float())
 
     if c.use_speaker_embedding:
         if c.use_external_speaker_embedding_file:
-            speaker_embeddings = data[8]
+            speaker_embeddings = data[9]
             speaker_ids = None
         else:
             speaker_ids = [speaker_mapping[speaker_name] for speaker_name in speaker_names]
@@ -107,6 +133,12 @@ def format_data(data):
         speaker_embeddings = None
         speaker_ids = None
 
+    if c.use_language_embedding:
+        language_ids = [language_mapping[language_name] for language_name in language_names]
+        language_ids = torch.LongTensor(language_ids)
+    else:
+        language_ids = None
+
     # set stop targets view, we predict a single stop token per iteration.
     stop_targets = stop_targets.view(text_input.shape[0], stop_targets.size(1) // c.r, -1)
     stop_targets = (stop_targets.sum(2) > 0.0).unsqueeze(2).float().squeeze(2)
@@ -123,6 +155,8 @@ def format_data(data):
             speaker_ids = speaker_ids.cuda(non_blocking=True)
         if speaker_embeddings is not None:
             speaker_embeddings = speaker_embeddings.cuda(non_blocking=True)
+        if language_ids is not None:
+            language_ids = language_ids.cuda(non_blocking=True)
 
     return (
         text_input,
@@ -132,12 +166,26 @@ def format_data(data):
         linear_input,
         stop_targets,
         speaker_ids,
+        language_ids,
         speaker_embeddings,
         max_text_length,
         max_spec_length,
     )
 
 
+def extract_parameters(test_sentence):
+    splited = test_sentence.split('|')
+    if len(splited) == 1: # No language or speaker info
+        return (splited[0], None, None)
+    if len(splited) == 2: # No language info
+        sentence, speaker = splited
+        return (sentence, speaker_mapping[speaker], None)
+    if len(splited) == 3:
+        sentence, speaker, language = splited
+        return (sentence, speaker_mapping[speaker], language_mapping[language])
+    raise RuntimeError("Invalid line was given in the test sentence file.")
+
+
 def train(data_loader, model, criterion, optimizer, optimizer_st, scheduler, ap, global_step, epoch, scaler, scaler_st):
     model.train()
     epoch_time = 0
@@ -160,6 +208,7 @@ def train(data_loader, model, criterion, optimizer, optimizer_st, scheduler, ap,
             linear_input,
             stop_targets,
             speaker_ids,
+            language_ids,
             speaker_embeddings,
             max_text_length,
             max_spec_length,
@@ -186,22 +235,25 @@ def train(data_loader, model, criterion, optimizer, optimizer_st, scheduler, ap,
                     stop_tokens,
                     decoder_backward_output,
                     alignments_backward,
+                    speaker_prediction,
                 ) = model(
                     text_input,
                     text_lengths,
                     mel_input,
                     mel_lengths,
                     speaker_ids=speaker_ids,
                     speaker_embeddings=speaker_embeddings,
+                    language_ids=language_ids,
                 )
             else:
-                decoder_output, postnet_output, alignments, stop_tokens = model(
+                decoder_output, postnet_output, alignments, stop_tokens, speaker_prediction = model(
                     text_input,
                     text_lengths,
                     mel_input,
                     mel_lengths,
                     speaker_ids=speaker_ids,
                     speaker_embeddings=speaker_embeddings,
+                    language_ids=language_ids,
                 )
                 decoder_backward_output = None
                 alignments_backward = None
@@ -228,6 +280,8 @@ def train(data_loader, model, criterion, optimizer, optimizer_st, scheduler, ap,
                 alignment_lengths,
                 alignments_backward,
                 text_lengths,
+                speaker_prediction,
+                speaker_ids,
             )
 
         # check nan loss
@@ -405,6 +459,7 @@ def evaluate(data_loader, model, criterion, ap, global_step, epoch):
                 linear_input,
                 stop_targets,
                 speaker_ids,
+                language_ids,
                 speaker_embeddings,
                 _,
                 _,
@@ -420,12 +475,23 @@ def evaluate(data_loader, model, criterion, ap, global_step, epoch):
                     stop_tokens,
                     decoder_backward_output,
                     alignments_backward,
+                    speaker_prediction,
                 ) = model(
-                    text_input, text_lengths, mel_input, speaker_ids=speaker_ids, speaker_embeddings=speaker_embeddings
+                    text_input,
+                    text_lengths,
+                    mel_input,
+                    speaker_ids=speaker_ids,
+                    language_ids=language_ids,
+                    speaker_embeddings=speaker_embeddings,
                 )
             else:
-                decoder_output, postnet_output, alignments, stop_tokens = model(
-                    text_input, text_lengths, mel_input, speaker_ids=speaker_ids, speaker_embeddings=speaker_embeddings
+                decoder_output, postnet_output, alignments, stop_tokens, speaker_prediction = model(
+                    text_input,
+                    text_lengths,
+                    mel_input,
+                    speaker_ids=speaker_ids,
+                    language_ids=language_ids,
+                    speaker_embeddings=speaker_embeddings,
                 )
                 decoder_backward_output = None
                 alignments_backward = None
@@ -452,6 +518,8 @@ def evaluate(data_loader, model, criterion, ap, global_step, epoch):
                 alignment_lengths,
                 alignments_backward,
                 text_lengths,
+                speaker_prediction,
+                speaker_ids,
             )
 
             # step time
@@ -536,12 +604,14 @@ def evaluate(data_loader, model, criterion, ap, global_step, epoch):
         test_audios = {}
         test_figures = {}
         print(" | > Synthesizing test sentences")
-        speaker_id = 0 if c.use_speaker_embedding else None
+        # Those defaults are used if speaker and language are not defined in the test_sentences_file
+        speaker_id = 5 if c.use_speaker_embedding else None
         speaker_embedding = (
-            speaker_mapping[list(speaker_mapping.keys())[randrange(len(speaker_mapping) - 1)]]["embedding"]
+            speaker_mapping[list(speaker_mapping.keys())[random.randrange(len(speaker_mapping) - 1)]]["embedding"]
             if c.use_external_speaker_embedding_file and c.use_speaker_embedding
             else None
         )
+        language_id = 0 if c.use_language_embedding else None
         style_wav = c.get("gst_style_input")
         if style_wav is None and c.use_gst:
             # inicialize GST with zero dict.
@@ -552,13 +622,16 @@ def evaluate(data_loader, model, criterion, ap, global_step, epoch):
         style_wav = c.get("gst_style_input")
         for idx, test_sentence in enumerate(test_sentences):
             try:
+                test_sentence, speaker_id, language_id = extract_parameters(test_sentence)
                 wav, alignment, decoder_output, postnet_output, stop_tokens, _ = synthesis(
                     model,
                     test_sentence,
                     c,
                     use_cuda,
                     ap,
                     speaker_id=speaker_id,
+                    language_id=language_id,
+                    language_mapping=language_mapping,
                     speaker_embedding=speaker_embedding,
                     style_wav=style_wav,
                     truncated=False,
@@ -584,7 +657,7 @@ def evaluate(data_loader, model, criterion, ap, global_step, epoch):
 
 def main(args):  # pylint: disable=redefined-outer-name
     # pylint: disable=global-variable-undefined
-    global meta_data_train, meta_data_eval, speaker_mapping, symbols, phonemes, model_characters
+    global meta_data_train, meta_data_eval, speaker_mapping, symbols, phonemes, model_characters, language_mapping
     # Audio processor
     ap = AudioProcessor(**c.audio)
 
@@ -600,6 +673,7 @@ def main(args):  # pylint: disable=redefined-outer-name
 
     # load data instances
     meta_data_train, meta_data_eval = load_meta_data(c.datasets)
+    # meta_data_train = random.sample(meta_data_train, len(meta_data_train)//64)  # to speedup train phase for dev purposes
 
     # set the portion of the data used for training
     if "train_portion" in c.keys():
@@ -609,8 +683,9 @@ def main(args):  # pylint: disable=redefined-outer-name
 
     # parse speakers
     num_speakers, speaker_embedding_dim, speaker_mapping = parse_speakers(c, args, meta_data_train, OUT_PATH)
+    num_langs, language_embedding_dim, language_mapping = parse_languages(c, args, meta_data_train, OUT_PATH)
 
-    model = setup_model(num_chars, num_speakers, c, speaker_embedding_dim)
+    model = setup_model(num_chars, num_speakers, num_langs, c, speaker_embedding_dim, language_embedding_dim)
 
     # scalers for mixed precision training
     scaler = torch.cuda.amp.GradScaler() if c.mixed_precision else None

TTS/tts/configs/config.json

@@ -159,6 +159,10 @@
         "gst_style_tokens": 10,
         "gst_use_speaker_embedding": false
 	},
+    "reversal_classifier": false,
+    "reversal_classifier_dim": 256,
+    "reversal_classifier_w": 0.125,
+    "reversal_gradient_clipping": 0.25,
 
     // DATASETS
     "datasets":   // List of datasets. They all merged and they get different speaker_ids.

TTS/tts/datasets/TTSDataset.py

@@ -80,7 +80,7 @@ def __init__(
             print("\n > DataLoader initialization")
             print(" | > Use phonemes: {}".format(self.use_phonemes))
             if use_phonemes:
-                print("   | > phoneme language: {}".format(phoneme_language))
+                print("   | > Default phoneme language: {}".format(phoneme_language))
             print(" | > Number of instances : {}".format(len(self.items)))
 
     def load_wav(self, filename):
@@ -133,10 +133,10 @@ def _load_or_generate_phoneme_sequence(
     def load_data(self, idx):
         item = self.items[idx]
 
-        if len(item) == 4:
-            text, wav_file, speaker_name, attn_file = item
+        if len(item) == 5:
+            text, wav_file, speaker_name, language_name, attn_file = item
         else:
-            text, wav_file, speaker_name = item
+            text, wav_file, speaker_name, language_name = item
             attn = None
 
         wav = np.asarray(self.load_wav(wav_file), dtype=np.float32)
@@ -153,7 +153,7 @@ def load_data(self, idx):
                     self.phoneme_cache_path,
                     self.enable_eos_bos,
                     self.cleaners,
-                    self.phoneme_language,
+                    language_name if language_name else self.phoneme_language,
                     self.tp,
                     self.add_blank,
                 )
@@ -181,6 +181,7 @@ def load_data(self, idx):
             "attn": attn,
             "item_idx": self.items[idx][1],
             "speaker_name": speaker_name,
+            "language_name": language_name,
             "wav_file_name": os.path.basename(wav_file),
         }
         return sample
@@ -294,6 +295,9 @@ def collate_fn(self, batch):
             text = [batch[idx]["text"] for idx in ids_sorted_decreasing]
 
             speaker_name = [batch[idx]["speaker_name"] for idx in ids_sorted_decreasing]
+
+            language_name = [batch[idx]["language_name"] for idx in ids_sorted_decreasing]
+
             # get speaker embeddings
             if self.speaker_mapping is not None:
                 wav_files_names = [batch[idx]["wav_file_name"] for idx in ids_sorted_decreasing]
@@ -360,6 +364,7 @@ def collate_fn(self, batch):
                 mel,
                 mel_lengths,
                 stop_targets,
+                language_name,
                 item_idxs,
                 speaker_embedding,
                 attns,