TalkNet_Training_Offline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# TalkNet Training (offline)\n",
    "\n",
    "To train a 22KHz TalkNet, run the cells below and follow the instructions.\n",
    "\n",
    "The notebook will automatically resume training any models from the last saved checkpoint. If you're resuming from a new session, always re-run step 1 first."
   ],
   "metadata": {
    "id": "Gss5Ox_RNiba"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 1:** Configure training data paths. Create the following and change the paths below:\n",
    "* A dataset of .wav files, packaged as a .zip or .tar file\n",
    "* Training and validation filelists, in LJSpeech format with relative paths (note: ARPABET transcripts are not supported)\n",
    "* An output path for checkpoints"
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "### CHANGE THESE ###\n",
    "\n",
    "dataset = \"/absolute/path_to_dataset.zip\"\n",
    "train_filelist = \"/absolute/train_filelist.txt\"\n",
    "val_filelist = \"/absolute/val_filelist.txt\"\n",
    "output_dir = \"/absolute/name_of_character\"\n",
    "\n",
    "### ------------ ###\n",
    "\n",
    "import os\n",
    "import gdown\n",
    "import platform\n",
    "\n",
    "assert os.path.exists(dataset), \"Cannot find dataset\"\n",
    "assert os.path.exists(train_filelist), \"Cannot find training filelist\"\n",
    "assert os.path.exists(val_filelist), \"Cannot find validation filelist\"\n",
    "if not os.path.exists(output_dir):\n",
    "   os.makedirs(output_dir)\n",
    "\n",
    "if os.getcwd() != output_dir:\n",
    "    cwd = os.getcwd()\n",
    "else:\n",
    "    os.chdir(cwd)\n",
    "    \n",
    "# Download pre-trained models\n",
    "if not os.path.exists(\"talknet_spect.nemo\"):\n",
    "    print(\"Downloading pre-trained models...\")\n",
    "    zip_path = \"TalkNet_Pretrained.zip\"\n",
    "    if not os.path.exists(zip_path) or os.stat(zip_path).st_size < 100:\n",
    "        d = 'https://drive.google.com/uc?id='\n",
    "        gdown.download(d+\"19wSym9mNEnmzLS9XdPlfNAW9_u-mP1hR\", zip_path, quiet=False)\n",
    "    print(platform.system())\n",
    "    if \"Windows\" in platform.system():\n",
    "        !tar -xf {zip_path}\n",
    "    else:\n",
    "        !unzip -qo {zip_path}\n",
    "    os.remove(zip_path)\n",
    "\n",
    "print(\"OK\")"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "nfSawDUD5tqv"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 2:** Dataset processing, part 1.\n",
    "\n",
    "If this step fails, try the following:\n",
    "* Make sure your filelists are correct. They should have relative \n",
    "paths that match the contents of the archive."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import os\n",
    "import shutil\n",
    "import sys\n",
    "import json\n",
    "import nemo\n",
    "import torch\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "#from pysptk import sptk\n",
    "from pathlib import Path\n",
    "from tqdm.notebook import tqdm\n",
    "import ffmpeg\n",
    "\n",
    "def fix_transcripts(inpath):\n",
    "    found_arpabet = False\n",
    "    found_grapheme = False\n",
    "    with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
    "        lines = f.readlines()\n",
    "    with open(inpath, \"w\", encoding=\"utf8\") as f:\n",
    "        for l in lines:\n",
    "            if l.strip() == \"\":\n",
    "                continue\n",
    "            if \"{\" in l:\n",
    "                if not found_arpabet:\n",
    "                    print(\"Warning: Skipping ARPABET lines (not supported).\")\n",
    "                    found_arpabet = True\n",
    "            else:\n",
    "                f.write(l)\n",
    "                found_grapheme = True\n",
    "    assert found_grapheme, \"No non-ARPABET lines found in \" + inpath\n",
    "\n",
    "def generate_json(inpath, outpath):\n",
    "    output = \"\"\n",
    "    sample_rate = 22050\n",
    "    with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
    "        for l in f.readlines():\n",
    "            lpath = l.split(\"|\")[0].strip()\n",
    "            if lpath[:5] != \"wavs/\":\n",
    "                lpath = \"wavs/\" + lpath\n",
    "            lpath = os.path.join(output_dir, \"_training_data\", lpath)\n",
    "            size = os.stat(lpath).st_size\n",
    "            x = {\n",
    "                \"audio_filepath\": lpath,\n",
    "                \"duration\": size / (sample_rate * 2),\n",
    "                \"text\": l.split(\"|\")[1].strip(),\n",
    "            }\n",
    "            output += json.dumps(x) + \"\\n\"\n",
    "        with open(outpath, \"w\", encoding=\"utf8\") as w:\n",
    "            w.write(output)\n",
    "\n",
    "def convert_to_22k(inpath):\n",
    "    if inpath.strip()[-4:].lower() != \".wav\":\n",
    "        print(\"Warning: \" + inpath.strip() + \" is not a .wav file!\")\n",
    "        return\n",
    "    ffmpeg.input(inpath).output(\n",
    "        inpath + \"_22k.wav\",\n",
    "        ar=\"22050\",\n",
    "        ac=\"1\",\n",
    "        acodec=\"pcm_s16le\",\n",
    "        map_metadata=\"-1\",\n",
    "        fflags=\"+bitexact\",\n",
    "    ).overwrite_output().run(quiet=True)\n",
    "    os.remove(inpath)\n",
    "    os.rename(inpath + \"_22k.wav\", inpath)\n",
    "\n",
    "# Extract dataset\n",
    "os.chdir(output_dir)\n",
    "\n",
    "if os.path.exists(\"_training_data\"):\n",
    "    shutil.rmtree(\"_training_data\")\n",
    "os.mkdir(\"_training_data\")\n",
    "os.chdir(\"_training_data\")\n",
    "os.mkdir(\"wavs\")\n",
    "os.chdir(\"wavs\")\n",
    "if dataset[-4:] == \".zip\":\n",
    "    if \"Windows\" in platform.system():\n",
    "        !tar -xf \"{dataset}\"\n",
    "    else:\n",
    "        !unzip -qo \"{dataset}\"\n",
    "elif dataset[-4:] == \".tar\":\n",
    "    !tar -xf \"{dataset}\"\n",
    "else:\n",
    "    raise Exception(\"Unknown extension for dataset\")\n",
    "if os.path.exists(os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\")):\n",
    "    shutil.move(\n",
    "        os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\"), \n",
    "        os.path.join(output_dir, \"_training_data\", \"tempwavs\")\n",
    "    )\n",
    "    shutil.rmtree(os.path.join(output_dir, \"_training_data\", \"wavs\"))\n",
    "    shutil.move(\n",
    "        os.path.join(output_dir, \"_training_data\", \"tempwavs\"), \n",
    "        os.path.join(output_dir, \"_training_data\", \"wavs\")\n",
    "    )\n",
    "# Filelist for preprocessing\n",
    "os.chdir(output_dir)\n",
    "shutil.copy(train_filelist, \"trainfiles.txt\")\n",
    "shutil.copy(val_filelist, \"valfiles.txt\")\n",
    "fix_transcripts(\"trainfiles.txt\")\n",
    "fix_transcripts(\"valfiles.txt\")\n",
    "seen_files = []\n",
    "with open(\"trainfiles.txt\") as f:\n",
    "    t = f.read().split(\"\\n\")\n",
    "with open(\"valfiles.txt\") as f:\n",
    "    v = f.read().split(\"\\n\")\n",
    "    all_filelist = t[:] + v[:]\n",
    "with open(\"allfiles.txt\", \"w\") as f:\n",
    "    for x in all_filelist:\n",
    "        if x.strip() == \"\":\n",
    "            continue\n",
    "        if x.split(\"|\")[0] not in seen_files:\n",
    "            seen_files.append(x.split(\"|\")[0])\n",
    "            f.write(x.strip() + \"\\n\")\n",
    "\n",
    "# Ensure audio is 22k\n",
    "print(\"Converting audio...\")\n",
    "for r, _, f in os.walk(os.path.join(\"_training_data\", \"wavs\")):\n",
    "    for name in tqdm(f):\n",
    "        convert_to_22k(os.path.join(r, name))\n",
    "\n",
    "# Convert to JSON\n",
    "generate_json(\"trainfiles.txt\", \"trainfiles.json\")\n",
    "generate_json(\"valfiles.txt\", \"valfiles.json\")\n",
    "generate_json(\"allfiles.txt\", \"allfiles.json\")\n",
    "\n",
    "print(\"OK\")"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "bxFr3Fdi_kOC"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 3:** Dataset processing, part 2. This takes a while, but\n",
    "you only have to run this once per dataset (results are saved to Drive).\n",
    "\n",
    "If this step fails, try the following:\n",
    "* Make sure your dataset only contains WAV files."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "# Extract phoneme duration\n",
    "\n",
    "import json\n",
    "from nemo.collections.asr.models import EncDecCTCModel\n",
    "asr_model = EncDecCTCModel.from_pretrained(model_name=\"asr_talknet_aligner\").cpu().eval()\n",
    "\n",
    "def forward_extractor(tokens, log_probs, blank):\n",
    "    \"\"\"Computes states f and p.\"\"\"\n",
    "    n, m = len(tokens), log_probs.shape[0]\n",
    "    # `f[s, t]` -- max sum of log probs for `s` first codes\n",
    "    # with `t` first timesteps with ending in `tokens[s]`.\n",
    "    f = np.empty((n + 1, m + 1), dtype=float)\n",
    "    f.fill(-(10 ** 9))\n",
    "    p = np.empty((n + 1, m + 1), dtype=int)\n",
    "    f[0, 0] = 0.0  # Start\n",
    "    for s in range(1, n + 1):\n",
    "        c = tokens[s - 1]\n",
    "        for t in range((s + 1) // 2, m + 1):\n",
    "            f[s, t] = log_probs[t - 1, c]\n",
    "            # Option #1: prev char is equal to current one.\n",
    "            if s == 1 or c == blank or c == tokens[s - 3]:\n",
    "                options = f[s : (s - 2 if s > 1 else None) : -1, t - 1]\n",
    "            else:  # Is not equal to current one.\n",
    "                options = f[s : (s - 3 if s > 2 else None) : -1, t - 1]\n",
    "            f[s, t] += np.max(options)\n",
    "            p[s, t] = np.argmax(options)\n",
    "    return f, p\n",
    "\n",
    "\n",
    "def backward_extractor(f, p):\n",
    "    \"\"\"Computes durs from f and p.\"\"\"\n",
    "    n, m = f.shape\n",
    "    n -= 1\n",
    "    m -= 1\n",
    "    durs = np.zeros(n, dtype=int)\n",
    "    if f[-1, -1] >= f[-2, -1]:\n",
    "        s, t = n, m\n",
    "    else:\n",
    "        s, t = n - 1, m\n",
    "    while s > 0:\n",
    "        durs[s - 1] += 1\n",
    "        s -= p[s, t]\n",
    "        t -= 1\n",
    "    assert durs.shape[0] == n\n",
    "    assert np.sum(durs) == m\n",
    "    assert np.all(durs[1::2] > 0)\n",
    "    return durs\n",
    "\n",
    "def preprocess_tokens(tokens, blank):\n",
    "    new_tokens = [blank]\n",
    "    for c in tokens:\n",
    "        new_tokens.extend([c, blank])\n",
    "    tokens = new_tokens\n",
    "    return tokens\n",
    "\n",
    "data_config = {\n",
    "    'manifest_filepath': \"allfiles.json\",\n",
    "    'sample_rate': 22050,\n",
    "    'labels': asr_model.decoder.vocabulary,\n",
    "    'batch_size': 1,\n",
    "}\n",
    "\n",
    "parser = nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset.make_vocab(\n",
    "    notation='phonemes', punct=True, spaces=True, stresses=False, add_blank_at=\"last\"\n",
    ")\n",
    "\n",
    "dataset = nemo.collections.asr.data.audio_to_text._AudioTextDataset(\n",
    "    manifest_filepath=data_config['manifest_filepath'], sample_rate=data_config['sample_rate'], parser=parser,\n",
    ")\n",
    "\n",
    "dl = torch.utils.data.DataLoader(\n",
    "    dataset=dataset, batch_size=data_config['batch_size'], collate_fn=dataset.collate_fn, shuffle=False,\n",
    ")\n",
    "\n",
    "blank_id = asr_model.decoder.num_classes_with_blank - 1\n",
    "\n",
    "if os.path.exists(os.path.join(output_dir, \"durations.pt\")):\n",
    "    print(\"durations.pt already exists; skipping\")\n",
    "else:\n",
    "    dur_data = {}\n",
    "    for sample_idx, test_sample in tqdm(enumerate(dl), total=len(dl)):\n",
    "        log_probs, _, greedy_predictions = asr_model(\n",
    "            input_signal=test_sample[0], input_signal_length=test_sample[1]\n",
    "        )\n",
    "\n",
    "        log_probs = log_probs[0].cpu().detach().numpy()\n",
    "        seq_ids = test_sample[2][0].cpu().detach().numpy()\n",
    "\n",
    "        target_tokens = preprocess_tokens(seq_ids, blank_id)\n",
    "\n",
    "        f, p = forward_extractor(target_tokens, log_probs, blank_id)\n",
    "        durs = backward_extractor(f, p)\n",
    "\n",
    "        dur_key = Path(dl.dataset.collection[sample_idx].audio_file).stem\n",
    "        dur_data[dur_key] = {\n",
    "            'blanks': torch.tensor(durs[::2], dtype=torch.long).cpu().detach(), \n",
    "            'tokens': torch.tensor(durs[1::2], dtype=torch.long).cpu().detach()\n",
    "        }\n",
    "\n",
    "        del test_sample\n",
    "\n",
    "    torch.save(dur_data, os.path.join(output_dir, \"durations.pt\"))\n",
    "\n",
    "#Extract F0 (pitch)\n",
    "import crepe\n",
    "from scipy.io import wavfile\n",
    "\n",
    "def crepe_f0(audio_file, hop_length=256):\n",
    "    sr, audio = wavfile.read(audio_file)\n",
    "    audio_x = np.arange(0, len(audio)) / 22050.0\n",
    "    time, frequency, confidence, activation = crepe.predict(audio, sr, viterbi=True)\n",
    "\n",
    "    x = np.arange(0, len(audio), hop_length) / 22050.0\n",
    "    freq_interp = np.interp(x, time, frequency)\n",
    "    conf_interp = np.interp(x, time, confidence)\n",
    "    audio_interp = np.interp(x, audio_x, np.absolute(audio)) / 32768.0\n",
    "    weights = [0.5, 0.25, 0.25]\n",
    "    audio_smooth = np.convolve(audio_interp, np.array(weights)[::-1], \"same\")\n",
    "\n",
    "    conf_threshold = 0.25\n",
    "    audio_threshold = 0.0005\n",
    "    for i in range(len(freq_interp)):\n",
    "        if conf_interp[i] < conf_threshold:\n",
    "            freq_interp[i] = 0.0\n",
    "        if audio_smooth[i] < audio_threshold:\n",
    "            freq_interp[i] = 0.0\n",
    "\n",
    "    # Hack to make f0 and mel lengths equal\n",
    "    if len(audio) % hop_length == 0:\n",
    "        freq_interp = np.pad(freq_interp, pad_width=[0, 1])\n",
    "    return torch.from_numpy(freq_interp.astype(np.float32))\n",
    "\n",
    "if os.path.exists(os.path.join(output_dir, \"f0s.pt\")):\n",
    "    print(\"f0s.pt already exists; skipping\")\n",
    "else:\n",
    "    f0_data = {}\n",
    "    with open(\"allfiles.json\") as f:\n",
    "        for i, l in enumerate(f.readlines()):\n",
    "            print(str(i))\n",
    "            audio_path = json.loads(l)[\"audio_filepath\"]\n",
    "            f0_data[Path(audio_path).stem] = crepe_f0(audio_path)\n",
    "\n",
    "    # calculate f0 stats (mean & std) only for train set\n",
    "    with open(\"trainfiles.json\") as f:\n",
    "        train_ids = {Path(json.loads(l)[\"audio_filepath\"]).stem for l in f}\n",
    "    all_f0 = torch.cat([f0[f0 >= 1e-5] for f0_id, f0 in f0_data.items() if f0_id in train_ids])\n",
    "\n",
    "    F0_MEAN, F0_STD = all_f0.mean().item(), all_f0.std().item()        \n",
    "    print(\"F0_MEAN: \" + str(F0_MEAN) + \", F0_STD: \" + str(F0_STD))\n",
    "    torch.save(f0_data, os.path.join(output_dir, \"f0s.pt\"))\n",
    "    with open(os.path.join(output_dir, \"f0_info.json\"), \"w\") as f:\n",
    "        f.write(json.dumps({\"FO_MEAN\": F0_MEAN, \"F0_STD\": F0_STD}))\n",
    "\n",
    "print(\"OK\")"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "sos9vsxPkIN7"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 4:** Train duration predictor.\n",
    "\n",
    "If CUDA runs out of memory, try the following:\n",
    "* Click on Kernel -> Restart, re-run step 1, and try again.\n",
    "* If that doesn't help, reduce the batch size (default 64)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "batch_size = 64\n",
    "\n",
    "import os\n",
    "from hydra.experimental import compose, initialize\n",
    "from hydra.core.global_hydra import GlobalHydra\n",
    "from omegaconf import OmegaConf\n",
    "import pytorch_lightning as pl\n",
    "from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
    "from nemo.collections.tts.models import TalkNetDursModel\n",
    "from nemo.core.config import hydra_runner\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "epochs = 20\n",
    "learning_rate = 1e-3\n",
    "min_learning_rate = 3e-6\n",
    "load_checkpoints = True\n",
    "\n",
    "def train(cfg):\n",
    "    cfg.sample_rate = 22050\n",
    "    cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
    "    cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
    "    cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
    "    cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
    "    cfg.trainer.accelerator = \"dp\"\n",
    "    cfg.trainer.max_epochs = epochs\n",
    "    cfg.trainer.check_val_every_n_epoch = 5\n",
    "    cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
    "    cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
    "    cfg.model.optim.lr = learning_rate\n",
    "    cfg.model.optim.sched.min_lr = min_learning_rate\n",
    "    cfg.exp_manager.exp_dir = output_dir\n",
    "\n",
    "    # Find checkpoints\n",
    "    ckpt_path = \"\"\n",
    "    if load_checkpoints:\n",
    "      path0 = os.path.join(output_dir, \"TalkNetDurs\")\n",
    "      if os.path.exists(path0):\n",
    "          path1 = sorted(os.listdir(path0))\n",
    "          for i in range(len(path1)):\n",
    "              path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
    "              if os.path.exists(path2):\n",
    "                  match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
    "                  if len(match) > 0:\n",
    "                      ckpt_path = os.path.join(path2, match[0])\n",
    "                      print(\"Resuming training from \" + match[0])\n",
    "                      break\n",
    "    \n",
    "    if ckpt_path != \"\":\n",
    "        trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
    "        model = TalkNetDursModel(cfg=cfg.model, trainer=trainer)\n",
    "    else:\n",
    "        warmstart_path = os.path.join(cwd, \"talknet_durs.nemo\")\n",
    "        trainer = pl.Trainer(**cfg.trainer)\n",
    "        model = TalkNetDursModel.restore_from(warmstart_path, override_config_path=cfg)\n",
    "        model.set_trainer(trainer)\n",
    "        model.setup_training_data(cfg.model.train_ds)\n",
    "        model.setup_validation_data(cfg.model.validation_ds)\n",
    "        model.setup_optimization(cfg.model.optim)\n",
    "        print(\"Warm-starting from \" + warmstart_path)\n",
    "    exp_manager(trainer, cfg.get('exp_manager', None))\n",
    "    trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()])  # noqa\n",
    "    trainer.fit(model)\n",
    "\n",
    "os.chdir(cwd)\n",
    "GlobalHydra().clear()\n",
    "initialize(config_path=\"conf\")\n",
    "cfg = compose(config_name=\"talknet-durs\")\n",
    "train(cfg)\n"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "nM7-bMpKO7U2"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 5:** Train pitch predictor.\n",
    "\n",
    "If CUDA runs out of memory, try the following:\n",
    "* Click on Kernel -> Restart, re-run step 1, and try again.\n",
    "* If that doesn't help, reduce the batch size (default 64)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "batch_size = 64\n",
    "epochs = 50\n",
    "\n",
    "import json\n",
    "\n",
    "with open(os.path.join(output_dir, \"f0_info.json\"), \"r\") as f:\n",
    "    f0_info = json.load(f)\n",
    "    f0_mean = f0_info[\"FO_MEAN\"]\n",
    "    f0_std = f0_info[\"F0_STD\"]\n",
    "\n",
    "learning_rate = 1e-3\n",
    "min_learning_rate = 3e-6\n",
    "load_checkpoints = True\n",
    "\n",
    "import os\n",
    "from hydra.experimental import compose, initialize\n",
    "from hydra.core.global_hydra import GlobalHydra\n",
    "from omegaconf import OmegaConf\n",
    "import pytorch_lightning as pl\n",
    "from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
    "from nemo.collections.tts.models import TalkNetPitchModel\n",
    "from nemo.core.config import hydra_runner\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "def train(cfg):\n",
    "    cfg.sample_rate = 22050\n",
    "    cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
    "    cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
    "    cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
    "    cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
    "    cfg.trainer.accelerator = \"dp\"\n",
    "    cfg.trainer.max_epochs = epochs\n",
    "    cfg.trainer.check_val_every_n_epoch = 5\n",
    "    cfg.model.f0_mean=f0_mean\n",
    "    cfg.model.f0_std=f0_std\n",
    "    cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
    "    cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
    "    cfg.model.optim.lr = learning_rate\n",
    "    cfg.model.optim.sched.min_lr = min_learning_rate\n",
    "    cfg.exp_manager.exp_dir = output_dir\n",
    "\n",
    "    # Find checkpoints\n",
    "    ckpt_path = \"\"\n",
    "    if load_checkpoints:\n",
    "      path0 = os.path.join(output_dir, \"TalkNetPitch\")\n",
    "      if os.path.exists(path0):\n",
    "          path1 = sorted(os.listdir(path0))\n",
    "          for i in range(len(path1)):\n",
    "              path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
    "              if os.path.exists(path2):\n",
    "                  match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
    "                  if len(match) > 0:\n",
    "                      ckpt_path = os.path.join(path2, match[0])\n",
    "                      print(\"Resuming training from \" + match[0])\n",
    "                      break\n",
    "    \n",
    "    if ckpt_path != \"\":\n",
    "        trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
    "        model = TalkNetPitchModel(cfg=cfg.model, trainer=trainer)\n",
    "    else:\n",
    "        warmstart_path = os.path.join(cwd, \"talknet_pitch.nemo\")\n",
    "        trainer = pl.Trainer(**cfg.trainer)\n",
    "        model = TalkNetPitchModel.restore_from(warmstart_path, override_config_path=cfg)\n",
    "        model.set_trainer(trainer)\n",
    "        model.setup_training_data(cfg.model.train_ds)\n",
    "        model.setup_validation_data(cfg.model.validation_ds)\n",
    "        model.setup_optimization(cfg.model.optim)\n",
    "        print(\"Warm-starting from \" + warmstart_path)\n",
    "    exp_manager(trainer, cfg.get('exp_manager', None))\n",
    "    trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()])  # noqa\n",
    "    trainer.fit(model)\n",
    "\n",
    "os.chdir(cwd)\n",
    "GlobalHydra().clear()\n",
    "initialize(config_path=\"conf\")\n",
    "cfg = compose(config_name=\"talknet-pitch\")\n",
    "train(cfg)"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "JLfm00NuJfon"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 6:** Train spectrogram generator. 200+ epochs are recommended. \n",
    "\n",
    "This is the slowest of the three models to train, and the hardest to\n",
    "get good results from. If your character sounds noisy or robotic,\n",
    "try improving the dataset, or adjusting the epochs and learning rate.\n",
    "\n",
    "If CUDA runs out of memory, try the following:\n",
    "* Click on Kernel -> Restart, re-run step 1, and try again.\n",
    "* If that doesn't help, reduce the batch size (default 32)."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "epochs = 200\n",
    "batch_size = 32\n",
    "\n",
    "# Advanced settings. You can probably leave these at their defaults (1e-3, 3e-6, empty, checked).\n",
    "learning_rate = 1e-3\n",
    "min_learning_rate = 3e-6\n",
    "pretrained_path = \"\"\n",
    "load_checkpoints = True\n",
    "\n",
    "import os\n",
    "from hydra.experimental import compose, initialize\n",
    "from hydra.core.global_hydra import GlobalHydra\n",
    "from omegaconf import OmegaConf\n",
    "import pytorch_lightning as pl\n",
    "from nemo.collections.common.callbacks import LogEpochTimeCallback\n",
    "from nemo.collections.tts.models import TalkNetSpectModel\n",
    "from nemo.core.config import hydra_runner\n",
    "from nemo.utils.exp_manager import exp_manager\n",
    "\n",
    "def train(cfg):\n",
    "    cfg.sample_rate = 22050\n",
    "    cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n",
    "    cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n",
    "    cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n",
    "    cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n",
    "    cfg.trainer.accelerator = \"dp\"\n",
    "    cfg.trainer.max_epochs = epochs\n",
    "    cfg.trainer.check_val_every_n_epoch = 5\n",
    "    cfg.model.train_ds.dataloader_params.batch_size = batch_size\n",
    "    cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n",
    "    cfg.model.optim.lr = learning_rate\n",
    "    cfg.model.optim.sched.min_lr = min_learning_rate\n",
    "    cfg.exp_manager.exp_dir = output_dir\n",
    "\n",
    "    # Find checkpoints\n",
    "    ckpt_path = \"\"\n",
    "    if load_checkpoints:\n",
    "      path0 = os.path.join(output_dir, \"TalkNetSpect\")\n",
    "      if os.path.exists(path0):\n",
    "          path1 = sorted(os.listdir(path0))\n",
    "          for i in range(len(path1)):\n",
    "              path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
    "              if os.path.exists(path2):\n",
    "                  match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n",
    "                  if len(match) > 0:\n",
    "                      ckpt_path = os.path.join(path2, match[0])\n",
    "                      print(\"Resuming training from \" + match[0])\n",
    "                      break\n",
    "    \n",
    "    if ckpt_path != \"\":\n",
    "        trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n",
    "        model = TalkNetSpectModel(cfg=cfg.model, trainer=trainer)\n",
    "    else:\n",
    "        if pretrained_path != \"\":\n",
    "            warmstart_path = pretrained_path\n",
    "        else:\n",
    "            warmstart_path = os.path.join(cwd, \"talknet_spect.nemo\")\n",
    "        trainer = pl.Trainer(**cfg.trainer)\n",
    "        model = TalkNetSpectModel.restore_from(warmstart_path, override_config_path=cfg)\n",
    "        model.set_trainer(trainer)\n",
    "        model.setup_training_data(cfg.model.train_ds)\n",
    "        model.setup_validation_data(cfg.model.validation_ds)\n",
    "        model.setup_optimization(cfg.model.optim)\n",
    "        print(\"Warm-starting from \" + warmstart_path)\n",
    "    exp_manager(trainer, cfg.get('exp_manager', None))\n",
    "    trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()])  # noqa\n",
    "    trainer.fit(model)\n",
    "\n",
    "os.chdir(cwd)\n",
    "GlobalHydra().clear()\n",
    "initialize(config_path=\"conf\")\n",
    "cfg = compose(config_name=\"talknet-spect\")\n",
    "train(cfg)"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "N9hh4WPHbCcn"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 7:** Generate GTA spectrograms. This will help HiFi-GAN learn what your TalkNet model sounds like.\n",
    "\n",
    "If this step fails, make sure you've finished training the spectrogram generator."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import sys\n",
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from nemo.collections.tts.models import TalkNetSpectModel\n",
    "import shutil\n",
    "\n",
    "def fix_paths(inpath):\n",
    "    output = \"\"\n",
    "    with open(inpath, \"r\", encoding=\"utf8\") as f:\n",
    "        for l in f.readlines():\n",
    "            if l[:5].lower() != \"wavs/\":\n",
    "                output += \"wavs/\" + l\n",
    "            else:\n",
    "                output += l\n",
    "    with open(inpath, \"w\", encoding=\"utf8\") as w:\n",
    "        w.write(output)\n",
    "\n",
    "shutil.copyfile(train_filelist, os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n",
    "shutil.copyfile(val_filelist, os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n",
    "fix_paths(os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n",
    "fix_paths(os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n",
    "fix_paths(os.path.join(output_dir, \"allfiles.txt\"))\n",
    "\n",
    "os.chdir(cwd)\n",
    "outdir = os.path.join(output_dir, \"_training_data\", \"wavs\")\n",
    "if not os.path.exists(outdir):\n",
    "    os.mkdir(outdir)\n",
    "\n",
    "model_path = \"\"\n",
    "path0 = os.path.join(output_dir, \"TalkNetSpect\")\n",
    "if os.path.exists(path0):\n",
    "    path1 = sorted(os.listdir(path0))\n",
    "    for i in range(len(path1)):\n",
    "        path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
    "        if os.path.exists(path2):\n",
    "            match = [x for x in os.listdir(path2) if \"TalkNetSpect.nemo\" in x]\n",
    "            if len(match) > 0:\n",
    "                model_path = os.path.join(path2, match[0])\n",
    "                break\n",
    "assert model_path != \"\", \"TalkNetSpect.nemo not found\"\n",
    "\n",
    "dur_path = os.path.join(output_dir, \"durations.pt\")\n",
    "f0_path = os.path.join(output_dir, \"f0s.pt\")\n",
    "\n",
    "model = TalkNetSpectModel.restore_from(model_path)\n",
    "model.eval()\n",
    "with open(os.path.join(output_dir, \"allfiles.txt\"), \"r\", encoding=\"utf-8\") as f:\n",
    "    dataset = f.readlines()\n",
    "durs = torch.load(dur_path)\n",
    "f0s = torch.load(f0_path)\n",
    "\n",
    "for x in tqdm(dataset):\n",
    "    x_name = os.path.splitext(os.path.basename(x.split(\"|\")[0].strip()))[0]\n",
    "    x_tokens = model.parse(text=x.split(\"|\")[1].strip())\n",
    "    x_durs = (\n",
    "        torch.stack(\n",
    "            (\n",
    "                durs[x_name][\"blanks\"],\n",
    "                torch.cat((durs[x_name][\"tokens\"], torch.zeros(1).int())),\n",
    "            ),\n",
    "            dim=1,\n",
    "        )\n",
    "        .view(-1)[:-1]\n",
    "        .view(1, -1)\n",
    "        .to(\"cuda:0\")\n",
    "    )\n",
    "    x_f0s = f0s[x_name].view(1, -1).to(\"cuda:0\")\n",
    "    x_spect = model.force_spectrogram(tokens=x_tokens, durs=x_durs, f0=x_f0s)\n",
    "    rel_path = os.path.splitext(x.split(\"|\")[0].strip())[0][5:]\n",
    "    abs_dir = os.path.join(outdir, os.path.dirname(rel_path))\n",
    "    if abs_dir != \"\" and not os.path.exists(abs_dir):\n",
    "        os.makedirs(abs_dir, exist_ok=True)\n",
    "    np.save(os.path.join(outdir, rel_path + \".npy\"), x_spect.detach().cpu().numpy())\n"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "Ajfyfz2p9Ior"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 8:** Train HiFi-GAN. 2,000+ steps are recommended.\n",
    "Stop this cell to finish training the model.\n",
    "\n",
    "* Click on Kernel -> Restart, re-run step 1, and try again.\n",
    "If this step still fails to start, make sure step 7 finished successfully.\n",
    "\n",
    "Note: If the training process starts at step 2500000, delete the HiFiGAN folder and try again."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "import gdown\n",
    "d = 'https://drive.google.com/uc?id='\n",
    "\n",
    "os.chdir(os.path.join(cwd, \"hifi-gan\"))\n",
    "\n",
    "if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\")):\n",
    "    os.makedirs(os.path.join(output_dir, \"HiFiGAN\"))\n",
    "if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\")):\n",
    "    print(\"Downloading universal model...\")\n",
    "    gdown.download(d+\"1qpgI41wNXFcH-iKq1Y42JlBC9j0je8PW\", os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), quiet=False)\n",
    "    gdown.download(d+\"1O63eHZR9t1haCdRHQcEgMfMNxiOciSru\", os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\"), quiet=False)\n",
    "    start_from_universal = \"--warm_start True \"\n",
    "else:\n",
    "    start_from_universal = \"\"\n",
    "\n",
    "hifi_train = os.path.join(cwd, \"hifi-gan\", \"training.txt\")\n",
    "hifi_val = os.path.join(cwd, \"hifi-gan\", \"validation.txt\")\n",
    "hifi_wavs = os.path.join(output_dir, \"_training_data\")\n",
    "\n",
    "!python train.py --fine_tuning True --config config_v1b.json \\\n",
    "{start_from_universal} \\\n",
    "--checkpoint_interval 250 --checkpoint_path \"{os.path.join(output_dir, 'HiFiGAN')}\" \\\n",
    "--input_training_file \"{hifi_train}\" \\\n",
    "--input_validation_file \"{hifi_val}\" \\\n",
    "--input_wavs_dir \"{hifi_wavs}\"\n"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "yVBjGhRB9hUJ"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "**Step 9:** Package the models. They'll be saved to the output directory as [character_name]_TalkNet.zip.\n",
    "\n",
    "When done, upload it to Google Drive, with permissions set to \"Anyone with the link\". \n",
    "You can then use it with TalkNet by selecting \"Custom model\" as your character. \n",
    "\n",
    "This cell will also delete the training checkpoints and logs.\n",
    "That should free up roughly 2 GB of space.\n",
    "If you wish to keep them, set delete_checkpoints to False."
   ],
   "metadata": {}
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "source": [
    "character_name = \"Character\"\n",
    "delete_checkpoints = True\n",
    "\n",
    "import shutil\n",
    "from zipfile import ZipFile\n",
    "\n",
    "def find_talknet(model_dir):\n",
    "    ckpt_path = \"\"\n",
    "    path0 = os.path.join(output_dir, model_dir)\n",
    "    if os.path.exists(path0):\n",
    "        path1 = sorted(os.listdir(path0))\n",
    "        for i in range(len(path1)):\n",
    "            path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n",
    "            if os.path.exists(path2):\n",
    "                match = [x for x in os.listdir(path2) if \".nemo\" in x]\n",
    "                if len(match) > 0:\n",
    "                    ckpt_path = os.path.join(path2, match[0])\n",
    "                    break\n",
    "    assert ckpt_path != \"\", \"Couldn't find \" + model_dir\n",
    "    return ckpt_path\n",
    "\n",
    "durs_path = find_talknet(\"TalkNetDurs\")\n",
    "pitch_path = find_talknet(\"TalkNetPitch\")\n",
    "spect_path = find_talknet(\"TalkNetSpect\")\n",
    "assert os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\")), \"Couldn't find HiFi-GAN\"\n",
    "\n",
    "zip = ZipFile(os.path.join(output_dir, character_name + \"_TalkNet.zip\"), 'w')\n",
    "zip.write(durs_path, \"TalkNetDurs.nemo\")\n",
    "zip.write(pitch_path, \"TalkNetPitch.nemo\")\n",
    "zip.write(spect_path, \"TalkNetSpect.nemo\")\n",
    "zip.write(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), \"hifiganmodel\")\n",
    "zip.write(os.path.join(output_dir, \"HiFiGAN\", \"config.json\"), \"config.json\")\n",
    "zip.write(os.path.join(output_dir, \"f0_info.json\"), \"f0_info.json\")\n",
    "zip.close()\n",
    "print(\"Archived model to \" + os.path.join(output_dir, character_name + \"_TalkNet.zip\"))\n",
    "\n",
    "if delete_checkpoints:\n",
    "    shutil.rmtree((os.path.join(output_dir, \"TalkNetDurs\")))\n",
    "    shutil.rmtree((os.path.join(output_dir, \"TalkNetPitch\")))\n",
    "    shutil.rmtree((os.path.join(output_dir, \"TalkNetSpect\")))\n",
    "    shutil.rmtree((os.path.join(output_dir, \"HiFiGAN\")))\n"
   ],
   "outputs": [],
   "metadata": {
    "cellView": "form",
    "id": "5OtwfrTT-blU"
   }
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "TalkNet_Training_Offline.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}