Demo/Inference_LJSpeech.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9adb7bd1",
   "metadata": {},
   "source": [
    "# StyleTTS 2 Demo (LJSpeech)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6108384d",
   "metadata": {},
   "source": [
    "### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96e173bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.manual_seed(0)\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "import random\n",
    "random.seed(0)\n",
    "\n",
    "import numpy as np\n",
    "np.random.seed(0)"
   ]
  },
  {
   "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 time\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",
    "from text_utils import TextCleaner\n",
    "textclenaer = TextCleaner()\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": "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": "48e7b644",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = yaml.safe_load(open(\"Models/LJSpeech/config.yml\"))\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",
    "# load BERT model\n",
    "from Utils.PLBERT.util import load_plbert\n",
    "BERT_path = config.get('PLBERT_dir', False)\n",
    "plbert = load_plbert(BERT_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffc18cf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = build_model(recursive_munch(config['model_params']), text_aligner, pitch_extractor, plbert)\n",
    "_ = [model[key].eval() for key in model]\n",
    "_ = [model[key].to(device) for key in model]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64529d5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "params_whole = torch.load(\"Models/LJSpeech/epoch_2nd_00100.pth\", map_location='cpu')\n",
    "params = params_whole['net']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "895d9706",
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in model:\n",
    "    if key in params:\n",
    "        print('%s loaded' % key)\n",
    "        try:\n",
    "            model[key].load_state_dict(params[key])\n",
    "        except:\n",
    "            from collections import OrderedDict\n",
    "            state_dict = params[key]\n",
    "            new_state_dict = OrderedDict()\n",
    "            for k, v in state_dict.items():\n",
    "                name = k[7:] # remove `module.`\n",
    "                new_state_dict[name] = v\n",
    "            # load params\n",
    "            model[key].load_state_dict(new_state_dict, strict=False)\n",
    "#             except:\n",
    "#                 _load(params[key], model[key])\n",
    "_ = [model[key].eval() for key in model]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a59db2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e30985ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = DiffusionSampler(\n",
    "    model.diffusion.diffusion,\n",
    "    sampler=ADPM2Sampler(),\n",
    "    sigma_schedule=KarrasSchedule(sigma_min=0.0001, sigma_max=3.0, rho=9.0), # empirical parameters\n",
    "    clamp=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b803110e",
   "metadata": {},
   "source": [
    "### Synthesize speech"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24655f46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# synthesize a text\n",
    "text = ''' StyleTTS 2 is a text-to-speech model that leverages style diffusion and adversarial training with large speech language models to achieve human-level text-to-speech synthesis. '''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca57469c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inference(text, noise, diffusion_steps=5, embedding_scale=1):\n",
    "    text = text.strip()\n",
    "    text = text.replace('\"', '')\n",
    "    ps = global_phonemizer.phonemize([text])\n",
    "    ps = word_tokenize(ps[0])\n",
    "    ps = ' '.join(ps)\n",
    "\n",
    "    tokens = textclenaer(ps)\n",
    "    tokens.insert(0, 0)\n",
    "    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(tokens.device)\n",
    "        text_mask = length_to_mask(input_lengths).to(tokens.device)\n",
    "\n",
    "        t_en = model.text_encoder(tokens, input_lengths, text_mask)\n",
    "        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())\n",
    "        d_en = model.bert_encoder(bert_dur).transpose(-1, -2) \n",
    "\n",
    "        s_pred = sampler(noise, \n",
    "              embedding=bert_dur[0].unsqueeze(0), num_steps=diffusion_steps,\n",
    "              embedding_scale=embedding_scale).squeeze(0)\n",
    "\n",
    "        s = s_pred[:, 128:]\n",
    "        ref = s_pred[:, :128]\n",
    "\n",
    "        d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)\n",
    "\n",
    "        x, _ = model.predictor.lstm(d)\n",
    "        duration = model.predictor.duration_proj(x)\n",
    "        duration = torch.sigmoid(duration).sum(axis=-1)\n",
    "        pred_dur = torch.round(duration.squeeze()).clamp(min=1)\n",
    "\n",
    "        pred_dur[-1] += 5\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",
    "        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)\n",
    "        out = model.decoder((t_en @ pred_aln_trg.unsqueeze(0).to(device)), \n",
    "                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))\n",
    "        \n",
    "    return out.squeeze().cpu().numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d438ef4f",
   "metadata": {},
   "source": [
    "#### Basic synthesis (5 diffusion steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d7f7d5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "noise = torch.randn(1,1,256).to(device)\n",
    "wav = inference(text, noise, diffusion_steps=5, embedding_scale=1)\n",
    "rtf = (time.time() - start) / (len(wav) / 24000)\n",
    "print(f\"RTF = {rtf:5f}\")\n",
    "import IPython.display as ipd\n",
    "display(ipd.Audio(wav, rate=24000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d5d9df0",
   "metadata": {},
   "source": [
    "#### With higher diffusion steps (more diverse)\n",
    "Since the sampler is ancestral, the higher the stpes, the more diverse the samples are, with the cost of slower synthesis speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a10129fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "noise = torch.randn(1,1,256).to(device)\n",
    "wav = inference(text, noise, diffusion_steps=10, embedding_scale=1)\n",
    "rtf = (time.time() - start) / (len(wav) / 24000)\n",
    "print(f\"RTF = {rtf:5f}\")\n",
    "import IPython.display as ipd\n",
    "display(ipd.Audio(wav, rate=24000))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1877ea15",
   "metadata": {},
   "source": [
    "### Speech expressiveness\n",
    "The following section recreates the samples shown in [Section 6](https://styletts2.github.io/#emo) of the demo page."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c4777b7",
   "metadata": {},
   "source": [
    "#### With embedding_scale=1\n",
    "This is the classifier-free guidance scale. The higher the scale, the more conditional the style is to the input text and hence more emotional. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c29ea2f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = {}\n",
    "texts['Happy'] = \"We are happy to invite you to join us on a journey to the past, where we will visit the most amazing monuments ever built by human hands.\"\n",
    "texts['Sad'] = \"I am sorry to say that we have suffered a severe setback in our efforts to restore prosperity and confidence.\"\n",
    "texts['Angry'] = \"The field of astronomy is a joke! Its theories are based on flawed observations and biased interpretations!\"\n",
    "texts['Surprised'] = \"I can't believe it! You mean to tell me that you have discovered a new species of bacteria in this pond?\"\n",
    "\n",
    "for k,v in texts.items():\n",
    "    noise = torch.randn(1,1,256).to(device)\n",
    "    wav = inference(v, noise, diffusion_steps=10, embedding_scale=1)\n",
    "    print(k + \": \")\n",
    "    display(ipd.Audio(wav, rate=24000, normalize=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c89499f",
   "metadata": {},
   "source": [
    "#### With embedding_scale=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f73be3aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts = {}\n",
    "texts['Happy'] = \"We are happy to invite you to join us on a journey to the past, where we will visit the most amazing monuments ever built by human hands.\"\n",
    "texts['Sad'] = \"I am sorry to say that we have suffered a severe setback in our efforts to restore prosperity and confidence.\"\n",
    "texts['Angry'] = \"The field of astronomy is a joke! Its theories are based on flawed observations and biased interpretations!\"\n",
    "texts['Surprised'] = \"I can't believe it! You mean to tell me that you have discovered a new species of bacteria in this pond?\"\n",
    "\n",
    "for k,v in texts.items():\n",
    "    noise = torch.randn(1,1,256).to(device)\n",
    "    wav = inference(v, noise, diffusion_steps=10, embedding_scale=2) # embedding_scale=2 for more pronounced emotion\n",
    "    print(k + \": \")\n",
    "    display(ipd.Audio(wav, rate=24000, normalize=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9320da63",
   "metadata": {},
   "source": [
    "### Long-form generation\n",
    "This section includes basic implementation of Algorithm 1 in the paper for consistent longform audio generation. The example passage is taken from [Section 5](https://styletts2.github.io/#long) of the demo page. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdd4db51",
   "metadata": {},
   "outputs": [],
   "source": [
    "passage = '''If the supply of fruit is greater than the family needs, it may be made a source of income by sending the fresh fruit to the market if there is one near enough, or by preserving, canning, and making jelly for sale. To make such an enterprise a success the fruit and work must be first class. There is magic in the word \"Homemade,\" when the product appeals to the eye and the palate; but many careless and incompetent people have found to their sorrow that this word has not magic enough to float inferior goods on the market. As a rule large canning and preserving establishments are clean and have the best appliances, and they employ chemists and skilled labor. The home product must be very good to compete with the attractive goods that are sent out from such establishments. Yet for first-class homemade products there is a market in all large cities. All first-class grocers have customers who purchase such goods.'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebb941c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def LFinference(text, s_prev, noise, alpha=0.7, diffusion_steps=5, embedding_scale=1):\n",
    "    text = text.strip()\n",
    "    text = text.replace('\"', '')\n",
    "    ps = global_phonemizer.phonemize([text])\n",
    "    ps = word_tokenize(ps[0])\n",
    "    ps = ' '.join(ps)\n",
    "\n",
    "    tokens = textclenaer(ps)\n",
    "    tokens.insert(0, 0)\n",
    "    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(tokens.device)\n",
    "        text_mask = length_to_mask(input_lengths).to(tokens.device)\n",
    "\n",
    "        t_en = model.text_encoder(tokens, input_lengths, text_mask)\n",
    "        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())\n",
    "        d_en = model.bert_encoder(bert_dur).transpose(-1, -2) \n",
    "\n",
    "        s_pred = sampler(noise, \n",
    "              embedding=bert_dur[0].unsqueeze(0), num_steps=diffusion_steps,\n",
    "              embedding_scale=embedding_scale).squeeze(0)\n",
    "        \n",
    "        if s_prev is not None:\n",
    "            # convex combination of previous and current style\n",
    "            s_pred = alpha * s_prev + (1 - alpha) * s_pred\n",
    "        \n",
    "        s = s_pred[:, 128:]\n",
    "        ref = s_pred[:, :128]\n",
    "\n",
    "        d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)\n",
    "\n",
    "        x, _ = model.predictor.lstm(d)\n",
    "        duration = model.predictor.duration_proj(x)\n",
    "        duration = torch.sigmoid(duration).sum(axis=-1)\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",
    "        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)\n",
    "        out = model.decoder((t_en @ pred_aln_trg.unsqueeze(0).to(device)), \n",
    "                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))\n",
    "        \n",
    "    return out.squeeze().cpu().numpy(), s_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ca0ef2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences = passage.split('.') # simple split by comma\n",
    "wavs = []\n",
    "s_prev = None\n",
    "for text in sentences:\n",
    "    if text.strip() == \"\": continue\n",
    "    text += '.' # add it back\n",
    "    noise = torch.randn(1,1,256).to(device)\n",
    "    wav, s_prev = LFinference(text, s_prev, noise, alpha=0.7, diffusion_steps=10, embedding_scale=1.5)\n",
    "    wavs.append(wav)\n",
    "display(ipd.Audio(np.concatenate(wavs), rate=24000, normalize=False))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "NLP",
   "language": "python",
   "name": "nlp"
  },
  "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
}