Demo/Inference_LJSpeech.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9adb7bd1",
   "metadata": {},
   "source": [
    "# StyleTTS Demo (LJSpeech)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6108384d",
   "metadata": {},
   "source": [
    "### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da84c60f",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a3ddcc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load packages\n",
    "import random\n",
    "import yaml\n",
    "from munch import Munch\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import torchaudio\n",
    "import librosa\n",
    "from nltk.tokenize import word_tokenize\n",
    "\n",
    "from models import *\n",
    "from utils import *\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbdc04c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a173af4",
   "metadata": {},
   "outputs": [],
   "source": [
    "_pad = \"$\"\n",
    "_punctuation = ';:,.!?¡¿—…\"«»“” '\n",
    "_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'\n",
    "_letters_ipa = \"ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ\"\n",
    "\n",
    "\n",
    "# Export all symbols:\n",
    "symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)\n",
    "\n",
    "dicts = {}\n",
    "for i in range(len((symbols))):\n",
    "    dicts[symbols[i]] = i\n",
    "\n",
    "class TextCleaner:\n",
    "    def __init__(self, dummy=None):\n",
    "        self.word_index_dictionary = dicts\n",
    "    def __call__(self, text):\n",
    "        indexes = []\n",
    "        for char in text:\n",
    "            try:\n",
    "                indexes.append(self.word_index_dictionary[char])\n",
    "            except KeyError:\n",
    "                print(char)\n",
    "        return indexes\n",
    "\n",
    "textclenaer = TextCleaner()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00ee05e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_mel = torchaudio.transforms.MelSpectrogram(\n",
    "    n_mels=80, n_fft=2048, win_length=1200, hop_length=300)\n",
    "mean, std = -4, 4\n",
    "\n",
    "def length_to_mask(lengths):\n",
    "    mask = torch.arange(lengths.max()).unsqueeze(0).expand(lengths.shape[0], -1).type_as(lengths)\n",
    "    mask = torch.gt(mask+1, lengths.unsqueeze(1))\n",
    "    return mask\n",
    "\n",
    "def preprocess(wave):\n",
    "    wave_tensor = torch.from_numpy(wave).float()\n",
    "    mel_tensor = to_mel(wave_tensor)\n",
    "    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std\n",
    "    return mel_tensor\n",
    "\n",
    "def compute_style(ref_dicts):\n",
    "    reference_embeddings = {}\n",
    "    for key, path in ref_dicts.items():\n",
    "        wave, sr = librosa.load(path, sr=24000)\n",
    "        audio, index = librosa.effects.trim(wave, top_db=30)\n",
    "        if sr != 24000:\n",
    "            audio = librosa.resample(audio, sr, 24000)\n",
    "        mel_tensor = preprocess(audio).to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            ref = model.style_encoder(mel_tensor.unsqueeze(1))\n",
    "        reference_embeddings[key] = (ref.squeeze(1), audio)\n",
    "    \n",
    "    return reference_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b9cecbe",
   "metadata": {},
   "source": [
    "### Load models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64fc4c0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load phonemizer\n",
    "import phonemizer\n",
    "global_phonemizer = phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True,  with_stress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54cfbe48",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load hifi-gan\n",
    "\n",
    "import sys\n",
    "sys.path.insert(0, \"./Demo/hifi-gan\")\n",
    "\n",
    "import glob\n",
    "import os\n",
    "import argparse\n",
    "import json\n",
    "import torch\n",
    "from scipy.io.wavfile import write\n",
    "from attrdict import AttrDict\n",
    "from vocoder import Generator\n",
    "import librosa\n",
    "import numpy as np\n",
    "import torchaudio\n",
    "\n",
    "h = None\n",
    "\n",
    "def load_checkpoint(filepath, device):\n",
    "    assert os.path.isfile(filepath)\n",
    "    print(\"Loading '{}'\".format(filepath))\n",
    "    checkpoint_dict = torch.load(filepath, map_location=device)\n",
    "    print(\"Complete.\")\n",
    "    return checkpoint_dict\n",
    "\n",
    "def scan_checkpoint(cp_dir, prefix):\n",
    "    pattern = os.path.join(cp_dir, prefix + '*')\n",
    "    cp_list = glob.glob(pattern)\n",
    "    if len(cp_list) == 0:\n",
    "        return ''\n",
    "    return sorted(cp_list)[-1]\n",
    "\n",
    "cp_g = scan_checkpoint(\"Vocoder/\", 'g_')\n",
    "\n",
    "config_file = os.path.join(os.path.split(cp_g)[0], 'config.json')\n",
    "with open(config_file) as f:\n",
    "    data = f.read()\n",
    "json_config = json.loads(data)\n",
    "h = AttrDict(json_config)\n",
    "\n",
    "device = torch.device(device)\n",
    "generator = Generator(h).to(device)\n",
    "\n",
    "state_dict_g = load_checkpoint(cp_g, device)\n",
    "generator.load_state_dict(state_dict_g['generator'])\n",
    "generator.eval()\n",
    "generator.remove_weight_norm()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02fb18a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load StyleTTS\n",
    "model_path = \"./Models/LJSpeech/epoch_2nd_00180.pth\"\n",
    "model_config_path = \"./Models/LJSpeech/config.yml\"\n",
    "\n",
    "config = yaml.safe_load(open(model_config_path))\n",
    "\n",
    "# load pretrained ASR model\n",
    "ASR_config = config.get('ASR_config', False)\n",
    "ASR_path = config.get('ASR_path', False)\n",
    "text_aligner = load_ASR_models(ASR_path, ASR_config)\n",
    "\n",
    "# load pretrained F0 model\n",
    "F0_path = config.get('F0_path', False)\n",
    "pitch_extractor = load_F0_models(F0_path)\n",
    "\n",
    "model = build_model(Munch(config['model_params']), text_aligner, pitch_extractor)\n",
    "\n",
    "params = torch.load(model_path, map_location='cpu')\n",
    "params = params['net']\n",
    "for key in model:\n",
    "    if key in params:\n",
    "        if not \"discriminator\" in key:\n",
    "            print('%s loaded' % key)\n",
    "            model[key].load_state_dict(params[key])\n",
    "_ = [model[key].eval() for key in model]\n",
    "_ = [model[key].to(device) for key in model]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b803110e",
   "metadata": {},
   "source": [
    "### Synthesize speech"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30e8ff2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get first 3 training sample as references\n",
    "\n",
    "train_path = config.get('train_data', None)\n",
    "val_path = config.get('val_data', None)\n",
    "train_list, val_list = get_data_path_list(train_path, val_path)\n",
    "\n",
    "ref_dicts = {}\n",
    "for j in range(3):\n",
    "    filename = train_list[j].split('|')[0]\n",
    "    name = filename.split('/')[-1].replace('.wav', '')\n",
    "    ref_dicts[name] = filename\n",
    "    \n",
    "reference_embeddings = compute_style(ref_dicts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24655f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# synthesize a text\n",
    "text = ''' StyleTTS is a style-based generative model for parallel TTS that can synthesize diverse speech with natural prosody from a reference speech utterance. '''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43e9f635",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tokenize\n",
    "ps = global_phonemizer.phonemize([text])\n",
    "ps = word_tokenize(ps[0])\n",
    "ps = ' '.join(ps)\n",
    "tokens = textclenaer(ps)\n",
    "tokens.insert(0, 0)\n",
    "tokens.append(0)\n",
    "tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca57469c",
   "metadata": {},
   "outputs": [],
   "source": [
    "converted_samples = {}\n",
    "\n",
    "with torch.no_grad():\n",
    "    input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)\n",
    "    m = length_to_mask(input_lengths).to(device)\n",
    "    t_en = model.text_encoder(tokens, input_lengths, m)\n",
    "        \n",
    "    for key, (ref, _) in reference_embeddings.items():\n",
    "        \n",
    "        s = ref.squeeze(1)\n",
    "        style = s\n",
    "        \n",
    "        d = model.predictor.text_encoder(t_en, style, input_lengths, m)\n",
    "\n",
    "        x, _ = model.predictor.lstm(d)\n",
    "        duration = model.predictor.duration_proj(x)\n",
    "        pred_dur = torch.round(duration.squeeze()).clamp(min=1)\n",
    "        \n",
    "        pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))\n",
    "        c_frame = 0\n",
    "        for i in range(pred_aln_trg.size(0)):\n",
    "            pred_aln_trg[i, c_frame:c_frame + int(pred_dur[i].data)] = 1\n",
    "            c_frame += int(pred_dur[i].data)\n",
    "\n",
    "        # encode prosody\n",
    "        en = (d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device))\n",
    "        style = s.expand(en.shape[0], en.shape[1], -1)\n",
    "\n",
    "        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)\n",
    "\n",
    "        out = model.decoder((t_en @ pred_aln_trg.unsqueeze(0).to(device)), \n",
    "                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))\n",
    "\n",
    "\n",
    "        c = out.squeeze()\n",
    "        y_g_hat = generator(c.unsqueeze(0))\n",
    "        y_out = y_g_hat.squeeze().cpu().numpy()\n",
    "\n",
    "        c = out.squeeze()\n",
    "        y_g_hat = generator(c.unsqueeze(0))\n",
    "        y_out = y_g_hat.squeeze()\n",
    "        \n",
    "        converted_samples[key] = y_out.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d7f7d5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import IPython.display as ipd\n",
    "for key, wave in converted_samples.items():\n",
    "    print('Synthesized: %s' % key)\n",
    "    display(ipd.Audio(wave, rate=24000))\n",
    "    try:\n",
    "        print('Reference: %s' % key)\n",
    "        display(ipd.Audio(reference_embeddings[key][-1], rate=24000))\n",
    "    except:\n",
    "        continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74fe14d9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a97c5e82",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}