TalkNet_Training_A100.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/justinjohn0306/TalkNET-colab/blob/main/TalkNet_Training_A100.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gss5Ox_RNiba"
      },
      "source": [
        "# TalkNet Training\n",
        "Last updated: 2023-05-17\n",
        "\n",
        "To train a 22KHz TalkNet, run the cells below and follow the instructions.\n",
        "\n",
        "This will take a while, and you might have to do it in multiple Colab sessions. The notebook will automatically resume training any models from the last saved checkpoint. If you're resuming from a new session, always re-run steps 1 through 5 first.\n",
        "\n",
        "##**IMPORTANT:**\n",
        "Your Trash folder on Drive will fill up with old checkpoints\n",
        "as you train the various models. Keep an eye on your Drive storage, and empty the trash if it starts to become full.\n",
        "\n",
        "- Fixes by ``justinjohn0306`` and ``Tapiocapioca#6641``"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LCqXqFgP2ri0",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 1:** Check which GPU you've been allocated.\n",
        "\n",
        "#@markdown You want a P100, V100, T4, V100 or A100.\n",
        "!nvidia-smi -L\n",
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j9TI-Q6m3qlx",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 2:** Mount Google Drive.\n",
        "from google.colab import drive\n",
        "drive.mount(\"/content/drive\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nfSawDUD5tqv",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 3:** Configure training data paths. Upload the following to your Drive and change the paths below:\n",
        "#@markdown * A dataset of .wav files, packaged as a .zip or .tar file\n",
        "#@markdown * Training and validation filelists, in LJSpeech format with relative paths (note: ARPABET transcripts are not supported)\n",
        "#@markdown * An output path for checkpoints\n",
        "\n",
        "import os\n",
        "\n",
        "dataset = \"/content/drive/My Drive/path_to_dataset.zip\" #@param {type:\"string\"}\n",
        "train_filelist = \"/content/drive/My Drive/train_filelist.txt\" #@param {type:\"string\"}\n",
        "val_filelist = \"/content/drive/My Drive/val_filelist.txt\" #@param {type:\"string\"}\n",
        "output_dir = \"/content/drive/My Drive/talknet/name_of_character\" #@param {type:\"string\"}\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",
        "print(\"OK\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        " **Note: The runtime will crash after step 4, which is intentional.** **Just re-run step 3 and proceed to step 5 and further.**"
      ],
      "metadata": {
        "id": "d02YMYEnqsGJ"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "teF-Ut8Z7Gjp",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 4:** Download NVIDIA NeMo.\n",
        "!pip install pip==24.0\n",
        "!pip uninstall torch torchvision torchaudio -y\n",
        "!pip uninstall gdown -y\n",
        "!pip install git+https://github.com/wkentaro/gdown.git\n",
        "import os\n",
        "import time\n",
        "import gdown\n",
        "\n",
        "os.chdir('/content')\n",
        "!apt-get install sox libsndfile1 ffmpeg\n",
        "!pip install \"cython<3.0.0\" && pip install --no-build-isolation \"pyyaml<6.0\" && pip install --no-build-isolation \"pysptk==0.2.0\"\n",
        "!pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117\n",
        "!pip install wget unidecode tensorflow==2.9 tensorboardX pysptk frozendict torch_stft pytorch-lightning==1.3.8 kaldiio pydub pyannote.audio g2p_en pesq pystoi crepe ffmpeg-python\n",
        "!python -m pip install git+https://github.com/SortAnon/NeMo.git\n",
        "!git clone -q https://github.com/SortAnon/hifi-gan.git\n",
        "!pip install --pre torchtext==0.6.0 --no-deps --quiet\n",
        "!pip install -q numpy --upgrade torchmetrics==0.6 omegaconf==2.2.3 hmmlearn==0.2.7 crepe==0.0.12 tensorboard==2.9 librosa==0.9.1 protobuf==3.20.0 torch_stft==0.1.4\n",
        "\n",
        "\n",
        "!mkdir -p conf && cd conf \\\n",
        "&& wget https://raw.githubusercontent.com/SortAnon/NeMo/main/examples/tts/conf/talknet-durs.yaml \\\n",
        "&& wget https://raw.githubusercontent.com/SortAnon/NeMo/main/examples/tts/conf/talknet-pitch.yaml \\\n",
        "&& wget https://raw.githubusercontent.com/SortAnon/NeMo/main/examples/tts/conf/talknet-spect.yaml \\\n",
        "&& cd ..\n",
        "\n",
        "# Download pre-trained models\n",
        "zip_path = \"tts_en_talknet_1.0.0rc1.zip\"\n",
        "for i in range(10):\n",
        "    if not os.path.exists(zip_path) or os.stat(zip_path).st_size < 100:\n",
        "        gdown.download(\n",
        "            \"https://drive.google.com/uc?id=19wSym9mNEnmzLS9XdPlfNAW9_u-mP1hR\",\n",
        "            zip_path,\n",
        "            quiet=False,\n",
        "        )\n",
        "!unzip -qo {zip_path}\n",
        "\n",
        "# restart the runtime\n",
        "os.kill(os.getpid(), 9)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bxFr3Fdi_kOC",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 5:** Dataset processing, part 1.\n",
        "\n",
        "#@markdown If this step fails, try the following:\n",
        "#@markdown * Make sure your filelists are correct. They should have relative\n",
        "#@markdown paths that match the contents of the archive.\n",
        "\n",
        "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",
        "            size = os.stat(\n",
        "                os.path.join(os.path.dirname(inpath), lpath)\n",
        "            ).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('/content')\n",
        "if os.path.exists(\"/content/wavs\"):\n",
        "    shutil.rmtree(\"/content/wavs\")\n",
        "os.mkdir(\"wavs\")\n",
        "os.chdir(\"wavs\")\n",
        "if dataset[-4:] == \".zip\":\n",
        "    !unzip -q \"{dataset}\"\n",
        "elif dataset[-4:] == \".tar\":\n",
        "    !tar -xf \"{dataset}\"\n",
        "else:\n",
        "    raise Exception(\"Unknown extension for dataset\")\n",
        "if os.path.exists(\"/content/wavs/wavs\"):\n",
        "    shutil.move(\"/content/wavs/wavs\", \"/content/tempwavs\")\n",
        "    shutil.rmtree(\"/content/wavs\")\n",
        "    shutil.move(\"/content/tempwavs\", \"/content/wavs\")\n",
        "\n",
        "# Filelist for preprocessing\n",
        "os.chdir('/content')\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\", encoding=\"utf-8\") as f:\n",
        "    t = f.read().split(\"\\n\")\n",
        "with open(\"valfiles.txt\", encoding=\"utf-8\") as f:\n",
        "    v = f.read().split(\"\\n\")\n",
        "    all_filelist = t[:] + v[:]\n",
        "with open(\"/content/allfiles.txt\", \"w\", encoding=\"utf-8\") 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(\"/content/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\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sos9vsxPkIN7",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 6:** Dataset processing, part 2. This takes a while, but\n",
        "#@markdown you only have to run this once per dataset (results are saved to Drive).\n",
        "\n",
        "#@markdown If this step fails, try the following:\n",
        "#@markdown * Make sure your dataset only contains WAV files.\n",
        "\n",
        "# Extract phoneme duration\n",
        "\n",
        "import nltk\n",
        "nltk.download('averaged_perceptron_tagger_eng')\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}))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nM7-bMpKO7U2",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 7:** Train duration predictor.\n",
        "\n",
        "#@markdown If CUDA runs out of memory, try the following:\n",
        "#@markdown * Click on Runtime -> Restart runtime, re-run step 3, and try again.\n",
        "#@markdown * If that doesn't help, reduce the batch size (default 64).\n",
        "batch_size = 64 #@param {type:\"integer\"}\n",
        "\n",
        "epochs = 20\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 TalkNetDursModel\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 = \"trainfiles.json\"\n",
        "    cfg.validation_datasets = \"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 = \"/content/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",
        "GlobalHydra().clear()\n",
        "initialize(config_path=\"conf\")\n",
        "cfg = compose(config_name=\"talknet-durs\")\n",
        "train(cfg)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JLfm00NuJfon",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 8:** Train pitch predictor.\n",
        "\n",
        "#@markdown If CUDA runs out of memory, try the following:\n",
        "#@markdown * Click on Runtime -> Restart runtime, re-run step 3, and try again.\n",
        "#@markdown * If that doesn't help, reduce the batch size (default 64).\n",
        "batch_size = 64 #@param {type:\"integer\"}\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 = \"trainfiles.json\"\n",
        "    cfg.validation_datasets = \"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 = \"/content/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",
        "GlobalHydra().clear()\n",
        "initialize(config_path=\"conf\")\n",
        "cfg = compose(config_name=\"talknet-pitch\")\n",
        "train(cfg)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N9hh4WPHbCcn",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 9:** Train spectrogram generator. 200+ epochs are recommended.\n",
        "\n",
        "#@markdown This is the slowest of the three models to train, and the hardest to\n",
        "#@markdown get good results from. If your character sounds noisy or robotic,\n",
        "#@markdown try improving the dataset, or adjusting the epochs and learning rate.\n",
        "\n",
        "epochs = 200 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown If CUDA runs out of memory, try the following:\n",
        "#@markdown * Click on Runtime -> Restart runtime, re-run step 3, and try again.\n",
        "#@markdown * If that doesn't help, reduce the batch size (default 32).\n",
        "batch_size = 32 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Advanced settings. You can probably leave these at their defaults (1e-3, 3e-6, empty, checked).\n",
        "learning_rate = 1e-3 #@param {type:\"number\"}\n",
        "min_learning_rate = 3e-6 #@param {type:\"number\"}\n",
        "pretrained_path = \"\" #@param {type:\"string\"}\n",
        "load_checkpoints = True #@param {type:\"boolean\"}\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 = \"trainfiles.json\"\n",
        "    cfg.validation_datasets = \"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 = \"/content/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",
        "GlobalHydra().clear()\n",
        "initialize(config_path=\"conf\")\n",
        "cfg = compose(config_name=\"talknet-spect\")\n",
        "train(cfg)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ajfyfz2p9Ior",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 10:** Generate GTA spectrograms. This will help HiFi-GAN learn what your TalkNet model sounds like.\n",
        "\n",
        "#@markdown If this step fails, make sure you've finished training the spectrogram generator.\n",
        "\n",
        "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, \"/content/hifi-gan/training.txt\")\n",
        "shutil.copyfile(val_filelist, \"/content/hifi-gan/validation.txt\")\n",
        "fix_paths(\"/content/hifi-gan/training.txt\")\n",
        "fix_paths(\"/content/hifi-gan/validation.txt\")\n",
        "fix_paths(\"/content/allfiles.txt\")\n",
        "\n",
        "os.chdir('/content')\n",
        "indir = \"wavs\"\n",
        "outdir = \"hifi-gan/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(\"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"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yVBjGhRB9hUJ",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 11:** Train HiFi-GAN. 2,000+ steps are recommended.\n",
        "#@markdown Stop this cell to finish training the model.\n",
        "\n",
        "#@markdown If CUDA runs out of memory, click on Runtime -> Restart runtime, re-run step 3, and try again.\n",
        "#@markdown If this step still fails to start, make sure step 10 finished successfully.\n",
        "\n",
        "#@markdown Note: If the training process starts at step 2500000, delete the HiFiGAN folder and try again.\n",
        "\n",
        "import gdown\n",
        "d = 'https://drive.google.com/uc?id='\n",
        "\n",
        "os.chdir('/content/hifi-gan')\n",
        "assert os.path.exists(\"wavs\"), \"Spectrogram folder not found\"\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",
        "!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 \"/content/hifi-gan/training.txt\" \\\n",
        "--input_validation_file \"/content/hifi-gan/validation.txt\" \\\n",
        "--input_wavs_dir \"..\" --input_mels_dir \"wavs\"\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5OtwfrTT-blU",
        "cellView": "form"
      },
      "source": [
        "#@markdown **Step 12:** Package the models. They'll be saved to the output directory as [character_name]_TalkNet.zip.\n",
        "\n",
        "character_name = \"Character\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown When done, generate a Drive share link, with permissions set to \"Anyone with the link\".\n",
        "#@markdown You can then use it with the [Controllable TalkNet notebook](https://colab.research.google.com/drive/1aj6Jk8cpRw7SsN3JSYCv57CrR6s0gYPB)\n",
        "#@markdown by selecting \"Custom model\" as your character.\n",
        "\n",
        "#@markdown This cell will also move the training checkpoints and logs to the trash.\n",
        "#@markdown That should free up roughly 2 GB of space on your Drive (remember to empty your trash).\n",
        "#@markdown If you wish to keep them, uncheck this box.\n",
        "\n",
        "delete_checkpoints = True #@param {type:\"boolean\"}\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"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}