{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87f5579b-56f4-4ddc-9bc5-f073697b3628",
   "metadata": {},
   "outputs": [],
   "source": [
    "#vocoder\n",
    "import os\n",
    "\n",
    "from TTS.trainer import Trainer, TrainingArgs\n",
    "from TTS.utils.audio import AudioProcessor\n",
    "from TTS.vocoder.configs import HifiganConfig, MultibandMelganConfig\n",
    "from TTS.vocoder.datasets.preprocess import load_wav_data\n",
    "from TTS.vocoder.models.gan import GAN\n",
    "\n",
    "\n",
    "\n",
    "output_path =os.path.join(\"./model/vocoder/mbmelgan/\")\n",
    "\n",
    "config = MultibandMelganConfig(\n",
    "    batch_size=256,\n",
    "    eval_batch_size=16,\n",
    "    num_loader_workers=4,\n",
    "    num_eval_loader_workers=4,\n",
    "    run_eval=True,\n",
    "    test_delay_epochs=5,\n",
    "    epochs=1000,\n",
    "    seq_len=8192,\n",
    "    pad_short=2000,\n",
    "    use_noise_augment=True,\n",
    "    eval_split_size=10,\n",
    "    print_step=25,\n",
    "    print_eval=False,\n",
    "    mixed_precision=False,\n",
    "    lr_gen=1e-4,\n",
    "    lr_disc=1e-4,\n",
    "    data_path=\"\",\n",
    "    output_path=output_path,\n",
    "    steps_to_start_discriminator=0\n",
    ")\n",
    "\n",
    "\n",
    "# init audio processor\n",
    "ap = AudioProcessor(**config.audio.to_dict())\n",
    "\n",
    "# load training samples\n",
    "eval_samples, train_samples=[[],[]]\n",
    "for path in [\"./resample/halfLife/wavs\"]*50: \n",
    "    eval_samples_temp, train_samples_temp = load_wav_data(path, config.eval_split_size)\n",
    "    eval_samples+=eval_samples_temp\n",
    "    train_samples+=train_samples_temp\n",
    "    \n",
    "eval_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8710bba6-271a-4670-8484-7e934f35a47a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " > Generator Model: multiband_melgan_generator\n",
      " > Discriminator Model: melgan_multiscale_discriminator\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "fatal: not a git repository (or any of the parent directories): .git\n",
      "fatal: not a git repository (or any of the parent directories): .git\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " > Using CUDA:  True\n",
      " > Number of GPUs:  1\n",
      " > Restoring from checkpoint_150000.pth.tar ...\n",
      " > Restoring Model...\n",
      " > Restoring Optimizer...\n",
      " > Model restored from step 150000\n",
      "\n",
      " > Model has 6894446 parameters\n",
      " > Restoring best loss from  ...\n",
      " > Starting with loaded last best loss 1.6186144163734035.\n",
      "\n",
      "\u001b[4m\u001b[1m > EPOCH: 0/1000\u001b[0m\n",
      " --> ./model/vocoder/mbmelgan/coqui_tts-January-03-2022_10+17PM-0000000\n",
      "\n",
      "\u001b[1m > TRAINING (2022-01-03 22:17:40) \u001b[0m\n",
      "\n",
      "\u001b[1m   --> STEP: 0/48 -- GLOBAL_STEP: 150000\u001b[0m\n",
      "     | > G_stft_loss_mg: 1.18554  (1.18554)\n",
      "     | > G_stft_loss_sc: 0.51856  (0.51856)\n",
      "     | > G_subband_stft_loss_mg: 1.14449  (1.14449)\n",
      "     | > G_subband_stft_loss_sc: 0.54307  (0.54307)\n",
      "     | > G_mse_fake_loss: 0.38303  (0.38303)\n",
      "     | > G_gen_loss: 0.85205  (0.85205)\n",
      "     | > G_adv_loss: 0.95758  (0.95758)\n",
      "     | > loss_0: 1.80963  (1.80963)\n",
      "     | > grad_norm_0: 9910.39160  (9910.39160)\n",
      "     | > D_mse_gan_loss: 0.57532  (0.57532)\n",
      "     | > D_mse_gan_real_loss: 0.31135  (0.31135)\n",
      "     | > D_mse_gan_fake_loss: 0.23920  (0.23920)\n",
      "     | > loss_1: 0.57532  (0.57532)\n",
      "     | > grad_norm_1: 44.54868  (44.54868)\n",
      "     | > current_lr_0: 0.00010 \n",
      "     | > current_lr_1: 0.00010 \n",
      "     | > step_time: 1.85050  (1.85048)\n",
      "     | > loader_time: 8.34370  (8.34372)\n",
      "\n",
      "\n",
      " > CHECKPOINT : ./model/vocoder/mbmelgan/coqui_tts-January-03-2022_10+17PM-0000000/checkpoint_150000.pth.tar\n",
      "\n",
      "\u001b[1m   --> STEP: 25/48 -- GLOBAL_STEP: 150025\u001b[0m\n",
      "     | > G_stft_loss_mg: 1.39181  (1.35757)\n",
      "     | > G_stft_loss_sc: 0.89431  (0.71446)\n",
      "     | > G_subband_stft_loss_mg: 1.36580  (1.31034)\n",
      "     | > G_subband_stft_loss_sc: 0.90438  (0.72538)\n",
      "     | > G_mse_fake_loss: 0.91134  (0.77881)\n",
      "     | > G_gen_loss: 1.14306  (1.03602)\n",
      "     | > G_adv_loss: 2.27836  (1.94704)\n",
      "     | > loss_0: 3.42142  (2.98305)\n",
      "     | > grad_norm_0: 1361.07996  (1400.05530)\n",
      "     | > D_mse_gan_loss: 0.18951  (0.34498)\n",
      "     | > D_mse_gan_real_loss: 0.15920  (0.20557)\n",
      "     | > D_mse_gan_fake_loss: 0.06108  (0.15397)\n",
      "     | > loss_1: 0.18951  (0.34498)\n",
      "     | > grad_norm_1: 5.91609  (12.26910)\n",
      "     | > current_lr_0: 0.00010 \n",
      "     | > current_lr_1: 0.00010 \n",
      "     | > step_time: 1.57640  (1.58514)\n",
      "     | > loader_time: 0.02500  (0.45780)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# init model\n",
    "model = GAN(config)\n",
    "\n",
    "# init the trainer and 🚀\n",
    "trainer = Trainer(\n",
    "    TrainingArgs(restore_path=\"model/vocoder/mbmelgan/coqui_tts-December-14-2021_05+54PM-0000000/checkpoint_150000.pth.tar\"),\n",
    "    config,\n",
    "    output_path,\n",
    "    model=model,\n",
    "    train_samples=train_samples,\n",
    "    eval_samples=eval_samples,\n",
    "    training_assets={\"audio_processor\": ap},\n",
    ")\n",
    "trainer.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f427d13a-1e2a-4ab7-b146-951eb488b0ce",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:tf2]",
   "language": "python",
   "name": "conda-env-tf2-py"
  },
  "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
