notebooks/synthesis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "148f4bc0-c28e-4670-9a5e-4c7928ab8992",
   "metadata": {},
   "outputs": [],
   "source": [
    "%env CUDA_VISIBLE_DEVICES=\"0\"\n",
    "import sys\n",
    "sys.path.append('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d5876c0-b47e-4c80-9e9c-62550f81b64e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt\n",
    "from pathlib import Path\n",
    "\n",
    "import IPython.display as ipd\n",
    "import numpy as np\n",
    "import soundfile as sf\n",
    "import torch\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "from pflow.models.pflow_tts import pflowTTS\n",
    "from pflow.text import sequence_to_text, text_to_sequence\n",
    "from pflow.utils.model import denormalize\n",
    "from pflow.utils.utils import get_user_data_dir, intersperse\n",
    "\n",
    "from pflow.hifigan.config import v1\n",
    "from pflow.hifigan.denoiser import Denoiser\n",
    "from pflow.hifigan.env import AttrDict\n",
    "from pflow.hifigan.models import Generator as HiFiGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1a30306-588c-4f22-8d9b-e2676880b0e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "# This allows for real time code changes being reflected in the notebook, no need to restart the kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a312856b-01a9-4d75-a4c8-4666dffa0692",
   "metadata": {},
   "outputs": [],
   "source": [
    "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23b3bbf9",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d3f3db9",
   "metadata": {},
   "outputs": [],
   "source": [
    "get_user_data_dir()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88f3b3c3-d014-443b-84eb-e143cdec3e21",
   "metadata": {},
   "source": [
    "## Filepaths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7640a4c1-44ce-447c-a8ff-45012fb7bddd",
   "metadata": {},
   "outputs": [],
   "source": [
    "PFLOW_CHECKPOINT = \"\" #fill in the path to the pflow checkpoint\n",
    "HIFIGAN_CHECKPOINT = get_user_data_dir() / \"hifigan_T2_v1\"\n",
    "OUTPUT_FOLDER = \"synth_output\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6477a3a9-71f2-4d2f-bb86-bdf3e31c2461",
   "metadata": {},
   "source": [
    "## Load TTS model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26a16230-04ba-4825-a844-2fb5ab945e24",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model(checkpoint_path):\n",
    "    model = pflowTTS.load_from_checkpoint(checkpoint_path, map_location=device)\n",
    "    model.eval()\n",
    "    return model\n",
    "count_params = lambda x: f\"{sum(p.numel() for p in x.parameters()):,}\"\n",
    "\n",
    "\n",
    "model = load_model(PFLOW_CHECKPOINT)\n",
    "print(f\"Model loaded! Parameter count: {count_params(model)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_vocoder(checkpoint_path):\n",
    "    h = AttrDict(v1)\n",
    "    hifigan = HiFiGAN(h).to(device)\n",
    "    hifigan.load_state_dict(torch.load(checkpoint_path, map_location=device)['generator'])\n",
    "    _ = hifigan.eval()\n",
    "    hifigan.remove_weight_norm()\n",
    "    return hifigan\n",
    "\n",
    "vocoder = load_vocoder(HIFIGAN_CHECKPOINT)\n",
    "denoiser = Denoiser(vocoder, mode='zeros')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cbc2ba0-09ff-40e2-9e60-6b77b534f9fb",
   "metadata": {},
   "source": [
    "### Helper functions to synthesise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "880a1879-24fd-4757-849c-850339120796",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.inference_mode()\n",
    "def process_text(text: str):\n",
    "    x = torch.tensor(intersperse(text_to_sequence(text, ['english_cleaners2']), 0),dtype=torch.long, device=device)[None]\n",
    "    x_lengths = torch.tensor([x.shape[-1]],dtype=torch.long, device=device)\n",
    "    x_phones = sequence_to_text(x.squeeze(0).tolist())\n",
    "    return {\n",
    "        'x_orig': text,\n",
    "        'x': x,\n",
    "        'x_lengths': x_lengths,\n",
    "        'x_phones': x_phones\n",
    "    }\n",
    "\n",
    "\n",
    "@torch.inference_mode()\n",
    "def synthesise(text, prompt):\n",
    "    text_processed = process_text(text)\n",
    "    start_t = dt.datetime.now()\n",
    "    output = model.synthesise(\n",
    "        text_processed['x'], \n",
    "        text_processed['x_lengths'],\n",
    "        n_timesteps=n_timesteps,\n",
    "        temperature=temperature,\n",
    "        length_scale=length_scale,\n",
    "        prompt=prompt\n",
    "    )\n",
    "    # merge everything to one dict    \n",
    "    output.update({'start_t': start_t, **text_processed})\n",
    "    return output\n",
    "\n",
    "@torch.inference_mode()\n",
    "def to_waveform(mel, vocoder):\n",
    "    audio = vocoder(mel).clamp(-1, 1)\n",
    "    audio = denoiser(audio.squeeze(0), strength=0.00025).cpu().squeeze()\n",
    "    return audio.cpu().squeeze()\n",
    "    \n",
    "# def save_to_folder(filename: str, output: dict, folder: str):\n",
    "#     folder = Path(folder)\n",
    "#     folder.mkdir(exist_ok=True, parents=True)\n",
    "#     np.save(folder / f'{filename}', output['mel'].cpu().numpy())\n",
    "#     sf.write(folder / f'{filename}.wav', output['waveform'], 22050, 'PCM_24')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78f857e3-2ef7-4c86-b776-596c4d3cf875",
   "metadata": {},
   "source": [
    "## Setup text to synthesise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e0a9acd-0845-4192-ba09-b9683e28a3ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = [\n",
    "\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9da9e2d-99b9-4c6f-8a08-c828e2cba121",
   "metadata": {},
   "source": [
    "### Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d216e5-4895-4da8-9d24-9e61021d2556",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Number of ODE Solver steps\n",
    "n_timesteps = 10\n",
    "\n",
    "## Changes to the speaking rate\n",
    "length_scale=1.0\n",
    "\n",
    "## Sampling temperature\n",
    "temperature = 0.667"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "import glob\n",
    "wav_files = glob.glob(\"<path/to/LJSpeech-1.1/wavs>/*.wav\") ## fill in the path to the LJSpeech-1.1 dataset\n",
    "wav, sr = torchaudio.load(wav_files[0])\n",
    "from pflow.data.text_mel_datamodule import mel_spectrogram\n",
    "mel = mel_spectrogram(\n",
    "            wav,\n",
    "            1024,\n",
    "            80,\n",
    "            22050,\n",
    "            256,\n",
    "            1024,\n",
    "            0,\n",
    "            8000,\n",
    "            center=False,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = mel #load a mel spectrogram from a file and paste it here; check dimensions [batch, channels, time]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b93aac89-c7f8-4975-8510-4e763c9689f4",
   "metadata": {},
   "source": [
    "## Synthesis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a227963-aa12-43b9-a706-1168b6fc0ba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs, rtfs = [], []\n",
    "rtfs_w = []\n",
    "for i, text in enumerate(tqdm(texts)):\n",
    "    prompt = prompt[:,:,:264]\n",
    "    from pflow.utils.model import normalize\n",
    "    prompt = normalize(prompt, model.mel_mean, model.mel_std)\n",
    "    output = synthesise(text, prompt) #, torch.tensor([15], device=device, dtype=torch.long).unsqueeze(0))\n",
    "    output['waveform'] = to_waveform(output['mel'], vocoder)\n",
    "\n",
    "    # Compute Real Time Factor (RTF) with HiFi-GAN\n",
    "    t = (dt.datetime.now() - output['start_t']).total_seconds()\n",
    "    rtf_w = t * 22050 / (output['waveform'].shape[-1])\n",
    "\n",
    "    ## Pretty print\n",
    "    print(f\"{'*' * 53}\")\n",
    "    print(f\"Input text - {i}\")\n",
    "    print(f\"{'-' * 53}\")\n",
    "    print(output['x_orig'])\n",
    "    print(f\"{'*' * 53}\")\n",
    "    print(f\"Phonetised text - {i}\")\n",
    "    print(f\"{'-' * 53}\")\n",
    "    print(output['x_phones'])\n",
    "    print(f\"{'*' * 53}\")\n",
    "    print(f\"RTF:\\t\\t{output['rtf']:.6f}\")\n",
    "    print(f\"RTF Waveform:\\t{rtf_w:.6f}\")\n",
    "    rtfs.append(output['rtf'])\n",
    "    rtfs_w.append(rtf_w)\n",
    "\n",
    "    # Display the synthesised waveform\n",
    "    ipd.display(ipd.Audio(output['waveform'], rate=22050))\n",
    "\n",
    "    ## Save the generated waveform\n",
    "#     save_to_folder(i, output, OUTPUT_FOLDER)\n",
    "\n",
    "print(f\"Number of ODE steps: {n_timesteps}\")\n",
    "print(f\"Mean RTF:\\t\\t\\t\\t{np.mean(rtfs):.6f} ± {np.std(rtfs):.6f}\")\n",
    "print(f\"Mean RTF Waveform (incl. vocoder):\\t{np.mean(rtfs_w):.6f} ± {np.std(rtfs_w):.6f}\")"
   ]
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}