{ "cells": [ { "cell_type": "markdown", "id": "2610c4d9", "metadata": {}, "source": [ "## infer" ] }, { "cell_type": "code", "execution_count": null, "id": "191bb314", "metadata": {}, "outputs": [], "source": [ "\n", "import os\n", "\n", "import sys\n", "sys.path.append(\"/path/to/dune_codec\")\n", "from inference_utils_darya import *\n", "\n", "from transformers import AutoTokenizer\n", "tokenizer_path = \"KavirLabs/Darya_Tokenizer_bpe\"\n", "tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, add_bos_token=True, add_eos_token=True)\n", "\n", "device = 'cuda'\n", "\n", "\n", "import sys\n", "sys.path.append(\"dune_codec\")\n", "\n", "\n", "\n", "device = 'cuda'\n", "from codec.audio_processing.dune_codec import load_dune_audio_tokenizer\n", "dac_model = load_dune_audio_tokenizer(\"Respair/dune_codec\", \n", " device=device)\n", "\n", "\n", "CONFIG_PATH = \"config_transformer.json\"\n", "CKPT_DIR = \"exp/darya_fix/checkpoints/step_0080000\"\n", "\n", "\n", "\n", "tts_model, tts_cfg = load_model_from_checkpoint(CONFIG_PATH, CKPT_DIR)\n", "\n", "\n", "\n", "duration_model = load_duration_model(\n", " \"ckpt/duration_model/step_0075000.pt\",\n", " vocab_size=4096,\n", " latent_dim=52,\n", " hidden_dim=256,\n", " n_text_layer=8,\n", " n_cross_layer=8,\n", " n_head=8,\n", " output_dim=378,\n", " device=device,\n", ")\n", "\n", "\n", "silence_pack = torch.load(\"silence_fsq.pt\", map_location=\"cpu\", weights_only=True)\n", "\n", "silence_fsq = silence_pack[\"latents\"]\n", "silence_codec_rate_hz = silence_pack.get(\"codec_rate_hz\", 12.5)" ] }, { "cell_type": "code", "execution_count": null, "id": "e777d100", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "utterances: [' Oh my god, 🤣 ... I am actually losing my mind right now! Holy shit bro 🤣 You have absolutely no idea how perfect this is. That is hands-down the most amazing thing I have ever seen in my life!']\n", "duration: [163]\n", "rtf: 0.1307852395703592\n" ] }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "execution_count": 59, "metadata": {}, "output_type": "execute_result" } ], "source": [ "\n", "import re\n", "import time\n", "import torch\n", "\n", "\n", "text_list = [\" write your input text here. uh, maybe another speaker perhaps.\"]\n", "utterances = [text_list] if isinstance(text_list, str) else text_list\n", "\n", "td = Extractor(tokenizer, tts_cfg, device, duration_model, speaker_model=None)\n", "\n", "\n", "text_ids, text_mask, _, _ = td.get_tokens(\n", " utterances,\n", " utterances,\n", ")\n", "\n", "\n", "speed = 1.15\n", "\n", "# duration computed per utterance\n", "duration_items = []\n", "\n", "for utterance in utterances:\n", " duration_text = strip_tags(utterance)\n", "\n", " _, _, duration_text_ids_i, duration_text_mask_i = td.get_tokens(\n", " utterance,\n", " duration_text,\n", " )\n", "\n", " duration_i = td.get_duration(\n", " duration_text_ids_i,\n", " duration_text_mask_i,\n", " speed=speed,\n", " )\n", "\n", " duration_items.append(duration_i.squeeze(0))\n", "\n", "duration = torch.stack(duration_items, dim=0).to(device=device, dtype=torch.long)\n", "\n", "n_frame_per_class = 1\n", "max_duration = int(tts_cfg[\"data\"][\"max_audio_seconds\"] * tts_cfg[\"data\"][\"codec_rate_hz\"])\n", "\n", "torch.cuda.synchronize()\n", "t0 = time.perf_counter()\n", "latents = sample_euler_reducio_spk(\n", " tts_model,\n", " text_ids,\n", " text_mask,\n", " duration=duration,\n", " cond_latents=None,\n", " cond_latent_mask=None,\n", " duration_model=duration_model,\n", " n_frame_per_class=n_frame_per_class,\n", " max_duration=max_duration,\n", " steps=16,\n", " cfg=3.,\n", " seed=None,\n", ")\n", "\n", "silence_seconds = 0.3\n", "silence_frames = int(round(silence_seconds * 12.5))\n", "\n", "sil = silence_fsq[:silence_frames].to(\n", " device=latents.device,\n", " dtype=latents.dtype,\n", ")\n", "\n", "latents_merged = torch.cat(\n", " [\n", " x\n", " for i in range(latents.shape[0])\n", " for x in (\n", " [latents[i, :int(duration[i].item())], sil]\n", " if i < latents.shape[0] - 1\n", " else [latents[i, :int(duration[i].item())]]\n", " )\n", " ],\n", " dim=0,\n", ")\n", "\n", "audio_cropped = decode_audio(\n", " dac_model,\n", " latents_merged.transpose(0, 1),\n", ")\n", "torch.cuda.synchronize()\n", "total_time = time.perf_counter() - t0\n", "\n", "audio_duration = audio_cropped.shape[-1] / 44100\n", "rtf = total_time / max(audio_duration, 1e-6)\n", "\n", "print(\"utterances:\", utterances)\n", "print(\"duration:\", duration.tolist())\n", "print(\"rtf:\", rtf)\n", "\n", "Sawt(audio_cropped, rate=44100, normalize=False)" ] }, { "cell_type": "code", "execution_count": null, "id": "e590c2a7", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "utterances: [' Oh my god 🤣 I am actually losing my mind right now! Holy shit bro 🤣 You have absolutely no idea how perfect this is. That is hands-down the most amazing thing I have ever seen in my life.']\n", "duration: [163]\n", "rtf: 0.12585385544532682\n" ] }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "execution_count": 55, "metadata": {}, "output_type": "execute_result" } ], "source": [ "\n", "Sawt(audio_cropped, rate=44100, normalize=False)" ] }, { "cell_type": "markdown", "id": "0b733fd3", "metadata": {}, "source": [ "## prompt - infer" ] }, { "cell_type": "code", "execution_count": null, "id": "ecf6fd42", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " They were not! Their memory serves as an example to us all! The courageous fallen, the anguished fallen! Their lives have meaning because we, the living, refuse to forget them!\n", "fixed pre_q: torch.Size([1, 169, 52])\n" ] } ], "source": [ "AUDIO_PATH = \"/path/to/prompt_audio.mp3\"\n", "import librosa\n", "\n", "\n", "wav = librosa.load(AUDIO_PATH, sr=22050)[0]\n", "\n", "prompt_tensor = torch.from_numpy(wav)\n", "# prompt_text = decode_ids(beam_text[0])\n", "prompt_text = ' prompt transcription here.'\n", "print(prompt_text)\n", "\n", "\n", "\n", "@torch.no_grad()\n", "def extract_prequant_latents(dac_model, wav: torch.Tensor, device=\"cuda\"):\n", " wav = wav.float()\n", "\n", " audio = wav.unsqueeze(0).to(device)\n", " audio_len = torch.tensor([wav.shape[0]], dtype=torch.long, device=device)\n", "\n", " encoder = dac_model.codec.audio_encoder\n", " pre_q, latent_len = encoder(audio=audio, audio_len=audio_len)\n", " \n", " pre_q = pre_q.transpose(1, 2)\n", "\n", " valid_len = int(latent_len[0].item())\n", " pre_q = pre_q[:, :valid_len, :]\n", " return pre_q[0]\n", "\n", "prompt_lats = extract_prequant_latents(dac_model, prompt_tensor, device=device)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "72213f55", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "used_frames=324 | generated_frames=155 | Total: 3.074s | Audio: 12.40s | RTF: 0.2479\n" ] }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import time\n", "import torch\n", "import re\n", "\n", "text = \"Some text, tags should only appear in the prompt transcription and not here. unless it is cross lingual, otherwise the accent will leak.\"\n", "\n", "td = Extractor(tokenizer, tts_cfg, device, duration_model, None)\n", "\n", "full_text = prompt_text + \"\" + text\n", "duration_text = re.sub(r\"<[^>]+>\\s*\", \"\", full_text).strip()\n", "\n", "text_ids, text_mask, duration_text_ids, duration_text_mask = td.get_tokens(\n", " full_text,\n", " duration_text,\n", ")\n", "\n", "prompt_lats_b = prompt_lats.unsqueeze(0)\n", "\n", "prompt_mask = torch.ones(\n", " 1,\n", " prompt_lats.shape[0],\n", " device=device,\n", " dtype=torch.bool,\n", ")\n", "\n", "n_frame_per_class = 1\n", "max_duration = td.max_duration\n", "speed = 1.\n", "\n", "T_prompt = prompt_lats.shape[0]\n", "\n", "duration = td.get_duration(\n", " duration_text_ids,\n", " duration_text_mask,\n", " speed=speed,\n", " cond_latents=prompt_lats_b,\n", " cond_latent_mask=prompt_mask,\n", ")\n", "\n", "torch.cuda.synchronize()\n", "t0 = time.perf_counter()\n", "\n", "output = sample_euler_reducio(\n", " tts_model,\n", " text_ids=text_ids,\n", " text_mask=text_mask,\n", " duration=duration,\n", " cond_latents=prompt_lats_b,\n", " cond_latent_mask=prompt_mask,\n", " duration_model=duration_model,\n", " n_frame_per_class=n_frame_per_class,\n", " max_duration=max_duration,\n", " steps=32,\n", " cfg=3.0,\n", " seed=None,\n", ")\n", "\n", "continuation_lats = output[:, T_prompt:int(duration[0].item()), :]\n", "\n", "audio_cropped = decode_audio(\n", " dac_model,\n", " continuation_lats.transpose(1, 2),\n", ")\n", "\n", "torch.cuda.synchronize()\n", "total_time = time.perf_counter() - t0\n", "\n", "audio_duration = audio_cropped.shape[-1] / 44100\n", "rtf = total_time / max(audio_duration, 1e-6)\n", "\n", "print(\n", " f\"used_frames={int(duration[0].item())} | \"\n", " f\"generated_frames={continuation_lats.shape[1]} | \"\n", " f\"Total: {total_time:.3f}s | Audio: {audio_duration:.2f}s | RTF: {rtf:.4f}\"\n", ")\n", "\n", "Sawt(audio_cropped, rate=44100)" ] }, { "cell_type": "code", "execution_count": null, "id": "8c65caf0", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "used_frames=212 | generated_frames=144 | Total: 3.506s | Audio: 11.52s | RTF: 0.3043\n" ] }, { "data": { "text/html": [ "\n", " \n", " " ], "text/plain": [ "" ] }, "execution_count": 102, "metadata": {}, "output_type": "execute_result" } ], "source": [ "\n", "\n", "Sawt(audio_cropped, rate=44100)" ] }, { "cell_type": "markdown", "id": "7be3cf21", "metadata": {}, "source": [ "# Style Conditioned Path\n", "\n", "Here you can prompt the model another way.
\n", "The benefits of this approach are :\n", "- no need to load an ASR or provide prompt transcriptions\n", "- context length of the model is preserved, because the prompt is a fixed global vector and does not eat the context away. \n", "- You can use cross-lingual prompts, regardless of whether the model has seen that language or not.\n", "- mix various clips by doing a convex combination of their vectors, creating new voices.\n", "\n", "the downside is that the model is limited to the capacity of Titanet, which is trained on a limited amount of English-only data.
\n", "The checkpoint is also not trained and optimized from scratch with this approach, so it can impact the acoustic quality.\n", "\n", "If you want the maximum stability / quality and did not like what you see below, feel free to go back to the infilling path." ] }, { "cell_type": "markdown", "id": "21c099fb", "metadata": {}, "source": [ "## startup" ] }, { "cell_type": "code", "execution_count": null, "id": "7bc879a3", "metadata": {}, "outputs": [], "source": [ "\n", "# ---- paths ----\n", "CONFIG_PATH = \"config_transformer.json\"\n", "CKPT_DIR = \"exp/darya_fix_2ndstage/checkpoints/step_0070000\" # load the model trained with a speaker encoder\n", "\n", "tts_model, tts_cfg = load_model_from_checkpoint(CONFIG_PATH, CKPT_DIR, use_speaker_conditioning=True)\n", "\n", "import nemo.collections.asr as nemo_asr\n", "speaker_model = nemo_asr.models.EncDecSpeakerLabelModel.from_pretrained(\"nvidia/speakerverification_en_titanet_large\")" ] }, { "cell_type": "markdown", "id": "20824931", "metadata": {}, "source": [ "## infer" ] }, { "cell_type": "markdown", "id": "ae7d19c8", "metadata": {}, "source": [ "### long form or simple\n", "\n", "at least the way I am formulating it.\n", "\n", "- NOTE: since Titanet speaker embeds are pooled global tensors, they were not made to handle multi speakers. try to keep it limited to single speaker inputs" ] }, { "cell_type": "code", "execution_count": null, "id": "70ea7f92", "metadata": {}, "outputs": [], "source": [ "\"\"\"long form one\"\"\"\n", "\n", "import re\n", "import time\n", "import torch\n", "\n", "torch.cuda.synchronize()\n", "t0 = time.perf_counter()\n", "\n", "\n", "text_list = [\n", "\n", " \" some text\"\n", "]\n", "\n", "# normalize single string/list into utterance list\n", "utterances = [text_list] if isinstance(text_list, str) else text_list\n", "\n", "td = Extractor(tokenizer, tts_cfg, device, duration_model, speaker_model)\n", "\n", "# main TTS tokens\n", "text_ids, text_mask, _, _ = td.get_tokens(\n", " utterances,\n", " utterances,\n", ")\n", "\n", "\n", "speed = 1.\n", "\n", "emb = td.get_speaker_embedding(\"/path/to/prompt_audio.mp3\")\n", "# duration computed per utterance\n", "duration_items = []\n", "\n", "for utterance in utterances:\n", " duration_text = strip_tags(utterance)\n", "\n", " _, _, duration_text_ids_i, duration_text_mask_i = td.get_tokens(\n", " utterance,\n", " duration_text,\n", " )\n", "\n", " duration_i = td.get_duration(\n", " duration_text_ids_i,\n", " duration_text_mask_i,\n", " speed=speed,\n", " )\n", "\n", " duration_items.append(duration_i.squeeze(0))\n", "\n", "duration = torch.stack(duration_items, dim=0).to(device=device, dtype=torch.long)\n", "\n", "n_frame_per_class = 1\n", "max_duration = int(tts_cfg[\"data\"][\"max_audio_seconds\"] * tts_cfg[\"data\"][\"codec_rate_hz\"])\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e3602bda", "metadata": {}, "outputs": [], "source": [ "\n", "latents = sample_euler_reducio_spk(\n", " tts_model,\n", " text_ids,\n", " text_mask,\n", " duration=duration,\n", " cond_latents=None,\n", " cond_latent_mask=None,\n", " duration_model=duration_model,\n", " n_frame_per_class=n_frame_per_class,\n", " max_duration=max_duration,\n", " speaker_emb=emb,\n", " speaker_adaln_scale=2.,\n", " steps=32,\n", " cfg=2.,\n", " speaker_cfg=1.5, # less like the speaker\n", " seed=42,\n", ")\n", "\n", "silence_seconds = 0.3\n", "silence_frames = int(round(silence_seconds * 12.5))\n", "\n", "sil = silence_fsq[:silence_frames].to(\n", " device=latents.device,\n", " dtype=latents.dtype,\n", ")\n", "\n", "latents_merged = torch.cat(\n", " [\n", " x\n", " for i in range(latents.shape[0])\n", " for x in (\n", " [latents[i, :int(duration[i].item())], sil]\n", " if i < latents.shape[0] - 1\n", " else [latents[i, :int(duration[i].item())]]\n", " )\n", " ],\n", " dim=0,\n", ")\n", "\n", "audio_cropped = decode_audio(\n", " dac_model,\n", " latents_merged.transpose(0, 1),\n", ")\n", "torch.cuda.synchronize()\n", "total_time = time.perf_counter() - t0\n", "\n", "audio_duration = audio_cropped.shape[-1] / 44100\n", "rtf = total_time / max(audio_duration, 1e-6)\n", "\n", "print(\"utterances:\", utterances)\n", "print(\"duration:\", duration.tolist())\n", "print(\"rtf:\", rtf)\n", "\n", "Sawt(audio_cropped, rate=44100, normalize=False)" ] }, { "cell_type": "markdown", "id": "8d51bd9b", "metadata": {}, "source": [ "### long form sequential\n", "\n", "each generation is conditioned on the one that came before it.\n", "in production, each chunk can and should be played after it is decoded.\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "e9c19a0b", "metadata": {}, "outputs": [], "source": [ "import re\n", "import time\n", "import torch\n", "import numpy as np\n", "\n", "torch.cuda.synchronize()\n", "t0 = time.perf_counter()\n", "\n", "# emb0 = td.get_speaker_embedding(audio)\n", "emb0 = None\n", "emb = emb0\n", "\n", "audio_chunks = []\n", "\n", "for i, utterance in enumerate(text_list):\n", " duration_text = strip_tags(utterance)\n", "\n", " text_ids, text_mask, duration_text_ids, duration_text_mask = td.get_tokens(\n", " utterance,\n", " duration_text,\n", " )\n", "\n", " duration = td.get_duration(\n", " duration_text_ids,\n", " duration_text_mask,\n", " speed=1.,\n", " )\n", "\n", " latents = sample_euler_reducio_spk(\n", " tts_model,\n", " text_ids,\n", " text_mask,\n", " duration=duration,\n", " cond_latents=None,\n", " cond_latent_mask=None,\n", " duration_model=duration_model,\n", " n_frame_per_class=n_frame_per_class,\n", " max_duration=max_duration,\n", " speaker_emb=emb,\n", " speaker_adaln_scale=1.,\n", " steps=32,\n", " cfg=3.,\n", " speaker_cfg=None,\n", " seed=None,\n", " )\n", "\n", " dur_i = int(duration[0].item())\n", " latents_i = latents[0, :dur_i]\n", "\n", " audio_cropped = decode_audio(\n", " dac_model,\n", " latents_i.transpose(0, 1),\n", " )\n", "\n", " if isinstance(audio_cropped, torch.Tensor):\n", " audio_np = audio_cropped.detach().float().cpu().squeeze().numpy()\n", " else:\n", " audio_np = np.asarray(audio_cropped, dtype=np.float32).squeeze()\n", "\n", " audio_chunks.append(audio_np)\n", "\n", " emb_prev = td.get_speaker_embedding(audio_np)\n", "\n", " if emb0 is None:\n", " emb0 = emb_prev\n", " emb = emb_prev\n", " else:\n", " emb = 0.25 * emb0 + 0.75 * emb_prev\n", "\n", "audio_cropped = np.concatenate(audio_chunks, axis=-1)\n", "\n", "torch.cuda.synchronize()\n", "total_time = time.perf_counter() - t0\n", "\n", "audio_duration = audio_cropped.shape[-1] / 44100\n", "rtf = total_time / max(audio_duration, 1e-6)\n", "\n", "print(\"text_list:\", text_list)\n", "print(\"rtf:\", rtf)\n", "\n", "Sawt(audio_cropped, rate=44100, normalize=False)" ] }, { "cell_type": "code", "execution_count": null, "id": "176f7e17", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "e82158c9", "metadata": {}, "source": [ "# SPEECH EDIT" ] }, { "cell_type": "code", "execution_count": null, "id": "1e1101d4", "metadata": {}, "outputs": [], "source": [ "AUDIO_TO_EDIT = \"/path/to/audio_to_edit.mp3\"\n", "\n", "target_text = ' please put your target text here with the corrected segments.'\n", "\n", "wav, wav_sr = librosa.load(AUDIO_TO_EDIT, sr=22050)\n", "wav_tensor = torch.from_numpy(wav)\n", "\n", "original_lats = extract_prequant_latents(dac_model, wav_tensor, device=device)\n", "\n", "codec_rate_hz = tts_cfg[\"data\"][\"codec_rate_hz\"]\n", "\n", "parts_to_edit = [\n", " [0, 2],\n", " [7, 9], \n", "]\n", "\n", "extra_duration = [\n", " 0.5, # add room \n", "]\n", "\n", "cond, keep_mask, span_mask, valid, edit_ranges = build_latent_edit_condition(\n", " original_lats,\n", " wav_num_samples=len(wav),\n", " wav_sr=wav_sr,\n", " parts_to_edit=parts_to_edit,\n", " extra_duration=extra_duration,\n", " duration_scale=1.0,\n", " padding_sec=0.5,\n", " min_edit_sec=0.15,\n", " codec_rate_hz=codec_rate_hz,\n", " device=device,\n", " dtype=tts_model.dtype,\n", ")\n", "\n", "td = Extractor(tokenizer, tts_cfg, device, duration_model, None)\n", "\n", "duration_text = re.sub(r\"<[^>]+>\\s*\", \"\", target_text).strip()\n", "\n", "text_ids, text_mask, _, _ = td.get_tokens(\n", " target_text,\n", " duration_text,\n", ")\n", "\n", "edited_lats = sample_euler_reducio_edit(\n", " tts_model,\n", " text_ids=text_ids,\n", " text_mask=text_mask,\n", " cond=cond,\n", " keep_mask=keep_mask,\n", " valid_mask=valid,\n", " steps=32,\n", " cfg=2.,\n", " seed=None,\n", ")\n", "\n", "audio_edited = decode_audio(\n", " dac_model,\n", " edited_lats.transpose(1, 2),\n", ")\n", "\n", "Sawt(audio_edited, rate=44100)" ] } ], "metadata": { "kernelspec": { "display_name": "respair", "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.11.6" } }, "nbformat": 4, "nbformat_minor": 5 }