02_FastPitch_finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ea801208-e8e3-47d4-b099-ac2e952e4b4a",
   "metadata": {
    "id": "ef75d1d5"
   },
   "source": [
    "## Finetuning FastPitch\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": "code",
   "execution_count": 1,
   "id": "321fdfb6-0ce1-46d9-a48f-77f071e2feb3",
   "metadata": {
    "id": "LggELooctXCT",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from pathlib import Path\n",
    "\n",
    "import wandb\n",
    "\n",
    "import torch\n",
    "import pandas as pd\n",
    "import IPython.display as ipd\n",
    "from matplotlib.pyplot import imshow\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a80c8066-a8eb-48c3-bd0d-dbda23094bdb",
   "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": 2,
   "id": "b12bff96-7ed5-4a06-8e70-3fd7e40b9eec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SPEAKER_ID = 9017\n",
    "MODEL_NAME = \"tts_en_fastpitch\"\n",
    "\n",
    "WANDB_PROJECT = \"tts-workshop\"\n",
    "WANDB_ENTITY = \"capecape\" # replace with your wandb username or team"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8b7bc11-2d34-4f4a-99c6-fb8e125bc8a1",
   "metadata": {},
   "source": [
    "Let's download the pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a5f3d986-3b44-4e3a-a0da-88e4a5fa0a9c",
   "metadata": {
    "id": "LggELooctXCT",
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 20:13:40 optimizers:55] Apex was not found. Using the lamb or fused_adam optimizer will error out.\n",
      "[NeMo W 2022-12-07 20:13:40 experimental:27] Module <class 'nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer'> is experimental, not ready for production and is not fully supported. Use at your own risk.\n",
      "[NeMo W 2022-12-07 20:13:40 experimental:27] Module <class 'nemo.collections.tts.models.radtts.RadTTSModel'> is experimental, not ready for production and is not fully supported. Use at your own risk.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:13:40 cloud:56] Found existing object /home/tcapelle/.cache/torch/NeMo/NeMo_1.14.0rc0/tts_en_fastpitch_align/b7d086a07b5126c12d5077d9a641a38c/tts_en_fastpitch_align.nemo.\n",
      "[NeMo I 2022-12-07 20:13:40 cloud:62] Re-using file from: /home/tcapelle/.cache/torch/NeMo/NeMo_1.14.0rc0/tts_en_fastpitch_align/b7d086a07b5126c12d5077d9a641a38c/tts_en_fastpitch_align.nemo\n",
      "[NeMo I 2022-12-07 20:13:40 common:912] Instantiating model from pre-trained checkpoint\n",
      "[NeMo I 2022-12-07 20:13:43 tokenize_and_classify:87] Creating ClassifyFst grammars.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 20:14:13 experimental:27] Module <class 'nemo_text_processing.g2p.modules.IPAG2P'> is experimental, not ready for production and is not fully supported. Use at your own risk.\n",
      "[NeMo W 2022-12-07 20:14:13 modules:95] apply_to_oov_word=None, This means that some of words will remain unchanged if they are not handled by any of the rules in self.parse_one_word(). This may be intended if phonemes and chars are both valid inputs, otherwise, you may see unexpected deletions in your input.\n",
      "[NeMo W 2022-12-07 20:14:13 modelPT:142] If you intend to do training or fine-tuning, please call the ModelPT.setup_training_data() method and provide a valid configuration file to setup the train data loader.\n",
      "    Train config : \n",
      "    dataset:\n",
      "      _target_: nemo.collections.tts.torch.data.TTSDataset\n",
      "      manifest_filepath: /ws/LJSpeech/nvidia_ljspeech_train_clean_ngc.json\n",
      "      sample_rate: 22050\n",
      "      sup_data_path: /raid/LJSpeech/supplementary\n",
      "      sup_data_types:\n",
      "      - align_prior_matrix\n",
      "      - pitch\n",
      "      n_fft: 1024\n",
      "      win_length: 1024\n",
      "      hop_length: 256\n",
      "      window: hann\n",
      "      n_mels: 80\n",
      "      lowfreq: 0\n",
      "      highfreq: 8000\n",
      "      max_duration: null\n",
      "      min_duration: 0.1\n",
      "      ignore_file: null\n",
      "      trim: false\n",
      "      pitch_fmin: 65.40639132514966\n",
      "      pitch_fmax: 2093.004522404789\n",
      "      pitch_norm: true\n",
      "      pitch_mean: 212.35873413085938\n",
      "      pitch_std: 68.52806091308594\n",
      "      use_beta_binomial_interpolator: true\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: true\n",
      "      batch_size: 24\n",
      "      num_workers: 0\n",
      "    \n",
      "[NeMo W 2022-12-07 20:14:13 modelPT:149] If you intend to do validation, please call the ModelPT.setup_validation_data() or ModelPT.setup_multiple_validation_data() method and provide a valid configuration file to setup the validation data loader(s). \n",
      "    Validation config : \n",
      "    dataset:\n",
      "      _target_: nemo.collections.tts.torch.data.TTSDataset\n",
      "      manifest_filepath: /ws/LJSpeech/nvidia_ljspeech_val_clean_ngc.json\n",
      "      sample_rate: 22050\n",
      "      sup_data_path: /raid/LJSpeech/supplementary\n",
      "      sup_data_types:\n",
      "      - align_prior_matrix\n",
      "      - pitch\n",
      "      n_fft: 1024\n",
      "      win_length: 1024\n",
      "      hop_length: 256\n",
      "      window: hann\n",
      "      n_mels: 80\n",
      "      lowfreq: 0\n",
      "      highfreq: 8000\n",
      "      max_duration: null\n",
      "      min_duration: null\n",
      "      ignore_file: null\n",
      "      trim: false\n",
      "      pitch_fmin: 65.40639132514966\n",
      "      pitch_fmax: 2093.004522404789\n",
      "      pitch_norm: true\n",
      "      pitch_mean: 212.35873413085938\n",
      "      pitch_std: 68.52806091308594\n",
      "      use_beta_binomial_interpolator: true\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: false\n",
      "      batch_size: 24\n",
      "      num_workers: 0\n",
      "    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:14:14 features:267] PADDING: 1\n",
      "[NeMo I 2022-12-07 20:14:16 save_restore_connector:243] Model FastPitchModel was successfully restored from /home/tcapelle/.cache/torch/NeMo/NeMo_1.14.0rc0/tts_en_fastpitch_align/b7d086a07b5126c12d5077d9a641a38c/tts_en_fastpitch_align.nemo.\n"
     ]
    }
   ],
   "source": [
    "from nemo.collections.tts.models import FastPitchModel\n",
    "FastPitchModel.from_pretrained(MODEL_NAME);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92c62728-0b7f-404b-831d-50d80bd6dcab",
   "metadata": {
    "id": "lhhg2wBNtW0r"
   },
   "source": [
    "We will copy the model to our current working dir for simplicity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0ed70f0b-6651-4997-b755-e63b93987301",
   "metadata": {
    "id": "LggELooctXCT",
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Copying /home/tcapelle/.cache/torch/NeMo/NeMo_1.14.0rc0/tts_en_fastpitch_align/b7d086a07b5126c12d5077d9a641a38c/tts_en_fastpitch_align.nemo to ./tts_en_fastpitch_align.nemo\n",
      "Copying /home/tcapelle/.cache/torch/NeMo/NeMo_1.12.0/tts_en_fastpitch_align/26d7e09971f1d611e24df90c7a9d9b38/tts_en_fastpitch_align.nemo to ./tts_en_fastpitch_align.nemo\n"
     ]
    }
   ],
   "source": [
    "for nemo_file in (Path.home()/\".cache/torch/NeMo/\").glob(f\"**/{MODEL_NAME}_align.nemo\"):\n",
    "    print(f\"Copying {nemo_file} to ./{nemo_file.name}\")\n",
    "    Path(f\"./{nemo_file.name}\").write_bytes(nemo_file.read_bytes())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13dac038-c59c-46c3-94e3-15419d0666e2",
   "metadata": {
    "id": "12b5511c"
   },
   "source": [
    "## Train Time!\n",
    "\n",
    "We can now train our model with the following command:\n",
    "\n",
    "**NOTE: This will take about 30 minutes on colab's T4 GPUs.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "486620ac-f9c9-46e6-841b-89ae6aa483a0",
   "metadata": {},
   "source": [
    "### Using Hydra in Jupyter..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c7e2d1c7-5f4d-402a-8ccd-199dc423c4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytorch_lightning as pl\n",
    "\n",
    "from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "from hydra import compose, initialize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "efc80d7f-1bd8-4bfc-b7fc-a6e946bb2dd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "with initialize(config_path=\"./conf\"):\n",
    "    cfg = compose(config_name=\"fastpitch_align_v1.05.yaml\")\n",
    "    \n",
    "    # or use the nice hydra API 🤣\n",
    "    # cfg = compose(config_name=\"fastpitch_align_v1.05.yaml\",\n",
    "    #               overrides=['train_dataset=9017_manifest_train_local.json',\n",
    "    #                          'validation_datasets=9017_manifest_valid_local.json',\n",
    "    #                          '+exp_manager.create_wandb_logger=True',\n",
    "    #                          f'+exp_manager.wandb_logger_kwargs=\\{project:{WANDB_PROJECT}, job_type:training, log_model:True\\}',\n",
    "    #                          'model.pitch_mean=152.3',\n",
    "    #                          'model.pitch_std=64.0',\n",
    "    #                          'model.pitch_fmin=30',\n",
    "    #                          'model.pitch_fmax=512',\n",
    "    #                          'trainer.max_steps=10',\n",
    "    #                         ]\n",
    "    #              )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4bb565f6-39d8-40a3-a7e0-bcfddf6eba79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# datasets\n",
    "cfg.train_dataset = \"9017_manifest_train_local.json\"\n",
    "cfg.validation_datasets = \"9017_manifest_valid_local.json\"\n",
    "\n",
    "# speaker params\n",
    "cfg.model.pitch_mean = 152.3\n",
    "cfg.model.pitch_std = 64.0\n",
    "cfg.model.pitch_fmin = 30\n",
    "cfg.model.pitch_fmax = 512\n",
    "\n",
    "#wandb\n",
    "cfg.exp_manager.create_wandb_logger = True\n",
    "cfg.exp_manager.wandb_logger_kwargs = {\"project\":WANDB_PROJECT, \"job_type\":\"training\", \"log_model\":True}\n",
    "\n",
    "# iterate fast\n",
    "cfg.trainer.max_steps = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "da88d9f1-5922-4b23-8e53-2a22f2e9193e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using 16bit native Automatic Mixed Precision (AMP)\n",
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:14:18 exp_manager:343] Experiments will be logged at exp_dir/FastPitch/2022-12-07_20-14-18\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcapecape\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "wandb version 0.13.6 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.13.5"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>./exp_dir/wandb/run-20221207_201420-2022-12-07_20-14-18</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href=\"https://wandb.ai/capecape/tts-workshop/runs/2022-12-07_20-14-18\" target=\"_blank\">2022-12-07_20-14-18</a></strong> to <a href=\"https://wandb.ai/capecape/tts-workshop\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://wandb.me/run\" target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:14:25 exp_manager:733] WandBLogger has been set up\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 20:14:25 exp_manager:988] The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to 1000. Please ensure that max_steps will run for at least 10 epochs to ensure that checkpointing will not error out.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PosixPath('exp_dir/FastPitch/2022-12-07_20-14-18')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = pl.Trainer(**cfg.trainer, enable_progress_bar=False)\n",
    "exp_manager(trainer, cfg.get(\"exp_manager\", None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6949d243-5ade-43dd-9798-627ac340cfa1",
   "metadata": {},
   "source": [
    "We can create model lineage with the split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "547be895-24eb-4d37-a3a3-24acff10472b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Artifact QXJ0aWZhY3Q6Mjk0NzgzMDEx>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wandb_artifact=f'{WANDB_ENTITY}/{WANDB_PROJECT}/9017_5_split:v0'\n",
    "wandb.use_artifact(wandb_artifact)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4719385d-e2ac-4d20-a175-47f395c29c8e",
   "metadata": {},
   "source": [
    "Configure the logging and load the pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b0211a13-f2a1-495e-b549-edabd26d1933",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:14:28 tokenize_and_classify:87] Creating ClassifyFst grammars.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 20:14:58 modules:95] apply_to_oov_word=None, This means that some of words will remain unchanged if they are not handled by any of the rules in self.parse_one_word(). This may be intended if phonemes and chars are both valid inputs, otherwise, you may see unexpected deletions in your input.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:14:58 data:217] Loading dataset from 9017_manifest_train_local.json.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "76it [00:00, 84508.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:14:58 data:254] Loaded dataset with 76 files.\n",
      "[NeMo I 2022-12-07 20:14:58 data:256] Dataset contains 0.08 hours.\n",
      "[NeMo I 2022-12-07 20:14:58 data:358] Pruned 0 files. Final dataset contains 76 files\n",
      "[NeMo I 2022-12-07 20:14:58 data:360] Pruned 0.00 hours. Final dataset contains 0.08 hours.\n",
      "[NeMo I 2022-12-07 20:14:58 data:217] Loading dataset from 9017_manifest_valid_local.json.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "2it [00:00, 4350.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:14:58 data:254] Loaded dataset with 2 files.\n",
      "[NeMo I 2022-12-07 20:14:58 data:256] Dataset contains 0.00 hours.\n",
      "[NeMo I 2022-12-07 20:14:58 data:358] Pruned 0 files. Final dataset contains 2 files\n",
      "[NeMo I 2022-12-07 20:14:58 data:360] Pruned 0.00 hours. Final dataset contains 0.00 hours.\n",
      "[NeMo I 2022-12-07 20:14:58 features:267] PADDING: 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:15:02 tokenize_and_classify:87] Creating ClassifyFst grammars.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 20:15:34 modules:95] apply_to_oov_word=None, This means that some of words will remain unchanged if they are not handled by any of the rules in self.parse_one_word(). This may be intended if phonemes and chars are both valid inputs, otherwise, you may see unexpected deletions in your input.\n",
      "[NeMo W 2022-12-07 20:15:34 modelPT:142] If you intend to do training or fine-tuning, please call the ModelPT.setup_training_data() method and provide a valid configuration file to setup the train data loader.\n",
      "    Train config : \n",
      "    dataset:\n",
      "      _target_: nemo.collections.tts.torch.data.TTSDataset\n",
      "      manifest_filepath: /ws/LJSpeech/nvidia_ljspeech_train_clean_ngc.json\n",
      "      sample_rate: 22050\n",
      "      sup_data_path: /raid/LJSpeech/supplementary\n",
      "      sup_data_types:\n",
      "      - align_prior_matrix\n",
      "      - pitch\n",
      "      n_fft: 1024\n",
      "      win_length: 1024\n",
      "      hop_length: 256\n",
      "      window: hann\n",
      "      n_mels: 80\n",
      "      lowfreq: 0\n",
      "      highfreq: 8000\n",
      "      max_duration: null\n",
      "      min_duration: 0.1\n",
      "      ignore_file: null\n",
      "      trim: false\n",
      "      pitch_fmin: 65.40639132514966\n",
      "      pitch_fmax: 2093.004522404789\n",
      "      pitch_norm: true\n",
      "      pitch_mean: 212.35873413085938\n",
      "      pitch_std: 68.52806091308594\n",
      "      use_beta_binomial_interpolator: true\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: true\n",
      "      batch_size: 24\n",
      "      num_workers: 0\n",
      "    \n",
      "[NeMo W 2022-12-07 20:15:34 modelPT:149] If you intend to do validation, please call the ModelPT.setup_validation_data() or ModelPT.setup_multiple_validation_data() method and provide a valid configuration file to setup the validation data loader(s). \n",
      "    Validation config : \n",
      "    dataset:\n",
      "      _target_: nemo.collections.tts.torch.data.TTSDataset\n",
      "      manifest_filepath: /ws/LJSpeech/nvidia_ljspeech_val_clean_ngc.json\n",
      "      sample_rate: 22050\n",
      "      sup_data_path: /raid/LJSpeech/supplementary\n",
      "      sup_data_types:\n",
      "      - align_prior_matrix\n",
      "      - pitch\n",
      "      n_fft: 1024\n",
      "      win_length: 1024\n",
      "      hop_length: 256\n",
      "      window: hann\n",
      "      n_mels: 80\n",
      "      lowfreq: 0\n",
      "      highfreq: 8000\n",
      "      max_duration: null\n",
      "      min_duration: null\n",
      "      ignore_file: null\n",
      "      trim: false\n",
      "      pitch_fmin: 65.40639132514966\n",
      "      pitch_fmax: 2093.004522404789\n",
      "      pitch_norm: true\n",
      "      pitch_mean: 212.35873413085938\n",
      "      pitch_std: 68.52806091308594\n",
      "      use_beta_binomial_interpolator: true\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: false\n",
      "      batch_size: 24\n",
      "      num_workers: 0\n",
      "    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:15:35 features:267] PADDING: 1\n",
      "[NeMo I 2022-12-07 20:15:36 save_restore_connector:243] Model FastPitchModel was successfully restored from /home/tcapelle/wandb/nvidia-workshop/tts_en_fastpitch_align.nemo.\n",
      "[NeMo I 2022-12-07 20:15:36 modelPT:1067] Model checkpoint restored from nemo file with path : `./tts_en_fastpitch_align.nemo`\n"
     ]
    }
   ],
   "source": [
    "model = FastPitchModel(cfg=cfg.model, trainer=trainer)\n",
    "model.maybe_init_from_pretrained_checkpoint(cfg=cfg)\n",
    "lr_logger = pl.callbacks.LearningRateMonitor()\n",
    "epoch_time_logger = LogEpochTimeCallback()\n",
    "trainer.callbacks.extend([lr_logger, epoch_time_logger])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94925fd4-ecb1-4b92-a4a7-11fc66dbe3ef",
   "metadata": {},
   "source": [
    "time to train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "777109c4-134b-4f6c-9896-679cf304066f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 20:15:36 modelPT:602] Optimizer config = Adam (\n",
      "    Parameter Group 0\n",
      "        amsgrad: False\n",
      "        betas: [0.9, 0.999]\n",
      "        capturable: False\n",
      "        differentiable: False\n",
      "        eps: 1e-08\n",
      "        foreach: None\n",
      "        fused: False\n",
      "        lr: 5e-05\n",
      "        maximize: False\n",
      "        weight_decay: 1e-06\n",
      "    )\n",
      "[NeMo I 2022-12-07 20:15:36 lr_scheduler:772] Scheduler not initialized as no `sched` config supplied to setup_optimizer()\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "  | Name                | Type                              | Params\n",
      "--------------------------------------------------------------------------\n",
      "0 | mel_loss_fn         | MelLoss                           | 0     \n",
      "1 | pitch_loss_fn       | PitchLoss                         | 0     \n",
      "2 | duration_loss_fn    | DurationLoss                      | 0     \n",
      "3 | energy_loss_fn      | EnergyLoss                        | 0     \n",
      "4 | aligner             | AlignmentEncoder                  | 1.0 M \n",
      "5 | forward_sum_loss_fn | ForwardSumLoss                    | 0     \n",
      "6 | bin_loss_fn         | BinLoss                           | 0     \n",
      "7 | preprocessor        | AudioToMelSpectrogramPreprocessor | 0     \n",
      "8 | fastpitch           | FastPitchModule                   | 45.8 M\n",
      "--------------------------------------------------------------------------\n",
      "45.8 M    Trainable params\n",
      "0         Non-trainable params\n",
      "45.8 M    Total params\n",
      "91.518    Total estimated model params size (MB)\n",
      "Epoch 9, global step 50: 'val_loss' reached 1.82647 (best 1.82647), saving model to '/home/tcapelle/wandb/nvidia-workshop/exp_dir/FastPitch/2022-12-07_20-14-18/checkpoints/FastPitch.ckpt' as top 3\n",
      "Epoch 19, global step 100: 'val_loss' reached 1.93638 (best 1.82647), saving model to '/home/tcapelle/wandb/nvidia-workshop/exp_dir/FastPitch/2022-12-07_20-14-18/checkpoints/FastPitch-v1.ckpt' as top 3\n",
      "Epoch 29, global step 150: 'val_loss' reached 1.94588 (best 1.82647), saving model to '/home/tcapelle/wandb/nvidia-workshop/exp_dir/FastPitch/2022-12-07_20-14-18/checkpoints/FastPitch-v2.ckpt' as top 3\n",
      "Epoch 39, global step 200: 'val_loss' reached 1.93850 (best 1.82647), saving model to '/home/tcapelle/wandb/nvidia-workshop/exp_dir/FastPitch/2022-12-07_20-14-18/checkpoints/FastPitch-v2.ckpt' as top 3\n",
      "Epoch 49, global step 250: 'val_loss' was not in top 3\n",
      "Epoch 59, global step 300: 'val_loss' was not in top 3\n",
      "Epoch 69, global step 350: 'val_loss' was not in top 3\n",
      "Epoch 79, global step 400: 'val_loss' was not in top 3\n",
      "Epoch 89, global step 450: 'val_loss' was not in top 3\n",
      "Epoch 129, global step 650: 'val_loss' was not in top 3\n",
      "Epoch 139, global step 700: 'val_loss' was not in top 3\n",
      "Epoch 149, global step 750: 'val_loss' was not in top 3\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a0508af-5149-42ee-a742-fac0512dddf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19ffce2a-6489-46fc-a542-fbeccdff650e",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Running as a python program on a shell\n",
    "\n",
    "> I recommend to run this command on a terminal instead of on the notebooks itself"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "041c8852-f35c-43b3-b8b1-82bdd54e86e1",
   "metadata": {
    "id": "reY1LV4lwWoq",
    "tags": []
   },
   "source": [
    "```bash\n",
    "python fastpitch_finetune.py --config-name=fastpitch_align_v1.05.yaml \\\n",
    "  train_dataset=9017_manifest_train_local.json \\\n",
    "  validation_datasets=9017_manifest_valid_local.json \\\n",
    "  exp_manager.exp_dir=./fastpitch_finetune \\\n",
    "  +exp_manager.create_wandb_logger=True \\\n",
    "  +wandb_artifact='capecape/tts-workshop/9017_5_split:v0' \\\n",
    "  +exp_manager.wandb_logger_kwargs='{project:tts-workshop, job_type:training, log_model:True}' \\\n",
    "  +init_from_nemo_model=./tts_en_fastpitch_align.nemo \\\n",
    "  model.pitch_mean=152.3 model.pitch_std=64.0 \\\n",
    "  model.pitch_fmin=30 model.pitch_fmax=512 \\\n",
    "  trainer.max_steps=10 \\\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2a64ef3-3382-4e4a-b23e-0a1f857f9b96",
   "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=./6097_manifest_train_dur_5_mins_local.json \n",
    "  validation_datasets=./6097_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/lj_speech.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_6097_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",
    "  * fmin and fmax are hyperparameters to librosa's pyin function. We recommend tweaking these per speaker.\n",
    "  * After fmin and fmax are defined, pitch mean and std can be easily extracted.\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": "9497baa7-e87a-47e6-a4f9-3e631a67f9e1",
   "metadata": {
    "id": "c3bdf1ed"
   },
   "source": [
    "## Synthesize Samples from Finetuned Checkpoints"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0c3b08d-c201-4ba5-b34e-b325eafe11f6",
   "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": "e3e0db1a-0f9a-476f-926a-0f2307de2f65",
   "metadata": {
    "id": "886c91dc"
   },
   "outputs": [],
   "source": [
    "from nemo.collections.tts.models import HifiGanModel\n",
    "from nemo.collections.tts.models import FastPitchModel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "001113fa-2365-446b-943e-fb4daf83b545",
   "metadata": {},
   "source": [
    "we will load a pretrained model to generate the voice from the spectrogram (we will later fine tune this model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28e92db0-cf97-49c1-8632-1a5fe772a3cd",
   "metadata": {
    "id": "886c91dc"
   },
   "outputs": [],
   "source": [
    "vocoder = HifiGanModel.from_pretrained(\"tts_hifigan\")\n",
    "vocoder = vocoder.eval().cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60440ec1-26d6-49ff-991b-c1b7809b6213",
   "metadata": {},
   "source": [
    "We can grab the fine tuned models from the `wandb` artifact:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aede82a5-5be1-4a34-abb9-209c911c4def",
   "metadata": {},
   "outputs": [],
   "source": [
    "# which split we are using\n",
    "validation_split_artifact = f'{WANDB_ENTITY}/{WANDB_PROJECT}/9017_5_split:v0'\n",
    "\n",
    "# which model\n",
    "model_artifact = f'{WANDB_ENTITY}/{WANDB_PROJECT}/model-2022-12-07_17-56-27:v2'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec083108-7bc4-48c2-8896-7ad58be68926",
   "metadata": {},
   "source": [
    "Let's log the model predictions to `W&B`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5531df3a-f71e-478a-8f9f-5933589b0913",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(entity=WANDB_ENTITY, project=WANDB_PROJECT, job_type=\"fastptich_validation\")\n",
    "\n",
    "split_artifact = wandb.use_artifact(validation_split_artifact)\n",
    "split_artifact_dir = split_artifact.download()\n",
    "\n",
    "model_artifact = wandb.use_artifact(model_artifact, type='model')\n",
    "model_artifact_dir = model_artifact.download()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5bb5a72-b5ac-49fb-a384-520705786999",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ls(path): return list(Path(path).iterdir())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df519bc1-3b42-4ef6-aef7-378c8b767000",
   "metadata": {},
   "outputs": [],
   "source": [
    "ls(split_artifact_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03d421e5-9ecb-4dce-a53a-97270785c644",
   "metadata": {},
   "outputs": [],
   "source": [
    "ls(model_artifact_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48b3158e-2897-4f11-a7b1-765e4e88b04d",
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d3e2c2d-d164-4725-b159-2bbf268182fb",
   "metadata": {
    "id": "8901f88b"
   },
   "outputs": [],
   "source": [
    "last_ckpt = str(ls(model_artifact_dir)[0])\n",
    "print(last_ckpt)\n",
    "\n",
    "spec_model = FastPitchModel.load_from_checkpoint(last_ckpt)\n",
    "spec_model.eval().cuda();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6ec1b18-11e1-419e-8e35-da71de9ff752",
   "metadata": {},
   "source": [
    "### View results in the notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9be1a2ba-ddad-4df0-a510-76fa6b23d5bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_df = pd.read_json(Path(split_artifact_dir)/f\"{SPEAKER_ID}_manifest_valid_local.json\", lines=True)\n",
    "valid_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3851c08c-e72d-49af-b51e-56902165a922",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_audio(text, speaker_id):\n",
    "    \"Generate MEL and Synth Audio\"\n",
    "    spec, audio = infer(spec_model, vocoder, text, speaker=speaker_id)\n",
    "    return spec, audio.flatten()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba6810c3-2d73-485e-9730-52ee64867bdb",
   "metadata": {
    "id": "8901f88b"
   },
   "outputs": [],
   "source": [
    "new_speaker_id = SPEAKER_ID\n",
    "duration_mins = 5\n",
    "mixing = False\n",
    "original_speaker_id = \"ljspeech\"\n",
    "\n",
    "for _, val_record in valid_df.iterrows():\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 = generate_audio(val_record[\"text\"], new_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": "3753d143-b2c8-420c-9d8f-31cb59c03546",
   "metadata": {},
   "source": [
    "### The same on a wandb.Table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4306977-f04d-46da-92ea-a7201eb4d6c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "table = wandb.Table(columns=['Text', 'Real validation audio', f'Audio Speaker {new_speaker_id}', 'Spec'])\n",
    "\n",
    "sample_rate=22050\n",
    "\n",
    "for _, val_record in valid_df.iterrows():\n",
    "    speaker_spec, speaker_audio = generate_audio(val_record['text'], speaker_id=new_speaker_id)\n",
    "    row = [val_record[\"text_no_preprocessing\"],\n",
    "           wandb.Audio(val_record['audio_filepath'], sample_rate=sample_rate), \n",
    "           wandb.Audio(speaker_audio.flatten(), sample_rate=sample_rate),\n",
    "           wandb.Image(speaker_spec)]\n",
    "    table.add_data(*row)\n",
    "\n",
    "wandb.log({\"fastpitch_predictions\": table})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2bdd02-66e8-43c7-9343-012f7e157d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  }
 ],
 "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}