Train_VITS_restore_path.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Train VITS model using restore path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    In this notebook I restore from a model and start training from that checkpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from trainer import Trainer, TrainerArgs\n",
    "\n",
    "from TTS.tts.configs.shared_configs import BaseDatasetConfig\n",
    "from TTS.tts.configs.vits_config import VitsConfig\n",
    "from TTS.tts.datasets import load_tts_samples\n",
    "from TTS.tts.models.vits import Vits, VitsArgs\n",
    "from TTS.tts.utils.speakers import SpeakerManager\n",
    "from TTS.tts.utils.text.tokenizer import TTSTokenizer\n",
    "from TTS.utils.audio import AudioProcessor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = \"TTS_train\"\n",
    "#print(os.path.join(output_path, \"LJSpeech-1.1/\"))\n",
    "dataset_config = BaseDatasetConfig(\n",
    "    name=\"ljspeech\", meta_file_train=\"metadata.csv\", path=os.path.join(output_path, \"LJSpeech-1.1/\")\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Removed a few config, because the VITs model has roughly 83 million parameters\n",
    "# batch_group_size=5\n",
    "# num_loader_workers=8, -- changed to 4\n",
    "# eval_batch_size=16,\n",
    "\n",
    "config = VitsConfig(\n",
    "    run_name=\"vits_ljspeech_finetune\",\n",
    "    batch_size=64,\n",
    "    eval_batch_size=8,\n",
    "    num_loader_workers=8,\n",
    "    num_eval_loader_workers=4,\n",
    "    run_eval=False,\n",
    "    test_delay_epochs=-1,\n",
    "    epochs=100,\n",
    "    text_cleaner=\"english_cleaners\",\n",
    "    use_phonemes=True,\n",
    "    phoneme_language=\"en-us\",\n",
    "    phoneme_cache_path=os.path.join(output_path, \"phoneme_cache\"),\n",
    "    compute_input_seq_cache=True,\n",
    "    print_step=25,\n",
    "    print_eval=True,\n",
    "    mixed_precision=True,\n",
    "    output_path=output_path,\n",
    "    max_text_len=200,\n",
    "    datasets=[dataset_config],\n",
    "    cudnn_benchmark=False,\n",
    "    save_step=1000,\n",
    "    batch_group_size=2,\n",
    "    lr_gen=0.0001,\n",
    "    lr_disc=0.0001,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# INITIALIZE THE AUDIO PROCESSOR\n",
    "# Audio processor is used for feature extraction and audio I/O.\n",
    "# It mainly serves to the dataloader and the training loggers.\n",
    "ap = AudioProcessor.init_from_config(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# INITIALIZE THE TOKENIZER\n",
    "# Tokenizer is used to convert text to sequences of token IDs.\n",
    "# config is updated with the default characters if not defined in the config.\n",
    "tokenizer, config = TTSTokenizer.init_from_config(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Train Eval split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# LOAD DATA SAMPLES\n",
    "# Each sample is a list of ```[text, audio_file_path, speaker_name]```\n",
    "# You can define your custom sample loader returning the list of samples.\n",
    "# Or define your custom formatter and pass it to the `load_tts_samples`.\n",
    "# Check `TTS.tts.datasets.load_tts_samples` for more details.\n",
    "train_samples, eval_samples = load_tts_samples(\n",
    "    dataset_config,\n",
    "    eval_split=True,\n",
    "    eval_split_max_size=config.eval_split_max_size,\n",
    "    eval_split_size=config.eval_split_size,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Train: Get set go"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CONFIG_PATH=\"./VITS_MODEL_LJSPEECH/config.json\"\n",
    "config.save_json(CONFIG_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "MODEL_PATH = \"./VITS_MODEL_LJSPEECH/model_file.pth\"\n",
    "command = f\"python ./TTS/TTS/bin/train_tts.py --config_path {CONFIG_PATH} --restore_path {MODEL_PATH}\"\n",
    "os.system(command)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Tensorboard "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install tensorboard\n",
    "!tensorboard --logdir=TTS_train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# USe this cell to see if cuda memory does not remain stuck\n",
    "import torch\n",
    "print(torch.cuda.memory_reserved())\n",
    "print(torch.cuda.memory_allocated())\n",
    "torch.cuda.empty_cache()\n",
    "print(torch.cuda.memory_reserved())\n",
    "print(torch.cuda.memory_allocated())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.is_available()\n",
    "print(torch.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Test the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob, os\n",
    "output_path = \"TTS_train\"\n",
    "ckpt = output_path + \"best_model.pth\"\n",
    "config = output_path + \"config.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tts --text \"Is it recognize speech or wreck a nice beach?\" \\\n",
    "      --model_path $ckpt \\\n",
    "      --config_path $config \\\n",
    "      --out_path test1.wav"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Listen to the synthesized speech"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "IPython.display.Audio(\"test1.wav\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.10 ('py38tts': venv)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5058910fdc281ff0cea44e05d7ce1c0555f0d90c4a43084b3046c094e190c0de"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}