Latest_Tacotron2_and_HiFi_GAN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# <b><font color=\"pink\" size=\"+2\"> **FakeYou Tacotron 2 GPU Synthesis Colab**\n",
        "___\n",
        "\n",
        "Last update: 07/09/2023\n",
        "\n",
        "___\n",
        "\n",
        "### [Tacotron 2](https://github.com/NVIDIA/tacotron2)\n",
        "<a href=\"https://fakeyou.com/\"><img src=\"https://fakeyou.com/fakeyou/FakeYou-Logo.png\" alt=\"FakeYou Logo. Click here to go to the official website.\"></a>\n",
        "\n",
        "---\n",
        "\n",
        "## 🤝**Credits**🤝\n",
        "Code originally created by `CookieGalaxy#8351` for the Pony Preservation Project\n",
        "\n",
        "Colab notebook prepared by `justinjohn-03`."
      ],
      "metadata": {
        "id": "ONw4bDQAuBBw"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GHIBEHtW-eHZ",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Tacotron2 GPU Synthesizer\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "\n",
        "#@markdown **The \"tacotron_id\" is where you can put a link to your trained tacotron2 model from Google Drive.**\n",
        "\n",
        "\n",
        "#@markdown If the audio sounds too artificial, you can lower the superres_strength\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown Config:\n",
        "\n",
        "#@markdown Restart the runtime to apply any changes.\n",
        "\n",
        "#Add new characters here.\n",
        "import logging\n",
        "\n",
        "logging.getLogger('matplotlib').setLevel(logging.WARNING)\n",
        "logging.getLogger('numba').setLevel(logging.WARNING)\n",
        "logging.getLogger('librosa').setLevel(logging.WARNING)\n",
        "\n",
        "\n",
        "#Universal HiFi-GAN (has some robotic noise): 1qpgI41wNXFcH-iKq1Y42JlBC9j0je8PW\n",
        "tacotron_id = \"\" #@param {type:\"string\"}\n",
        "\n",
        "hifigan_id = \"universal\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown _leave blank or enter \"universal\" for universal model_\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "if tacotron_id != \"\":\n",
        "    TACOTRON2_ID = tacotron_id\n",
        "else:\n",
        "    raise Exception(\"No ID provided.\")\n",
        "\n",
        "if hifigan_id in {\"\", \"universal\"}:\n",
        "    HIFIGAN_ID = \"universal\"\n",
        "    print(\"Using universal Hifi-Gan model.\")\n",
        "else:\n",
        "    HIFIGAN_ID = hifigan_id\n",
        "\n",
        "# Check if Initialized\n",
        "try:\n",
        "    initialized\n",
        "except NameError:\n",
        "    print(\"Setting up, please wait.\\n\")\n",
        "    !pip install tqdm -q\n",
        "    from tqdm.notebook import tqdm\n",
        "    with tqdm(total=5, leave=False) as pbar:\n",
        "        import os\n",
        "        from os.path import exists, join, basename, splitext\n",
        "        !pip install resampy\n",
        "        !pip install git+https://github.com/IAHispano/gdown.git\n",
        "        git_repo_url = 'https://github.com/justinjohn0306/TTS-TT2.git'\n",
        "        project_name = splitext(basename(git_repo_url))[0]\n",
        "        if not exists(project_name):\n",
        "            # clone and install\n",
        "            !git clone -q --recursive {git_repo_url}\n",
        "            !git clone -q --recursive https://github.com/justinjohn0306/hifi-gan\n",
        "            !pip install -q unidecode\n",
        "        pbar.update(1) # downloaded TT2 and HiFi-GAN\n",
        "        import sys\n",
        "        sys.path.append('hifi-gan')\n",
        "        sys.path.append(project_name)\n",
        "        import time\n",
        "        import matplotlib\n",
        "        import matplotlib.pylab as plt\n",
        "        import gdown\n",
        "        d = 'https://drive.google.com/uc?id='\n",
        "\n",
        "        %matplotlib inline\n",
        "        import IPython.display as ipd\n",
        "        import numpy as np\n",
        "        import torch\n",
        "        import json\n",
        "        from hparams import create_hparams\n",
        "        from model import Tacotron2\n",
        "        from layers import TacotronSTFT\n",
        "        from audio_processing import griffin_lim\n",
        "        from text import text_to_sequence\n",
        "        from env import AttrDict\n",
        "        from meldataset import mel_spectrogram, MAX_WAV_VALUE\n",
        "        from models import Generator\n",
        "        from denoiser import Denoiser\n",
        "        import resampy\n",
        "        import scipy.signal\n",
        "\n",
        "        pbar.update(1) # initialized Dependancies\n",
        "\n",
        "        graph_width = 900\n",
        "        graph_height = 360\n",
        "        def plot_data(data, figsize=(int(graph_width/100), int(graph_height/100))):\n",
        "            %matplotlib inline\n",
        "            fig, axes = plt.subplots(1, len(data), figsize=figsize)\n",
        "            for i in range(len(data)):\n",
        "                axes[i].imshow(data[i], aspect='auto', origin='lower',\n",
        "                            interpolation='none', cmap='inferno')\n",
        "            fig.canvas.draw()\n",
        "            plt.show()\n",
        "\n",
        "        # Setup Pronounciation Dictionary\n",
        "        !wget 'https://github.com/justinjohn0306/FakeYou-Tacotron2-Notebook/releases/download/CMU_dict/merged.dict.txt'\n",
        "        thisdict = {}\n",
        "        for line in reversed((open('merged.dict.txt', \"r\").read()).splitlines()):\n",
        "            thisdict[(line.split(\" \",1))[0]] = (line.split(\" \",1))[1].strip()\n",
        "\n",
        "        pbar.update(1) # Downloaded and Set up Pronounciation Dictionary\n",
        "\n",
        "        def ARPA(text, punctuation=r\"!?,.;\", EOS_Token=True):\n",
        "            out = ''\n",
        "            for word_ in text.split(\" \"):\n",
        "                word=word_; end_chars = ''\n",
        "                while any(elem in word for elem in punctuation) and len(word) > 1:\n",
        "                    if word[-1] in punctuation: end_chars = word[-1] + end_chars; word = word[:-1]\n",
        "                    else: break\n",
        "                try:\n",
        "                    word_arpa = thisdict[word.upper()]\n",
        "                    word = \"{\" + str(word_arpa) + \"}\"\n",
        "                except KeyError: pass\n",
        "                out = (out + \" \" + word + end_chars).strip()\n",
        "            if EOS_Token and out[-1] != \";\": out += \";\"\n",
        "            return out\n",
        "\n",
        "        def get_hifigan(MODEL_ID, conf_name):\n",
        "            # Download HiFi-GAN\n",
        "            hifigan_pretrained_model = 'hifimodel_' + conf_name\n",
        "            #gdown.download(d+MODEL_ID, hifigan_pretrained_model, quiet=False)\n",
        "\n",
        "            if MODEL_ID == 1:\n",
        "              !wget \"https://github.com/justinjohn0306/tacotron2/releases/download/assets/Superres_Twilight_33000\" -O $hifigan_pretrained_model\n",
        "            elif MODEL_ID == \"universal\":\n",
        "              !wget \"https://github.com/justinjohn0306/tacotron2/releases/download/assets/g_02500000\" -O $hifigan_pretrained_model\n",
        "            else:\n",
        "              gdown.download(d+MODEL_ID, hifigan_pretrained_model, quiet=False)\n",
        "\n",
        "            # Load HiFi-GAN\n",
        "            conf = os.path.join(\"hifi-gan\", conf_name + \".json\")\n",
        "            with open(conf) as f:\n",
        "                json_config = json.loads(f.read())\n",
        "            h = AttrDict(json_config)\n",
        "            torch.manual_seed(h.seed)\n",
        "            hifigan = Generator(h).to(torch.device(\"cuda\"))\n",
        "            state_dict_g = torch.load(hifigan_pretrained_model, map_location=torch.device(\"cuda\"))\n",
        "            hifigan.load_state_dict(state_dict_g[\"generator\"])\n",
        "            hifigan.eval()\n",
        "            hifigan.remove_weight_norm()\n",
        "            denoiser = Denoiser(hifigan, mode=\"normal\")\n",
        "            return hifigan, h, denoiser\n",
        "\n",
        "        # Download character HiFi-GAN\n",
        "        hifigan, h, denoiser = get_hifigan(HIFIGAN_ID, \"config_v1\")\n",
        "        # Download super-resolution HiFi-GAN\n",
        "        hifigan_sr, h2, denoiser_sr = get_hifigan(1, \"config_32k\")\n",
        "        pbar.update(1) # Downloaded and Set up HiFi-GAN\n",
        "\n",
        "        def has_MMI(STATE_DICT):\n",
        "            return any(True for x in STATE_DICT.keys() if \"mi.\" in x)\n",
        "\n",
        "        def get_Tactron2(MODEL_ID):\n",
        "            # Download Tacotron2\n",
        "            tacotron2_pretrained_model = 'MLPTTS'\n",
        "            gdown.download(d+MODEL_ID, tacotron2_pretrained_model, quiet=False)\n",
        "            if not exists(tacotron2_pretrained_model):\n",
        "                raise Exception(\"Tacotron2 model failed to download!\")\n",
        "            # Load Tacotron2 and Config\n",
        "            hparams = create_hparams()\n",
        "            hparams.sampling_rate = 22050\n",
        "            hparams.max_decoder_steps = 3000 # Max Duration\n",
        "            hparams.gate_threshold = 0.25 # Model must be 25% sure the clip is over before ending generation\n",
        "            model = Tacotron2(hparams)\n",
        "            state_dict = torch.load(tacotron2_pretrained_model)['state_dict']\n",
        "            if has_MMI(state_dict):\n",
        "                raise Exception(\"ERROR: This notebook does not currently support MMI models.\")\n",
        "            model.load_state_dict(state_dict)\n",
        "            _ = model.cuda().eval().half()\n",
        "            return model, hparams\n",
        "\n",
        "        model, hparams = get_Tactron2(TACOTRON2_ID)\n",
        "        previous_tt2_id = TACOTRON2_ID\n",
        "\n",
        "        pbar.update(1) # Downloaded and Set up Tacotron2\n",
        "\n",
        "        # Extra Info\n",
        "        def end_to_end_infer(text, pronounciation_dictionary, show_graphs):\n",
        "            for i in [x for x in text.split(\"\\n\") if len(x)]:\n",
        "                if not pronounciation_dictionary:\n",
        "                    if i[-1] != \";\": i=i+\";\"\n",
        "                else: i = ARPA(i)\n",
        "                with torch.no_grad(): # save VRAM by not including gradients\n",
        "                    sequence = np.array(text_to_sequence(i, ['english_cleaners']))[None, :]\n",
        "                    sequence = torch.autograd.Variable(torch.from_numpy(sequence)).cuda().long()\n",
        "                    mel_outputs, mel_outputs_postnet, _, alignments = model.inference(sequence)\n",
        "                    if show_graphs:\n",
        "                        plot_data((mel_outputs_postnet.float().data.cpu().numpy()[0],\n",
        "                                alignments.float().data.cpu().numpy()[0].T))\n",
        "                    y_g_hat = hifigan(mel_outputs_postnet.float())\n",
        "                    audio = y_g_hat.squeeze()\n",
        "                    audio = audio * MAX_WAV_VALUE\n",
        "                    audio_denoised = denoiser(audio.view(1, -1), strength=35)[:, 0]\n",
        "\n",
        "                    # Resample to 32k\n",
        "                    audio_denoised = audio_denoised.cpu().numpy().reshape(-1)\n",
        "\n",
        "                    normalize = (MAX_WAV_VALUE / np.max(np.abs(audio_denoised))) ** 0.9\n",
        "                    audio_denoised = audio_denoised * normalize\n",
        "                    wave = resampy.resample(\n",
        "                        audio_denoised,\n",
        "                        h.sampling_rate,\n",
        "                        h2.sampling_rate,\n",
        "                        filter=\"sinc_window\",\n",
        "                        window=scipy.signal.windows.hann,\n",
        "                        num_zeros=8,\n",
        "                    )\n",
        "                    wave_out = wave.astype(np.int16)\n",
        "\n",
        "                    # HiFi-GAN super-resolution\n",
        "                    wave = wave / MAX_WAV_VALUE\n",
        "                    wave = torch.FloatTensor(wave).to(torch.device(\"cuda\"))\n",
        "                    new_mel = mel_spectrogram(\n",
        "                        wave.unsqueeze(0),\n",
        "                        h2.n_fft,\n",
        "                        h2.num_mels,\n",
        "                        h2.sampling_rate,\n",
        "                        h2.hop_size,\n",
        "                        h2.win_size,\n",
        "                        h2.fmin,\n",
        "                        h2.fmax,\n",
        "                    )\n",
        "                    y_g_hat2 = hifigan_sr(new_mel)\n",
        "                    audio2 = y_g_hat2.squeeze()\n",
        "                    audio2 = audio2 * MAX_WAV_VALUE\n",
        "                    audio2_denoised = denoiser(audio2.view(1, -1), strength=35)[:, 0]\n",
        "\n",
        "                    # High-pass filter, mixing and denormalizing\n",
        "                    audio2_denoised = audio2_denoised.cpu().numpy().reshape(-1)\n",
        "                    b = scipy.signal.firwin(\n",
        "                        101, cutoff=10500, fs=h2.sampling_rate, pass_zero=False\n",
        "                    )\n",
        "                    y = scipy.signal.lfilter(b, [1.0], audio2_denoised)\n",
        "                    y *= superres_strength\n",
        "                    y_out = y.astype(np.int16)\n",
        "                    y_padded = np.zeros(wave_out.shape)\n",
        "                    y_padded[: y_out.shape[0]] = y_out\n",
        "                    sr_mix = wave_out + y_padded\n",
        "                    sr_mix = sr_mix / normalize\n",
        "\n",
        "                    print(\"\")\n",
        "                    ipd.display(ipd.Audio(sr_mix.astype(np.int16), rate=h2.sampling_rate))\n",
        "    from IPython.display import clear_output\n",
        "    clear_output()\n",
        "    initialized = \"Ready\"\n",
        "\n",
        "if previous_tt2_id != TACOTRON2_ID:\n",
        "    print(\"Updating Models\")\n",
        "    model, hparams = get_Tactron2(TACOTRON2_ID)\n",
        "    hifigan, h, denoiser = get_hifigan(HIFIGAN_ID, \"config_v1\")\n",
        "    previous_tt2_id = TACOTRON2_ID\n",
        "\n",
        "pronounciation_dictionary = False #@param {type:\"boolean\"}\n",
        "# disables automatic ARPAbet conversion, useful for inputting your own ARPAbet pronounciations or just for testing\n",
        "show_graphs = True #@param {type:\"boolean\"}\n",
        "max_duration = 20 #@param {type:\"integer\"}\n",
        "model.decoder.max_decoder_steps = max_duration * 80\n",
        "stop_threshold = 0.5 #@param {type:\"number\"}\n",
        "model.decoder.gate_threshold = stop_threshold\n",
        "superres_strength = 10 #@param {type:\"number\"}\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "print(f\"Current Config:\\npronounciation_dictionary: {pronounciation_dictionary}\\nshow_graphs: {show_graphs}\\nmax_duration (in seconds): {max_duration}\\nstop_threshold: {stop_threshold}\\nsuperres_strength: {superres_strength}\\n\\n\")\n",
        "\n",
        "time.sleep(1)\n",
        "print(\"Enter/Paste your text.\")\n",
        "contents = []\n",
        "while True:\n",
        "    try:\n",
        "        print(\"-\"*50)\n",
        "        line = input()\n",
        "        if line == \"\":\n",
        "            continue\n",
        "        end_to_end_infer(line, not pronounciation_dictionary, show_graphs)\n",
        "    except EOFError:\n",
        "        break\n",
        "    except KeyboardInterrupt:\n",
        "        print(\"Stopping...\")\n",
        "        break"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}