t5tts_inference_multiturndialogues.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "466ccdc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemo.collections.tts.models import T5TTS_Model\n",
    "from nemo.collections.tts.data.text_to_speech_dataset import T5TTSDataset, DatasetSample\n",
    "from omegaconf.omegaconf import OmegaConf, open_dict\n",
    "import torch\n",
    "import os\n",
    "import soundfile as sf\n",
    "from IPython.display import display, Audio\n",
    "import os\n",
    "import numpy as np\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"]=\"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6659ae78",
   "metadata": {},
   "source": [
    "## Set checkpoint and other file paths on your machine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04445f11",
   "metadata": {},
   "outputs": [],
   "source": [
    "hparams_file = \"/datap/misc/duplexcheckpoints/blackwell/duplex_blackwell_medium_decoder_noexpresso_onlyphonemeFT_hparams.yaml\"\n",
    "checkpoint_file = \"/datap/misc/duplexcheckpoints/blackwell/duplex_blackwell_medium_decoder_withTC_fromroycheckpoint_lowsestvalloss.ckpt\"\n",
    "codecmodel_path = \"/datap/misc/checkpoints/AudioCodec_21Hz_no_eliz.nemo\"\n",
    "out_dir = \"/datap/misc/t5tts_inference_notebook_samples\"\n",
    "if not os.path.exists(out_dir):\n",
    "    os.makedirs(out_dir)\n",
    "\n",
    "dummy_audio_filepath = os.path.join(out_dir, \"dummy_audio.wav\")\n",
    "sf.write(dummy_audio_filepath, np.zeros(22050 * 3), 22050)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01e24b14",
   "metadata": {},
   "source": [
    "## Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87bf66f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_cfg = OmegaConf.load(hparams_file).cfg\n",
    "\n",
    "with open_dict(model_cfg):\n",
    "    model_cfg.codecmodel_path = codecmodel_path\n",
    "    if hasattr(model_cfg, 'text_tokenizer'):\n",
    "        # Backward compatibility for models trained with absolute paths in text_tokenizer\n",
    "        model_cfg.text_tokenizer.g2p.phoneme_dict = \"scripts/tts_dataset_files/ipa_cmudict-0.7b_nv23.01.txt\"\n",
    "        model_cfg.text_tokenizer.g2p.heteronyms = \"scripts/tts_dataset_files/heteronyms-052722\"\n",
    "        model_cfg.text_tokenizer.g2p.phoneme_probability = 1.0\n",
    "    model_cfg.train_ds = None\n",
    "    model_cfg.validation_ds = None\n",
    "\n",
    "\n",
    "model = T5TTS_Model(cfg=model_cfg)\n",
    "# Load weights from checkpoint file\n",
    "print(\"Loading weights from checkpoint\")\n",
    "ckpt = torch.load(checkpoint_file)\n",
    "model.load_state_dict(ckpt['state_dict'])\n",
    "print(\"Loaded weights.\")\n",
    "\n",
    "if model_cfg.t5_decoder.pos_emb == \"learnable\":\n",
    "    if (model_cfg.t5_decoder.use_flash_self_attention) is False and (model_cfg.t5_decoder.use_flash_self_attention is False):\n",
    "        print(\"Using kv cache for inference.\")\n",
    "        model.use_kv_cache_for_inference = True\n",
    "\n",
    "model.cuda()\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843167ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = T5TTSDataset(\n",
    "    dataset_meta={},\n",
    "    sample_rate=model_cfg.sample_rate,\n",
    "    min_duration=0.5,\n",
    "    max_duration=20,\n",
    "    codec_model_downsample_factor=model_cfg.codec_model_downsample_factor,\n",
    "    bos_id=model.bos_id,\n",
    "    eos_id=model.eos_id,\n",
    "    context_audio_bos_id=model.context_audio_bos_id,\n",
    "    context_audio_eos_id=model.context_audio_eos_id,\n",
    "    audio_bos_id=model.audio_bos_id,\n",
    "    audio_eos_id=model.audio_eos_id,\n",
    "    num_audio_codebooks=model_cfg.num_audio_codebooks,\n",
    "    prior_scaling_factor=None,\n",
    "    load_cached_codes_if_available=True,\n",
    "    dataset_type='test',\n",
    "    tokenizer_config=None,\n",
    "    load_16khz_audio=model.model_type == 'single_encoder_sv_tts',\n",
    "    use_text_conditioning_tokenizer=model.use_text_conditioning_encoder,\n",
    "    pad_context_text_to_max_duration=model.pad_context_text_to_max_duration,\n",
    "    context_duration_min=model.cfg.get('context_duration_min', 5.0),\n",
    "    context_duration_max=model.cfg.get('context_duration_max', 5.0),\n",
    ")\n",
    "test_dataset.text_tokenizer, test_dataset.text_conditioning_tokenizer = model._setup_tokenizers(model.cfg, mode='test')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd78348",
   "metadata": {},
   "source": [
    "### Set dialogues to generate\n",
    "\n",
    "Each item in the list can be a single-turn or a multi-turn dialogue.\n",
    "\n",
    "[SPK-BWL-B-F] is the speaker tag for female speaker and [SPK-BWL-B-M] is the speaker tag for male speaker.\n",
    "\n",
    "ChatGPT prompt that can generate something like this:\n",
    "\n",
    "```\n",
    "Generate dialogues for a 30 second podcast about <TOPIC>.\n",
    "The conversation should be between a male and female speaker formatted as follows:\n",
    "[SPK-BWL-B-F] Sentence by a female speaker\n",
    "[SPK-BWL-B-M] Sentence by a male speaker\n",
    "where [SPK-BWL-B-F] and [SPK-BWL-B-M] indicate speaker tags. Dont have any quotation marks in the text, and if there are any numbers spell them out. Basically, keep the text normalized suitable for a TTS model. Keep the conversation fun and engaging with the speakers talking and responding to each other. \n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5e32a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "dialogues = [\"[SPK-BWL-B-F] Have you been keeping up with multi-turn TTS models? They are so fascinating [SPK-BWL-B-M] Absolutely. The way they handle conversations feels almost human.\",\n",
    "\"[SPK-BWL-B-F] I know. It is amazing how they can remember context across multiple exchanges [SPK-BWL-B-M] Exactly. It is not just about saying words anymore. They actually make the responses flow naturally.\",\n",
    "\"[SPK-BWL-B-F] And it is not just for assistants or chatbots. I heard they are being used for things like audiobooks and interactive storytelling.\",\n",
    "\"[SPK-BWL-B-M] That is true. Plus, they can even adjust tone to match emotions. Imagine hearing a story where the narrator sounds genuinely excited during a twist.\",\n",
    "\"[SPK-BWL-B-F] Or even a bit sarcastic during a funny moment. That makes everything feel more real.\",\n",
    "\"[SPK-BWL-B-M] For sure. It is like conversations with a machine are finally catching up to the way we actually talk.\",\n",
    "\"[SPK-BWL-B-F] The future of tech keeps surprising me. Every day, there is something new to explore.\",]\n",
    "dialogues = [d for d in dialogues if len(d) > 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa182540",
   "metadata": {},
   "source": [
    "### Generation\n",
    "\n",
    "Below code generates 4 samples for each item in the dialogues list. Then asks, which one you like the best (index 0,1,2 or 3) and adds that to the already generated dialogue. You may modify the below code to automate and just select the first generation if you dont want to do this manually. After every dialogue it also plays the combined dialogues until that item."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "168f157d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydub import AudioSegment\n",
    "\n",
    "def play_combined_audio(audio_files):\n",
    "    combined_audio = AudioSegment.empty()\n",
    "    for file in audio_files:\n",
    "        audio = AudioSegment.from_file(file)\n",
    "        combined_audio += audio\n",
    "    output_path = os.path.join(out_dir, \"combined_audio.wav\")\n",
    "    combined_audio.export(output_path, format=\"wav\")\n",
    "    display(Audio(output_path))\n",
    "    return output_path\n",
    "    \n",
    "\n",
    "\n",
    "context_path = None\n",
    "generated_audios = []\n",
    "for didx, dialogue in enumerate(dialogues):\n",
    "    audio_dir = \"/\"\n",
    "    entry = {\n",
    "        \"audio_filepath\": dummy_audio_filepath,\n",
    "        \"duration\": 3.0,\n",
    "        \"text\": dialogue,\n",
    "        \"speaker\": \"dummy\",\n",
    "    }\n",
    "    if didx == 0:\n",
    "        entry[\"context_text\"] = \"MIXED SPEECH TTS\"\n",
    "    else:\n",
    "        entry['context_audio_filepath'] = context_path\n",
    "        entry['context_audio_duration'] = 5.0\n",
    "        \n",
    "        \n",
    "        \n",
    "    data_sample = DatasetSample(\n",
    "        dataset_name=\"sample\",\n",
    "        manifest_entry=entry,\n",
    "        audio_dir=audio_dir,\n",
    "        feature_dir=audio_dir,\n",
    "        text=entry['text'],\n",
    "        speaker=None,\n",
    "        speaker_index=0,\n",
    "        tokenizer_names=[\"english_phoneme\"]\n",
    "    )\n",
    "    test_dataset.data_samples = [data_sample]\n",
    "\n",
    "    test_data_loader = torch.utils.data.DataLoader(\n",
    "        test_dataset,\n",
    "        batch_size=1,\n",
    "        collate_fn=test_dataset.collate_fn,\n",
    "        num_workers=0,\n",
    "        shuffle=False\n",
    "    )\n",
    "    \n",
    "    \n",
    "    item_idx = 0\n",
    "    for bidx, batch in enumerate(test_data_loader):\n",
    "        print(\"Processing batch {} out of {}\".format(bidx, len(test_data_loader)))\n",
    "        model.t5_decoder.reset_cache(use_cache=True)\n",
    "        batch_cuda ={}\n",
    "        for key in batch:\n",
    "            if isinstance(batch[key], torch.Tensor):\n",
    "                batch_cuda[key] = batch[key].cuda()\n",
    "            else:\n",
    "                batch_cuda[key] = batch[key]\n",
    "        import time\n",
    "        \n",
    "        candidates = []\n",
    "        for try_idx in range(4):\n",
    "            st = time.time()\n",
    "            predicted_audio, predicted_audio_lens, _, _ = model.infer_batch(\n",
    "                batch_cuda, \n",
    "                max_decoder_steps=500, \n",
    "                temperature=0.6, \n",
    "                topk=80, \n",
    "                use_cfg=True, \n",
    "                cfg_scale=1.6\n",
    "            )\n",
    "            print(\"generation time\", time.time() - st)\n",
    "            for idx in range(predicted_audio.size(0)):\n",
    "                predicted_audio_np = predicted_audio[idx].float().detach().cpu().numpy()\n",
    "                predicted_audio_np = predicted_audio_np[:predicted_audio_lens[idx]]\n",
    "                audio_path = os.path.join(out_dir, f\"predicted_audio_{try_idx}_{didx}_{item_idx}.wav\")\n",
    "                sf.write(audio_path, predicted_audio_np, model.cfg.sample_rate)\n",
    "                print(\"Dialogue:\", item_idx, \"Candidate: \", try_idx)\n",
    "                display(Audio(audio_path))\n",
    "                candidates.append(audio_path)\n",
    "        \n",
    "        user_input = input(\"Enter Candidate number that sounds the best:\").strip().lower()\n",
    "        selected_audio_idx = int(user_input)\n",
    "        audio_path = candidates[selected_audio_idx]\n",
    "        item_idx += 1\n",
    "        generated_audios.append(audio_path)\n",
    "        print(\"Podcast generated until now:\")\n",
    "        combined_audio_path = play_combined_audio(generated_audios)\n",
    "        last_gen_audio = AudioSegment.from_file(combined_audio_path)\n",
    "        last_5_seconds = last_gen_audio[-5000:]  # Duration is in milliseconds\n",
    "        context_path = os.path.join(out_dir, \"context.wav\")\n",
    "        last_5_seconds.export(context_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1368c380",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943a791f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}