TTS/utils/synthesis.py

import io
import time
import librosa
import torch
import numpy as np
from .text import text_to_sequence, phoneme_to_sequence, sequence_to_phoneme
from .visual import visualize
from matplotlib import pylab as plt


def text_to_seqvec(text, CONFIG, use_cuda):
    text_cleaner = [CONFIG.text_cleaner]
    # text ot phonemes to sequence vector
    if CONFIG.use_phonemes:
        seq = np.asarray(
            phoneme_to_sequence(text, text_cleaner, CONFIG.phoneme_language,
                                CONFIG.enable_eos_bos_chars),
            dtype=np.int32)
    else:
        seq = np.asarray(text_to_sequence(text, text_cleaner), dtype=np.int32)
    # torch tensor
    chars_var = torch.from_numpy(seq).unsqueeze(0)
    if use_cuda:
        chars_var = chars_var.cuda()
    return chars_var.long()


def compute_style_mel(style_wav, ap, use_cuda):
    print(style_wav)
    style_mel = torch.FloatTensor(ap.melspectrogram(
        ap.load_wav(style_wav))).unsqueeze(0)
    if use_cuda:
        return style_mel.cuda()
    else:
        return style_mel


def run_model(model, inputs, CONFIG, truncated, style_mel=None):
    if CONFIG.model == "TacotronGST" and style_mel is not None:
        decoder_output, postnet_output, alignments, stop_tokens = model.inference(
            inputs, style_mel)
    else:
        if truncated:
            decoder_output, postnet_output, alignments, stop_tokens = model.inference_truncated(
                inputs)
        else:
            decoder_output, postnet_output, alignments, stop_tokens = model.inference(
                inputs)
    return decoder_output, postnet_output, alignments, stop_tokens


def parse_outputs(postnet_output, decoder_output, alignments):
    postnet_output = postnet_output[0].data.cpu().numpy()
    decoder_output = decoder_output[0].data.cpu().numpy()
    alignment = alignments[0].cpu().data.numpy()
    return postnet_output, decoder_output, alignment


def trim_silence(wav):
    return wav[:ap.find_endpoint(wav)]


def inv_spectrogram(postnet_output, ap, CONFIG):
    if CONFIG.model in ["Tacotron", "TacotronGST"]:
        wav = ap.inv_spectrogram(postnet_output.T)
    else:
        wav = ap.inv_mel_spectrogram(postnet_output.T)
    return wav


def synthesis(model,
              text,
              speaker_id,
              CONFIG,
              use_cuda,
              ap,
              style_wav=None,
              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
            speaker_id (int): id of speaker
            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.
            style_wav (str): Uses for style embedding of GST.
            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.
    """
    # GST processing
    style_mel = None
    if CONFIG.model == "TacotronGST" and style_wav is not None:
        style_mel = compute_style_mel(style_wav, ap, use_cuda)
    # preprocess the given text
    inputs = text_to_seqvec(text, CONFIG, use_cuda)
    speaker_id = speaker_id_var = torch.from_numpy(speaker_id).unsqueeze(0)
    if use_cuda:
        speaker_id.cuda()
    # synthesize voice
    decoder_output, postnet_output, alignments, stop_tokens = run_model(
        model, inputs, CONFIG, truncated, style_mel)
    # convert outputs to numpy
    postnet_output, decoder_output, alignment = parse_outputs(
        postnet_output, decoder_output, alignments)
    # plot results
    wav = inv_spectrogram(postnet_output, ap, CONFIG)
    # trim silence
    if trim_silence:
        wav = trim_silence(wav)
    return wav, alignment, decoder_output, postnet_output, stop_tokens