Demo/inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HwaQq4GRU_Nw"
   },
   "source": [
    "# StarGANv2-VC Demo (VCTK 20 Speakers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hCpoXuZeGKAn"
   },
   "source": [
    "### Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 24923,
     "status": "ok",
     "timestamp": 1613984920200,
     "user": {
      "displayName": "Yinghao Li",
      "photoUrl": "",
      "userId": "12798981472803960591"
     },
     "user_tz": 300
    },
    "id": "3on9IjGhVGTP",
    "outputId": "63a799f8-564d-48c2-fb0f-e66c0cd9fdb8"
   },
   "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",
    "\n",
    "from Utils.ASR.models import ASRCNN\n",
    "from Utils.JDC.model import JDCNet\n",
    "from models import Generator, MappingNetwork, StyleEncoder\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source: http://speech.ee.ntu.edu.tw/~jjery2243542/resource/model/is18/en_speaker_used.txt\n",
    "# Source: https://github.com/jjery2243542/voice_conversion\n",
    "\n",
    "speakers = [225,228,229,230,231,233,236,239,240,244,226,227,232,243,254,256,258,259,270,273]\n",
    "\n",
    "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 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 build_model(model_params={}):\n",
    "    args = Munch(model_params)\n",
    "    generator = Generator(args.dim_in, args.style_dim, args.max_conv_dim, w_hpf=args.w_hpf, F0_channel=args.F0_channel)\n",
    "    mapping_network = MappingNetwork(args.latent_dim, args.style_dim, args.num_domains, hidden_dim=args.max_conv_dim)\n",
    "    style_encoder = StyleEncoder(args.dim_in, args.style_dim, args.num_domains, args.max_conv_dim)\n",
    "    \n",
    "    nets_ema = Munch(generator=generator,\n",
    "                     mapping_network=mapping_network,\n",
    "                     style_encoder=style_encoder)\n",
    "\n",
    "    return nets_ema\n",
    "\n",
    "def compute_style(speaker_dicts):\n",
    "    reference_embeddings = {}\n",
    "    for key, (path, speaker) in speaker_dicts.items():\n",
    "        if path == \"\":\n",
    "            label = torch.LongTensor([speaker]).to('cuda')\n",
    "            latent_dim = starganv2.mapping_network.shared[0].in_features\n",
    "            ref = starganv2.mapping_network(torch.randn(1, latent_dim).to('cuda'), label)\n",
    "        else:\n",
    "            wave, sr = librosa.load(path, sr=24000)\n",
    "            audio, index = librosa.effects.trim(wave, top_db=30)\n",
    "            if sr != 24000:\n",
    "                wave = librosa.resample(wave, sr, 24000)\n",
    "            mel_tensor = preprocess(wave).to('cuda')\n",
    "\n",
    "            with torch.no_grad():\n",
    "                label = torch.LongTensor([speaker])\n",
    "                ref = starganv2.style_encoder(mel_tensor.unsqueeze(1), label)\n",
    "        reference_embeddings[key] = (ref, label)\n",
    "    \n",
    "    return reference_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load F0 model\n",
    "\n",
    "F0_model = JDCNet(num_class=1, seq_len=192)\n",
    "params = torch.load(\"Utils/JDC/bst.t7\")['net']\n",
    "F0_model.load_state_dict(params)\n",
    "_ = F0_model.eval()\n",
    "F0_model = F0_model.to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 43003,
     "status": "ok",
     "timestamp": 1613984938321,
     "user": {
      "displayName": "Yinghao Li",
      "photoUrl": "",
      "userId": "12798981472803960591"
     },
     "user_tz": 300
    },
    "id": "NZA3ot-oF5t-"
   },
   "outputs": [],
   "source": [
    "# load vocoder\n",
    "from parallel_wavegan.utils import load_model\n",
    "vocoder = load_model(\"Vocoder/checkpoint-400000steps.pkl\").to('cuda').eval()\n",
    "vocoder.remove_weight_norm()\n",
    "_ = vocoder.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 24462,
     "status": "ok",
     "timestamp": 1613985522414,
     "user": {
      "displayName": "Yinghao Li",
      "photoUrl": "",
      "userId": "12798981472803960591"
     },
     "user_tz": 300
    },
    "id": "Ou4367LCyefA",
    "outputId": "19c61f6f-f39a-43b9-9275-09418c2aebb4"
   },
   "outputs": [],
   "source": [
    "# load starganv2\n",
    "\n",
    "model_path = 'Models/VCTK20/epoch_00150.pth'\n",
    "\n",
    "with open('Models/VCTK20/config.yml') as f:\n",
    "    starganv2_config = yaml.safe_load(f)\n",
    "starganv2 = build_model(model_params=starganv2_config[\"model_params\"])\n",
    "params = torch.load(model_path, map_location='cpu')\n",
    "params = params['model_ema']\n",
    "_ = [starganv2[key].load_state_dict(params[key]) for key in starganv2]\n",
    "_ = [starganv2[key].eval() for key in starganv2]\n",
    "starganv2.style_encoder = starganv2.style_encoder.to('cuda')\n",
    "starganv2.mapping_network = starganv2.mapping_network.to('cuda')\n",
    "starganv2.generator = starganv2.generator.to('cuda')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load input wave\n",
    "selected_speakers = [273, 259, 258, 243, 254, 244, 236, 233, 230, 228]\n",
    "k = random.choice(selected_speakers)\n",
    "wav_path = 'Demo/VCTK-corpus/p' + str(k) + '/p' + str(k) + '_023.wav'\n",
    "audio, source_sr = librosa.load(wav_path, sr=24000)\n",
    "audio = audio / np.max(np.abs(audio))\n",
    "audio.dtype = np.float32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Convert by style encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# with reference, using style encoder\n",
    "speaker_dicts = {}\n",
    "for s in selected_speakers:\n",
    "    k = s\n",
    "    speaker_dicts['p' + str(s)] = ('Demo/VCTK-corpus/p' + str(k) + '/p' + str(k) + '_023.wav', speakers.index(s))\n",
    "\n",
    "reference_embeddings = compute_style(speaker_dicts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 333
    },
    "executionInfo": {
     "elapsed": 1424,
     "status": "ok",
     "timestamp": 1613986299525,
     "user": {
      "displayName": "Yinghao Li",
      "photoUrl": "",
      "userId": "12798981472803960591"
     },
     "user_tz": 300
    },
    "id": "T5tahObUyN-d",
    "outputId": "f4f38742-2235-4f59-cb2a-5008912cd870",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# conversion \n",
    "import time\n",
    "start = time.time()\n",
    "    \n",
    "source = preprocess(audio).to('cuda:0')\n",
    "keys = []\n",
    "converted_samples = {}\n",
    "reconstructed_samples = {}\n",
    "converted_mels = {}\n",
    "\n",
    "for key, (ref, _) in reference_embeddings.items():\n",
    "    with torch.no_grad():\n",
    "        f0_feat = F0_model.get_feature_GAN(source.unsqueeze(1))\n",
    "        out = starganv2.generator(source.unsqueeze(1), ref, F0=f0_feat)\n",
    "        \n",
    "        c = out.transpose(-1, -2).squeeze().to('cuda')\n",
    "        y_out = vocoder.inference(c)\n",
    "        y_out = y_out.view(-1).cpu()\n",
    "\n",
    "        if key not in speaker_dicts or speaker_dicts[key][0] == \"\":\n",
    "            recon = None\n",
    "        else:\n",
    "            wave, sr = librosa.load(speaker_dicts[key][0], sr=24000)\n",
    "            mel = preprocess(wave)\n",
    "            c = mel.transpose(-1, -2).squeeze().to('cuda')\n",
    "            recon = vocoder.inference(c)\n",
    "            recon = recon.view(-1).cpu().numpy()\n",
    "\n",
    "    converted_samples[key] = y_out.numpy()\n",
    "    reconstructed_samples[key] = recon\n",
    "\n",
    "    converted_mels[key] = out\n",
    "    \n",
    "    keys.append(key)\n",
    "end = time.time()\n",
    "print('total processing time: %.3f sec' % (end - start) )\n",
    "\n",
    "import IPython.display as ipd\n",
    "for key, wave in converted_samples.items():\n",
    "    print('Converted: %s' % key)\n",
    "    display(ipd.Audio(wave, rate=24000))\n",
    "    print('Reference (vocoder): %s' % key)\n",
    "    if reconstructed_samples[key] is not None:\n",
    "        display(ipd.Audio(reconstructed_samples[key], rate=24000))\n",
    "\n",
    "print('Original (vocoder):')\n",
    "wave, sr = librosa.load(wav_path, sr=24000)\n",
    "mel = preprocess(wave)\n",
    "c = mel.transpose(-1, -2).squeeze().to('cuda')\n",
    "with torch.no_grad():\n",
    "    recon = vocoder.inference(c)\n",
    "    recon = recon.view(-1).cpu().numpy()\n",
    "display(ipd.Audio(recon, rate=24000))\n",
    "print('Original:')\n",
    "display(ipd.Audio(wav_path, rate=24000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SWh3o9hvGvJt"
   },
   "source": [
    "#### Convert by mapping network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# no reference, using mapping network\n",
    "speaker_dicts = {}\n",
    "selected_speakers = [273, 259, 258, 243, 254, 244, 236, 233, 230, 228]\n",
    "for s in selected_speakers:\n",
    "    k = s\n",
    "    speaker_dicts['p' + str(s)] = ('', speakers.index(s))\n",
    "\n",
    "reference_embeddings = compute_style(speaker_dicts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# conversion \n",
    "import time\n",
    "start = time.time()\n",
    "    \n",
    "source = preprocess(audio).to('cuda:0')\n",
    "keys = []\n",
    "converted_samples = {}\n",
    "reconstructed_samples = {}\n",
    "converted_mels = {}\n",
    "\n",
    "for key, (ref, _) in reference_embeddings.items():\n",
    "    with torch.no_grad():\n",
    "        f0_feat = F0_model.get_feature_GAN(source.unsqueeze(1))\n",
    "        out = starganv2.generator(source.unsqueeze(1), ref, F0=f0_feat)\n",
    "        \n",
    "        c = out.transpose(-1, -2).squeeze().to('cuda')\n",
    "        y_out = vocoder.inference(c)\n",
    "        y_out = y_out.view(-1).cpu()\n",
    "\n",
    "        if key not in speaker_dicts or speaker_dicts[key][0] == \"\":\n",
    "            recon = None\n",
    "        else:\n",
    "            wave, sr = librosa.load(speaker_dicts[key][0], sr=24000)\n",
    "            mel = preprocess(wave)\n",
    "            c = mel.transpose(-1, -2).squeeze().to('cuda')\n",
    "            recon = vocoder.inference(c)\n",
    "            recon = recon.view(-1).cpu().numpy()\n",
    "\n",
    "    converted_samples[key] = y_out.numpy()\n",
    "    reconstructed_samples[key] = recon\n",
    "\n",
    "    converted_mels[key] = out\n",
    "    \n",
    "    keys.append(key)\n",
    "end = time.time()\n",
    "print('total processing time: %.3f sec' % (end - start) )\n",
    "\n",
    "import IPython.display as ipd\n",
    "for key, wave in converted_samples.items():\n",
    "    print('Converted: %s' % key)\n",
    "    display(ipd.Audio(wave, rate=24000))\n",
    "    print('Reference (vocoder): %s' % key)\n",
    "    if reconstructed_samples[key] is not None:\n",
    "        display(ipd.Audio(reconstructed_samples[key], rate=24000))\n",
    "\n",
    "print('Original (vocoder):')\n",
    "wave, sr = librosa.load(wav_path, sr=24000)\n",
    "mel = preprocess(wave)\n",
    "c = mel.transpose(-1, -2).squeeze().to('cuda')\n",
    "with torch.no_grad():\n",
    "    recon = vocoder.inference(c)\n",
    "    recon = recon.view(-1).cpu().numpy()\n",
    "display(ipd.Audio(recon, rate=24000))\n",
    "print('Original:')\n",
    "display(ipd.Audio(wav_path, rate=24000))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "hCpoXuZeGKAn"
   ],
   "name": "Starganv2_vc.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}