{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tacotron 2 inference code \n",
    "Edit the variables **checkpoint_path** and **text** to match yours and run the entire code to generate plots of mel outputs, alignments and audio synthesis from the generated mel-spectrogram using Griffin-Lim."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Import libraries and setup matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [ ]
    }
   ],
   "source": [
    "import matplotlib\n",
    "matplotlib.use(\"Agg\")\n",
    "import matplotlib.pylab as plt\n",
    "%matplotlib inline\n",
    "import IPython.display as ipd\n",
    "\n",
    "import sys\n",
    "sys.path.append('waveglow/')\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from hparams import create_hparams\n",
    "from model import Tacotron2\n",
    "from layers import TacotronSTFT, STFT\n",
    "from audio_processing import griffin_lim\n",
    "from train import load_model\n",
    "from text import text_to_sequence\n",
    "from denoiser import Denoiser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_data(data, figsize=(16, 4)):\n",
    "    fig, axes = plt.subplots(1, len(data), figsize=figsize)\n",
    "    for i in range(len(data)):\n",
    "        axes[i].imshow(data[i], aspect='auto', origin='bottom', \n",
    "                       interpolation='none')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup hparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams = create_hparams()\n",
    "hparams.sampling_rate = 22050"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load model from checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = \"tacotron2_statedict\"\n",
    "model = load_model(hparams)\n",
    "model.load_state_dict(torch.load(checkpoint_path)['state_dict'])\n",
    "_ = model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load WaveGlow for mel2audio synthesis and denoiser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [ ]
    }
   ],
   "source": [
    "waveglow_path = 'waveglow_old.pt'\n",
    "waveglow = torch.load(waveglow_path)['model']\n",
    "waveglow.cuda()\n",
    "denoiser = Denoiser(waveglow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare text input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"Waveglow is really awesome!\"\n",
    "sequence = np.array(text_to_sequence(text, ['english_cleaners']))[None, :]\n",
    "sequence = torch.autograd.Variable(\n",
    "    torch.from_numpy(sequence)).cuda().long()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Decode text input and plot results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84492 characters]\n",
      "text/plain": [
       "<Figure size 1152x288 with 3 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "mel_outputs, mel_outputs_postnet, _, alignments = model.inference(sequence)\n",
    "plot_data((mel_outputs.data.cpu().numpy()[0],\n",
    "           mel_outputs_postnet.data.cpu().numpy()[0],\n",
    "           alignments.data.cpu().numpy()[0].T))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Synthesize audio from spectrogram using WaveGlow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 109968 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    audio = waveglow.infer(mel_outputs_postnet, sigma=0.666)\n",
    "ipd.Audio(audio[0].data.cpu().numpy(), rate=hparams.sampling_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### (Optional) Remove WaveGlow bias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 109968 characters]\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "audio_denoised = denoiser(audio, strength=0.01)[:, 0]\n",
    "ipd.Audio(audio_denoised.cpu().numpy(), rate=hparams.sampling_rate) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
