tts-finetune-nemo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Text to Speech Finetuning using NeMo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NeMo Toolkit  is a python based AI toolkit for training and customizing purpose-built pre-trained AI models with your own data.\n",
    "\n",
    "Transfer learning extracts learned features from an existing neural network to a new one. Transfer learning is often used when creating a large training dataset is not feasible.\n",
    "\n",
    "Developers, researchers and software partners building intelligent AI apps and services, can bring their own data to fine-tune pre-trained models instead of going through the hassle of training from scratch."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see this in action with a use case for Speech Synthesis!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text to Speech"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Text to Speech (TTS) is often the last step in building a Conversational AI model. A TTS model converts text into audible speech. The main objective is to synthesize reasonable and natural speech for given text. Since there are no universal standard to measure quality of synthesized speech, you will need to listen to some inferred speech to tell whether a TTS model is well trained.\n",
    "\n",
    "In this tutorial we will look at two models: [FastPitch](https://arxiv.org/pdf/2006.06873.pdf) for spectrogram generation and [HiFiGAN](https://arxiv.org/pdf/2010.05646.pdf) as vocoder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## Let's Dig in: TTS using NeMo\n",
    "\n",
    "This notebook assumes that you are already familiar with TTS Training using NeMo, as described in the [text-to-speech-training](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tts/FastPitch_MixerTTS_Training.ipynb) notebook, and that you have a pretrained TTS model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After [installing NeMo](https://github.com/NVIDIA/NeMo#installation), the next step is to setup the paths to save data and results. NeMo can be used with docker containers or virtual environments.\n",
    "\n",
    "Replace the variables FIXME with the required paths enclosed in \"\" as a string.\n",
    "\n",
    "`IMPORTANT NOTE:` Here, we map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container. Make sure this tutorial is in the NeMo folder."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installation of packages and importing of files\n",
    "\n",
    "We will first install all necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install numba>=0.53\n",
    "! pip install librosa\n",
    "! pip install soundfile\n",
    "! pip install tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the following packages only if you want to export your models to `.riva` format else you can skip it.\n",
    "We will now install the packages NeMo and nemo2riva. nemo2riva is available on [ngc](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/riva/resources/riva_quickstart/files?version=2.8.1). Make sure you install NGC CLI first before running the following commands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install nvidia-pyindex\n",
    "!pip install nemo_toolkit['all']==1.23.0\n",
    "!ngc registry resource download-version \"nvidia/riva/riva_quickstart:2.8.1\"\n",
    "!pip install \"riva_quickstart_v2.8.1/nemo2riva-2.8.1-py3-none-any.whl\"\n",
    "!pip install protobuf==3.20.0\n",
    "# Installing pynini separately\n",
    "!wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/nemo_text_processing/install_pynini.sh \\\n",
    "bash install_pynini.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now import all the relevant files from NeMo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/dataset_processing/tts/ljspeech/get_data.py\n",
    "    \n",
    "! wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/dataset_processing/tts/extract_sup_data.py\n",
    "! mkdir -p ljspeech && cd ljspeech \\\n",
    "&& wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/dataset_processing/tts/ljspeech/ds_conf/ds_for_fastpitch_align.yaml \\\n",
    "&& cd ..\n",
    "    \n",
    "# additional files\n",
    "!wget https://raw.githubusercontent.com/nvidia/NeMo/main/examples/tts/fastpitch_finetune.py\n",
    "\n",
    "!mkdir -p conf \\\n",
    "&& cd conf \\\n",
    "&& wget https://raw.githubusercontent.com/nvidia/NeMo/main/examples/tts/conf/fastpitch_align_v1.05.yaml \\\n",
    "&& cd ..\n",
    "\n",
    "!mkdir -p tts_dataset_files && cd tts_dataset_files \\\n",
    "&& wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/tts_dataset_files/cmudict-0.7b_nv22.10 \\\n",
    "&& wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/tts_dataset_files/heteronyms-052722 \\\n",
    "&& wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/nemo_text_processing/text_normalization/en/data/whitelist/tts.tsv \\\n",
    "&& cd ..\n",
    "            \n",
    "! wget https://raw.githubusercontent.com/NVIDIA/NeMo/main/scripts/dataset_processing/tts/generate_mels.py\n",
    "    \n",
    "! wget https://raw.githubusercontent.com/nvidia/NeMo/main/examples/tts/hifigan_finetune.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set Relevant Paths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the NeMo Docker.\n",
    "\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "# The data is saved here\n",
    "DATA_DIR = FIXME\n",
    "RESULTS_DIR = FIXME\n",
    "\n",
    "! mkdir -p {DATA_DIR}\n",
    "! mkdir -p {RESULTS_DIR}\n",
    "\n",
    "os.environ[\"DATA_DIR\"] = DATA_DIR\n",
    "os.environ[\"RESULTS_DIR\"] = RESULTS_DIR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data\n",
    "\n",
    "For the rest of this notebook, it is assumed that you have:\n",
    "\n",
    " - Pretrained FastPitch and HiFiGAN models that were trained on LJSpeech sampled at 22kHz\n",
    " \n",
    "In the case that you are not using a TTS model trained on LJSpeech at the correct sampling rate. Please ensure that you have the original data, including wav files and a .json manifest file. If you have a TTS model but not at 22kHz, please ensure that you set the correct sampling rate, and fft parameters.\n",
    "\n",
    "For the rest of this notebook, we will be using a subset of audio samples from the [Hi-Fi TTS dataset](https://vctk-subset.s3.amazonaws.com/vctk_subset.tar.gz) adding up to about one minute. This dataset is for demo purposes only. For a good quality model, we recommend at least 30 minutes of audio. If you want to record your own dataset, you can follow the [Guidelines to Record a TTS Dataset at Home](https://github.com/nvidia-riva/tutorials/blob/main/tts-dataset-recording-at-home.md). Sample scripts to download and preprocess datasets supported by NeMo can be found [here](https://github.com/NVIDIA/NeMo/tree/main/scripts/dataset_processing/tts).\n",
    "\n",
    "Let's first download and pre-process the original LJSpeech dataset and set variables that point to this as the original data's `.json` file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pre-Processing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This step downloads audio to text file lists from NVIDIA for LJSpeech and generates the manifest files. If you use your own dataset, you have to generate three files: `ljs_audio_text_train_manifest.json`, `ljs_audio_text_val_manifest.json`, `ljs_audio_text_test_manifest.json` yourself. Those files correspond to your train / val / test split. For each text file, the number of rows should be equal to number of samples in this split and each row for a single speaker dataset should be like:\n",
    "\n",
    "```\n",
    "{\"audio_filepath\": \"path_to_audio_file\", \"text\": \"text_of_the_audio\", \"duration\": duration_of_the_audio}\n",
    "```\n",
    "\n",
    "In case of multi-speaker dataset\n",
    "\n",
    "```\n",
    "{\"audio_filepath\": \"path_to_audio_file\", \"text\": \"text_of_the_audio\", \"duration\": duration_of_the_audio, \"speaker\": speaker_id}\n",
    "```\n",
    "\n",
    "An example row is:\n",
    "\n",
    "```\n",
    "{\"audio_filepath\": \"actressinhighlife_01_bowen_0001.flac\", \"text\": \"the pleasant season did my heart employ\", \"duration\": 2.4}\n",
    "```\n",
    "\n",
    "We will now download the audio and the manifest files then convert them to the above format, also normalize the text. These steps for LJSpeech can be found in NeMo [`scripts/dataset_processing/tts/ljspeech/get_data.py`](https://github.com/NVIDIA/NeMo/blob/main/scripts/dataset_processing/tts/ljspeech/get_data.py). Be patient, this step is expected to take some time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! python get_data.py --data-root {DATA_DIR}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "original_data_json = os.path.join(os.environ[\"DATA_DIR\"], \"LJSpeech-1.1/train_manifest.json\")\n",
    "os.environ[\"original_data_json\"] = original_data_json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now download the Hi-Fi TTS audio samples, and place the data in the `DATA_DIR`. Create a manifest file named `manifest.json` and copy the contents of both `dev.json` and `train.json` into it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Name of the untarred Hi-Fi TTS audio samples directory.\n",
    "finetune_data_name = FIX_ME\n",
    "# Absolute path of finetuning dataset from the perspective of NeMo container\n",
    "finetune_data_path = os.path.join(os.environ[\"DATA_DIR\"], finetune_data_name)\n",
    "\n",
    "os.environ[\"finetune_data_name\"] = finetune_data_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that you have downloaded the data, let's make sure that the audio clips and sample at the same sampling frequency as the clips used to train the pretrained model. For the course of this notebook, NVIDIA recommends using a model trained on the LJSpeech dataset. The sampling rate for this model is 22.05kHz."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import soundfile\n",
    "import librosa\n",
    "import json\n",
    "import os\n",
    "\n",
    "def resample_audio(input_file_path, output_path, target_sampling_rate=22050):\n",
    "    \"\"\"Resample a single audio file.\n",
    "    \n",
    "    Args:\n",
    "        input_file_path (str): Path to the input audio file.\n",
    "        output_path (str): Path to the output audio file.\n",
    "        target_sampling_rate (int): Sampling rate for output audio file.\n",
    "        \n",
    "    Returns:\n",
    "        No explicit returns\n",
    "    \"\"\"\n",
    "    if not input_file_path.endswith(\".wav\"):\n",
    "        raise NotImplementedError(\"Loading only implemented for wav files.\")\n",
    "    if not os.path.exists(input_file_path):\n",
    "        raise FileNotFoundError(f\"Cannot file input file at {input_file_path}\")\n",
    "    audio, sampling_rate = librosa.load(\n",
    "      input_file_path,\n",
    "      sr=target_sampling_rate\n",
    "    )\n",
    "    # Filterring out empty audio files.\n",
    "    if librosa.get_duration(y=audio, sr=sampling_rate) == 0:\n",
    "        print(f\"0 duration audio file encountered at {input_file_path}\")\n",
    "        return None\n",
    "    filename = os.path.basename(input_file_path)\n",
    "    if not os.path.exists(output_path):\n",
    "        os.makedirs(output_path)\n",
    "    soundfile.write(\n",
    "        os.path.join(output_path, filename),\n",
    "        audio,\n",
    "        samplerate=target_sampling_rate,\n",
    "        format=\"wav\"\n",
    "    )\n",
    "    return filename"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm.notebook import tqdm\n",
    "\n",
    "relative_path = f\"{finetune_data_name}/clips_resampled\"\n",
    "resampled_manifest_file = os.path.join(\n",
    "    os.environ[\"DATA_DIR\"],\n",
    "    f\"{finetune_data_name}/manifest_resampled.json\"\n",
    ")\n",
    "input_manifest_file = os.path.join(\n",
    "    os.environ[\"DATA_DIR\"],\n",
    "    f\"{finetune_data_name}/manifest.json\"\n",
    ")\n",
    "sampling_rate = 22050\n",
    "output_path = os.path.join(os.environ[\"DATA_DIR\"], relative_path)\n",
    "\n",
    "# Resampling the audio clip.\n",
    "with open(input_manifest_file, \"r\") as finetune_file:\n",
    "    with open(resampled_manifest_file, \"w\") as resampled_file:\n",
    "        for line in tqdm(finetune_file.readlines()):\n",
    "            data = json.loads(line)\n",
    "            filename = resample_audio(\n",
    "                os.path.join(\n",
    "                    os.environ[\"DATA_DIR\"],\n",
    "                    finetune_data_name,\n",
    "                    data[\"audio_filepath\"]\n",
    "                ),\n",
    "                output_path,\n",
    "                target_sampling_rate=sampling_rate\n",
    "            )\n",
    "            if not filename:\n",
    "                print(\"Skipping clip {} from training dataset\")\n",
    "                continue\n",
    "            data[\"audio_filepath\"] = os.path.join(\n",
    "                os.environ[\"DATA_DIR\"],\n",
    "                relative_path, filename\n",
    "            )\n",
    "            resampled_file.write(f\"{json.dumps(data)}\\n\")\n",
    "\n",
    "assert resampled_file.closed, \"Output file wasn't closed properly\"\n",
    "assert finetune_file.closed, \"Input file wasn't closed properly\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Splitting the dataset to train and val set.\n",
    "! cat $finetune_data_path/manifest_resampled.json | tail -n 2 > $finetune_data_path/manifest_val.json\n",
    "! cat $finetune_data_path/manifest_resampled.json | head -n -2 > $finetune_data_path/manifest_train.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "finetune_data_json = os.path.join(os.environ[\"DATA_DIR\"], f'{finetune_data_name}/manifest_train.json')\n",
    "os.environ[\"finetune_data_json\"] = finetune_data_json\n",
    "os.environ[\"finetune_val_data_json\"] = os.path.join(os.environ[\"DATA_DIR\"], f'{finetune_data_name}/manifest_val.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first step is to create a json that contains data from both the original data and the finetuning data. Since, the original data is much larger than the finetuning data, we merge the finetuning data with a sample of the original data. We can do this using the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import json\n",
    "\n",
    "def json_reader(filename):\n",
    "    with open(filename) as f:\n",
    "        for line in f:\n",
    "            yield json.loads(line)\n",
    "            \n",
    "            \n",
    "def json_writer(file, json_objects):\n",
    "    with open(file, \"w\") as f:\n",
    "        for jsonobj in json_objects:\n",
    "            jsonstr = json.dumps(jsonobj)\n",
    "            f.write(jsonstr + \"\\n\")\n",
    "            \n",
    "            \n",
    "def dataset_merge(original_manifest, finetune_manifest, num_records_original=50):\n",
    "    original_ds = list(json_reader(original_manifest))\n",
    "    finetune_ds = list(json_reader(finetune_manifest))\n",
    "    original_ds = random.sample(original_ds, num_records_original)\n",
    "    merged_ds = original_ds + finetune_ds\n",
    "    random.shuffle(merged_ds)\n",
    "    return merged_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "merged_ds = dataset_merge(os.environ[\"original_data_json\"], os.environ[\"finetune_data_json\"])\n",
    "\n",
    "os.environ[\"merged_data_json\"] = f\"{DATA_DIR}/{finetune_data_name}/merged_train.json\"\n",
    "json_writer(os.environ[\"merged_data_json\"], merged_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting Pitch Statistics\n",
    "\n",
    "Training Fastpitch requires you to set 2 values for pitch extraction:\n",
    "  - `avg`: The average used to normalize the pitch\n",
    "  - `std`: The std deviation used to normalize the pitch\n",
    "\n",
    "We can compute pitch for the training data using [`scripts/dataset_processing/tts/extract_sup_data.py`](https://github.com/NVIDIA/NeMo/blob/main/scripts/dataset_processing/tts/extract_sup_data.py) and extract pitch statistics using the NeMo script [`scripts/dataset_processing/tts/compute_speaker_stats.py`](https://github.com/NVIDIA/NeMo/blob/main/scripts/dataset_processing/tts/compute_speaker_stats.py), We have already downloaded the files earlier in the tutorial. Let's use it to get `pitch_mean` and `pitch_std`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we will extract the pitch supplementary data using `extract_sup_data.py` file. This file works with a yaml config file `ds_for_fastpitch_align`, which we downloaded above. To make this work for your dataset simply change the `manifest_path` to your manifest path. The argument `sup_data_path` determines where the supplementary data is stored."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sup_data_path = f'{finetune_data_path}/sup_data_path'\n",
    "pitch_stats_path = f'{finetune_data_path}/pitch_stats.json'\n",
    "\n",
    "# The script extract_sup_data.py writes the pitch mean and pitch std in the commandline. We will parse it to get the pitch mean and std\n",
    "cmd_str_list = !python extract_sup_data.py --config-path \"ljspeech\" manifest_filepath={os.environ[\"merged_data_json\"]} sup_data_path={sup_data_path}\n",
    "# Select only the line that contains PITCH_MEAN\n",
    "cmd_str = [c for c in cmd_str_list if \"PITCH_MEAN\" in c][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract pitch mean and std from the commandline\n",
    "pitch_mean_str = cmd_str.split(',')[0]\n",
    "pitch_mean = float(pitch_mean_str.split('=')[1])\n",
    "pitch_std_str = cmd_str.split(',')[1]\n",
    "pitch_std = float(pitch_std_str.split('=')[1])\n",
    "pitch_mean, pitch_std"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setting the `pitch_fmean` and `pitch_fmax` based on the results from the cell above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"pitch_mean\"] = str(pitch_mean)\n",
    "os.environ[\"pitch_std\"] = str(pitch_std)\n",
    "\n",
    "print(f\"pitch mean: {pitch_mean}\")\n",
    "print(f\"pitch std: {pitch_std}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finetuning\n",
    "\n",
    "We are now ready to finetune our TTS pipeline. In order to do so, you need to finetune FastPitch. For best results, you need to finetune HiFiGAN as well.\n",
    "\n",
    "Here we are using pretrained checkpoints from NGC, [FastPitch](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/tts_en_fastpitch) and [HiFiGAN](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/tts_hifigan)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Finetuning FastPitch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will need some additional files from NeMo to run finetuning on FastPitch, we have downloaded them earlier in the tutorial. In NeMo you can find the `fastpitch_finetuning.py` script and the config in [`examples`](https://github.com/NVIDIA/NeMo/tree/main/examples/tts) section."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!(python fastpitch_finetune.py --config-name=fastpitch_align_v1.05.yaml \\\n",
    "  train_dataset={os.environ[\"merged_data_json\"]} \\\n",
    "  validation_datasets={os.environ[\"finetune_val_data_json\"]} \\\n",
    "  sup_data_path={sup_data_path} \\\n",
    "  phoneme_dict_path=tts_dataset_files/cmudict-0.7b_nv22.10 \\\n",
    "  heteronyms_path=tts_dataset_files/heteronyms-052722 \\\n",
    "  whitelist_path=tts_dataset_files/tts.tsv \\\n",
    "  exp_manager.exp_dir={os.environ[\"RESULTS_DIR\"]} \\\n",
    "  +init_from_pretrained_model=\"tts_en_fastpitch\" \\\n",
    "  +trainer.max_steps=1000 \\\n",
    "  ~trainer.max_epochs \\\n",
    "  trainer.check_val_every_n_epoch=10 \\\n",
    "  model.train_ds.dataloader_params.batch_size=24 \\\n",
    "  model.validation_ds.dataloader_params.batch_size=24 \\\n",
    "  model.n_speakers=1 \\\n",
    "  model.pitch_mean={os.environ[\"pitch_mean\"]} model.pitch_std={os.environ[\"pitch_std\"]} \\\n",
    "  model.optim.lr=2e-4 \\\n",
    "  ~model.optim.sched \\\n",
    "  model.optim.name=adam \\\n",
    "  trainer.devices=1 \\\n",
    "  trainer.strategy=null \\\n",
    "  +model.text_tokenizer.add_blank_at=true \\\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a closer look at the training command:\n",
    "\n",
    "* `--config-name=fastpitch_align_v1.05.yaml`\n",
    "  * We first tell the script what config file to use.\n",
    "\n",
    "* `train_dataset=./9017_manifest_train_dur_5_mins_local.json \n",
    "  validation_datasets=./9017_manifest_dev_ns_all_local.json \n",
    "  sup_data_path=./fastpitch_sup_data`\n",
    "  * We tell the script what manifest files to train and eval on, as well as where supplementary data is located (or will be calculated and saved during training if not provided).\n",
    "  \n",
    "* `phoneme_dict_path=tts_dataset_files/cmudict-0.7b_nv22.10 \n",
    "heteronyms_path=tts_dataset_files/heteronyms-052722\n",
    "whitelist_path=tts_dataset_files/tts.tsv \n",
    "`\n",
    "  * We tell the script where `phoneme_dict_path`, `heteronyms-052722` and `whitelist_path` are located. These are the additional files we downloaded earlier, and are used in preprocessing the data.\n",
    "  \n",
    "* `exp_manager.exp_dir=./ljspeech_to_9017_no_mixing_5_mins`\n",
    "  * Where we want to save our log files, tensorboard file, checkpoints, and more.\n",
    "\n",
    "* `+init_from_nemo_model=./tts_en_fastpitch_align.nemo`\n",
    "  * We tell the script what checkpoint to finetune from.\n",
    "\n",
    "* `+trainer.max_steps=1000 ~trainer.max_epochs trainer.check_val_every_n_epoch=25`\n",
    "  * For this experiment, we tell the script to train for 1000 training steps/iterations rather than specifying a number of epochs to run. Since the config file specifies `max_epochs` instead, we need to remove that using `~trainer.max_epochs`.\n",
    "\n",
    "* `model.train_ds.dataloader_params.batch_size=24 model.validation_ds.dataloader_params.batch_size=24`\n",
    "  * Set batch sizes for the training and validation data loaders.\n",
    "\n",
    "* `model.n_speakers=1`\n",
    "  * The number of speakers in the data. There is only 1 for now, but we will revisit this parameter later in the notebook.\n",
    "\n",
    "* `model.pitch_mean=152.3 model.pitch_std=64.0 model.pitch_fmin=30 model.pitch_fmax=512`\n",
    "  * For the new speaker, we need to define new pitch hyperparameters for better audio quality.\n",
    "  * These parameters work for speaker 9017 from the Hi-Fi TTS dataset.\n",
    "  * If you are using a custom dataset, running the script `python <NeMo_base>/scripts/dataset_processing/tts/extract_sup_data.py manifest_filepath=<your_manifest_path>` will precalculate supplementary data and print these pitch stats.\n",
    "  * fmin and fmax are hyperparameters to librosa's pyin function. We recommend tweaking these only if the speaker is in a noisy environment, such that background noise isn't predicted to be speech.\n",
    "\n",
    "* `model.optim.lr=2e-4 ~model.optim.sched model.optim.name=adam`\n",
    "  * For fine-tuning, we lower the learning rate.\n",
    "  * We use a fixed learning rate of 2e-4.\n",
    "  * We switch from the lamb optimizer to the adam optimizer.\n",
    "\n",
    "* `trainer.devices=1 trainer.strategy=null`\n",
    "  * For this notebook, we default to 1 gpu which means that we do not need ddp.\n",
    "  * If you have the compute resources, feel free to scale this up to the number of free gpus you have available.\n",
    "  * Please remove the `trainer.strategy=null` section if you intend on multi-gpu training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Finetuning HiFiGAN\n",
    "\n",
    "In order to get the best audio from HiFiGAN, we need to finetune it:\n",
    "  - on the new speaker\n",
    "  - using mel spectrograms from our finetuned FastPitch Model\n",
    "\n",
    "Let's first generate mels from our FastPitch model, and save it to a new .json manifest for use with HiFiGAN. We can generate the mels using [`generate_mels.py`](https://github.com/NVIDIA/NeMo/blob/main/scripts/dataset_processing/tts/generate_mels.py) file from NeMo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fastpitch_checkpoint = FIXME\n",
    "mel_dir = f\"{finetune_data_path}/mels\"\n",
    "! mkdir -p mel_dir\n",
    "\n",
    "!(python generate_mels.py \\\n",
    "  --fastpitch-model-ckpt {fastpitch_checkpoint} \\\n",
    "  --input-json-manifests {os.environ[\"merged_data_json\"]} \\\n",
    "  --output-json-manifest-root {mel_dir} \\\n",
    " )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Finetuning HiFiGAN\n",
    "\n",
    "Now let's finetune hifigan. Finetuning HiFiGAN can be done in NeMo using the script [examples/tts/hifigan_finetune.py](https://github.com/NVIDIA/NeMo/blob/main/examples/tts/hifigan_finetune.py) and configs present in [examples/tts/conf/hifigan](https://github.com/NVIDIA/NeMo/tree/main/examples/tts/conf/hifigan).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a small validation dataset for HiFiGAN finetuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hifigan_full_ds = f\"{finetune_data_path}/mels/merged_full_mel.json\"\n",
    "hifigan_train_ds = f\"{finetune_data_path}/mels/merged_train_mel.json\"\n",
    "hifigan_val_ds = f\"{finetune_data_path}/mels/merged_val_mel.json\"\n",
    "\n",
    "! cat {hifigan_train_ds} > {hifigan_full_ds}\n",
    "! cat {hifigan_full_ds} | tail -n 2 > {hifigan_val_ds}\n",
    "! cat {hifigan_full_ds} | head -n -2 > {hifigan_train_ds}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the following command for HiFiGAN finetuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!(python examples/tts/hifigan_finetune.py \\\n",
    "--config-name=hifigan.yaml \\\n",
    "model.train_ds.dataloader_params.batch_size=32 \\\n",
    "model.max_steps=1000 \\\n",
    "model.optim.lr=0.00001 \\\n",
    "~model.optim.sched \\\n",
    "train_dataset={hifigan_train_ds} \\\n",
    "validation_datasets={hifigan_val_ds} \\\n",
    "exp_manager.exp_dir={os.environ[\"RESULTS_DIR\"]} \\\n",
    "+init_from_pretrained_model=tts_hifigan \\\n",
    "trainer.check_val_every_n_epoch=10 \\\n",
    "model/train_ds=train_ds_finetune \\\n",
    "model/validation_ds=val_ds_finetune)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TTS Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As aforementioned, since there are no universal standard to measure quality of synthesized speech, you will need to listen to some inferred speech to tell whether a TTS model is well trained. Therefore, we do not provide `evaluate` functionality in NeMo Toolkit for TTS but only provide `infer` functionality."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generate spectrogram and audio\n",
    "\n",
    "The first step for inference is generating spectrogram. That's a numpy array (saved as `.npy` file) for a sentence which can be converted to voice by a vocoder. We use FastPitch we just trained to generate spectrogram\n",
    "\n",
    "Please update the `hifigan_checkpoint` variable with the path to the HiFiGAN checkpoint you want to use.\n",
    "\n",
    "Let's load the two models, FastPitch and HiFiGAN, for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.tts.models import FastPitchModel, HifiGanModel\n",
    "\n",
    "hifigan_checkpoint = FIXME\n",
    "vocoder = HifiGanModel.load_from_checkpoint(hifigan_checkpoint)\n",
    "vocoder = vocoder.eval().cuda()\n",
    "spec_model = FastPitchModel.load_from_checkpoint(fastpitch_checkpoint)\n",
    "spec_model.eval().cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a helper method to do inference given a string input. In case of multi-speaker inference the same method can\n",
    "be used by passing the speaker ID as a parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def infer(spec_gen_model, vocoder_model, str_input, speaker=None):\n",
    "    \"\"\"\n",
    "    Synthesizes spectrogram and audio from a text string given a spectrogram synthesis and vocoder model.\n",
    "    \n",
    "    Args:\n",
    "        spec_gen_model: Spectrogram generator model (FastPitch in our case)\n",
    "        vocoder_model: Vocoder model (HiFiGAN in our case)\n",
    "        str_input: Text input for the synthesis\n",
    "        speaker: Speaker ID\n",
    "    \n",
    "    Returns:\n",
    "        spectrogram and waveform of the synthesized audio.\n",
    "    \"\"\"\n",
    "    with torch.no_grad():\n",
    "        parsed = spec_gen_model.parse(str_input)\n",
    "        if speaker is not None:\n",
    "            speaker = torch.tensor([speaker]).long().to(device=spec_gen_model.device)\n",
    "        spectrogram = spec_gen_model.generate_spectrogram(tokens=parsed, speaker=speaker)\n",
    "        audio = vocoder_model.convert_spectrogram_to_audio(spec=spectrogram)\n",
    "        \n",
    "    if spectrogram is not None:\n",
    "        if isinstance(spectrogram, torch.Tensor):\n",
    "            spectrogram = spectrogram.to('cpu').numpy()\n",
    "        if len(spectrogram.shape) == 3:\n",
    "            spectrogram = spectrogram[0]\n",
    "    if isinstance(audio, torch.Tensor):\n",
    "        audio = audio.to('cpu').numpy()\n",
    "    return spectrogram, audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "# Path to test manifest file (.json)\n",
    "test_records_path = FIXME\n",
    "test_records = list(json_reader(test_records_path))\n",
    "new_speaker_id = FIXME\n",
    "\n",
    "for test_record in test_records:\n",
    "    print(\"Real validation audio\")\n",
    "    ipd.display(ipd.Audio(test_record['audio_filepath'], rate=22050))\n",
    "    duration_mins = test_record['duration']\n",
    "    if 'speaker' in test_record:\n",
    "        speaker_id = test_record['speaker']\n",
    "    else:\n",
    "        speaker_id = new_speaker_id\n",
    "    print(f\"SYNTHESIZED | Duration: {duration_mins} mins | Text: {test_record['text']}\")\n",
    "    spec, audio = infer(spec_model, vocoder, test_record['text'], speaker=speaker_id)\n",
    "    ipd.display(ipd.Audio(audio, rate=22050))\n",
    "    %matplotlib inline\n",
    "    imshow(spec, origin=\"lower\", aspect=\"auto\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Debug\n",
    "\n",
    "The data provided is only meant to be a sample to understand how finetuning works in NeMo. In order to generate better speech quality, we recommend recording at least 30 mins of audio, and increasing the number of finetuning steps from the current `trainer.max_steps=1000` to `trainer.max_steps=5000` for both models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TTS model export\n",
    "\n",
    "You can also export your model in a format that can deployed using Nvidia Riva, a highly performant application framework for multi-modal conversational AI services using GPUs!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Export to RIVA\n",
    "\n",
    "Executing the snippets in the cells below, allows you to generate a `.riva` model file for the spectrogram generator and vocoder models that were trained the preceding cells. These models are required to generate a complete Text-To-Speech pipeline.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Convert to riva.\n",
    "\n",
    "Convert the downloaded model to .riva format, we will use encryption key=nemotoriva. Change this while generating .riva models for production."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hifigan_nemo_file_path = FIXME\n",
    "hifigan_riva_file_path = hifigan_nemo_file_path[:-5]+\".riva\"\n",
    "fastpitch_nemo_file_path = FIXME\n",
    "fastpitch_riva_file_path = fastpitch_nemo_file_path[:-5]+\".riva\"\n",
    "\n",
    "!nemo2riva --out {fastpitch_riva_file_path} --key=nemotoriva {fastpitch_nemo_file_path}\n",
    "!nemo2riva --out {hifigan_riva_file_path} --key=nemotoriva {hifigan_nemo_file_path}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What's Next ?\n",
    "\n",
    " You could use NeMo to build custom models for your own applications, and deploy them to Nvidia Riva! To try deploying these models to RIVA, use the [tts-deploy.ipynb](https://github.com/nvidia-riva/tutorials/blob/stable/tts-deploy.ipynb) as a quick sample."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "741d73fab70d7eb29e7b56260ebaa567f0620f4d2780830ca385f600e5120e14"
  },
  "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}