Merge branch 'dev-tacotron2'

utils/audio.py

@@ -1,5 +1,6 @@
 import os
 import librosa
+import soundfile as sf
 import pickle
 import copy
 import numpy as np
@@ -172,6 +173,14 @@ class AudioProcessor(object):
         else:
             return self._griffin_lim(S**self.power)
 
+    def out_linear_to_mel(self, linear_spec):
+        S = self._denormalize(linear_spec)
+        S = self._db_to_amp(S + self.ref_level_db)  
+        S = self._linear_to_mel(np.abs(S))
+        S = self._amp_to_db(S) - self.ref_level_db
+        mel = self._normalize(S)
+        return mel
+
     def _griffin_lim(self, S):
         angles = np.exp(2j * np.pi * np.random.rand(*S.shape))
         S_complex = np.abs(S).astype(np.complex)
@@ -230,7 +239,7 @@ class AudioProcessor(object):
     #     return np.sign(signal) * magnitude
 
     def load_wav(self, filename, encode=False):
-        x, sr = librosa.load(filename, sr=self.sample_rate)
+        x, sr = sf.read(filename)
         if self.do_trim_silence:
             x = self.trim_silence(x)
         # sr, x = io.wavfile.read(filename)

utils/generic_utils.py

@@ -8,6 +8,7 @@ import datetime
 import json
 import torch
 import subprocess
+import importlib
 import numpy as np
 from collections import OrderedDict
 from torch.autograd import Variable
@@ -31,6 +32,12 @@ def load_config(config_path):
     return config
 
 
+def get_git_branch():
+    out = subprocess.check_output(["git", "branch"]).decode("utf8")
+    current = next(line for line in out.split("\n") if line.startswith("*"))
+    return current.replace("* ", "")
+
+
 def get_commit_hash():
     """https://stackoverflow.com/questions/14989858/get-the-current-git-hash-in-a-python-script"""
     # try:
@@ -48,10 +55,10 @@ def get_commit_hash():
 def create_experiment_folder(root_path, model_name, debug):
     """ Create a folder with the current date and time """
     date_str = datetime.datetime.now().strftime("%B-%d-%Y_%I+%M%p")
-    if debug:
-        commit_hash = 'debug'
-    else:
-        commit_hash = get_commit_hash()
+    # if debug:
+    # commit_hash = 'debug'
+    # else:
+    commit_hash = get_commit_hash()
     output_folder = os.path.join(
         root_path, model_name + '-' + date_str + '-' + commit_hash)
     os.makedirs(output_folder, exist_ok=True)
@@ -71,10 +78,19 @@ def remove_experiment_folder(experiment_path):
         print(" ! Run is kept in {}".format(experiment_path))
 
 
-def copy_config_file(config_file, path):
+def copy_config_file(config_file, out_path, new_fields):
     config_name = os.path.basename(config_file)
-    out_path = os.path.join(path, config_name)
-    shutil.copyfile(config_file, out_path)
+    config_lines = open(config_file, "r").readlines()
+    # add extra information fields
+    for key, value in new_fields.items():
+        if type(value) == str:
+            new_line = '"{}":"{}",\n'.format(key, value)
+        else:
+            new_line = '"{}":{},\n'.format(key, value)
+        config_lines.insert(1, new_line)
+    config_out_file = open(out_path, "w")
+    config_out_file.writelines(config_lines)
+    config_out_file.close()
 
 
 def _trim_model_state_dict(state_dict):
