tutorials/tts/FastPitch_Finetuning.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "id": "b5c36beb",
            "metadata": {
                "id": "8d0bbac2"
            },
            "source": [
                "# Finetuning FastPitch for a new speaker\n",
                "\n",
                "In this tutorial, we will finetune a single speaker FastPitch (with alignment) model on 5 mins of a new speaker's data. We will finetune the model parameters only on the new speaker's text and speech pairs (though see the section at the end to learn more about mixing speaker data).\n",
                "\n",
                "We will download the training data, then generate and run a training command to finetune Fastpitch on 5 mins of data, and synthesize the audio from the trained checkpoint.\n",
                "\n",
                "A final section will describe approaches to improve audio quality past this notebook."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "4881d33e",
            "metadata": {
                "id": "nGw0CBaAtmQ6"
            },
            "source": [
                "## License\n",
                "\n",
                "> Copyright (c) 2021, 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,
            "id": "df5e057b",
            "metadata": {
                "id": "U7bOoIgLttRC"
            },
            "outputs": [],
            "source": [
                "\"\"\"\n",
                "You can either run this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\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",
                "\"\"\"\n",
                "BRANCH = 'r1.17.0'\n",
                "# # If you're using Google Colab and not running locally, uncomment and run this cell.\n",
                "# !apt-get install sox libsndfile1 ffmpeg\n",
                "# !pip install wget text-unidecode \n",
                "# !python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "c4f10958",
            "metadata": {
                "id": "2502cf61"
            },
            "source": [
                "## Downloading data"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "69e622c9",
            "metadata": {
                "id": "81fa2c02"
            },
            "source": [
                "For our tutorial, we will use a small part of the Hi-Fi Multi-Speaker English TTS (Hi-Fi TTS) dataset. You can read more about dataset [here](https://arxiv.org/abs/2104.01497). We will use speaker 9017 as the target speaker, and only a 5-minute subset of audio will be used for this fine-tuning example. We additionally resample audio to 22050 kHz."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "067c0a6c",
            "metadata": {
                "id": "VIFgqxLOpxha"
            },
            "outputs": [],
            "source": [
                "!wget https://multilangaudiosamples.s3.us-east-2.amazonaws.com/9017_5_mins.tar.gz\n",
                "!tar -xzf 9017_5_mins.tar.gz"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "b67737e5",
            "metadata": {
                "id": "gSQqq0fBqy8K"
            },
            "source": [
                "Looking at `manifest.json`, we see a standard NeMo json that contains the filepath, text, and duration. Please note that our `manifest.json` contains the relative path.\n",
                "\n",
                "Let's make sure that the entries look something like this:\n",
                "\n",
                "```\n",
                "{\"audio_filepath\": \"audio/dartagnan03part1_027_dumas_0047.wav\", \"text\": \"yes monsieur\", \"duration\": 1.04, \"text_no_preprocessing\": \"Yes, monsieur.\", \"text_normalized\": \"Yes, monsieur.\"}\n",
                "```"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "526ae2e9",
            "metadata": {},
            "outputs": [],
            "source": [
                "!head -n 1 ./9017_5_mins/manifest.json"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "3076f65a",
            "metadata": {},
            "source": [
                "Let's take 2 samples from the dataset and split it off into a validation set. Then, split all other samples into the training set.\n",
                "\n",
                "As mentioned, since the paths in the manifest are relative, we also create a symbolic link to the audio folder such that `audio/` goes to the correct directory."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "8d7946ae",
            "metadata": {
                "id": "B8gVfp5SsuDd"
            },
            "outputs": [],
            "source": [
                "!cat ./9017_5_mins/manifest.json | tail -n 2 > ./9017_manifest_dev_ns_all_local.json\n",
                "!cat ./9017_5_mins/manifest.json | head -n -2 > ./9017_manifest_train_dur_5_mins_local.json\n",
                "!ln -s ./9017_5_mins/audio audio"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "bf1f9c46",
            "metadata": {
                "id": "lhhg2wBNtW0r"
            },
            "source": [
                "Let's also download the pretrained checkpoint that we want to finetune from. NeMo will save checkpoints to `~/.cache`, so let's move that to our current directory. \n",
                "\n",
                "*Note: please, check that `home_path` refers to your home folder. Otherwise, change it manually.*"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "ed1898b9",
            "metadata": {},
            "outputs": [],
            "source": [
                "home_path = !(echo $HOME)\n",
                "home_path = home_path[0]\n",
                "print(home_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "dbaae91a",
            "metadata": {
                "id": "LggELooctXCT",
                "scrolled": true
            },
            "outputs": [],
            "source": [
                "import os\n",
                "import json\n",
                "\n",
                "import torch\n",
                "import IPython.display as ipd\n",
                "from matplotlib.pyplot import imshow\n",
                "from matplotlib import pyplot as plt\n",
                "\n",
                "from nemo.collections.tts.models import FastPitchModel\n",
                "FastPitchModel.from_pretrained(\"tts_en_fastpitch\")\n",
                "\n",
                "from pathlib import Path\n",
                "nemo_files = [p for p in Path(f\"{home_path}/.cache/torch/NeMo/\").glob(\"**/tts_en_fastpitch_align.nemo\")]\n",
                "print(f\"Copying {nemo_files[0]} to ./\")\n",
                "Path(\"./tts_en_fastpitch_align.nemo\").write_bytes(nemo_files[0].read_bytes())"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "49aa5048",
            "metadata": {
                "id": "6c8b13b8"
            },
            "source": [
                "To finetune the FastPitch model on the above created filelists, we use the `examples/tts/fastpitch_finetune.py` script to train the models with the `fastpitch_align_v1.05.yaml` configuration.\n",
                "\n",
                "Let's grab those files."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "200c7b26",
            "metadata": {
                "id": "3zg2H-32dNBU"
            },
            "outputs": [],
            "source": [
                "!wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/fastpitch_finetune.py\n",
                "\n",
                "!mkdir -p conf \\\n",
                "&& cd conf \\\n",
                "&& wget https://raw.githubusercontent.com/nvidia/NeMo/$BRANCH/examples/tts/conf/fastpitch_align_v1.05.yaml \\\n",
                "&& cd .."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "5415162b",
            "metadata": {},
            "source": [
                "We also need some additional files (see `FastPitch_MixerTTS_Training.ipynb` tutorial for more details) for training. Let's download these, too."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "20374059",
            "metadata": {},
            "outputs": [],
            "source": [
                "# 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/cmudict-0.7b_nv22.10 \\\n",
                "&& wget https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/scripts/tts_dataset_files/heteronyms-052722 \\\n",
                "&& cd .."
            ]
        },
        {
            "cell_type": "markdown",
            "id": "779af190",
            "metadata": {
                "id": "ef75d1d5"
            },
            "source": [
                "## Finetuning FastPitch"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "094c3383",
            "metadata": {
                "id": "12b5511c"
            },
            "source": [
                "We can now train our model with the following command:\n",
                "\n",
                "**NOTE: This will take about 50 minutes on colab's K80 GPUs.**"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "1e69d923",
            "metadata": {
                "id": "reY1LV4lwWoq"
            },
            "outputs": [],
            "source": [
                "# TODO(oktai15): remove +model.text_tokenizer.add_blank_at=true when we update FastPitch checkpoint\n",
                "!(python fastpitch_finetune.py --config-name=fastpitch_align_v1.05.yaml \\\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",
                "  phoneme_dict_path=tts_dataset_files/cmudict-0.7b_nv22.10 \\\n",
                "  heteronyms_path=tts_dataset_files/heteronyms-052722 \\\n",
                "  exp_manager.exp_dir=./ljspeech_to_9017_no_mixing_5_mins \\\n",
                "  +init_from_nemo_model=./tts_en_fastpitch_align.nemo \\\n",
                "  +trainer.max_steps=1000 ~trainer.max_epochs \\\n",
                "  trainer.check_val_every_n_epoch=25 \\\n",
                "  model.train_ds.dataloader_params.batch_size=24 model.validation_ds.dataloader_params.batch_size=24 \\\n",
                "  model.n_speakers=1 model.pitch_mean=152.3 model.pitch_std=64.0 \\\n",
                "  model.pitch_fmin=30 model.pitch_fmax=512 model.optim.lr=2e-4 \\\n",
                "  ~model.optim.sched model.optim.name=adam trainer.devices=1 trainer.strategy=null \\\n",
                "  +model.text_tokenizer.add_blank_at=true \\\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "c67a1086",
            "metadata": {
                "id": "j2svKvd1eMhf"
            },
            "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",
                "`\n",
                "  * We tell the script where `phoneme_dict_path` and `heteronyms-052722` 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",
            "id": "86675c74",
            "metadata": {
                "id": "c3bdf1ed"
            },
            "source": [
                "## Synthesize Samples from Finetuned Checkpoints"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "908ad3bb",
            "metadata": {
                "id": "f2b46325"
            },
            "source": [
                "Once we have finetuned our FastPitch model, we can synthesize the audio samples for given text using the following inference steps. We use a HiFi-GAN vocoder trained on LJSpeech.\n",
                "\n",
                "We define some helper functions as well."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "f5e85644",
            "metadata": {
                "id": "886c91dc"
            },
            "outputs": [],
            "source": [
                "from nemo.collections.tts.models import HifiGanModel\n",
                "from nemo.collections.tts.models import FastPitchModel\n",
                "\n",
                "vocoder = HifiGanModel.from_pretrained(\"tts_en_hifigan\")\n",
                "vocoder = vocoder.eval().cuda()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "559a4333",
            "metadata": {
                "id": "0a4c986f"
            },
            "outputs": [],
            "source": [
                "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\n",
                "\n",
                "def get_best_ckpt_from_last_run(\n",
                "        base_dir, \n",
                "        new_speaker_id, \n",
                "        duration_mins, \n",
                "        mixing_enabled, \n",
                "        original_speaker_id, \n",
                "        model_name=\"FastPitch\"\n",
                "    ):    \n",
                "    mixing = \"no_mixing\" if not mixing_enabled else \"mixing\"\n",
                "    \n",
                "    d = f\"{original_speaker_id}_to_{new_speaker_id}_{mixing}_{duration_mins}_mins\"\n",
                "    \n",
                "    exp_dirs = list([i for i in (Path(base_dir) / d / model_name).iterdir() if i.is_dir()])\n",
                "    last_exp_dir = sorted(exp_dirs)[-1]\n",
                "    \n",
                "    last_checkpoint_dir = last_exp_dir / \"checkpoints\"\n",
                "    \n",
                "    last_ckpt = list(last_checkpoint_dir.glob('*-last.ckpt'))\n",
                "\n",
                "    if len(last_ckpt) == 0:\n",
                "        raise ValueError(f\"There is no last checkpoint in {last_checkpoint_dir}.\")\n",
                "    \n",
                "    return str(last_ckpt[0])"
            ]
        },
        {
            "cell_type": "markdown",
            "id": "3d07c870",
            "metadata": {
                "id": "0153bd5a"
            },
            "source": [
                "Specify the speaker ID, duration of the dataset in minutes, and speaker mixing variables to find the relevant checkpoint (for example, we've saved our model in `ljspeech_to_9017_no_mixing_5_mins/`) and compare the synthesized audio with validation samples of the new speaker.\n",
                "\n",
                "The mixing variable is False for now, but we include some code to handle multiple speakers for reference."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "5ad71372",
            "metadata": {
                "id": "8901f88b",
                "scrolled": false
            },
            "outputs": [],
            "source": [
                "new_speaker_id = 9017\n",
                "duration_mins = 5\n",
                "mixing = False\n",
                "original_speaker_id = \"ljspeech\"\n",
                "\n",
                "last_ckpt = get_best_ckpt_from_last_run(\"./\", new_speaker_id, duration_mins, mixing, original_speaker_id)\n",
                "print(last_ckpt)\n",
                "\n",
                "spec_model = FastPitchModel.load_from_checkpoint(last_ckpt)\n",
                "spec_model.eval().cuda()\n",
                "\n",
                "# Only need to set speaker_id if there is more than one speaker\n",
                "speaker_id = None\n",
                "if mixing:\n",
                "    speaker_id = 1\n",
                "\n",
                "num_val = 2  # Number of validation samples\n",
                "val_records = []\n",
                "with open(f\"{new_speaker_id}_manifest_dev_ns_all_local.json\", \"r\") as f:\n",
                "    for i, line in enumerate(f):\n",
                "        val_records.append(json.loads(line))\n",
                "        if len(val_records) >= num_val:\n",
                "            break\n",
                "            \n",
                "for val_record in val_records:\n",
                "    print(\"Real validation audio\")\n",
                "    ipd.display(ipd.Audio(val_record['audio_filepath'], rate=22050))\n",
                "    print(f\"SYNTHESIZED FOR -- Speaker: {new_speaker_id} | Dataset size: {duration_mins} mins | Mixing:{mixing} | Text: {val_record['text']}\")\n",
                "    spec, audio = infer(spec_model, vocoder, val_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",
            "id": "28454638",
            "metadata": {
                "id": "ge2s7s9-w3py"
            },
            "source": [
                "## Improving Speech Quality\n",
                "\n",
                "We see that from fine-tuning FastPitch, we were able to generate audio in a male voice but the audio quality is not as good as we expect. We recommend two steps to improve audio quality:\n",
                "\n",
                "* Finetuning HiFi-GAN\n",
                "* Adding more data\n",
                "\n",
                "**Note that both of these steps are outside the scope of the notebook due to the limited compute available on Colab, but the code is included below for you to use outside of this notebook.**\n",
                "\n",
                "### Finetuning HiFi-GAN\n",
                "From the synthesized samples, there might be audible audio crackling. To fix this, we need to finetune HiFi-GAN on the new speaker's data. HiFi-GAN shows improvement using **synthesized mel spectrograms**, so the first step is to generate mel spectrograms with our finetuned FastPitch model to use as input.\n",
                "\n",
                "The code below uses our finetuned model to generate synthesized mels for the training set we have been using. You will also need to do the same for the validation set (code should be very similar, just with paths changed)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "id": "b17d3760",
            "metadata": {},
            "outputs": [],
            "source": [
                "import json\n",
                "import numpy as np\n",
                "import torch\n",
                "import soundfile as sf\n",
                "\n",
                "from pathlib import Path\n",
                "\n",
                "from nemo.collections.tts.parts.utils.tts_dataset_utils import BetaBinomialInterpolator\n",
                "\n",
                "\n",
                "def load_wav(audio_file, target_sr=None):\n",
                "    with sf.SoundFile(audio_file, 'r') as f:\n",
                "        samples = f.read(dtype='float32')\n",
                "        sample_rate = f.samplerate\n",
                "        if target_sr is not None and target_sr != sample_rate:\n",
                "            samples = librosa.core.resample(samples, orig_sr=sample_rate, target_sr=target_sr)\n",
                "    return samples.transpose()\n",
                "\n",
                "# Get records from the training manifest\n",
                "manifest_path = \"./9017_manifest_train_dur_5_mins_local.json\"\n",
                "records = []\n",
                "with open(manifest_path, \"r\") as f:\n",
                "    for i, line in enumerate(f):\n",
                "        records.append(json.loads(line))\n",
                "\n",
                "beta_binomial_interpolator = BetaBinomialInterpolator()\n",
                "spec_model.eval()\n",
                "\n",
                "device = spec_model.device\n",
                "\n",
                "save_dir = Path(\"./9017_manifest_train_dur_5_mins_local_mels\")\n",
                "save_dir.mkdir(exist_ok=True, parents=True)\n",
                "\n",
                "# Generate a spectrograms (we need to use ground truth alignment for correct matching between audio and mels)\n",
                "for i, r in enumerate(records):\n",
                "    audio = load_wav(r[\"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",
                "    \n",
                "    # Again, our finetuned FastPitch model doesn't use multiple speakers,\n",
                "    # but we keep the code to support it here for reference\n",
                "    if spec_model.fastpitch.speaker_emb is not None and \"speaker\" in r:\n",
                "        speaker = torch.tensor([r['speaker']]).to(device)\n",
                "    else:\n",
                "        speaker = None\n",
                "    \n",
                "    with torch.no_grad():\n",
                "        if \"normalized_text\" in r:\n",
                "            text = spec_model.parse(r[\"normalized_text\"], normalize=False)\n",
                "        else:\n",
                "            text = spec_model.parse(r['text'])\n",
                "        \n",
                "        text_len = torch.tensor(text.shape[-1], dtype=torch.long, device=device).unsqueeze(0)\n",
                "    \n",
                "        spect, spect_len = spec_model.preprocessor(input_signal=audio, length=audio_len)\n",
                "\n",
                "        # Generate attention prior and spectrogram inputs for HiFi-GAN\n",
                "        attn_prior = torch.from_numpy(\n",
                "          beta_binomial_interpolator(spect_len.item(), text_len.item())\n",
                "        ).unsqueeze(0).to(text.device)\n",
                "            \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",
                "        save_path = save_dir / f\"mel_{i}.npy\"\n",
                "        np.save(save_path, spectrogram[0].to('cpu').numpy())\n",
                "        r[\"mel_filepath\"] = str(save_path)\n",
                "\n",
                "hifigan_manifest_path = \"hifigan_train_ft.json\"\n",
                "with open(hifigan_manifest_path, \"w\") as f:\n",
                "    for r in records:\n",
                "        f.write(json.dumps(r) + '\\n')\n",
                "# Please do the same for the validation json. Code is omitted."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "id": "843674e7",
            "metadata": {},
            "source": [
                "We can then finetune hifigan similarly to fastpitch using NeMo's [hifigan_finetune.py](https://github.com/NVIDIA/NeMo/blob/main/examples/tts/hifigan_finetune.py) and [hifigan.yaml](https://github.com/NVIDIA/NeMo/blob/main/examples/tts/conf/hifigan/hifigan.yaml):\n",
                "\n",
                "```bash\n",
                "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_ft.json \\\n",
                "validation_datasets=./hifigan_val_ft.json \\\n",
                "exp_manager.exp_dir=hifigan_ft \\\n",
                "+init_from_pretrained_model=tts_en_hifigan \\\n",
                "trainer.check_val_every_n_epoch=10 \\\n",
                "model/train_ds=train_ds_finetune \\\n",
                "model/validation_ds=val_ds_finetune\n",
                "```\n",
                "\n",
                "Like when finetuning FastPitch, we lower the learning rate and get rid of the optimizer schedule for finetuning. You will need to create `<your_hifigan_val_manifest>` and the synthesized mels corresponding to it.\n",
                "\n",
                "As mentioned, the above command is not runnable in Colab due to limited compute resources, but you are free to finetune HiFi-GAN on your own machines.\n",
                "\n",
                "### Adding more data\n",
                "We can add more data in two ways. They can be combined for the best effect:\n",
                "\n",
                "* **Add more training data from the new speaker**\n",
                "\n",
                "The entire notebook can be repeated from the top after a new JSON manifest is defined that includes the additional data. Modify your finetuning commands to point to the new manifest. Be sure to increase the number of steps as more data is added to both the FastPitch and HiFi-GAN finetuning.\n",
                "\n",
                "We recommend **1000 steps per minute of audio for fastpitch and 500 steps per minute of audio for HiFi-GAN**.\n",
                "\n",
                "* **Mix new speaker data with old speaker data**\n",
                "\n",
                "We recommend finetuning FastPitch (but not HiFi-GAN) using both old speaker data (LJSpeech in this notebook) and the new speaker data. In this case, please modify the JSON manifests when finetuning FastPitch to include speaker information by adding a `speaker` field to each entry:\n",
                "\n",
                "```\n",
                "{\"audio_filepath\": \"new_speaker.wav\", \"text\": \"sample\", \"duration\": 2.6, \"speaker\": 1}\n",
                "{\"audio_filepath\": \"old_speaker.wav\", \"text\": \"LJSpeech sample\", \"duration\": 2.6, \"speaker\": 0}\n",
                "```\n",
                "\n",
                "5 hours of data from the old speaker should be sufficient for training; it's up to you how much data from the old speaker to use in validation.\n",
                "\n",
                "For the training manifest, since we likely have less data from the new speaker, we need to ensure that the model sees a similar amount of new data and old data. We can do this by repeating samples from the new speaker until we have an equivalent number of samples from the old and new speaker. For each sample from the old speaker, please add a sample from the new speaker in the .json.\n",
                "\n",
                "As a toy example, if we use 4 samples of the old speaker and only 2 samples of the new speaker, we would want the order of samples in our manifest to look something like this:\n",
                "\n",
                "```\n",
                "old_speaker_sample_0\n",
                "new_speaker_sample_0\n",
                "old_speaker_sample_1\n",
                "new_speaker_sample_1\n",
                "old_speaker_sample_2\n",
                "new_speaker_sample_0  # Start repeat of new speaker samples\n",
                "old_speaker_sample_3\n",
                "new_speaker_sample_1\n",
                "```\n",
                "\n",
                "Once the manifests are created, we can modify the FastPitch training command to point to the new training and validation JSON files.\n",
                "\n",
                "We also need to update `model.n_speakers=1` to `model.n_speakers=2`, as well as update the `sup_data_types` specified in the config file to include `speaker_id` (`sup_data_types=[align_prior_matrix,pitch,speaker_id]`). Updating these two fields is very important--otherwise the model will not recognize that there is more than one speaker!\n",
                "\n",
                "Ensure the pitch statistics correspond to the new speaker rather than the old speaker for best results.\n",
                "\n",
                "**For HiFiGAN finetuning, the training should be done on the new speaker data.**"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3.9.15 ('ptl_venv')",
            "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.15"
        },
        "vscode": {
            "interpreter": {
                "hash": "f8a1d50fd7b1e17bd198f085b8ced031398c6134b0da7c4415c17601bbcccc4e"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 5
}