{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "import torchaudio\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from utils.task import load_checkpoint\n",
    "from utils.hparams import get_hparams_from_file\n",
    "from model.models import SynthesizerTrn\n",
    "from text.symbols import symbols\n",
    "from text import tokenizer\n",
    "\n",
    "\n",
    "def get_text(text: str, hps) -> torch.LongTensor:\n",
    "    text_norm = tokenizer(text, hps.data.text_cleaners, language=hps.data.language)\n",
    "    text_norm = torch.LongTensor(text_norm)\n",
    "    return text_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = \"filelists/madasr23_test.csv\"\n",
    "output_path = \"/path/to/output/directory\"\n",
    "data = pd.read_csv(dataset_path, sep=\"|\")\n",
    "print(data.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MADASR23 batch inference\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = \"custom_base\"\n",
    "hps = get_hparams_from_file(f\"./datasets/{model}/config.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_g = SynthesizerTrn(len(symbols), hps.data.n_mels if hps.data.use_mel else hps.data.n_fft // 2 + 1, hps.train.segment_size // hps.data.hop_length, n_speakers=hps.data.n_speakers, **hps.model).cuda()\n",
    "_ = net_g.eval()\n",
    "\n",
    "_ = load_checkpoint(f\"./datasets/{model}/logs/G_15000.pth\", net_g, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyDataset(Dataset):\n",
    "    def __init__(self, dataframe, hps):\n",
    "        self.data = dataframe\n",
    "        self.hps = hps\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        sid_idx = self.data[\"sid_idx\"][idx]\n",
    "        sid = self.data[\"sid\"][idx]\n",
    "        phonemes = self.data[\"phonemes\"][idx]\n",
    "        stn_tst = get_text(phonemes, self.hps)\n",
    "        return sid_idx, sid, stn_tst, idx\n",
    "\n",
    "\n",
    "# Initialize the dataset and data loader\n",
    "dataset = MyDataset(data, hps)\n",
    "data_loader = DataLoader(dataset, batch_size=1, num_workers=8)\n",
    "\n",
    "for sid_idx, spk_id, stn_tst, i in tqdm(data_loader):\n",
    "    sid_idx = int(sid_idx)\n",
    "    spk_id = int(spk_id)\n",
    "    i = int(i)\n",
    "    stn_tst = stn_tst[0]\n",
    "    with torch.no_grad():\n",
    "        x_tst = stn_tst.cuda().unsqueeze(0)\n",
    "        x_tst_lengths = torch.LongTensor([stn_tst.size(0)]).cuda()\n",
    "        sid = torch.LongTensor([sid_idx]).cuda()\n",
    "        audio = net_g.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=0.667,\n",
    "                            noise_scale_w=0.8, length_scale=1)[0][0].data.cpu()\n",
    "        torchaudio.save(f\"{output_path}/{spk_id}_{i}.wav\", audio,\n",
    "                        hps.data.sample_rate, bits_per_sample=hps.data.bits_per_sample)\n",
    "\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Voice Conversion\n",
    "\n",
    "TODO: Add batch inference for voice conversion\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
