voice_clone.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "59859b9a-f338-4d36-82dc-5cec9ca73676",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "177\n"
     ]
    }
   ],
   "source": [
    "#StyleTTS2 imports\n",
    "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)\n",
    "\n",
    "# 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\n",
    "\n",
    "import sounddevice as sd\n",
    "from scipy.io.wavfile import write\n",
    "\n",
    "import whisper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f72784e2-47d9-4f4d-8b63-73bfa5075efb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/elf/brego/dev/src/StyleTTS2/styleenv/lib/python3.11/site-packages/torch/cuda/__init__.py:138: UserWarning: CUDA initialization: CUDA driver initialization failed, you might not have a CUDA gpu. (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:108.)\n",
      "  return torch._C._cuda_getDeviceCount() > 0\n",
      "/home/elf/brego/dev/src/StyleTTS2/styleenv/lib/python3.11/site-packages/torch/nn/utils/weight_norm.py:30: UserWarning: torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\n",
      "  warnings.warn(\"torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\")\n",
      "/home/elf/brego/dev/src/StyleTTS2/styleenv/lib/python3.11/site-packages/torch/nn/modules/rnn.py:82: UserWarning: dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.2 and num_layers=1\n",
      "  warnings.warn(\"dropout option adds dropout after all but last \"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bert loaded\n",
      "bert_encoder loaded\n",
      "predictor loaded\n",
      "decoder loaded\n",
      "text_encoder loaded\n",
      "predictor_encoder loaded\n",
      "style_encoder loaded\n",
      "diffusion loaded\n",
      "text_aligner loaded\n",
      "pitch_extractor loaded\n",
      "mpd loaded\n",
      "msd loaded\n",
      "wd loaded\n"
     ]
    }
   ],
   "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(path):\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_s = model.style_encoder(mel_tensor.unsqueeze(1))\n",
    "        ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))\n",
    "\n",
    "    return torch.cat([ref_s, ref_p], dim=1)\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# load phonemizer\n",
    "import phonemizer\n",
    "global_phonemizer = phonemizer.backend.EspeakBackend(language='en-us', preserve_punctuation=True,  with_stress=True)\n",
    "\n",
    "config = yaml.safe_load(open(\"Models/LibriTTS/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)\n",
    "\n",
    "model_params = recursive_munch(config['model_params'])\n",
    "model = build_model(model_params, text_aligner, pitch_extractor, plbert)\n",
    "_ = [model[key].eval() for key in model]\n",
    "_ = [model[key].to(device) for key in model]\n",
    "\n",
    "params_whole = torch.load(\"Models/LibriTTS/epochs_2nd_00020.pth\", map_location='cpu')\n",
    "params = params_whole['net']\n",
    "\n",
    "\n",
    "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]\n",
    "\n",
    "from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule\n",
    "\n",
    "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",
    ")\n",
    "\n",
    "def inference(text, ref_s, alpha = 0.3, beta = 0.7, diffusion_steps=5, embedding_scale=1):\n",
    "    text = text.strip()\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 = torch.LongTensor(tokens).to(device).unsqueeze(0)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)\n",
    "        text_mask = length_to_mask(input_lengths).to(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 = torch.randn((1, 256)).unsqueeze(1).to(device), \n",
    "                                          embedding=bert_dur,\n",
    "                                          embedding_scale=embedding_scale,\n",
    "                                            features=ref_s, # reference from the same speaker as the embedding\n",
    "                                             num_steps=diffusion_steps).squeeze(1)\n",
    "\n",
    "\n",
    "        s = s_pred[:, 128:]\n",
    "        ref = s_pred[:, :128]\n",
    "\n",
    "        ref = alpha * ref + (1 - alpha)  * ref_s[:, :128]\n",
    "        s = beta * s + (1 - beta)  * ref_s[:, 128:]\n",
    "\n",
    "        d = model.predictor.text_encoder(d_en, \n",
    "                                         s, input_lengths, text_mask)\n",
    "\n",
    "        x, _ = model.predictor.lstm(d)\n",
    "        duration = model.predictor.duration_proj(x)\n",
    "\n",
    "        duration = torch.sigmoid(duration).sum(axis=-1)\n",
    "        pred_dur = torch.round(duration.squeeze()).clamp(min=1)\n",
    "\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",
    "        if model_params.decoder.type == \"hifigan\":\n",
    "            asr_new = torch.zeros_like(en)\n",
    "            asr_new[:, :, 0] = en[:, :, 0]\n",
    "            asr_new[:, :, 1:] = en[:, :, 0:-1]\n",
    "            en = asr_new\n",
    "\n",
    "        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)\n",
    "\n",
    "        asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))\n",
    "        if model_params.decoder.type == \"hifigan\":\n",
    "            asr_new = torch.zeros_like(asr)\n",
    "            asr_new[:, :, 0] = asr[:, :, 0]\n",
    "            asr_new[:, :, 1:] = asr[:, :, 0:-1]\n",
    "            asr = asr_new\n",
    "\n",
    "        out = model.decoder(asr, \n",
    "                                F0_pred, N_pred, ref.squeeze().unsqueeze(0))\n",
    "    \n",
    "        \n",
    "    return out.squeeze().cpu().numpy()[..., :-50] # weird pulse at the end of the model, need to be fixed later"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dbb5cfce-f54c-4120-8f65-714216afca01",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/elf/brego/dev/src/StyleTTS2/styleenv/lib/python3.11/site-packages/whisper/transcribe.py:115: UserWarning: FP16 is not supported on CPU; using FP32 instead\n",
      "  warnings.warn(\"FP16 is not supported on CPU; using FP32 instead\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "You said: \n",
      "\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "list index out of range",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 17\u001b[0m\n\u001b[1;32m     15\u001b[0m noise \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrandn(\u001b[38;5;241m1\u001b[39m,\u001b[38;5;241m1\u001b[39m,\u001b[38;5;241m256\u001b[39m)\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m     16\u001b[0m ref_s \u001b[38;5;241m=\u001b[39m compute_style(voice_path)\n\u001b[0;32m---> 17\u001b[0m wav \u001b[38;5;241m=\u001b[39m \u001b[43minference\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtranscribed_text\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mref_s\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43malpha\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbeta\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.5\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdiffusion_steps\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m10\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43membedding_scale\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m     18\u001b[0m rtf \u001b[38;5;241m=\u001b[39m (time\u001b[38;5;241m.\u001b[39mtime() \u001b[38;5;241m-\u001b[39m start) \u001b[38;5;241m/\u001b[39m (\u001b[38;5;28mlen\u001b[39m(wav) \u001b[38;5;241m/\u001b[39m \u001b[38;5;241m24000\u001b[39m)\n\u001b[1;32m     19\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mIPython\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdisplay\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mipd\u001b[39;00m\n",
      "Cell \u001b[0;32mIn[2], line 90\u001b[0m, in \u001b[0;36minference\u001b[0;34m(text, ref_s, alpha, beta, diffusion_steps, embedding_scale)\u001b[0m\n\u001b[1;32m     88\u001b[0m text \u001b[38;5;241m=\u001b[39m text\u001b[38;5;241m.\u001b[39mstrip()\n\u001b[1;32m     89\u001b[0m ps \u001b[38;5;241m=\u001b[39m global_phonemizer\u001b[38;5;241m.\u001b[39mphonemize([text])\n\u001b[0;32m---> 90\u001b[0m ps \u001b[38;5;241m=\u001b[39m word_tokenize(\u001b[43mps\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m)\n\u001b[1;32m     91\u001b[0m ps \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m \u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mjoin(ps)\n\u001b[1;32m     92\u001b[0m tokens \u001b[38;5;241m=\u001b[39m textclenaer(ps)\n",
      "\u001b[0;31mIndexError\u001b[0m: list index out of range"
     ]
    }
   ],
   "source": [
    "fs = 44100  # Sample rate\n",
    "seconds = 6  # Duration of recording\n",
    "myrecording = sd.rec(int(seconds * fs), samplerate=fs, channels=2)\n",
    "sd.wait()  # Wait until recording is finished\n",
    "recorded_audio_path = 'recorded_audio.wav'\n",
    "write(recorded_audio_path, fs, myrecording)  # Save as WAV file \n",
    "\n",
    "whisper_model = whisper.load_model(\"base\")\n",
    "result = whisper_model.transcribe(\"recorded_audio.wav\")\n",
    "transcribed_text = result[\"text\"]\n",
    "print(\"You said: \" + transcribed_text + \"\\n\")\n",
    "\n",
    "voice_path = \"Demo/reference_audio/James.wav\"\n",
    "start = time.time()\n",
    "noise = torch.randn(1,1,256).to(device)\n",
    "ref_s = compute_style(voice_path)\n",
    "wav = inference(transcribed_text, ref_s, alpha=0.1, beta=0.5, diffusion_steps=10, embedding_scale=1)\n",
    "rtf = (time.time() - start) / (len(wav) / 24000)\n",
    "import IPython.display as ipd\n",
    "print('Original')\n",
    "display(ipd.Audio(recorded_audio_path, rate=24000, normalize=False)) \n",
    "print('Synthesized:')\n",
    "display(ipd.Audio(wav, rate=24000, normalize=False))\n",
    "print('Reference Voice:')\n",
    "display(ipd.Audio(voice_path, rate=24000, normalize=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ce7927-4672-4754-8bcf-c38667d8e01f",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3a0691d-9fa3-4e69-be62-1a541c534381",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (styleenv)",
   "language": "python",
   "name": "styleenv"
  },
  "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}