{
 "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": 27,
   "metadata": {
    "collapsed": true
   },
   "outputs": [
   ],
   "source": [
    "import matplotlib\n",
    "matplotlib.use(\"Agg\")\n",
    "import matplotlib.pylab as plt\n",
    "import IPython.display as ipd\n",
    "\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\n",
    "from audio_processing import griffin_lim\n",
    "from train import load_model\n",
    "from text import text_to_sequence\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "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": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Parsing command line hparams: distributed_run=False,mask_padding=False\n"
     ]
    }
   ],
   "source": [
    "hparams = create_hparams(\"distributed_run=False,mask_padding=False\")\n",
    "hparams.sampling_rate = 22050\n",
    "hparams.filter_length = 1024\n",
    "hparams.hop_length = 256\n",
    "hparams.win_length = 1024"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load model from checkpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_path = \"/home/scratch.adlr-gcf/audio_denoising/runs/TTS-Tacotron2-LJS-MSE-DRC-NoMaskPadding-Unsorted-Distributed-22khz/checkpoint_15500\"\n",
    "model = load_model(hparams)\n",
    "model.load_state_dict(torch.load(checkpoint_path)['state_dict'])\n",
    "model = model.module\n",
    "_ = model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare text input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"This is an example of text to speech synthesis after 14 hours training.\"\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": 53,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 126872 characters]\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x7fff5b0fa780>"
      ]
     },
     "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": [
    "#### Load TacotronSTFT and convert mel-spectrogram to spectrogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "taco_stft = TacotronSTFT(\n",
    "    hparams.filter_length, hparams.hop_length, hparams.win_length, \n",
    "    sampling_rate=hparams.sampling_rate)\n",
    "mel_decompress = taco_stft.spectral_de_normalize(mel_outputs_postnet)\n",
    "mel_decompress = mel_decompress.transpose(1, 2).data.cpu()\n",
    "spec_from_mel_scaling = 1000\n",
    "spec_from_mel = torch.mm(mel_decompress[0], taco_stft.mel_basis)\n",
    "spec_from_mel = spec_from_mel.transpose(0, 1).unsqueeze(0)\n",
    "spec_from_mel = spec_from_mel * spec_from_mel_scaling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Synthesize audio from spectrogram using the Griffin-Lim algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
   ],
   "source": [
    "waveform = griffin_lim(torch.autograd.Variable(spec_from_mel[:, :, :-1]), \n",
    "                       taco_stft.stft_fn, 60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio controls=\"controls\" >\n",
       "                    <source src=\"data:audio/wav;base64,[encoded data, 303848 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": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ipd.Audio(waveform[0].data.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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
