tutorials/tts/FastPitch_GermanTTS_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2d65b905",
   "metadata": {},
   "source": [
    "# Modifying FastPitch to Train on a Non-English (German) Dataset\n",
    "\n",
    "This notebook is designed to provide a guide on how to train FastPitch on a German dataset from scratch as part of the TTS pipeline. It contains the following sections:\n",
    "  1. **Introduction**: FastPitch and HiFi-GAN in NeMo\n",
    "  2. **Dataset Preparation**: How to prepare German dataset for FastPitch\n",
    "  3. **Training**: Example of FastPitch training and evaluation\n",
    "  4. **Finetuning HiFi-GAN**: Improving speech quality by Finetuning HiFi-GAN on synthesized mel-spectrograms from FastPitch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42273266",
   "metadata": {},
   "source": [
    "# License\n",
    "\n",
    "> Copyright 2023 NVIDIA. All Rights Reserved.\n",
    "> \n",
    "> Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "> you may not use this file except in compliance with the License.\n",
    "> You may obtain a copy of the License at\n",
    "> \n",
    ">     http://www.apache.org/licenses/LICENSE-2.0\n",
    "> \n",
    "> Unless required by applicable law or agreed to in writing, software\n",
    "> distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "> WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "> See the License for the specific language governing permissions and\n",
    "> limitations under the License."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4c77c34",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL).\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator).\n",
    "4. Run this cell to set up dependencies.\n",
    "5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect.\n",
    "\"\"\"\n",
    "\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "\n",
    "## Install dependencies\n",
    "# !apt-get install sox libsndfile1 ffmpeg\n",
    "# !pip install wget text-unidecode matplotlib>=3.3.2\n",
    "\n",
    "## Install NeMo\n",
    "BRANCH = 'r1.17.0'\n",
    "# !python -m pip install \"git+https://github.com/NVIDIA/NeMo.git@${BRANCH}#egg=nemo_toolkit[all]\"\n",
    "\n",
    "## Install pynini\n",
    "# !wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/nemo_text_processing/install_pynini.sh\n",
    "# !bash install_pynini.sh\n",
    "\n",
    "\"\"\"\n",
    "Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
    "Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
    "that you want to use the \"Run All Cells\" (or similar) option.\n",
    "\"\"\"\n",
    "# exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f018747c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import nemo\n",
    "import torch\n",
    "import librosa\n",
    "import numpy as np\n",
    "\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "797bb19a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's download the files we need to run this tutorial\n",
    "!rm -rf NeMoGermanTTS\n",
    "!mkdir NeMoGermanTTS\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/thorsten_neutral/get_data.py\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/thorsten_neutral/ds_conf/ds_for_fastpitch_align.yaml\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/fastpitch.py\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/hifigan_finetune.py\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/extract_sup_data.py\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/generate_mels.py\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/conf/de/fastpitch_align_22050_grapheme.yaml\n",
    "!cd NeMoGermanTTS && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/conf/hifigan/hifigan.yaml\n",
    "!cd NeMoGermanTTS && mkdir -p model/train_ds && cd model/train_ds && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/conf/hifigan/model/train_ds/train_ds_finetune.yaml\n",
    "!cd NeMoGermanTTS && mkdir -p model/validation_ds && cd model/validation_ds && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/conf/hifigan/model/validation_ds/val_ds_finetune.yaml\n",
    "!cd NeMoGermanTTS && mkdir -p model/generator && cd model/generator && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/conf/hifigan/model/generator/v1.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9e63f12",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74b10183",
   "metadata": {},
   "source": [
    "### FastPitch\n",
    "\n",
    "FastPitch is non-autoregressive model for mel-spectrogram generation based on FastSpeech, conditioned on fundamental frequency contours. For more details about model, please refer to the original [paper](https://ieeexplore.ieee.org/abstract/document/9413889). Original [FastPitch model](https://ieeexplore.ieee.org/abstract/document/9413889) uses an external Tacotron 2 model trained on LJSpeech-1.1 to extract training alignments and estimate durations of input symbols. This implementation of FastPitch is based on [Deep Learning Examples](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechSynthesis/FastPitch), which uses an alignment mechanism proposed in [RAD-TTS](https://openreview.net/pdf?id=0NQwnnwAORi) and extended in [TTS Aligner](https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=9747707).\n",
    "\n",
    "For more information on training a basic FastPitch model, please refer to [FastPitch_MixerTTS_Training.ipynb](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/FastPitch_MixerTTS_Training.ipynb) tutorial.\n",
    "\n",
    "### HiFiGAN\n",
    "HiFiGAN is a generative adversarial network (GAN) model that generates audio from mel spectrograms. The generator uses transposed convolutions to upsample mel spectrograms to audio. For more details about the model, please refer to the original [paper](https://arxiv.org/abs/2010.05646). NeMo re-implementation of HiFi-GAN can be found [here](https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/tts/models/hifigan.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8b533b7",
   "metadata": {},
   "source": [
    "# Dataset Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89d94d7c",
   "metadata": {},
   "source": [
    "We will show example of preprocessing and training using OpenSLR's German Neutral TTS dataset ([link](https://www.openslr.org/95)). It is a free single german speaker dataset (> 23 hours) by Thorsten Müller (voice) and Dominik Kreutz (audio optimization) for tts training. \n",
    "\n",
    "In this section, we will cover:\n",
    "1. Downloading the dataset\n",
    "2. Extracting supplementary data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0863d60c",
   "metadata": {},
   "source": [
    "## 1. Download Thorsten's Neutral Dataset and Create Data Manifests \n",
    "\n",
    "We've created a script, `scripts/dataset_processing/tts/thorsten_neutral/get_data.py`, to download Thorsten's Datasets and generate train/val/test splits as JSON manifests with the following fields for each datapoint:\n",
    "1. `audio_filepath`: location of the wav file;\n",
    "2. `duration`: duration of the wav file;\n",
    "3. `text`: original text;\n",
    "4. `normalized_text`: normalized text through our text normalization pipline.\n",
    "    \n",
    "This script supports processing either of Thorsten's Neutral Datasets 21.02 or 22.10. In this tutorial, we only focus on the latest 22.10 version dataset. Please refer [thorsten-muller-s-german-neutral-tts-datasets](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/tts/datasets.html#thorsten-muller-s-german-neutral-tts-datasets) for more details about Thorsten's datasets. \n",
    "\n",
    "You can run the below command to obtain the final manifests, `train_manifest_text_normed.json`, `val_manifest_text_normed.json` and `test_manifest_text_normed.json`. **Note** that this script would take sometime (~2 hours) to dowload and normalize the entire dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3384c8fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Version 22.10\n",
    "!cd NeMoGermanTTS && python get_data.py \\\n",
    "    --data-root ../data_thorsten_2210 \\\n",
    "    --manifests-root ./ \\\n",
    "    --data-version \"22_10\" \\\n",
    "    --val-size 100 \\\n",
    "    --test-size 100 \\\n",
    "    --seed-for-ds-split 100 \\\n",
    "    --num-workers -1 \\\n",
    "    --normalize-text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a7cf49e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data_thorsten_2210 directory looks like\n",
    "!ls ./data_thorsten_2210/ThorstenVoice-Dataset-22_10/ThorstenVoice-Dataset_2022.10/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc6f120f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# manifests\n",
    "!ls ./NeMoGermanTTS/*_text_normed.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c42e260",
   "metadata": {},
   "source": [
    "## 2. Extracting Supplementary Data\n",
    "\n",
    "As mentioned in the [FastPitch and MixerTTS training tutorial](FastPitch_MixerTTS_Training.ipynb) - To accelerate and stabilize our training, we also need to extract pitch for every audio, estimate pitch statistics (mean, std, min, and max). To do this, all we need to do is iterate over our data one time, via `extract_sup_data.py` script.\n",
    "\n",
    "**Note**: This is an optional step, if skipped, it will be automatically executed within the first epoch of training FastPitch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d1a32b",
   "metadata": {},
   "source": [
    "The configuration remains the same as described in `scripts/dataset_processing/tts/thorsten_neutral/ds_conf/ds_for_fastpitch_align.yaml`. Note that there is no need to specify `whitelist_path` config anymore from NeMo Release 1.17.0 because it has been moved to a new dependency repo https://github.com/NVIDIA/NeMo-text-processing and it has been applied implicitly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9588c280",
   "metadata": {},
   "source": [
    "Run below command line and you will obtain the pitch normalization parameters estimated on the entire dataset. It will cost around 1.5 hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "808b5048",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd NeMoGermanTTS && python extract_sup_data.py \\\n",
    "        --config-path . \\\n",
    "        --config-name ds_for_fastpitch_align.yaml \\\n",
    "        manifest_filepath=train_manifest_text_normed.json \\\n",
    "        sup_data_path=sup_data \\\n",
    "        ++dataloader_params.num_workers=4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb2d7f2b",
   "metadata": {},
   "source": [
    "After running the above command line, you will observe a new folder `NeMoGermanTTS/sup_data/pitch` and printouts of pitch statistics below. Specify these values to the FastPitch training configurations. We will be there in the following section.\n",
    "```bash\n",
    "PITCH_MEAN=126.70349884033203, PITCH_STD=31.456707000732422\n",
    "PITCH_MIN=65.4063949584961, PITCH_MAX=1998.48779296875\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f4c41e9",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14c2ff60",
   "metadata": {},
   "source": [
    "Before we train our model, let's define model config. Most of the model config stays the same as defined here: `examples/tts/conf/de/fastpitch_align_22050_grapheme.yaml`, except that the `pitch_mean` and `pitch_std` should be updated with the values estimated by the above `extract_sup_data.py` script. Note that there is no need to specify `whitelist_path` config anymore from NeMo Release 1.17.0 because it has been moved to a new dependency repo https://github.com/NVIDIA/NeMo-text-processing and it has been applied implicitly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac938960",
   "metadata": {},
   "source": [
    "If you are using Weights and Biases, you may need to login first. More details [here](https://docs.wandb.ai/ref/cli/wandb-login)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab8b471",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "!wandb login #paste_wandb_apikey_here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4364261",
   "metadata": {},
   "source": [
    "Now we are ready for training our model! Let's try to train FastPitch. Copy and Paste the `PITCH_MEAN` and `PITCH_STD` from previous steps to overide `pitch_mean` and `pitch_std` configs below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e8e0cd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "!(cd NeMoGermanTTS && CUDA_VISIBLE_DEVICES=0 python fastpitch.py --config-path . --config-name fastpitch_align_22050_grapheme \\\n",
    "    model.train_ds.dataloader_params.batch_size=32 \\\n",
    "    model.validation_ds.dataloader_params.batch_size=32 \\\n",
    "    train_dataset=train_manifest_text_normed.json \\\n",
    "    validation_datasets=val_manifest_text_normed.json \\\n",
    "    sup_data_path=sup_data \\\n",
    "    exp_manager.exp_dir=resultGermanTTS \\\n",
    "    trainer.max_epochs=1 \\\n",
    "    trainer.check_val_every_n_epoch=1 \\\n",
    "    pitch_mean=126.70349884033203 \\\n",
    "    pitch_std=31.456707000732422 \\\n",
    "    +exp_manager.create_wandb_logger=true \\\n",
    "    +exp_manager.wandb_logger_kwargs.name=\"tutorial\" \\\n",
    "    +exp_manager.wandb_logger_kwargs.project=\"GermanTTS\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe0bfbb9",
   "metadata": {},
   "source": [
    "Note:\n",
    "1. We use `CUDA_VISIBLE_DEVICES=0` to limit training to single GPU.\n",
    "2. For debugging you may also add the following flags: `HYDRA_FULL_ERROR=1`, `CUDA_LAUNCH_BLOCKING=1`\n",
    "\n",
    "<b>Note</b>: We've limited the above run to 1 epoch only, so we can validate the implementation within the scope of this tutorial. We recommend around 5000 epochs when training FastPitch from scratch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "632d51b0",
   "metadata": {},
   "source": [
    "## Evaluating FastPitch + pretrained HiFi-GAN\n",
    "\n",
    "Let's evaluate the quality of the FastPitch model generated so far using a HiFi-GAN model pre-trained on English."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "053a00a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "from nemo.collections.tts.models import HifiGanModel, FastPitchModel\n",
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d82bc36b",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = \"das werden einmal \\u00fcber eine million kleiner buchenpflanzen werden.\" # text input to the model\n",
    "test_id = \"2b2b496ccc9b57130f559c4fd827825f\" # identifier for the audio corresponding to the test text\n",
    "data_path = \"data_thorsten_2210/ThorstenVoice-Dataset-22_10/ThorstenVoice-Dataset_2022.10/wavs/\" # path to dataset folder with wav files from original dataset\n",
    "seed = 1234"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87ff970d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_spec_fastpitch_ckpt(spec_gen_model, v_model, test):\n",
    "    with torch.no_grad():\n",
    "        torch.manual_seed(seed)\n",
    "        torch.cuda.manual_seed(seed)\n",
    "        torch.backends.cudnn.enabled = True\n",
    "        torch.backends.cudnn.benchmark = False\n",
    "        parsed = spec_gen_model.parse(str_input=test, normalize=True)\n",
    "        spectrogram = spec_gen_model.generate_spectrogram(tokens=parsed)\n",
    "        print(spectrogram.size())\n",
    "        audio = v_model.convert_spectrogram_to_audio(spec=spectrogram)\n",
    "\n",
    "    spectrogram = spectrogram.to('cpu').numpy()[0]\n",
    "    audio = audio.to('cpu').numpy()[0]\n",
    "    audio = audio / np.abs(audio).max()\n",
    "    return audio, spectrogram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70733af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load fastpitch and hifigan models\n",
    "import glob, os\n",
    "fastpitch_model_path = sorted(\n",
    "    glob.glob(\"NeMoGermanTTS/resultGermanTTS/FastPitch/*/checkpoints/FastPitch.nemo\"), \n",
    "    key=os.path.getmtime\n",
    ")[-1] # path_to_fastpitch_nemo_or_ckpt\n",
    "hfg_ngc = \"tts_en_lj_hifigan_ft_mixerttsx\" # NGC pretrained model name: https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/tts_en_lj_hifigan \n",
    "\n",
    "vocoder_model = HifiGanModel.from_pretrained(hfg_ngc, strict=False).eval().cuda()\n",
    "if \".nemo\" in fastpitch_model_path:\n",
    "    spec_gen_model = FastPitchModel.restore_from(fastpitch_model_path).eval().cuda()\n",
    "else:\n",
    "    spec_gen_model = FastPitchModel.load_from_checkpoint(checkpoint_path=fastpitch_model_path).eval().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27aa1827",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "audio, spectrogram = evaluate_spec_fastpitch_ckpt(spec_gen_model, vocoder_model, test)\n",
    "\n",
    "# visualize the spectrogram\n",
    "if spectrogram is not None:\n",
    "    imshow(spectrogram, origin=\"lower\")\n",
    "    plt.show()\n",
    "\n",
    "# audio\n",
    "print(\"original audio\")\n",
    "ipd.display(ipd.Audio(data_path+test_id+'.wav', rate=22050))\n",
    "print(\"predicted audio\")\n",
    "ipd.display(ipd.Audio(audio, rate=22050))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15006392",
   "metadata": {},
   "source": [
    "You would hear that the above synthesized audio quality is not as good as we expect. It would be improved after continuing to train 5000 epochs. But again,the quality is still not acceptable. A straightforward solution is to finetune the HiFi-GAN model following the tutorial [FastPitch_Finetuning.ipynb](FastPitch_Finetuning.ipynb). Let's try that out next!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7002b1e",
   "metadata": {},
   "source": [
    "# Finetuning HiFi-GAN\n",
    "\n",
    "Improving speech quality by finetuning HiFi-GAN on synthesized mel-spectrograms from FastPitch. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eccd357",
   "metadata": {},
   "source": [
    "## Generating synthetic mels\n",
    "\n",
    "To generate mel-spectrograms from FastPitch, we can use `generate_spectrogram` method defined in `nemo/collections/tts/models/fastpitch.py`. However, the resulting spectrogram may be different from ground truth mel spectrogram, as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5edb6fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_audio_filepath = \"data_thorsten_2210/ThorstenVoice-Dataset-22_10/ThorstenVoice-Dataset_2022.10/wavs/43ee28172fe1a9d1eebd77bc09f03e51.wav\"\n",
    "test_audio_text = \"daran beteiligten sich einhundertachtzig kommunen und hochschulen.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "769f270a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import imshow\n",
    "from nemo.collections.tts.models import FastPitchModel\n",
    "from matplotlib import pyplot as plt\n",
    "import librosa\n",
    "import librosa.display\n",
    "import torch\n",
    "import soundfile as sf\n",
    "import numpy as np\n",
    "from nemo.collections.tts.parts.utils.tts_dataset_utils import BetaBinomialInterpolator\n",
    "\n",
    "def load_wav(audio_file):\n",
    "    with sf.SoundFile(audio_file, 'r') as f:\n",
    "        samples = f.read(dtype='float32')\n",
    "    return samples.transpose()\n",
    "\n",
    "def plot_logspec(spec, axis=None):    \n",
    "    librosa.display.specshow(\n",
    "        librosa.amplitude_to_db(spec, ref=np.max),\n",
    "        y_axis='linear', \n",
    "        x_axis=\"time\",\n",
    "        fmin=0, \n",
    "        fmax=8000,\n",
    "        ax=axis\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aef1470b",
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_model = FastPitchModel.restore_from(fastpitch_model_path).eval().cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a840dcf",
   "metadata": {},
   "source": [
    "So we have 2 types of mel spectrograms that we can use for finetuning HiFi-GAN:\n",
    "\n",
    "### 1. Original mel spectrogram generated from original audio file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60d49e67",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"loading original melspec\")\n",
    "y, sr = librosa.load(test_audio_filepath)\n",
    "# change n_fft, win_length, hop_length parameters below based on your specific config file\n",
    "spectrogram2 = librosa.feature.melspectrogram(y=y, sr=sr, n_fft=1024, win_length=1024, hop_length=256)\n",
    "spectrogram = spectrogram2[ :80, :]\n",
    "print(\"spectrogram shape = \", spectrogram.shape)\n",
    "plot_logspec(spectrogram)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5641b7c4",
   "metadata": {},
   "source": [
    "### 2. Mel spectrogram predicted from FastPitch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f448efd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"loading fastpitch melspec via generate_spectrogram\")\n",
    "with torch.no_grad():\n",
    "    text = spec_model.parse(test_audio_text, normalize=False)\n",
    "    spectrogram = spec_model.generate_spectrogram(\n",
    "      tokens=text, \n",
    "      speaker=None,\n",
    "    )\n",
    "spectrogram = spectrogram.to('cpu').numpy()[0]\n",
    "plot_logspec(spectrogram)\n",
    "print(\"spectrogram shape = \", spectrogram.shape)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c240c79",
   "metadata": {},
   "source": [
    "**Note**: The above predicted spectrogram has the duration of 241 frames which is not equal to the ground truth 345 frames. In order to finetune HiFi-GAN we need mel spectrogram predicted from FastPitch with ground truth alignment and duration.\n",
    "\n",
    "### 2.1 Mel spectrogram predicted from FastPitch with groundtruth alignment and duration "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddd0eef5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"loading fastpitch melspec via forward method with groundtruth alignment and duration\")\n",
    "with torch.no_grad():\n",
    "    device = spec_model.device\n",
    "    beta_binomial_interpolator = BetaBinomialInterpolator()\n",
    "    text = spec_model.parse(test_audio_text, normalize=False)\n",
    "    text_len = torch.tensor(text.shape[-1], dtype=torch.long, device=device).unsqueeze(0)\n",
    "    audio = load_wav(test_audio_filepath)\n",
    "    audio = torch.from_numpy(audio).unsqueeze(0).to(device)\n",
    "    audio_len = torch.tensor(audio.shape[1], dtype=torch.long, device=device).unsqueeze(0)\n",
    "    spect, spect_len = spec_model.preprocessor(input_signal=audio, length=audio_len)\n",
    "    attn_prior = torch.from_numpy(\n",
    "      beta_binomial_interpolator(spect_len.item(), text_len.item())\n",
    "    ).unsqueeze(0).to(text.device)\n",
    "    spectrogram = spec_model.forward(\n",
    "      text=text, \n",
    "      input_lens=text_len, \n",
    "      spec=spect, \n",
    "      mel_lens=spect_len, \n",
    "      attn_prior=attn_prior,\n",
    "      speaker=None,\n",
    "    )[0]\n",
    "spectrogram = spectrogram.to('cpu').numpy()[0]\n",
    "print(\"spectrogram shape = \", spectrogram.shape)\n",
    "plot_logspec(spectrogram)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2c99521",
   "metadata": {},
   "source": [
    "In our experience, \n",
    "- Finetuning with #1 has artifacts from the original audio (noise) that get passed on as input to the vocoder resulting in artifacts in vocoder output in the form of noise.\n",
    "- <b> On the other hand, #2.1 (i.e. `Mel spectrogram predicted from FastPitch with groundtruth alignment and duration`) gives the best results because it enables HiFi-GAN to learn mel spectrograms generated by FastPitch as well as duration distributions closer to the real world (i.e. ground truth) durations. </b>\n",
    "\n",
    "From implementation perspective - we follow the same process described in [Finetuning FastPitch for a new speaker](FastPitch_Finetuning.ipynb) - i.e. take the latest checkpoint from FastPitch training and predict spectrograms for each of the input records in `train_manifest_text_normed.json`, `test_manifest_text_normed.json` and `val_manifest_text_normed.json`. NeMo provides an efficient script, [scripts/dataset_processing/tts/generate_mels.py](https://raw.githubusercontent.com/nvidia/NeMo/main/scripts/dataset_processing/tts/generate_mels.py), to generate Mel-spectrograms in the directory `NeMoGermanTTS/mels` and also create new JSON manifests with a suffix `_mel` by adding a new key `\"mel_filepath\"`. For example, `train_manifest_text_normed.json` corresponds to `train_manifest_text_normed_mel.json` saved in the same directory. You can run the following CLI to obtain the new JSON manifests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1dc3025",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd NeMoGermanTTS && python generate_mels.py \\\n",
    "    --cpu \\\n",
    "    --fastpitch-model-ckpt {fastpitch_model_path.split(\"/\", maxsplit=1)[1]} \\\n",
    "    --input-json-manifests train_manifest_text_normed.json val_manifest_text_normed.json test_manifest_text_normed.json \\\n",
    "    --output-json-manifest-root ./"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4cb8e4a",
   "metadata": {},
   "source": [
    "Revisiting how we implement #2.1 (i.e. Predicted mel spectrogram predicted from FastPitch with groundtruth alignment and duration):\n",
    "\n",
    "1. Notice above that we use audio from dataset (`audio` variable) to compute spectrogram length (`spect_len`):\n",
    "    ```python\n",
    "    spect, spect_len = spec_model.preprocessor(input_signal=audio, length=audio_len)\n",
    "    ```\n",
    "2. and groundtruth alignment (`attn_prior`).\n",
    "    ```python\n",
    "    attn_prior = torch.from_numpy(\n",
    "          beta_binomial_interpolator(spect_len.item(), text_len.item())\n",
    "        ).unsqueeze(0).to(text.device)\n",
    "    ```\n",
    "3. We use both of them to generate synthetic mel spectrogram via `spec_model.forward` method:\n",
    "    ```python\n",
    "    spectrogram = spec_model.forward(\n",
    "          text=text, \n",
    "          input_lens=text_len, \n",
    "          spec=spect, \n",
    "          mel_lens=spect_len, \n",
    "          attn_prior=attn_prior,\n",
    "          speaker=speaker,\n",
    "        )[0]\n",
    "    ```\n",
    "    \n",
    "Repeat the above script for train and validation datasets as well. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b0918cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example HiFi-GAN manifest:\n",
    "!head -n1 NeMoGermanTTS/train_manifest_text_normed_mel.json | jq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61cf4696",
   "metadata": {},
   "source": [
    "## Launch finetuning\n",
    "\n",
    "We will be re-using the existing HiFi-GAN config and HiFi-GAN pretrained on English."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "712df66f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd NeMoGermanTTS && CUDA_VISIBLE_DEVICES=0 python hifigan_finetune.py --config-path . --config-name hifigan.yaml \\\n",
    "    model.max_steps=10 \\\n",
    "    model.optim.lr=0.00001 \\\n",
    "    ~model.optim.sched \\\n",
    "    train_dataset=train_manifest_text_normed_mel.json \\\n",
    "    validation_datasets=val_manifest_text_normed_mel.json \\\n",
    "    exp_manager.exp_dir=resultGermanTTS \\\n",
    "    +init_from_pretrained_model={hfg_ngc} \\\n",
    "    +trainer.val_check_interval=5 \\\n",
    "    trainer.check_val_every_n_epoch=null \\\n",
    "    model/train_ds=train_ds_finetune \\\n",
    "    model/validation_ds=val_ds_finetune \\\n",
    "    exp_manager.create_wandb_logger=true \\\n",
    "    exp_manager.wandb_logger_kwargs.name=\"tutorial_2\" \\\n",
    "    exp_manager.wandb_logger_kwargs.project=\"GermanTTS\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1e18615",
   "metadata": {},
   "source": [
    "<b>Note</b>: We've limited the above run to 10 steps only, so we can validate the implementation within the scope of this tutorial. We recommend evaluating around every 50 steps HiFi-GAN until you get desired quality results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb4a8c7b",
   "metadata": {},
   "source": [
    "## Evaluating FastPitch and Finetuned HiFi-GAN\n",
    "\n",
    "Let's evaluate the quality of the FastPitch model generated so far using a HiFi-GAN model finetuned on predicted mels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd683d45",
   "metadata": {},
   "outputs": [],
   "source": [
    "hfg_path = sorted(glob.glob(\"NeMoGermanTTS/resultGermanTTS/HifiGan/*/checkpoints/HifiGan.nemo\"), key=os.path.getmtime)[-1]\n",
    "\n",
    "if \".nemo\" in hfg_path:\n",
    "    vocoder_model_pt = HifiGanModel.restore_from(hfg_path).eval().cuda()\n",
    "else:\n",
    "    vocoder_model_pt = HifiGanModel.load_from_checkpoint(checkpoint_path=hfg_path).eval().cuda()\n",
    "    \n",
    "if \".nemo\" in fastpitch_model_path:\n",
    "    spec_gen_model = FastPitchModel.restore_from(fastpitch_model_path).eval().cuda()\n",
    "else:\n",
    "    spec_gen_model = FastPitchModel.load_from_checkpoint(checkpoint_path=fastpitch_model_path).eval().cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a7b7aa2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "audio, spectrogram = evaluate_spec_fastpitch_ckpt(spec_gen_model, vocoder_model_pt, test)\n",
    "\n",
    "# visualize the spectrogram\n",
    "if spectrogram is not None:\n",
    "    imshow(spectrogram, origin=\"lower\")\n",
    "    plt.show()\n",
    "\n",
    "# audio\n",
    "print(\"original audio\")\n",
    "ipd.display(ipd.Audio(data_path+test_id+'.wav', rate=22050))\n",
    "print(\"predicted audio\")\n",
    "ipd.display(ipd.Audio(audio, rate=22050))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0700e6f4",
   "metadata": {},
   "source": [
    "That's it!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}