eval_audio.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "945496f2-8559-47ef-943e-14ffe110024b",
   "metadata": {},
   "source": [
    "# Audio Predictor Evaluation\n",
    "This notebook helps to see how audio predictor synthesize voice using existing phonemes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "97f69095-f2ce-4966-9544-04964e19eaa6",
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'prepare_textgrid_alignments' from 'utils.dataset' (/data/notebooks/supervoice/utils/dataset.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 19\u001b[0m\n\u001b[1;32m     17\u001b[0m \u001b[38;5;66;03m# Local\u001b[39;00m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmisc\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m plot_specgram, plot_waveform\n\u001b[0;32m---> 19\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdataset\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m prepare_textgrid_alignments\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01msupervoice\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01maudio\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m spectogram, load_mono_audio\n\u001b[1;32m     21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01msupervoice\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodel_audio\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AudioPredictor\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'prepare_textgrid_alignments' from 'utils.dataset' (/data/notebooks/supervoice/utils/dataset.py)"
     ]
    }
   ],
   "source": [
    "# Base\n",
    "import itertools\n",
    "from glob import glob\n",
    "from tqdm import tqdm\n",
    "import math\n",
    "import textgrid\n",
    "import random\n",
    "\n",
    "# ML\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DistributedSampler, DataLoader\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from IPython.display import Audio, display\n",
    "from vocos import Vocos\n",
    "\n",
    "# Local\n",
    "from utils.misc import plot_specgram, plot_waveform\n",
    "from utils.dataset import prepare_textgrid_alignments\n",
    "from supervoice.audio import spectogram, load_mono_audio\n",
    "from supervoice.model_audio import AudioPredictor\n",
    "from supervoice.tokenizer import Tokenizer\n",
    "from supervoice.tensors import drop_using_mask\n",
    "from supervoice.config import config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08046af5-d95f-4c9c-88ee-5d693d1b752b",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cpu\"\n",
    "\n",
    "# Vocoder\n",
    "vocoder = torch.hub.load(repo_or_dir='ex3ndr/supervoice-vocoder', model='bigvsan')\n",
    "vocoder = vocoder.to(device)\n",
    "def do_vocoder(src):\n",
    "    with torch.no_grad():\n",
    "        return vocoder.generate(src)\n",
    "\n",
    "# GPT\n",
    "phonemizator = torch.hub.load(repo_or_dir='ex3ndr/supervoice-gpt', model='phonemizer', force_reload = True)\n",
    "\n",
    "# Tokenizer\n",
    "tokenizer = Tokenizer(config)\n",
    "phoneme_duration = config.audio.hop_size / config.audio.sample_rate\n",
    "\n",
    "# Load index\n",
    "files = glob(\"datasets/eval-aligned/**/*.TextGrid\")\n",
    "files = [f[len(\"datasets/eval-aligned/\"):-len(\".TextGrid\")] for f in files]\n",
    "\n",
    "# Load textgrids\n",
    "tg = [textgrid.TextGrid.fromFile(\"datasets/eval-aligned/\" + f + \".TextGrid\") for f in files]\n",
    "\n",
    "# Load audio\n",
    "styles = [\"datasets/eval-prepared/\" + f + \".style.pt\" for f in files]\n",
    "files = [\"datasets/eval-prepared/\" + f + \".wav\" for f in files]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c07a148-f9d2-448d-b154-48a820c2c624",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_tokens(src, style, audio, condition = False):\n",
    "    aligned_phonemes = prepare_textgrid_alignments(src, style, audio.shape[1], phoneme_duration, not condition)\n",
    "    phonemes = []\n",
    "    styles = []\n",
    "    for t in aligned_phonemes:\n",
    "        for i in range(t[1]):\n",
    "            phonemes.append(t[0])\n",
    "            styles.append(0)\n",
    "    tokens = tokenizer(phonemes)\n",
    "    return tokens, phonemes, torch.tensor(styles)\n",
    "\n",
    "def audio_denormalize(src):\n",
    "    return (src * config.audio.norm_std) + config.audio.norm_mean\n",
    "\n",
    "def audio_normalize(src):\n",
    "    return (src - config.audio.norm_mean) / config.audio.norm_std\n",
    "\n",
    "def do_spectogram(src):\n",
    "    return spectogram(src, config.audio.n_fft, config.audio.n_mels, config.audio.hop_size, config.audio.win_size, config.audio.mel_norm, config.audio.mel_scale, config.audio.sample_rate)\n",
    "\n",
    "def empty_tokens(size):\n",
    "    return torch.zeros((size)).long().to(device)\n",
    "\n",
    "def empty_mask(size):\n",
    "    return torch.ones((size)).bool().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f989d707-50d9-444f-86bb-55a7a23881b4",
   "metadata": {},
   "source": [
    "### Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5e7a5b4-f033-4ed8-a0e5-c15052d002c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = AudioPredictor(config)\n",
    "model = model.to(device)\n",
    "checkpoint = torch.load(f'./output/audio_pitch.pt', map_location=\"cpu\")\n",
    "model.load_state_dict(checkpoint['model'])\n",
    "model.eval()\n",
    "print(\"Predictor at \", checkpoint['step'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a88f28b6-d8d9-4575-83fa-b8fbeae24d0f",
   "metadata": {},
   "source": [
    "### Load audio with phonemes\n",
    "This loads sample audio that we would be able to use to compare with synthesized one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ee03a91-3939-417a-b262-68d52810805e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load audio\n",
    "index = 300\n",
    "waveform = load_mono_audio(files[index], sample_rate = config.audio.sample_rate)\n",
    "style = torch.load(styles[index])\n",
    "spec = do_spectogram(waveform)\n",
    "spec = spec.to(device)\n",
    "tokens, _, token_styles = extract_tokens(tg[index], style, spec)\n",
    "tokens = tokens.to(device)\n",
    "token_styles = token_styles.to(device)\n",
    "spec = spec[:, :len(tokens)]\n",
    "\n",
    "# Source audio\n",
    "print(\"Source audio\")\n",
    "display(Audio(data=waveform, rate=config.audio.sample_rate))\n",
    "print(\"Re-synth audio\")\n",
    "display(Audio(data=do_vocoder(spec).cpu(), rate=config.audio.sample_rate))\n",
    "plot_specgram(spec.cpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fe85a57-bcd2-400d-a236-497c46ecae45",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load conditioning audio\n",
    "cond_index = 2107\n",
    "# filename = \"external_datasets/libritts-r-clean-100/60/121082/60_121082_000004_000001.wav\n",
    "cond_waveform = load_mono_audio(files[cond_index], sample_rate = config.audio.sample_rate)\n",
    "cond_style = torch.load(styles[cond_index])\n",
    "cond_spec = do_spectogram(cond_waveform)\n",
    "cond_tokens, _, cond_token_styles = extract_tokens(tg[cond_index], cond_style, cond_spec, True)\n",
    "cond_spec = cond_spec.to(device)\n",
    "cond_tokens = cond_tokens.to(device)\n",
    "cond_token_styles = cond_token_styles.to(device)\n",
    "cond_spec = cond_spec[:, :len(cond_tokens)]\n",
    "assert len(cond_tokens) == cond_spec.shape[1]\n",
    "\n",
    "# Source audio\n",
    "print(\"Conditioning audio\")\n",
    "display(Audio(data=cond_waveform, rate=config.audio.sample_rate))\n",
    "plot_specgram(cond_spec.cpu())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00d4ab09-4db7-4c9a-ab33-f75c977abd41",
   "metadata": {},
   "source": [
    "### Synthesize unconditioned voice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47d70e14-e316-44a1-9e8e-7a280f9ce97f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with torch.no_grad():\n",
    "#     sample_tokens = 250\n",
    "#     sample, trajectory = model.sample(\n",
    "#         tokens = empty_tokens(sample_tokens),\n",
    "#         audio = torch.zeros((sample_tokens, config.audio.n_mels)).to(device),\n",
    "#         mask = empty_mask(sample_tokens),\n",
    "#         steps = 4\n",
    "#     )\n",
    "#     sample = audio_denormalize(sample)\n",
    "#     trajectory = audio_denormalize(trajectory)\n",
    "\n",
    "# # Resynth\n",
    "# resynth = do_vocoder(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "\n",
    "# # Display audio\n",
    "# print(\"Synthesized without any input\")\n",
    "# display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "\n",
    "# # Display trajectory\n",
    "# # for i in reversed(range(len(trajectory))):\n",
    "# #     plot_specgram(trajectory[i].transpose(1,0).cpu())    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1154b1dc-bd77-4157-baf8-d39cd8bcd6d7",
   "metadata": {},
   "source": [
    "# Synthesize voice from phonemes\n",
    "This synthesizes random voice from existing phonemes (taken from another file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cefef11-80ce-403f-89fb-4e3fecbb87d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "\n",
    "    # Params\n",
    "    s_steps = 8\n",
    "    s_tokens = tokens\n",
    "    s_tokens_style = token_styles + 1\n",
    "    s_audio = torch.zeros((len(tokens), config.audio.n_mels)).to(device)\n",
    "    s_mask = empty_mask(len(tokens))\n",
    "\n",
    "    # Synthesize\n",
    "    sample, trajectory = model.sample(tokens = s_tokens, tokens_style = s_tokens_style, audio = s_audio, mask = s_mask, steps = s_steps, alpha = 0.5)\n",
    "    sample = audio_denormalize(sample)\n",
    "    trajectory = audio_denormalize(trajectory)\n",
    "\n",
    "# Resynth\n",
    "resynth = do_vocoder(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "\n",
    "# Display audio\n",
    "print(\"Source of phonemes\")\n",
    "display(Audio(data=waveform, rate=config.audio.sample_rate))\n",
    "print(\"Synthesized with phonemes provided\")\n",
    "display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "\n",
    "# Display trajectory\n",
    "# for i in reversed(range(len(trajectory))):\n",
    "#     plot_specgram(trajectory[i].transpose(1,0).cpu())    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692e5564-3d31-4c38-a7c3-39c1dcd5845e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "\n",
    "    # Params\n",
    "    s_steps = 8\n",
    "    s_tokens = tokens\n",
    "    s_tokens_style = torch.zeros(len(tokens)).long().to(device)\n",
    "    s_audio = torch.zeros((len(tokens), config.audio.n_mels)).to(device)\n",
    "    s_mask = empty_mask(len(tokens))\n",
    "\n",
    "    # Synthesize\n",
    "    sample, trajectory = model.sample(tokens = s_tokens, tokens_style = s_tokens_style, audio = s_audio, mask = s_mask, steps = s_steps, alpha = 0.5)\n",
    "    sample = audio_denormalize(sample)\n",
    "    trajectory = audio_denormalize(trajectory)\n",
    "\n",
    "# Resynth\n",
    "resynth = do_vocoder(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "\n",
    "# Display audio\n",
    "print(\"Source of phonemes\")\n",
    "display(Audio(data=waveform, rate=config.audio.sample_rate))\n",
    "print(\"Synthesized with phonemes provided without styles\")\n",
    "display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "\n",
    "# Display trajectory\n",
    "# for i in reversed(range(len(trajectory))):\n",
    "#     plot_specgram(trajectory[i].transpose(1,0).cpu())    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5be096c-9a23-428d-97b7-a0e3d0bca06f",
   "metadata": {},
   "source": [
    "### Restore segment\n",
    "Cut part of the audio and re-generate segment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0dbabd3-9872-4ab8-9c25-667ef7493a53",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with torch.no_grad():\n",
    "\n",
    "#     # Params\n",
    "#     s_steps = 4\n",
    "#     s_tokens = empty_tokens(len(tokens))\n",
    "#     s_audio = audio_normalize(spec.transpose(1,0))\n",
    "\n",
    "#     # Mask\n",
    "#     seq_len = len(tokens)\n",
    "#     s_mask = torch.zeros((seq_len)).bool().to(device)\n",
    "#     s_mask[math.floor(seq_len * 0.2): math.floor(seq_len * 0.6)] = True\n",
    "#     s_audio = drop_using_mask(s_audio, 0, s_mask)\n",
    "\n",
    "#     # Restore audio\n",
    "#     sample, trajectory = model.sample(tokens = s_tokens, audio = s_audio, mask = s_mask, steps = s_steps)\n",
    "#     sample = audio_denormalize(sample)\n",
    "#     trajectory = audio_denormalize(trajectory)\n",
    "\n",
    "# # Resynth\n",
    "# resynth = do_vocoder(sample.transpose(1,0).to(device)).detach().cpu()\n",
    "\n",
    "# # Display audio\n",
    "# print(\"Source\")\n",
    "# display(Audio(data=waveform, rate=config.audio.sample_rate))\n",
    "# print(\"Restored segment without tokens provided\")\n",
    "# display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "\n",
    "# # Display trajectory\n",
    "# plot_specgram(sample.transpose(1,0).cpu())\n",
    "# # for i in reversed(range(len(trajectory))):\n",
    "# #     plot_specgram(trajectory[i].transpose(1,0).cpu())    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe8cc1a0-602d-4707-887b-bea340573556",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "\n",
    "    # Params\n",
    "    s_steps = 4\n",
    "    s_tokens = tokens\n",
    "    s_tokens_style = token_styles + 1\n",
    "    s_audio = audio_normalize(spec.transpose(1,0))\n",
    "\n",
    "    # Mask\n",
    "    seq_len = len(tokens)\n",
    "    s_mask = torch.zeros((seq_len)).bool().to(device)\n",
    "    s_mask[math.floor(seq_len * 0.2): math.floor(seq_len * 0.6)] = True\n",
    "    s_audio = drop_using_mask(s_audio, 0, s_mask)\n",
    "\n",
    "    # Restore audio\n",
    "    sample, trajectory = model.sample(tokens = s_tokens, tokens_style = s_tokens_style, audio = s_audio, mask = s_mask, steps = s_steps)\n",
    "    sample = audio_denormalize(sample) * 1.2\n",
    "    trajectory = audio_denormalize(trajectory)\n",
    "\n",
    "# Resynth\n",
    "resynth = do_vocoder(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "\n",
    "# Display audio\n",
    "print(\"Source\")\n",
    "display(Audio(data=waveform, rate=config.audio.sample_rate))\n",
    "print(\"Restored segment with tokens provided\")\n",
    "display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "\n",
    "# Display trajectory\n",
    "plot_specgram(sample.transpose(1,0).cpu())\n",
    "# for i in reversed(range(len(trajectory))):\n",
    "#     plot_specgram(trajectory[i].transpose(1,0).cpu())    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a561cc3b-7790-4bcc-9121-c14e9980caac",
   "metadata": {},
   "source": [
    "### Conditioned TTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b84bf09f-3903-4cbf-99d4-0c4afed1aa2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def conditioned_tts(tokens, steps = 8, alpha = 0.5):\n",
    "    with torch.no_grad():\n",
    "\n",
    "        # Create tokens\n",
    "        s_audio = audio_normalize(torch.cat([cond_spec.transpose(1,0), torch.zeros(len(tokens) + 2, config.audio.n_mels, device=device)]))\n",
    "        s_tokens = torch.cat([\n",
    "            torch.tensor([tokenizer.begin_token_id], device=device), \n",
    "            cond_tokens[1:], \n",
    "            tokens, \n",
    "            torch.tensor([tokenizer.silence_token_id, tokenizer.end_token_id], device=device)])\n",
    "        s_tokens_style = torch.cat([\n",
    "            torch.tensor([0], device=device), \n",
    "            cond_token_styles[1:] + 1, \n",
    "            token_styles + 1, \n",
    "            torch.tensor([0, 0], device=device)])\n",
    "\n",
    "        # Create a mask\n",
    "        s_mask = torch.zeros((len(cond_tokens) + len(tokens) + 2)).bool().to(device)\n",
    "        s_mask[len(cond_tokens):len(s_mask)] = True\n",
    "\n",
    "        # Restore audio\n",
    "        sample, trajectory = model.sample(tokens = s_tokens, tokens_style = s_tokens_style, audio = s_audio, mask = s_mask, steps = s_steps, alpha = alpha)\n",
    "        sample = audio_denormalize(sample)\n",
    "        trajectory = audio_denormalize(trajectory)\n",
    "        sample = sample[len(cond_tokens):-2] * 1.2\n",
    "        trajectory = trajectory[:,len(cond_tokens):]\n",
    "        \n",
    "        return sample, trajectory\n",
    "\n",
    "sample, trajectory = conditioned_tts(tokens)\n",
    "\n",
    "# Resynth\n",
    "resynth = do_vocoder(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "\n",
    "# Display audio\n",
    "print(\"Source audio of phonemes and durations\")\n",
    "display(Audio(data=waveform, rate=config.audio.sample_rate))\n",
    "print(\"Synthesized voice\")\n",
    "display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "\n",
    "# Display trajectory\n",
    "# for i in reversed(range(len(trajectory))):\n",
    "#     plot_specgram(trajectory[i].transpose(1,0).cpu())    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34dea243-e50d-4a66-9017-b18c08b944ee",
   "metadata": {},
   "source": [
    "### Variativity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ffcd0f-a061-45f0-bd47-317a4dbf4c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in range(10):\n",
    "#     with torch.no_grad():\n",
    "#         sample, trajectory = model.sample(\n",
    "#             tokens = tokens, \n",
    "#             audio = torch.zeros((len(tokens), config.audio.n_mels)).to(device),  # Empty source audio\n",
    "#             mask = torch.ones((len(tokens))).bool().to(device), # Mask everything\n",
    "#             steps = 4\n",
    "#         )\n",
    "#         sample = audio_denormalize(sample) * 1.1\n",
    "#         resynth = do_vocoder(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "#         display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "#         # plot_specgram(sample.transpose(1,0).cpu())    \n",
    "\n",
    "# # Resynth\n",
    "# # resynth = generator(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "\n",
    "# # # Display audio\n",
    "# # display(Audio(data=waveform, rate=config.audio.sample_rate))\n",
    "# # display(Audio(data=resynth, rate=config.audio.sample_rate))\n",
    "\n",
    "# # # Display trajectory\n",
    "# # for i in reversed(range(len(trajectory))):\n",
    "# #     plot_specgram(trajectory[i].transpose(1,0).cpu())    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63326cfb-78df-4b7e-abff-ffd82009d4d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval = phonemizator(\"I'm Sorry Dave, I'm Afraid I Can't Do That.\", max_new_tokens = 1024)\n",
    "speedup = 1.1\n",
    "eval = [(char.replace('•', '<SIL>'), max(1, duration)) for char, duration in eval]\n",
    "print(eval)\n",
    "eval = [char for char, duration in eval for _ in range(round(duration * speedup))]\n",
    "\n",
    "# with torch.no_grad():\n",
    "\n",
    "#     # Params\n",
    "#     s_steps = 32\n",
    "#     s_tokens = torch.cat([torch.tensor([tokenizer.begin_token_id]), tokenizer(eval), torch.tensor([tokenizer.end_token_id])])\n",
    "#     s_tokens = s_tokens.to(device)\n",
    "#     s_audio = torch.zeros((len(s_tokens), config.audio.n_mels)).to(device)\n",
    "#     s_mask = empty_mask(len(s_tokens)).to(device)\n",
    "\n",
    "#     # Synthesize\n",
    "#     sample, trajectory = model.sample(tokens = s_tokens, audio = s_audio, mask = s_mask, steps = s_steps, alpha = 0.5)\n",
    "#     sample = audio_denormalize(sample) * 1.1\n",
    "#     trajectory = audio_denormalize(trajectory)\n",
    "\n",
    "sample, trajectory = conditioned_tts(tokenizer(eval).to(device), steps=4)\n",
    "\n",
    "# Resynth\n",
    "resynth = do_vocoder(sample.transpose(1,0).unsqueeze(0).to(device)).detach().cpu().squeeze(0)\n",
    "\n",
    "# Display audio\n",
    "print(\"Synthesized sound\")\n",
    "display(Audio(data=resynth, rate=config.audio.sample_rate))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8ce7943-4e11-4c0b-86c9-0edd02701bb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import pyworld as pw\n",
    "import torchaudio\n",
    "def load_mono_audio(src, sample_rate, device=None):\n",
    "\n",
    "    # Load audio\n",
    "    audio, sr = torchaudio.load(src)\n",
    "\n",
    "    # Move to device\n",
    "    if device is not None:\n",
    "        audio = audio.to(device)\n",
    "\n",
    "    # Resample\n",
    "    if sr != sample_rate:\n",
    "        audio = resampler(sr, sample_rate, device)(audio)\n",
    "        sr = sample_rate\n",
    "\n",
    "    # Convert to mono\n",
    "    if audio.shape[0] > 1:\n",
    "        audio = audio.mean(dim=0, keepdim=True)\n",
    "\n",
    "    # Convert to single dimension\n",
    "    audio = audio[0]\n",
    "\n",
    "    return audio\n",
    "    \n",
    "wf = load_mono_audio(\"./datasets/libritts-prepared/00000012/00001341.wav\", sample_rate = config.audio.sample_rate)\n",
    "\n",
    "sp = do_spectogram(wf)\n",
    "f0, t = pw.dio(wf.squeeze(0).numpy().astype('double'), 24000, frame_period=1000 * 256/24000)\n",
    "display(Audio(data=wf, rate=config.audio.sample_rate))\n",
    "_, axis = plt.subplots(1, 1, figsize=(20, 10))\n",
    "axis.imshow(sp.cpu(), cmap=\"viridis\", vmin=-10, vmax=0, origin=\"lower\", aspect=\"auto\")\n",
    "axis.set_title(\"Spectogram\")\n",
    "axis.plot(f0, color=\"red\")\n",
    "# Draw annotations\n",
    "for span in tg[210][0]:\n",
    "    span_start = math.floor(span.minTime / phoneme_duration)\n",
    "    span_end = math.floor(span.maxTime / phoneme_duration)\n",
    "    axis.axvspan(span_start, span_end, facecolor=\"None\", edgecolor=\"white\")\n",
    "    axis.annotate(span.mark, (span_start, 75), annotation_clip=True, color=\"white\")\n",
    "# for span in tg[210][1]:\n",
    "#     span_start = math.floor(span.minTime / phoneme_duration)\n",
    "#     span_end = math.floor(span.maxTime / phoneme_duration)\n",
    "#     axis.annotate(span.mark, (span_start, 70), annotation_clip=True, color=\"white\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aebf7a85-1ffa-4814-b7d1-21cc678f24f5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}