tutorials/tts/FastPitch_ChineseTTS_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bfe7f0e7",
   "metadata": {},
   "source": [
    "# Modifying FastPitch to Train on a Chinese and English Bilingual Dataset\n",
    "\n",
    "This notebook is designed to provide a guide on how to train FastPitch on a Chinese and English bilingual 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 Chinese dataset for FastPitch\n",
    "  3. **Training**: Example of FastPitch training and evaluation\n",
    "  4. **(TODO) Finetuning from LJSpeech Acoustic Model**: Improving English speech quality by finetuning LJ Speech pretrained model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4623c99",
   "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": "2406ae3b",
   "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",
    "# !pip install opencc-python-reimplemented\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": "d9fa8367",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import nemo\n",
    "import torch\n",
    "import librosa\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25f2b755",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's download the files we need to run this tutorial\n",
    "!mkdir -p NeMoChineseTTS\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/tts_dataset_files/zh/pinyin_dict_nv_22.10.txt\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/sfbilingual/get_data.py\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/sfbilingual/ds_conf/ds_for_fastpitch_align.yaml\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/fastpitch.py\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/hifigan_finetune.py\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/extract_sup_data.py\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/scripts/dataset_processing/tts/generate_mels.py\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/conf/zh/fastpitch_align_22050.yaml\n",
    "!cd NeMoChineseTTS && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/conf/hifigan/hifigan.yaml\n",
    "!cd NeMoChineseTTS && 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 NeMoChineseTTS && 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 NeMoChineseTTS && 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": "fe046c98",
   "metadata": {},
   "source": [
    "# Introduction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb151217",
   "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",
    "### HiFi-GAN\n",
    "HiFi-GAN 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": "85ac741a",
   "metadata": {},
   "source": [
    "# Dataset Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e37c09f",
   "metadata": {},
   "source": [
    "We will show example of preprocessing and training using SF Bilingual Speech TTS Dataset ([link](https://catalog.ngc.nvidia.com/orgs/nvidia/resources/sf_bilingual_speech_zh_en)). The dataset contains about 2,740 bilingual audio samples of a single female speaker and their corresponding text transcripts, each of them is an audio of around 5-6 seconds and have a total length of approximately 4.5 hours. The SF Bilingual Speech Dataset is published in NGC registry with CC BY-NC 4.0 license. Please review details from the above link.\n",
    "\n",
    "In this section, we will cover:\n",
    "1. Installing NGC Registry CLI\n",
    "2. Downloading SFSpeech Dataset\n",
    "3. Creating Data Manifests\n",
    "4. Phonemization\n",
    "5. Creating Supplementary Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "402e2494",
   "metadata": {},
   "source": [
    "## 1. Installing NGC Registry CLI\n",
    "You will need to install the [NGC registry CLI](https://docs.nvidia.com/ngc/ngc-overview/index.html#installing-ngc-registry-cli) to download SF Bilingual Speech Dataset. In general, you will need to,\n",
    "1. Log in to your enterprise account on the NGC website (https://ngc.nvidia.com).\n",
    "2. In the top right corner, click your user account icon and select **Setup**, then click **Downloads** under **CLI** from the Setup page.\n",
    "3. From the CLI Install page, click the **Windows**, **Linux**, or **macOS** tab, according to the platform from which you will be running NGC Registry CLI.\n",
    "4. Follow the instructions to install the CLI.\n",
    "5. Verify the installation by entering `ngc --version`. The output should be \"`NGC CLI x.y.z`\" where x.y.z indicates the version."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ae6f9d3",
   "metadata": {},
   "source": [
    "Below bash script demonstrate basic steps of NGC Registry CLI installation,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a70e841e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# https://ngc.nvidia.com/setup/installers/cli\n",
    "rm -rf ngc-cli\n",
    "wget --content-disposition \"https://ngc.nvidia.com/downloads/ngccli_linux.zip\" && unzip ngccli_linux.zip && chmod u+x ngc-cli/ngc\n",
    "find ngc-cli/ -type f -exec md5sum {} + | LC_ALL=C sort | md5sum -c ngc-cli.md5\n",
    "rm ngccli_linux.zip ngc-cli.md5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b11bd311",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"PATH\"] = f\"{os.getcwd()}/ngc-cli:{os.getenv('PATH', '')}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "305f7c0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo $PATH"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe4feb11",
   "metadata": {},
   "source": [
    "### Note: You must configure NGC CLI for your use by running `ngc config set` so that you can run the commands."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab0e2df7",
   "metadata": {},
   "source": [
    "Here is the example of configuring NGC CLI,\n",
    "``` bash\n",
    "$ ngc config set\n",
    "Enter API key [no-apikey]. Choices: [<VALID_APIKEY>, 'no-apikey']: <paste_your_API_key_here>\n",
    "Enter CLI output format type [ascii]. Choices: [ascii, csv, json]:\n",
    "Enter org [no-org]. Choices: ['<your_user_hashcode>']: <paste_your_user_hashcode_here> or leave it empty.\n",
    "Enter team [no-team]. Choices: ['no-team']:\n",
    "Enter ace [no-ace]. Choices: ['no-ace']:\n",
    "Successfully saved NGC configuration to /root/.ngc/config\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da7d6856",
   "metadata": {},
   "source": [
    "## 2. Downloading SFSpeech Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f11ba1e5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cd NeMoChineseTTS && mkdir DataChinese && \\\n",
    "    cd DataChinese && \\\n",
    "    ngc registry resource download-version \"nvidia/sf_bilingual_speech_zh_en:v1\" && \\\n",
    "    cd sf_bilingual_speech_zh_en_vv1 && \\\n",
    "    unzip SF_bilingual.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df3fb1ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# DataChineseTTS directory looks like\n",
    "!ls NeMoChineseTTS/DataChinese/sf_bilingual_speech_zh_en_vv1/SF_bilingual"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cbbe658",
   "metadata": {},
   "source": [
    "## 3. Creating Data Manifests \n",
    "\n",
    "We've created `scripts/dataset_processing/tts/sfbilingual/get_data.py` script that reads the `DataChinese/SF_bilingual/text_SF.txt` provided with the dataset and generates the following fields per 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",
    "Please refer to [sfspeech-chinese-english-bilingual-speech](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/tts/datasets.html#sfspeech-chinese-english-bilingual-speech) for more details about the SFSpeech dataset. \n",
    "\n",
    "You can run the below command to obtain the final manifests, `train_manifest.json`, `val_manifest.json` and `test_manifest.json`. This command splits 1% datapoints to validation set, 1% to test set, and the remaining 98% to training set. **Note** that this script would take sometime to process and normalize the entire dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15a8cdab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cd NeMoChineseTTS && python get_data.py \\\n",
    "        --data-root ./DataChinese/sf_bilingual_speech_zh_en_vv1/SF_bilingual/ \\\n",
    "        --manifests-path ./ \\\n",
    "        --val-size 0.005 \\\n",
    "        --test-size 0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2996e62b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# generated JSON manifests\n",
    "!ls NeMoChineseTTS/*.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40d575fd",
   "metadata": {},
   "source": [
    "## 4. Phonemization\n",
    "\n",
    "The pronunciation of a Chinese sentence can be represented as a string of phones. We would first convert a sentence into a pinyin sequences by using pypinyin library. Then we use a pre-defined pinyin-to-phoneme dict to convert them into phonemes. For English words in the sentences, we would directly use letters as input units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f58abf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"text: 我今天去了Apple Store, 买了一个iPhone。\")\n",
    "print(\"pinyin: 'wo3', 'jin1', 'tian1', 'qu4', 'le5', 'A', 'p', 'p', 'l', 'e', ' ', 'S', 't', 'o', 'r', 'e', ',', ' ', 'mai3', 'le5', 'yi2', 'ge4', 'i', 'P', 'h', 'o', 'n', 'e', '。'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fee48fa8",
   "metadata": {},
   "source": [
    "The original JSON dataset split generated from `get_data.py` only contains text/grapheme inputs. We recommend using phonemes as well to obtain better quality of synthesized audios. The tutorial uses Chinese phonemes and English letters as modeling unit by default."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93cb9b3a",
   "metadata": {},
   "source": [
    "## 5. 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.\n",
    "\n",
    "The configuration remains the same as described in `scripts/dataset_processing/tts/sfbilingual/ds_conf/ds_for_fastpitch_align.yaml`, except that `phoneme_dict_path` should point to `pinyin_dict_nv_22.10.txt` in this tutorial. 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": "code",
   "execution_count": null,
   "id": "039c7d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd NeMoChineseTTS && python extract_sup_data.py \\\n",
    "        --config-path . \\\n",
    "        --config-name ds_for_fastpitch_align.yaml \\\n",
    "        manifest_filepath=train_manifest.json \\\n",
    "        sup_data_path=sup_data \\\n",
    "        phoneme_dict_path=pinyin_dict_nv_22.10.txt \\\n",
    "        ++dataloader_params.num_workers=4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdf26d61",
   "metadata": {},
   "source": [
    "After running the above command line, you will observe a new folder `NeMoChineseTTS/sup_data/pitch` and printouts of pitch statistics like below. Specify these values to the FastPitch training configurations. We will be there in the following section.\n",
    "```bash\n",
    "PITCH_MEAN=226.75924682617188, PITCH_STD=58.773109436035156\n",
    "PITCH_MIN=65.4063949584961, PITCH_MAX=1986.977294921875\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e260190a",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6319a83d",
   "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/zh/fastpitch_align_22050.yaml`, except that in this tutorial,\n",
    "1. `phoneme_dict_path` should point to `pinyin_dict_nv_22.10.txt`;\n",
    "2. `pitch_mean` and `pitch_std` should be updated with the values estimated by the above `extract_sup_data.py` script.\n",
    "\n",
    "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": "ae8a960f",
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "!wandb login #paste_wandb_apikey_here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35f2f667",
   "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": "af763ead",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cd NeMoChineseTTS && CUDA_VISIBLE_DEVICES=0 python fastpitch.py --config-path . --config-name fastpitch_align_22050 \\\n",
    "    model.train_ds.dataloader_params.batch_size=32 \\\n",
    "    model.validation_ds.dataloader_params.batch_size=32 \\\n",
    "    train_dataset=train_manifest.json \\\n",
    "    validation_datasets=val_manifest.json \\\n",
    "    sup_data_path=sup_data \\\n",
    "    exp_manager.exp_dir=resultChineseTTS \\\n",
    "    trainer.max_epochs=1 \\\n",
    "    trainer.check_val_every_n_epoch=1 \\\n",
    "    pitch_mean=226.75924682617188 \\\n",
    "    pitch_std=58.773109436035156 \\\n",
    "    phoneme_dict_path=pinyin_dict_nv_22.10.txt \\\n",
    "    +exp_manager.create_wandb_logger=true \\\n",
    "    +exp_manager.wandb_logger_kwargs.name=\"tutorial\" \\\n",
    "    +exp_manager.wandb_logger_kwargs.project=\"ChineseTTS\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d478a420",
   "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",
    "**Note**: 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": "bb9375de",
   "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": "f28ecb25",
   "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": "93df3197",
   "metadata": {},
   "outputs": [],
   "source": [
    "test = \"GOOGLE也計畫讓社交網路技術成為ANDROID未來版本的要項。\"\n",
    "test_id = \"com_SF_ce73\"\n",
    "data_path = \"NeMoChineseTTS/DataChinese/sf_bilingual_speech_zh_en_vv1/SF_bilingual/wavs/\" # path to dataset folder with wav files from original dataset\n",
    "seed = 1234"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53434f70",
   "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": "77a88aa9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# load fastpitch and hifigan models\n",
    "import glob, os\n",
    "fastpitch_model_path = sorted(\n",
    "    glob.glob(\"NeMoChineseTTS/resultChineseTTS/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\" # pretrained hifigan from https://api.ngc.nvidia.com/v2/models/nvidia/nemo/tts_en_lj_hifigan/versions/1.6.0/files/tts_en_lj_hifigan_ft_mixerttsx.nemo\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": "86f7af16",
   "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(filename=data_path+test_id+'.wav', rate=22050))\n",
    "print(\"predicted audio\")\n",
    "ipd.display(ipd.Audio(audio, rate=22050))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c5e1d3c",
   "metadata": {},
   "source": [
    "You would hear that the above synthesized audio quality is pretty bad. It would be improved after continuing to train 1500 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). Lets try that out next!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7922073",
   "metadata": {},
   "source": [
    "# Finetuning HiFi-GAN\n",
    "\n",
    "Improving speech quality by Finetuning HiFi-GAN on synthesized mel-spectrograms from FastPitch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47584213",
   "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": "8ae3ee1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_audio_filepath = \"NeMoChineseTTS/DataChinese/sf_bilingual_speech_zh_en_vv1/SF_bilingual/wavs/com_SF_ce1.wav\"\n",
    "test_audio_text = \"NTHU對面有一條宵夜街。\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9278ac6d",
   "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": "de035be0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "spec_model = FastPitchModel.restore_from(fastpitch_model_path).eval().cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4310fe9",
   "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": "b9ee65a1",
   "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 = np.log(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": "432bd949",
   "metadata": {},
   "source": [
    "### 2. Mel spectrogram predicted from FastPitch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de034fb9",
   "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": "84a6d9f8",
   "metadata": {},
   "source": [
    "**Note**: The above spectrogram has the duration of 247 frames which is not equal to the ground truth 407 frames. In order to finetune HiFi-GAN we need mel spectrogram predicted from FastPitch with groundtruth alignment and duration.\n",
    "\n",
    "### 2.1 Mel spectrogram predicted from FastPitch with groundtruth alignment and duration "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d27682d9",
   "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": "0271d26f",
   "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.json`, `test_manifest.json` and `val_manifest.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 `NeMoChineseTTS/mels` and also create new JSON manifests with a suffix `_mel` by adding a new key `\"mel_filepath\"`. For example, `train_manifest.json` corresponds to `train_manifest_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": "37dd4579",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python NeMoChineseTTS/generate_mels.py \\\n",
    "    --cpu \\\n",
    "    --fastpitch-model-ckpt {fastpitch_model_path} \\\n",
    "    --input-json-manifests NeMoChineseTTS/train_manifest.json NeMoChineseTTS/val_manifest.json NeMoChineseTTS/test_manifest.json \\\n",
    "    --output-json-manifest-root NeMoChineseTTS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "729cbe29",
   "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": "3250e6b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example HiFi-GAN manifest:\n",
    "!head -n1 NeMoChineseTTS/train_manifest_mel.json | jq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee7c8af0",
   "metadata": {},
   "source": [
    "## Launch finetuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bbeb413",
   "metadata": {},
   "source": [
    "We will be re-using the existing HiFi-GAN config and HiFi-GAN pretrained on English."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f038d3e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!CUDA_VISIBLE_DEVICES=0 python NeMoChineseTTS/hifigan_finetune.py --config-path . --config-name hifigan.yaml \\\n",
    "    trainer.max_steps=10 \\\n",
    "    model.optim.lr=0.00001 \\\n",
    "    ~model.optim.sched \\\n",
    "    train_dataset=NeMoChineseTTS/train_manifest_mel.json \\\n",
    "    validation_datasets=NeMoChineseTTS/val_manifest_mel.json \\\n",
    "    exp_manager.exp_dir=NeMoChineseTTS/resultChineseTTS \\\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=\"ChineseTTS\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e660ce7",
   "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": "b4fc29f3",
   "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": "a805e88f",
   "metadata": {},
   "outputs": [],
   "source": [
    "hfg_path = sorted(glob.glob(\"NeMoChineseTTS/resultChineseTTS/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": "cdba9330",
   "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": "6b127508",
   "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"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}