tutorials/tts/Vits_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "htbJiaJjYQAD"
   },
   "source": [
    "# VITS Training\n",
    "\n",
    "This notebook is designed to provide a guide on how to train VITS as part of the TTS pipeline. It contains the following sections\n",
    "\n",
    "  1. VITS and NeMo - An introduction to the VITS model\n",
    "  2. LJSpeech - How to train VITS on LJSpeech"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wqPMTEXXYUP4"
   },
   "source": [
    "# License\n",
    "\n",
    "> Copyright 2023, NVIDIA CORPORATION & AFFILIATES. 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,
   "metadata": {
    "id": "SUkq9HAvYU7T"
   },
   "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 = 'main'\n",
    "# !python -m pip install \"git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\"\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": "markdown",
   "metadata": {
    "id": "ZivXzmq0YYLj"
   },
   "source": [
    "# VITS and NeMo\n",
    "\n",
    "VITS is a neural network that converts text characters into an audio sample. For more details on the model, please refer to Nvidia's [VITS Model Card](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/tts_en_lj_vits), or the original [paper](https://arxiv.org/abs/2106.06103).\n",
    "\n",
    "VITS like most NeMo models are defined as a LightningModule, allowing for easy training via PyTorch Lightning, and parameterized by a configuration, currently defined via a yaml file and loading using Hydra.\n",
    "\n",
    "Let's take a look using NeMo's pretrained model and how to use it to generate spectrograms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HEvdSU5WYZbj"
   },
   "outputs": [],
   "source": [
    "# Load the VITSModel\n",
    "from nemo.collections.tts.models import VitsModel\n",
    "from nemo.collections.tts.models.base import TextToWaveform\n",
    "\n",
    "# Let's see what pretrained models are available\n",
    "print(VitsModel.list_available_models())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3W8unatgYbUp"
   },
   "outputs": [],
   "source": [
    "# We can load the pre-trained model as follows\n",
    "model = VitsModel.from_pretrained(\"tts_en_lj_vits\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# VITS is a TextToWaveform\n",
    "assert isinstance(model, TextToWaveform)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TextToWaveform in NeMo have two helper functions:\n",
    "   1. ```python\n",
    "      parse(self, str_input: str),\n",
    "      ``` \n",
    "      which takes an English string and produces a token tensor\n",
    "\n",
    "\n",
    "\n",
    "   2. ```python\n",
    "      convert_text_to_waveform(self, *, tokens),\n",
    "      ```\n",
    "      which takes the token tensor and generates an audio sample\n",
    "Let's try it out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zZ90eCfdrNIf"
   },
   "source": [
    "# Training\n",
    "\n",
    "Now that we looked at the VITS model, let's see how to train a VITS Model\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7rHG-LERrPRY"
   },
   "outputs": [],
   "source": [
    "# NeMo's training scripts are stored inside the examples/ folder. Let's grab the vits.py file\n",
    "# as well as the vits.yaml file\n",
    "!wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/vits.py\n",
    "!(mkdir -p conf \\\n",
    "  && cd conf \\\n",
    "  && wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/tts/conf/vits.yaml \\\n",
    "  && cd ..)\n",
    "\n",
    "# additional files\n",
    "!mkdir -p tts_dataset_files && cd tts_dataset_files \\\n",
    "&& wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tts_dataset_files/ipa_cmudict-0.7b_nv23.01.txt \\\n",
    "&& wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tts_dataset_files/heteronyms-052722 \\\n",
    "&& cd ..\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Upv_LxBIsC51"
   },
   "source": [
    "Let's take a look at the vits.py file\n",
    "\n",
    "```python\n",
    "import lightning.pytorch as pl\n",
    "\n",
    "from nemo.collections.tts.models.vits import VitsModel\n",
    "from nemo.core.config import hydra_runner\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "# hydra_runner is a thin NeMo wrapper around Hydra\n",
    "# It looks for a config named vits.yaml inside the conf folder\n",
    "# Hydra parses the yaml and returns it as a Omegaconf DictConfig\n",
    "@hydra_runner(config_path=\"conf\", config_name=\"vits\")\n",
    "def main(cfg):\n",
    "    # Define the Lightning trainer\n",
    "    trainer = pl.Trainer(use_distributed_sampler=False, **cfg.trainer)\n",
    "    # exp_manager is a NeMo construct that helps with logging and checkpointing\n",
    "    exp_manager(trainer, cfg.get(\"exp_manager\", None))\n",
    "    # Define the VITS model, this will construct the model\n",
    "    model = VitsModel(cfg=cfg.model, trainer=trainer)\n",
    "    # Let's add a few more callbacks\n",
    "    trainer.callbacks.extend([pl.callbacks.LearningRateMonitor()])\n",
    "    # Call lightning trainer's fit() to train the model\n",
    "    trainer.fit(model)\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()  # noqa pylint: disable=no-value-for-parameter\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6nM-fZO-s75u"
   },
   "source": [
    "Let's take a look at the yaml config\n",
    "\n",
    "```yaml\n",
    "name: &name VITS\n",
    "\n",
    "train_dataset: ???\n",
    "validation_datasets: ???\n",
    "sup_data_path: null\n",
    "sup_data_types: null\n",
    "\n",
    "phoneme_dict_path: \"scripts/tts_dataset_files/ipa_cmudict-0.7b_nv23.01.txt\"\n",
    "heteronyms_path: \"scripts/tts_dataset_files/heteronyms-052722\"\n",
    "```\n",
    "\n",
    "The first part of the yaml defines dataset parameters used by VITS. Then in the head of 'model' section there are processing - related parameters. You can see\n",
    "that the sample rate is set to 22050 for LJSpeech. \n",
    "\n",
    "Looking at the yaml, there is `train_dataset: ???` and `validation_datasets: ???`. The ??? indicates to hydra that these values must be passed via the command line or the script will fail.\n",
    "\n",
    "Looking further down the yaml, we get to the pytorch lightning trainer parameters.\n",
    "\n",
    "```yaml\n",
    "trainer:\n",
    "  num_nodes: 1\n",
    "  devices: 2\n",
    "  accelerator: gpu\n",
    "  strategy: ddp_find_unused_parameters_true\n",
    "  precision: 32\n",
    "  max_epochs: -1\n",
    "  accumulate_grad_batches: 1\n",
    "  enable_checkpointing: false # Provided by exp_manager\n",
    "  logger: false # Provided by exp_manager\n",
    "  log_every_n_steps: 50\n",
    "  check_val_every_n_epoch: 1\n",
    "```\n",
    "\n",
    "These values can be changed either by editing the yaml or through the command line.\n",
    "\n",
    "Let's grab some simple audio data and test VITS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GnEzODcorugt"
   },
   "outputs": [],
   "source": [
    "!wget https://github.com/NVIDIA/NeMo/releases/download/v0.11.0/test_data.tar.gz \\\n",
    "&& mkdir -p tests/data \\\n",
    "&& tar xzf test_data.tar.gz -C tests/data\n",
    "\n",
    "# Just like ASR, the VITS require .json files to define the training and validation data.\n",
    "!cat tests/data/asr/an4_val.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have some sample data, we can try training VITS!\n",
    "\n",
    "Note that the sample data is not enough data to fully train a VITS model. The following code uses a toy dataset to illustrate how the pipeline for training would work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!(python vits.py \\\n",
    "  model.sample_rate=22050 \\\n",
    "  train_dataset=tests/data/asr/an4_train.json \\\n",
    "  validation_datasets=tests/data/asr/an4_val.json \\\n",
    "  phoneme_dict_path=tts_dataset_files/ipa_cmudict-0.7b_nv23.01.txt \\\n",
    "  heteronyms_path=tts_dataset_files/heteronyms-052722 \\\n",
    "  trainer.max_epochs=3 \\\n",
    "  trainer.accelerator='gpu' \\\n",
    "  trainer.strategy='ddp_find_unused_parameters_true' \\\n",
    "  trainer.check_val_every_n_epoch=1 \\\n",
    "  trainer.devices=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9erGDGZJ1H_p"
   },
   "source": [
    "# Training Data\n",
    "\n",
    "In order to train VITS, it is highly recommended to obtain high quality speech data with the following properties:\n",
    "  - Sampling rate of 22050Hz or higher\n",
    "  - Speech should contain a variety of speech phonemes\n",
    "  - Audio split into segments of 1-10 seconds\n",
    "  - Audio segments should not have silence at the beginning and end\n",
    "  - Audio segments should not contain long silences inside\n",
    "\n",
    "After obtaining the speech data and splitting into training, validation, and test sections, it is required to construct .json files to tell NeMo where to find these audio files.\n",
    "\n",
    "The .json files should adhere to the format required by the `nemo.collections.tts.data.dataset.TTSDataset` class. For example, here is a sample .json file\n",
    "\n",
    "```json\n",
    "{\"audio_filepath\": \"/path/to/audio1.wav\", \"text\": \"the transcription\", \"duration\": 0.82}\n",
    "{\"audio_filepath\": \"/path/to/audio2.wav\", \"text\": \"the other transcription\", \"duration\": 2.1}\n",
    "...\n",
    "```\n",
    "Please note that the duration is in seconds.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating VITS\n",
    "\n",
    "Let's evaluate the quality of the VITS model.\n",
    "\n",
    "VITS is end-to-end model, so we don't need any additional models to generate audios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://multilangaudiosamples.s3.us-east-2.amazonaws.com/LJ023-0089.wav"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt\n",
    "import IPython.display as ipd\n",
    "import numpy as np\n",
    "import torch\n",
    "import librosa\n",
    "import soundfile as sf\n",
    "\n",
    "target_sr = 22050\n",
    "\n",
    "audio_path = \"./LJ023-0089.wav\"\n",
    "text_raw = \"That is not only my accusation.\"\n",
    "\n",
    "\n",
    "audio_data, orig_sr = sf.read(audio_path)\n",
    "if orig_sr != target_sr:\n",
    "    audio_data = librosa.core.resample(audio_data, orig_sr=orig_sr, target_sr=target_sr)\n",
    "\n",
    "# Let's double-check that everything matches up!\n",
    "print(f\"Duration (s): {len(audio_data)/target_sr}\")\n",
    "print(\"Transcript:\", text_raw)\n",
    "ipd.Audio(audio_data, rate=target_sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = VitsModel.from_pretrained(\"tts_en_lj_vits\").cpu().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = model.parse(text_raw)\n",
    "audio_pred = model.convert_text_to_waveform(tokens=tokens).cpu().detach().numpy()\n",
    "\n",
    "print(\"predicted audio\")\n",
    "ipd.Audio(audio_pred, rate=target_sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_to_mel = model.audio_to_melspec_processor\n",
    "\n",
    "\n",
    "len_audio = torch.tensor(len(audio_data)).view(1, -1)\n",
    "\n",
    "spec_pred, _ = audio_to_mel(torch.tensor(audio_pred).view(1, -1), len_audio)\n",
    "spec_orig, _ = audio_to_mel(torch.tensor(audio_data).view(1, -1), len_audio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 2)\n",
    "\n",
    "ax[0].imshow(spec_orig[0][0].cpu().detach())\n",
    "ax[1].imshow(spec_pred[0][0].cpu().detach())\n",
    "\n",
    "ax[0].set_title('Original spectrogram')\n",
    "ax[1].set_title('Predicted spectrogram')\n",
    "fig.show()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Vits.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}