Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from hyper_parameters import tacotron_params as hparams | |
| from training import load_model | |
| from text import text_to_sequence | |
| from melgan.model.generator import Generator | |
| from melgan.utils.hparams import load_hparam | |
| import torch | |
| import numpy as np | |
| from matplotlib import pyplot as plt | |
| from matplotlib import gridspec | |
| torch.manual_seed(1234) | |
| MAX_WAV_VALUE = 32768.0 | |
| DESCRIPTION = """ | |
| This is a Tacotron2 model based on the NVIDIA's model plus three unsupervised Global Style Tokens (GST). | |
| The whole architecture has been trained from scratch with the LJSpeech dataset. In order to control the relevance | |
| of each style token, we configured the attention module as a single-head. | |
| Keep in mind that, for a better synthetic output, the sum of the three style weights should be around 1. A combination that sums less than 1 may work, but higher the | |
| generated speech may show more distortion and misspronunciations. | |
| """ | |
| # load trained tacotron2 + GST model: | |
| model = load_model(hparams) | |
| checkpoint_path = "trained_models/checkpoint_78000.model" | |
| model.load_state_dict(torch.load(checkpoint_path, map_location="cpu")['state_dict']) | |
| # model.to('cuda') | |
| _ = model.eval() | |
| # load pre trained MelGAN model for mel2audio: | |
| vocoder_checkpoint_path = "trained_models/nvidia_tacotron2_LJ11_epoch6400.pt" | |
| checkpoint = torch.load(vocoder_checkpoint_path, map_location="cpu") | |
| hp_melgan = load_hparam("melgan/config/default.yaml") | |
| vocoder_model = Generator(80) | |
| vocoder_model.load_state_dict(checkpoint['model_g']) | |
| # vocoder_model = vocoder_model.to('cuda') | |
| vocoder_model.eval(inference=False) | |
| def plot_spec_align(mel, align): | |
| fig_mel = plt.figure() | |
| ax_mel = fig_mel.add_subplot(111) | |
| ax_mel.imshow(mel) | |
| ax_mel.set_title('Mel-Scale Spectrogram', fontsize=20) | |
| fig_align = plt.figure() | |
| ax_align = fig_align.add_subplot(111) | |
| ax_align.imshow(align) | |
| ax_align.set_title('Alignment', fontsize=20) | |
| ''' | |
| grid_spec = gridspec.GridSpec(1, 1) | |
| ax = plt.subplot(grid_spec[0]) | |
| plt.imshow(mel) | |
| plt.axis('off') | |
| ax.set_title('Mel-Scale Spectrogram', fontsize=20) | |
| ax = plt.subplot(grid_spec[1]) | |
| plt.imshow(align) | |
| plt.axis('off') | |
| ax.set_title('Alignment', fontsize=20) | |
| plt.imshow(legend, interpolation='nearest') | |
| plt.grid('off') | |
| ''' | |
| return fig_mel, fig_align | |
| def synthesize(text, gst_1, gst_2, gst_3): | |
| sequence = np.array(text_to_sequence(text, ['english_cleaners']))[None, :] | |
| sequence = torch.from_numpy(sequence).to(device='cpu', dtype=torch.int64) | |
| # gst_head_scores = np.array([0.5, 0.15, 0.35]) # originally ([0.5, 0.15, 0.35]) | |
| gst_head_scores = np.array([gst_1, gst_2, gst_3]) | |
| gst_scores = torch.from_numpy(gst_head_scores).float() | |
| mel_outputs, mel_outputs_postnet, _, alignments = model.inference(sequence, gst_scores) | |
| # mel2wav inference: | |
| with torch.no_grad(): | |
| audio = vocoder_model.inference(mel_outputs_postnet) | |
| audio_numpy = audio.data.cpu().detach().numpy() | |
| # prepare plot for the output: | |
| mel_outputs_postnet = mel_outputs_postnet.squeeze().detach().numpy() | |
| alignments = alignments.squeeze().detach().numpy() | |
| fig_mel, fig_align = plot_spec_align(mel_outputs_postnet, alignments) | |
| return (22050, audio_numpy), fig_mel, fig_align | |
| iface = gr.Interface(fn=synthesize, inputs=[gr.Textbox(label="Input Text"), gr.Slider(0.2, 0.45, label="First style token weight:"), | |
| gr.Slider(0.2, 0.45, label="Second style token weight:"), gr.Slider(0.2, 0.45, label="Third style token weight:")], | |
| outputs=[gr.Audio(label="Generated Speech", type="numpy"), gr.Plot(label="Spectrogram"), gr.Plot(label="Alignments")], | |
| title="Single-Head Attention Tacotron2 with Style Tokens", description=DESCRIPTION) | |
| iface.launch() | |