03_HiFIGAN_finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "afdba47f-b132-45da-8830-db058476f38e",
   "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 SageMaker/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": 1,
   "id": "a22f9ada-c3cd-40f4-9c3c-c653f82bb4e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 18:13:06 optimizers:55] Apex was not found. Using the lamb or fused_adam optimizer will error out.\n",
      "[NeMo W 2022-12-07 18:13:07 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 18:13:07 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"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "import soundfile as sf\n",
    "\n",
    "import wandb\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "import pytorch_lightning as pl\n",
    "from hydra import compose, initialize\n",
    "\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "from nemo.collections.tts.models import FastPitchModel, HifiGanModel\n",
    "from nemo.collections.tts.torch.helpers import BetaBinomialInterpolator\n",
    "from nemo.collections.common.callbacks import LogEpochTimeCallback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "75ed6416-ccf8-4dc7-bc17-ff22a0be0082",
   "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": "1e819c66-e8f8-4040-b774-616e288ea21f",
   "metadata": {},
   "source": [
    "load the FastPitch finetuned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "719867de-37fd-4437-828e-ec03a7007b27",
   "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": "code",
   "execution_count": 4,
   "id": "76e70398-48d2-4544-bc09-45bde01c8595",
   "metadata": {},
   "outputs": [
    {
     "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>/home/tcapelle/wandb/nvidia-workshop/wandb/run-20221207_181313-32pmu5t7</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/32pmu5t7\" target=\"_blank\">volcanic-field-8</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": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m:   2 of 2 files downloaded.  \n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact model-2022-12-07_17-56-27:v2, 524.07MB. 1 files... \n",
      "\u001b[34m\u001b[1mwandb\u001b[0m:   1 of 1 files downloaded.  \n",
      "Done. 0:0:0.0\n"
     ]
    }
   ],
   "source": [
    "wandb.init(entity=WANDB_ENTITY, project=WANDB_PROJECT, job_type=\"hifigan_finetune\")\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": 5,
   "id": "785c418c-10f3-4e9e-8785-9878cfc9f0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ls(path): return list(Path(path).iterdir())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cb6f13a3-9182-473b-9706-776981a9a3d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PosixPath('artifacts/9017_5_split:v0/9017_manifest_valid_local.json'),\n",
       " PosixPath('artifacts/9017_5_split:v0/9017_manifest_train_local.json')]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ls(split_artifact_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0deb7e59-97ac-4f93-a0a5-7eecd42f8c10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PosixPath('artifacts/model-2022-12-07_17-56-27:v2/model.ckpt')]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ls(model_artifact_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "91bae020-82c8-4810-8620-83d44ae21002",
   "metadata": {
    "id": "8901f88b"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "artifacts/model-2022-12-07_17-56-27:v2/model.ckpt\n",
      "[NeMo I 2022-12-07 18:13:21 tokenize_and_classify:87] Creating ClassifyFst grammars.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 18:13:53 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 18:13:53 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 18:13:53 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: 9017_manifest_train_local.json\n",
      "      sample_rate: 22050\n",
      "      sup_data_path: ./fastpitch_sup_data\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: 30\n",
      "      pitch_fmax: 512\n",
      "      pitch_norm: true\n",
      "      pitch_mean: 152.3\n",
      "      pitch_std: 64.0\n",
      "      use_beta_binomial_interpolator: true\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: true\n",
      "      batch_size: 16\n",
      "      num_workers: 12\n",
      "      pin_memory: true\n",
      "    \n",
      "[NeMo W 2022-12-07 18:13:53 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: 9017_manifest_valid_local.json\n",
      "      sample_rate: 22050\n",
      "      sup_data_path: ./fastpitch_sup_data\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: 30\n",
      "      pitch_fmax: 512\n",
      "      pitch_norm: true\n",
      "      pitch_mean: 152.3\n",
      "      pitch_std: 64.0\n",
      "      use_beta_binomial_interpolator: true\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: false\n",
      "      batch_size: 16\n",
      "      num_workers: 8\n",
      "      pin_memory: true\n",
      "    \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 18:13:54 features:267] PADDING: 1\n"
     ]
    }
   ],
   "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": "code",
   "execution_count": 9,
   "id": "6e7dea70-ddb2-4e76-b6f2-1e2a913d6021",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_manifest = Path(split_artifact_dir)/f\"{SPEAKER_ID}_manifest_train_local.json\"\n",
    "valid_manifest = Path(split_artifact_dir)/f\"{SPEAKER_ID}_manifest_valid_local.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "23816772-df53-4bbe-91df-55d8ed7cd4f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "31831c8b-6946-4e53-885b-ae25f047ab64",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_manifest(manifest_path, hifigan_manifest_path=\"hifigan_train_ft.json\"):\n",
    "    \"Get records from the training manifest\"\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(manifest_path.name.split(\".\")[0] + \"_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",
    "    with open(hifigan_manifest_path, \"w\") as f:\n",
    "        for r in records:\n",
    "            f.write(json.dumps(r) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "565ad53a-f32a-44af-b8e6-38ccea3c2d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "process_manifest(train_manifest, \"hifigan_train_ft.json\")\n",
    "process_manifest(valid_manifest, \"hifigan_valid_ft.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02ca1f3b-b62b-4dd8-a9b3-fc0142058740",
   "metadata": {},
   "source": [
    "Just as before, we can fine tune using the underlying Pytorch Lightning code (or the `hifigan_finetune.py` script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "306d8dc0-70d0-45d4-830f-478b35ece019",
   "metadata": {},
   "outputs": [],
   "source": [
    "with initialize(config_path=\"./conf/hifigan\"):\n",
    "    cfg = compose(config_name=\"hifigan.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "fde818f1-11dc-465b-a251-c6302b60f86e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# datasets\n",
    "cfg.train_dataset = \"hifigan_train_ft.json\"\n",
    "cfg.validation_datasets = \"hifigan_valid_ft.json\"\n",
    "\n",
    "# model ckpt\n",
    "cfg.init_from_pretrained_model = \"tts_hifigan\"\n",
    "\n",
    "# iterate fast\n",
    "cfg.model.max_steps = 300\n",
    "cfg.model.optim.lr = 0.00001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1308cc6c-2b7d-4487-b2f2-d991a8b9a3b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "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 18:15:07 exp_manager:343] Experiments will be logged at exp_dir_hifi/HifiGan/2022-12-07_18-14-37\n",
      "[NeMo I 2022-12-07 18:15:07 exp_manager:733] WandBLogger has been set up\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 18:15:07 exp_manager:988] The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to 300. 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_hifi/HifiGan/2022-12-07_18-14-37')"
      ]
     },
     "execution_count": 20,
     "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": "code",
   "execution_count": 21,
   "id": "95e7bce2-4f2d-4ebb-ad54-bacd3accf813",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 18:15:09 data:930] Loading dataset from hifigan_train_ft.json.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "76it [00:00, 62307.88it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 18:15:09 data:954] Loaded dataset with 76 files.\n",
      "[NeMo I 2022-12-07 18:15:09 data:956] Dataset contains 0.08 hours.\n",
      "[NeMo I 2022-12-07 18:15:09 data:358] Pruned 2 files. Final dataset contains 74 files\n",
      "[NeMo I 2022-12-07 18:15:09 data:360] Pruned 0.00 hours. Final dataset contains 0.08 hours.\n",
      "[NeMo I 2022-12-07 18:15:09 data:930] Loading dataset from hifigan_valid_ft.json.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "2it [00:00, 3508.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 18:15:09 data:954] Loaded dataset with 2 files.\n",
      "[NeMo I 2022-12-07 18:15:09 data:956] Dataset contains 0.00 hours.\n",
      "[NeMo I 2022-12-07 18:15:09 data:358] Pruned 0 files. Final dataset contains 2 files\n",
      "[NeMo I 2022-12-07 18:15:09 data:360] Pruned 0.00 hours. Final dataset contains 0.00 hours.\n",
      "[NeMo I 2022-12-07 18:15:09 features:267] PADDING: 0\n",
      "[NeMo I 2022-12-07 18:15:09 features:275] STFT using exact pad\n",
      "[NeMo I 2022-12-07 18:15:09 features:267] PADDING: 0\n",
      "[NeMo I 2022-12-07 18:15:09 features:275] STFT using exact pad\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 18:15:10 cloud:56] Found existing object /home/tcapelle/.cache/torch/NeMo/NeMo_1.14.0rc0/tts_hifigan/e6da322f0f7e7dcf3f1900a9229a7e69/tts_hifigan.nemo.\n",
      "[NeMo I 2022-12-07 18:15:10 cloud:62] Re-using file from: /home/tcapelle/.cache/torch/NeMo/NeMo_1.14.0rc0/tts_hifigan/e6da322f0f7e7dcf3f1900a9229a7e69/tts_hifigan.nemo\n",
      "[NeMo I 2022-12-07 18:15:10 common:912] Instantiating model from pre-trained checkpoint\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 18:15: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.data.datalayers.MelAudioDataset\n",
      "      manifest_filepath: /home/fkreuk/data/train_finetune.txt\n",
      "      min_duration: 0.75\n",
      "      n_segments: 8192\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: true\n",
      "      batch_size: 64\n",
      "      num_workers: 4\n",
      "    \n",
      "[NeMo W 2022-12-07 18:15: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.data.datalayers.MelAudioDataset\n",
      "      manifest_filepath: /home/fkreuk/data/val_finetune.txt\n",
      "      min_duration: 3\n",
      "      n_segments: 66150\n",
      "    dataloader_params:\n",
      "      drop_last: false\n",
      "      shuffle: false\n",
      "      batch_size: 5\n",
      "      num_workers: 4\n",
      "    \n",
      "[NeMo W 2022-12-07 18:15:13 features:244] Using torch_stft is deprecated and has been removed. The values have been forcibly set to False for FilterbankFeatures and AudioToMelSpectrogramPreprocessor. Please set exact_pad to True as needed.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 18:15:13 features:267] PADDING: 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[NeMo W 2022-12-07 18:15:13 features:244] Using torch_stft is deprecated and has been removed. The values have been forcibly set to False for FilterbankFeatures and AudioToMelSpectrogramPreprocessor. Please set exact_pad to True as needed.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[NeMo I 2022-12-07 18:15:13 features:267] PADDING: 0\n",
      "[NeMo I 2022-12-07 18:15:14 save_restore_connector:243] Model HifiGanModel was successfully restored from /home/tcapelle/.cache/torch/NeMo/NeMo_1.14.0rc0/tts_hifigan/e6da322f0f7e7dcf3f1900a9229a7e69/tts_hifigan.nemo.\n",
      "[NeMo I 2022-12-07 18:15:14 modelPT:1115] Model checkpoint restored from pretrained checkpoint with name : `tts_hifigan`\n"
     ]
    }
   ],
   "source": [
    "model = HifiGanModel(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": "code",
   "execution_count": 22,
   "id": "42914e5f-f25a-4744-99f4-ba22f23e1a0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name                       | Type                     | Params\n",
      "------------------------------------------------------------------------\n",
      "0 | audio_to_melspec_precessor | FilterbankFeatures       | 0     \n",
      "1 | trg_melspec_fn             | FilterbankFeatures       | 0     \n",
      "2 | generator                  | Generator                | 13.9 M\n",
      "3 | mpd                        | MultiPeriodDiscriminator | 41.1 M\n",
      "4 | msd                        | MultiScaleDiscriminator  | 29.6 M\n",
      "5 | feature_loss               | FeatureMatchingLoss      | 0     \n",
      "6 | discriminator_loss         | DiscriminatorLoss        | 0     \n",
      "7 | generator_loss             | GeneratorLoss            | 0     \n",
      "------------------------------------------------------------------------\n",
      "84.7 M    Trainable params\n",
      "0         Non-trainable params\n",
      "84.7 M    Total params\n",
      "338.643   Total estimated model params size (MB)\n",
      "Epoch 9, global step 100: 'val_loss' reached 0.43742 (best 0.43742), saving model to '/home/tcapelle/wandb/nvidia-workshop/exp_dir_hifi/HifiGan/2022-12-07_18-14-37/checkpoints/HifiGan--val_loss=0.4374-epoch=9.ckpt' as top 3\n",
      "Epoch 19, global step 200: 'val_loss' reached 0.42872 (best 0.42872), saving model to '/home/tcapelle/wandb/nvidia-workshop/exp_dir_hifi/HifiGan/2022-12-07_18-14-37/checkpoints/HifiGan--val_loss=0.4287-epoch=19.ckpt' as top 3\n",
      "Epoch 29, global step 300: 'val_loss' reached 0.44745 (best 0.42872), saving model to '/home/tcapelle/wandb/nvidia-workshop/exp_dir_hifi/HifiGan/2022-12-07_18-14-37/checkpoints/HifiGan--val_loss=0.4474-epoch=29.ckpt' as top 3\n",
      "`Trainer.fit` stopped: `max_steps=300` reached.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n",
      "wandb: Network error (ConnectTimeout), entering retry loop.\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "49c00343-481a-4609-989c-5e68222133c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>d_loss</td><td>█▂▅▃▃▃▄▅█▂▂▁▄▂▄</td></tr><tr><td>d_loss_mpd</td><td>▆▄█▅▁▁▄▄▂▃▁▃▄▃▇</td></tr><tr><td>d_loss_msd</td><td>▆▂▂▂▃▄▃▄█▂▂▁▃▂▂</td></tr><tr><td>epoch</td><td>▁▁▂▃▃▃▃▄▅▅▅▅▆▇▇▇██</td></tr><tr><td>epoch_time</td><td>▁▁▁▁▁▁▁▁▁▇▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁█</td></tr><tr><td>g_loss</td><td>▆▄▂█▄▄▃▃▃▁▃▂▃▄▄</td></tr><tr><td>g_loss_fm_mpd</td><td>▆▄▁▃█▇▆▃▄▄▅▃▆▅▅</td></tr><tr><td>g_loss_fm_msd</td><td>█▄▃▅▇▅▅▅▆▄▃▁▆▃▅</td></tr><tr><td>g_loss_gen_mpd</td><td>▂▄▁▅▅▇▅▄▆▆██▅▇▆</td></tr><tr><td>g_loss_gen_msd</td><td>▆▆▃█▃▇▂▆▁▃▆▃█▂▄</td></tr><tr><td>global_step</td><td>▁▁▂▃▃▃▄▅▅▅▆▇▇▇█</td></tr><tr><td>lr</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>lr-AdamW</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>lr-AdamW-1</td><td>▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁</td></tr><tr><td>train_backward_timing</td><td>▁▂▂▂▁▃▇▇▃▁▃█▄▂▂</td></tr><tr><td>train_step_timing</td><td>▁▁▁▂▁▁▂█▂▁▂▃▁▁▂</td></tr><tr><td>trainer/global_step</td><td>▁▁▁▂▁▁▂▃▃▂▂▂▁▃▂▂▃▃▅▃▃▅▃▃▃▃▆▆▄▄▄▇▇▄▄▄█▅▅▅</td></tr><tr><td>val_loss</td><td>▄▁█</td></tr><tr><td>validation_step_timing</td><td>▁█▃</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>d_loss</td><td>3.42249</td></tr><tr><td>d_loss_mpd</td><td>1.99808</td></tr><tr><td>d_loss_msd</td><td>1.42442</td></tr><tr><td>epoch</td><td>29</td></tr><tr><td>epoch_time</td><td>9.63934</td></tr><tr><td>g_loss</td><td>42.46178</td></tr><tr><td>g_loss_fm_mpd</td><td>8.54431</td></tr><tr><td>g_loss_fm_msd</td><td>5.5821</td></tr><tr><td>g_loss_gen_mpd</td><td>2.38173</td></tr><tr><td>g_loss_gen_msd</td><td>1.69788</td></tr><tr><td>global_step</td><td>300.0</td></tr><tr><td>lr</td><td>1e-05</td></tr><tr><td>lr-AdamW</td><td>1e-05</td></tr><tr><td>lr-AdamW-1</td><td>1e-05</td></tr><tr><td>train_backward_timing</td><td>0.07299</td></tr><tr><td>train_step_timing</td><td>0.26561</td></tr><tr><td>trainer/global_step</td><td>300</td></tr><tr><td>val_loss</td><td>0.44745</td></tr><tr><td>validation_step_timing</td><td>1.29292</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Synced <strong style=\"color:#cdcd00\">volcanic-field-8</strong>: <a href=\"https://wandb.ai/capecape/tts-workshop/runs/32pmu5t7\" target=\"_blank\">https://wandb.ai/capecape/tts-workshop/runs/32pmu5t7</a><br/>Synced 6 W&B file(s), 56 media file(s), 4 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20221207_181313-32pmu5t7/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2027f049-4440-4d31-b35e-d6140acb9b8b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Running as a python program on a shell"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "520a067a-e302-44a4-b39f-44aada814dd2",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "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_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."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4b72a25-5320-44bd-bdef-2114c32f68d6",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Adding more data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b844aba-6435-459f-9151-c62dc9995de1",
   "metadata": {
    "tags": []
   },
   "source": [
    "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.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65d1794f-c038-4a66-9345-9356f30ebcfc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}