Add colab notebook
This commit is contained in:
James Betker
parent
54a946d0ae
commit
56b10cc54c
248
tortoise_tts.ipynb
Normal file
248
tortoise_tts.ipynb
Normal file
|
|
@ -0,0 +1,248 @@
|
|||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "tortoise-tts.ipynb",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "JrK20I32grP6"
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!git clone https://github.com/neonbjb/tortoise-tts.git\n",
|
||||
"%cd tortoise-tts\n",
|
||||
"!pip install -r requirements.txt"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Imports used through the rest of the notebook.\n",
|
||||
"import torch\n",
|
||||
"import torchaudio\n",
|
||||
"import torch.nn as nn\n",
|
||||
"import torch.nn.functional as F\n",
|
||||
"from tqdm import tqdm\n",
|
||||
"\n",
|
||||
"from utils.tokenizer import VoiceBpeTokenizer\n",
|
||||
"from models.discrete_diffusion_vocoder import DiscreteDiffusionVocoder\n",
|
||||
"from models.text_voice_clip import VoiceCLIP\n",
|
||||
"from models.dvae import DiscreteVAE\n",
|
||||
"from models.autoregressive import UnifiedVoice\n",
|
||||
"\n",
|
||||
"# These have some fairly interesting code that is hidden in the colab. Consider checking it out.\n",
|
||||
"from do_tts import download_models, load_discrete_vocoder_diffuser, load_conditioning, fix_autoregressive_output, do_spectrogram_diffusion"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Gen09NM4hONQ"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Download pretrained models and set up pretrained voice bank. Feel free to upload and add your own voices here.\n",
|
||||
"# To do so, upload two WAV files cropped to 5-10 seconds of someone speaking.\n",
|
||||
"download_models()\n",
|
||||
"preselected_cond_voices = {\n",
|
||||
" # Male voices\n",
|
||||
" 'dotrice': ['voices/dotrice/1.wav', 'voices/dotrice/2.wav'],\n",
|
||||
" 'harris': ['voices/harris/1.wav', 'voices/harris/2.wav'],\n",
|
||||
" 'lescault': ['voices/lescault/1.wav', 'voices/lescault/2.wav'],\n",
|
||||
" 'otto': ['voices/otto/1.wav', 'voices/otto/2.wav'],\n",
|
||||
" # Female voices\n",
|
||||
" 'atkins': ['voices/atkins/1.wav', 'voices/atkins/2.wav'],\n",
|
||||
" 'grace': ['voices/grace/1.wav', 'voices/grace/2.wav'],\n",
|
||||
" 'kennard': ['voices/kennard/1.wav', 'voices/kennard/2.wav'],\n",
|
||||
" 'mol': ['voices/mol/1.wav', 'voices/mol/2.wav'],\n",
|
||||
" }"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "SSleVnRAiEE2"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# This is the text that will be spoken.\n",
|
||||
"text = \"And took the other as just as fair, and having perhaps the better claim, because it was grassy and wanted wear.\"\n",
|
||||
"# This is the voice that will speak it.\n",
|
||||
"voice = 'atkins'\n",
|
||||
"# This is the number of samples we will generate from the DALLE-style model. More will produce better results, but will take longer to produce.\n",
|
||||
"# I don't recommend going less than 128.\n",
|
||||
"num_autoregressive_samples = 128"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "bt_aoxONjfL2"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Prepare data.\n",
|
||||
"tokenizer = VoiceBpeTokenizer()\n",
|
||||
"text = torch.IntTensor(tokenizer.encode(text)).unsqueeze(0).cuda()\n",
|
||||
"text = F.pad(text, (0,1)) # This may not be necessary.\n",
|
||||
"cond_paths = preselected_cond_voices[voice]\n",
|
||||
"conds = []\n",
|
||||
"for cond_path in cond_paths:\n",
|
||||
" c, cond_wav = load_conditioning(cond_path)\n",
|
||||
" conds.append(c)\n",
|
||||
"conds = torch.stack(conds, dim=1) # And just use the last cond_wav for the diffusion model."
|
||||
],
|
||||
"metadata": {
|
||||
"id": "KEXOKjIvn6NW"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Load the autoregressive model.\n",
|
||||
"autoregressive = UnifiedVoice(max_mel_tokens=300, max_text_tokens=200, max_conditioning_inputs=2, layers=30, model_dim=1024,\n",
|
||||
" heads=16, number_text_tokens=256, start_text_token=255, checkpointing=False, train_solo_embeddings=False).cuda().eval()\n",
|
||||
"autoregressive.load_state_dict(torch.load('.models/autoregressive.pth'))\n",
|
||||
"stop_mel_token = autoregressive.stop_mel_token"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "Z15xFT_uhP8v"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Perform inference with the autoregressive model, generating num_autoregressive_samples\n",
|
||||
"with torch.no_grad():\n",
|
||||
" samples = []\n",
|
||||
" for b in tqdm(range(num_autoregressive_samples // 16)):\n",
|
||||
" codes = autoregressive.inference_speech(conds, text, num_beams=1, repetition_penalty=1.0, do_sample=True, top_k=50, top_p=.95,\n",
|
||||
" temperature=.9, num_return_sequences=16, length_penalty=1)\n",
|
||||
" padding_needed = 250 - codes.shape[1]\n",
|
||||
" codes = F.pad(codes, (0, padding_needed), value=stop_mel_token)\n",
|
||||
" samples.append(codes)\n",
|
||||
"\n",
|
||||
"# Delete model weights to conserve memory.\n",
|
||||
"del autoregressive"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "xajqWiEik-j0"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Load the CLIP model.\n",
|
||||
"clip = VoiceCLIP(dim_text=512, dim_speech=512, dim_latent=512, num_text_tokens=256, text_enc_depth=8, text_seq_len=120, text_heads=8,\n",
|
||||
" num_speech_tokens=8192, speech_enc_depth=10, speech_heads=8, speech_seq_len=250).cuda().eval()\n",
|
||||
"clip.load_state_dict(torch.load('.models/clip.pth'))"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "KNgYSyuyliMs"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Use the CLIP model to select the best autoregressive output to match the given text.\n",
|
||||
"clip_results = []\n",
|
||||
"with torch.no_grad():\n",
|
||||
" for batch in samples:\n",
|
||||
" for i in range(batch.shape[0]):\n",
|
||||
" batch[i] = fix_autoregressive_output(batch[i], stop_mel_token)\n",
|
||||
" text = text[:, :120] # Ugly hack to fix the fact that I didn't train CLIP to handle long enough text.\n",
|
||||
" clip_results.append(clip(text.repeat(batch.shape[0], 1),\n",
|
||||
" torch.full((batch.shape[0],), fill_value=text.shape[1]-1, dtype=torch.long, device='cuda'),\n",
|
||||
" batch, torch.full((batch.shape[0],), fill_value=batch.shape[1]*1024, dtype=torch.long, device='cuda'),\n",
|
||||
" return_loss=False))\n",
|
||||
" clip_results = torch.cat(clip_results, dim=0)\n",
|
||||
" samples = torch.cat(samples, dim=0)\n",
|
||||
" best_results = samples[torch.topk(clip_results, k=1).indices]\n",
|
||||
"\n",
|
||||
"# Save samples to CPU memory, delete clip to conserve memory.\n",
|
||||
"samples = samples.cpu()\n",
|
||||
"del clip"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "DDXkM0lclp4U"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Load the DVAE and diffusion model.\n",
|
||||
"dvae = DiscreteVAE(positional_dims=1, channels=80, hidden_dim=512, num_resnet_blocks=3, codebook_dim=512, num_tokens=8192, num_layers=2,\n",
|
||||
" record_codes=True, kernel_size=3, use_transposed_convs=False).cuda().eval()\n",
|
||||
"dvae.load_state_dict(torch.load('.models/dvae.pth'), strict=False)\n",
|
||||
"diffusion = DiscreteDiffusionVocoder(model_channels=128, dvae_dim=80, channel_mult=[1, 1, 1.5, 2, 3, 4, 6, 8, 8, 8, 8], num_res_blocks=[1, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1],\n",
|
||||
" spectrogram_conditioning_resolutions=[2,512], attention_resolutions=[512,1024], num_heads=4, kernel_size=3, scale_factor=2,\n",
|
||||
" conditioning_inputs_provided=True, time_embed_dim_multiplier=4).cuda().eval()\n",
|
||||
"diffusion.load_state_dict(torch.load('.models/diffusion.pth'))\n",
|
||||
"diffuser = load_discrete_vocoder_diffuser(desired_diffusion_steps=100)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "97acSnBal8Q2"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Decode the (best) discrete sequence created by the autoregressive model.\n",
|
||||
"with torch.no_grad():\n",
|
||||
" for b in range(best_results.shape[0]):\n",
|
||||
" code = best_results[b].unsqueeze(0)\n",
|
||||
" wav = do_spectrogram_diffusion(diffusion, dvae, diffuser, code, cond_wav, spectrogram_compression_factor=256, mean=True)\n",
|
||||
" torchaudio.save(f'{voice}_{b}.wav', wav.squeeze(0).cpu(), 22050)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "HEDABTrdl_kM"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"# Listen to your text! (told you that'd take a long time..)\n",
|
||||
"from IPython.display import Audio\n",
|
||||
"Audio(data=wav.squeeze(0).cpu().numpy(), rate=22050)"
|
||||
],
|
||||
"metadata": {
|
||||
"id": "EyHmcdqBmSvf"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
}
|
||||
]
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user