{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Grad-TTS: A Diffusion Probabilistic Model for Text-to-Speech\r\n",
    "\r\n",
    "Official implementation of the Grad-TTS model based on Diffusion Probabilistic Models. For all details check out our paper accepted to ICML 2021 via [this](https://arxiv.org/abs/2105.06337) link.\r\n",
    "\r\n",
    "You can listen to the samples on our demo page via [this](https://grad-tts.github.io/) link.\r\n",
    "\r\n",
    "You can access Google Colab demo notebook via [this](https://colab.research.google.com/drive/1YNrXtkJQKcYDmIYJeyX8s5eXxB4zgpZI?usp=sharing) link.\r\n",
    "\r\n",
    "**Authors**: Vadim Popov\\*, Ivan Vovk\\*, Vladimir Gogoryan, Tasnima Sadekova, Mikhail Kudinov.\r\n",
    "\r\n",
    "<sup>\\*Equal contribution.</sup>\r\n",
    "\r\n",
    "**Note**: for fast synthesis prefer running inference on GPU device."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "%env CUDA_VISIBLE_DEVICES=0"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import argparse\r\n",
    "import json\r\n",
    "import datetime as dt\r\n",
    "import numpy as np\r\n",
    "import matplotlib.pyplot as plt\r\n",
    "import IPython.display as ipd\r\n",
    "from tqdm import tqdm\r\n",
    "from scipy.io.wavfile import write\r\n",
    "\r\n",
    "import torch\r\n",
    "\r\n",
    "# For Grad-TTS\r\n",
    "import params\r\n",
    "from model import GradTTS\r\n",
    "from text import text_to_sequence, cmudict\r\n",
    "from text.symbols import symbols\r\n",
    "from utils import intersperse\r\n",
    "\r\n",
    "# For HiFi-GAN\r\n",
    "import sys\r\n",
    "sys.path.append('./hifi-gan/')\r\n",
    "from env import AttrDict\r\n",
    "from models import Generator as HiFiGAN"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "N_SPKS = 247  # 247 for Libri-TTS model and 1 for single speaker (LJSpeech)"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "generator = GradTTS(len(symbols)+1, N_SPKS, params.spk_emb_dim,\r\n",
    "                    params.n_enc_channels, params.filter_channels,\r\n",
    "                    params.filter_channels_dp, params.n_heads, params.n_enc_layers,\r\n",
    "                    params.enc_kernel, params.enc_dropout, params.window_size,\r\n",
    "                    params.n_feats, params.dec_dim, params.beta_min, params.beta_max,\r\n",
    "                    pe_scale=1000)  # pe_scale=1 for `grad-tts-old.pt`\r\n",
    "generator.load_state_dict(torch.load('./checkpts/grad-tts-libri-tts.pt', map_location=lambda loc, storage: loc))\r\n",
    "_ = generator.cuda().eval()\r\n",
    "print(f'Number of parameters: {generator.nparams}')\r\n",
    "\r\n",
    "cmu = cmudict.CMUDict('./resources/cmu_dictionary')"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "with open('./checkpts/hifigan-config.json') as f:\r\n",
    "    h = AttrDict(json.load(f))\r\n",
    "hifigan = HiFiGAN(h)\r\n",
    "hifigan.load_state_dict(torch.load('./checkpts/hifigan.pt', map_location=lambda loc, storage: loc)['generator'])\r\n",
    "_ = hifigan.cuda().eval()\r\n",
    "hifigan.remove_weight_norm()\r\n",
    "%matplotlib inline"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "text = \"Here are the match lineups for the Colombia Haiti match.\""
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "x = torch.LongTensor(intersperse(text_to_sequence(text, dictionary=cmu), len(symbols))).cuda()[None]\r\n",
    "x_lengths = torch.LongTensor([x.shape[-1]]).cuda()\r\n",
    "x.shape, x_lengths"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "t = dt.datetime.now()\r\n",
    "y_enc, y_dec, attn = generator.forward(x, x_lengths, n_timesteps=50, temperature=1.3,\r\n",
    "                                       stoc=False, spk=None if N_SPKS==1 else torch.LongTensor([15]).cuda(),\r\n",
    "                                       length_scale=0.91)\r\n",
    "t = (dt.datetime.now() - t).total_seconds()\r\n",
    "print(f'Grad-TTS RTF: {t * 22050 / (y_dec.shape[-1] * 256)}')\r\n",
    "\r\n",
    "plt.figure(figsize=(15, 4))\r\n",
    "plt.subplot(1, 3, 1)\r\n",
    "plt.title('Encoder outputs')\r\n",
    "plt.imshow(y_enc.cpu().squeeze(), aspect='auto', origin='lower')\r\n",
    "plt.colorbar()\r\n",
    "plt.subplot(1, 3, 2)\r\n",
    "plt.title('Decoder outputs')\r\n",
    "plt.imshow(y_dec.cpu().squeeze(), aspect='auto', origin='lower')\r\n",
    "plt.colorbar()\r\n",
    "plt.subplot(1, 3, 3)\r\n",
    "plt.title('Alignment')\r\n",
    "plt.imshow(attn.cpu().squeeze(), aspect='auto', origin='lower');"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "with torch.no_grad():\r\n",
    "    audio = hifigan.forward(y_dec).cpu().squeeze().clamp(-1, 1)\r\n",
    "ipd.display(ipd.Audio(audio, rate=22050))"
   ],
   "outputs": [],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [],
   "outputs": [],
   "metadata": {}
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3.8.3 64-bit"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "metadata": {
   "interpreter": {
    "hash": "1c27759576147a09f82f75fe7e6da160ee29ac300de0ba196702adc9d307c9a1"
   }
  },
  "interpreter": {
   "hash": "1c27759576147a09f82f75fe7e6da160ee29ac300de0ba196702adc9d307c9a1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}