@@ -99,7 +115,6 @@ def save_checkpoint(model, optimizer, optimizer_st, model_loss, out_path,
     state = {
         'model': new_state_dict,
         'optimizer': optimizer.state_dict(),
-        'optimizer_st': optimizer_st.state_dict(),
         'step': current_step,
         'epoch': epoch,
         'linear_loss': model_loss,
@@ -191,7 +206,72 @@ def sequence_mask(sequence_length, max_len=None):
     seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
     if sequence_length.is_cuda:
         seq_range_expand = seq_range_expand.cuda()
-    seq_length_expand = (sequence_length.unsqueeze(1)
-                         .expand_as(seq_range_expand))
+    seq_length_expand = (
+        sequence_length.unsqueeze(1).expand_as(seq_range_expand))
     # B x T_max
     return seq_range_expand < seq_length_expand
+
+
+def set_init_dict(model_dict, checkpoint, c):
+    # Partial initialization: if there is a mismatch with new and old layer, it is skipped.
+    for k, v in checkpoint['model'].items():
+        if k not in model_dict:
+            print(" | > Layer missing in the model definition: {}".format(k))
+    # 1. filter out unnecessary keys
+    pretrained_dict = {
+        k: v
+        for k, v in checkpoint['model'].items() if k in model_dict
+    }
+    # 2. filter out different size layers
+    pretrained_dict = {
+        k: v
+        for k, v in pretrained_dict.items()
+        if v.numel() == model_dict[k].numel()
+    }
+    # 3. skip reinit layers
+    if c.reinit_layers is not None:
+        for reinit_layer_name in c.reinit_layers:
+            pretrained_dict = {
+                k: v
+                for k, v in pretrained_dict.items()
+                if reinit_layer_name not in k
+            }
+    # 4. overwrite entries in the existing state dict
+    model_dict.update(pretrained_dict)
+    print(" | > {} / {} layers are restored.".format(
+        len(pretrained_dict), len(model_dict)))
+    return model_dict
+
+
+def setup_model(num_chars, c):
+    print(" > Using model: {}".format(c.model))
+    MyModel = importlib.import_module('models.' + c.model.lower())
+    MyModel = getattr(MyModel, c.model)
+    if c.model.lower() == "tacotron":
+        model = MyModel(
+            num_chars=num_chars,
+            r=c.r,
+            linear_dim=1025,
+            mel_dim=80,
+            memory_size=c.memory_size,
+            attn_win=c.windowing,
+            attn_norm=c.attention_norm,
+            prenet_type=c.prenet_type,
+            prenet_dropout=c.prenet_dropout,
+            forward_attn=c.use_forward_attn,
+            trans_agent=c.transition_agent,
+            location_attn=c.location_attn,
+            separate_stopnet=c.separate_stopnet)
+    elif c.model.lower() == "tacotron2":
+        model = MyModel(
+            num_chars=num_chars,
+            r=c.r,
+            attn_win=c.windowing,
+            attn_norm=c.attention_norm,
+            prenet_type=c.prenet_type,
+            prenet_dropout=c.prenet_dropout,
+            forward_attn=c.use_forward_attn,
+            trans_agent=c.transition_agent,
+            location_attn=c.location_attn,
+            separate_stopnet=c.separate_stopnet)
+    return model

utils/synthesis.py

@@ -8,23 +8,49 @@ from .visual import visualize
 from matplotlib import pylab as plt
 
 
-def synthesis(m, s, CONFIG, use_cuda, ap):
-    """ Given the text, synthesising the audio """
+def synthesis(model, text, CONFIG, use_cuda, ap, truncated=False, enable_eos_bos_chars=False, trim_silence=False):
+    """Synthesize voice for the given text.
+
+        Args:
+            model (TTS.models): model to synthesize.
+            text (str): target text
+            CONFIG (dict): config dictionary to be loaded from config.json.
+            use_cuda (bool): enable cuda.
+            ap (TTS.utils.audio.AudioProcessor): audio processor to process
+                model outputs.
+            truncated (bool): keep model states after inference. It can be used
+                for continuous inference at long texts.
+            enable_eos_bos_chars (bool): enable special chars for end of sentence and start of sentence.
+            trim_silence (bool): trim silence after synthesis.
+    """
+    # preprocess the given text
     text_cleaner = [CONFIG.text_cleaner]
     if CONFIG.use_phonemes:
         seq = np.asarray(
-            phoneme_to_sequence(s, text_cleaner, CONFIG.phoneme_language),
+            phoneme_to_sequence(text, text_cleaner, CONFIG.phoneme_language, enable_eos_bos_chars),
             dtype=np.int32)
     else:
-        seq = np.asarray(text_to_sequence(s, text_cleaner), dtype=np.int32)
+        seq = np.asarray(text_to_sequence(text, text_cleaner), dtype=np.int32)
     chars_var = torch.from_numpy(seq).unsqueeze(0)
+    # synthesize voice
     if use_cuda:
         chars_var = chars_var.cuda()
-    mel_spec, linear_spec, alignments, stop_tokens = m.forward(
-        chars_var.long())
-    linear_spec = linear_spec[0].data.cpu().numpy()
-    mel_spec = mel_spec[0].data.cpu().numpy()
+    if truncated:
+        decoder_output, postnet_output, alignments, stop_tokens = model.inference_truncated(
+            chars_var.long())
+    else:
+        decoder_output, postnet_output, alignments, stop_tokens = model.inference(
+            chars_var.long())
+    # convert outputs to numpy
+    postnet_output = postnet_output[0].data.cpu().numpy()
+    decoder_output = decoder_output[0].data.cpu().numpy()
     alignment = alignments[0].cpu().data.numpy()
-    wav = ap.inv_spectrogram(linear_spec.T)
-    wav = wav[:ap.find_endpoint(wav)]
-    return wav, alignment, linear_spec, mel_spec, stop_tokens
+    # plot results
+    if CONFIG.model == "Tacotron":
+        wav = ap.inv_spectrogram(postnet_output.T)
+    else:
+        wav = ap.inv_mel_spectrogram(postnet_output.T)
+    # trim silence
+    if trim_silence:
+        wav = wav[:ap.find_endpoint(wav)]
+    return wav, alignment, decoder_output, postnet_output, stop_tokens

utils/text/__init__.py

@@ -17,7 +17,7 @@ _id_to_phonemes = {i: s for i, s in enumerate(phonemes)}
 _curly_re = re.compile(r'(.*?)\{(.+?)\}(.*)')
 
 # Regular expression matchinf punctuations, ignoring empty space
-pat = r'['+_phoneme_punctuations[:-1]+']+'
+pat = r'['+_phoneme_punctuations+']+'
 
 
 def text2phone(text, language):
@@ -28,32 +28,39 @@ def text2phone(text, language):
     #try:
     punctuations = re.findall(pat, text)
     ph = phonemize(text, separator=seperator, strip=False, njobs=1, backend='espeak', language=language)
+    ph = ph[:-1].strip() # skip the last empty character
     # Replace \n with matching punctuations.
     if len(punctuations) > 0:
-        for punct in punctuations[:-1]:
-             ph = ph.replace('| |\n', '|'+punct+'| |', 1)
-        try:
-             ph = ph[:-1] + punctuations[-1]
-        except:
-             print(text)
+        # if text ends with a punctuation.
+        if text[-1] == punctuations[-1]:
+            for punct in punctuations[:-1]:
+                ph = ph.replace('| |\n', '|'+punct+'| |', 1)
+            try:
+                ph = ph + punctuations[-1]
+            except:
+                print(text)
+        else:
+            for punct in punctuations:
+                ph = ph.replace('| |\n', '|'+punct+'| |', 1)
     return ph
 
 
-def phoneme_to_sequence(text, cleaner_names, language):
-    '''
-    TODO: This ignores punctuations
-    '''
-    sequence = []
+def phoneme_to_sequence(text, cleaner_names, language, enable_eos_bos=False):
+    if enable_eos_bos:
+        sequence = [_phonemes_to_id['^']]
+    else:
+        sequence = []
+    text = text.replace(":", "")
     clean_text = _clean_text(text, cleaner_names)
     phonemes = text2phone(clean_text, language)
-#     print(phonemes.replace('|', ''))
     if phonemes is None:
         print("!! After phoneme conversion the result is None. -- {} ".format(clean_text))
-    for phoneme in phonemes.split('|'):
-        # print(word, ' -- ', phonemes_text)
+    # iterate by skipping empty strings - NOTE: might be useful to keep it to have a better intonation.
+    for phoneme in filter(None, phonemes.split('|')):
         sequence += _phoneme_to_sequence(phoneme)
     # Append EOS char
-    sequence.append(_phonemes_to_id['~'])
+    if enable_eos_bos:
+        sequence.append(_phonemes_to_id['~'])
     return sequence
 
 
@@ -81,7 +88,6 @@ def text_to_sequence(text, cleaner_names):
         List of integers corresponding to the symbols in the text
     '''
     sequence = []
-
     # Check for curly braces and treat their contents as ARPAbet:
     while len(text):
         m = _curly_re.match(text)
@@ -92,9 +98,6 @@ def text_to_sequence(text, cleaner_names):
             _clean_text(m.group(1), cleaner_names))
         sequence += _arpabet_to_sequence(m.group(2))
         text = m.group(3)
-
-    # Append EOS token
-    sequence.append(_symbol_to_id['~'])
     return sequence
 
 
@@ -133,8 +136,8 @@ def _arpabet_to_sequence(text):
 
 
 def _should_keep_symbol(s):
-    return s in _symbol_to_id and s is not '_' and s is not '~'
+    return s in _symbol_to_id and s not in ['~', '^', '_']
 
 
 def _should_keep_phoneme(p):
-    return p in _phonemes_to_id and p is not '_' and p is not '~'
+    return p in _phonemes_to_id and p not in ['~', '^', '_']

utils/text/cleaners.py

@@ -56,7 +56,7 @@ def lowercase(text):
 
 
 def collapse_whitespace(text):
-    return re.sub(_whitespace_re, ' ', text)
+    return re.sub(_whitespace_re, ' ', text).strip()
 
 
 def convert_to_ascii(text):

utils/text/symbols.py

@@ -8,6 +8,7 @@ through Unidecode. For other data, you can modify _characters. See TRAINING_DATA
 
 _pad = '_'
 _eos = '~'
+_bos = '^'
 _characters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz!\'(),-.:;? '
 _punctuations = '!\'(),-.:;? '
 _phoneme_punctuations = '.!;:,?'
@@ -25,8 +26,12 @@ _phonemes = sorted(list(_vowels + _non_pulmonic_consonants + _pulmonic_consonant
 _arpabet = ['@' + s for s in _phonemes]
 
 # Export all symbols:
-symbols = [_pad, _eos] + list(_characters) + _arpabet
-phonemes = [_pad, _eos] + _phonemes + list(_punctuations)
+symbols = [_pad, _eos, _bos] + list(_characters) + _arpabet
+phonemes = [_pad, _eos, _bos] + list(_phonemes) + list(_punctuations)
+
+# Generate ALIEN language
+# from random import shuffle
+# shuffle(phonemes)
 
 if __name__ == '__main__':
     print(" > TTS symbols {}".format(len(symbols)))

utils/visual.py

@@ -30,22 +30,23 @@ def plot_spectrogram(linear_output, audio):
     return fig
 
 
-def visualize(alignment, spectrogram_postnet, stop_tokens, text, hop_length, CONFIG, spectrogram=None):
+def visualize(alignment, spectrogram_postnet, stop_tokens, text, hop_length, CONFIG, spectrogram=None, output_path=None):
     if spectrogram is not None:
         num_plot = 4
     else:
         num_plot = 3
 
     label_fontsize = 16
-    plt.figure(figsize=(16, 48))
+    fig = plt.figure(figsize=(8, 24))
 
     plt.subplot(num_plot, 1, 1)
     plt.imshow(alignment.T, aspect="auto", origin="lower", interpolation=None)
     plt.xlabel("Decoder timestamp", fontsize=label_fontsize)
     plt.ylabel("Encoder timestamp", fontsize=label_fontsize)
     if CONFIG.use_phonemes:
-        seq = phoneme_to_sequence(text, [CONFIG.text_cleaner], CONFIG.phoneme_language)
+        seq = phoneme_to_sequence(text, [CONFIG.text_cleaner], CONFIG.phoneme_language, CONFIG.enable_eos_bos_chars)
         text = sequence_to_phoneme(seq)
+        print(text)
     plt.yticks(range(len(text)), list(text))
     plt.colorbar()
     
@@ -69,3 +70,8 @@ def visualize(alignment, spectrogram_postnet, stop_tokens, text, hop_length, CON
         plt.ylabel("Hz", fontsize=label_fontsize)
         plt.tight_layout()
         plt.colorbar()
+    
+    if output_path:
+        print(output_path)
+        fig.savefig(output_path)
+        plt.close()