Ol_ControllableTalkNet(CPU).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Ol_ControllableTalkNet(CPU).ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gss5Ox_RNiba"
      },
      "source": [
        "# **CPU version of Controllable TalkNET**\n",
        "This colab change is made by **mega b#6696** and **justinjohn-03**\n",
        "# Original colab / code credits go to the PPP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JjKnX0ztK6Wc"
      },
      "source": [
        "# **Setup CPU**\n",
        "(run all)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "VqcXFy2qLH_g"
      },
      "source": [
        "#@markdown Download dependencies.\n",
        "%tensorflow_version 2.x\n",
        "import os\n",
        "\n",
        "custom_lists = [\n",
        "    #\"https://gist.githubusercontent.com/SortAnon/997cda157954a189259c9876fd804e53/raw/example_models.json\",\n",
        "]\n",
        "\n",
        "!apt-get install sox libsndfile1 ffmpeg\n",
        "!pip install tensorflow==2.4.1 dash==1.21.0 dash-bootstrap-components==0.13.0 jupyter-dash==0.4.0 psola wget unidecode pysptk frozendict torchvision==0.9.1 torchaudio==0.8.1 torchtext==0.9.1 torch_stft kaldiio pydub pyannote.audio g2p_en pesq pystoi crepe resampy ffmpeg-python torchcrepe einops taming-transformers-rom1504==0.0.6 tensorflow-hub gdown --upgrade\n",
        "!pip install numpy --upgrade\n",
        "!python -m pip install git+https://github.com/SortAnon/NeMo.git\n",
        "if not os.path.exists(\"hifi-gan\"):\n",
        "    !git clone -q --recursive https://github.com/SortAnon/hifi-gan\n",
        "!git clone -q https://github.com/justinjohn0306/ControllableTalkNet\n",
        "os.chdir(\"/content/ControllableTalkNet\")\n",
        "!git archive --output=./files.tar --format=tar HEAD\n",
        "os.chdir(\"/content\")\n",
        "!tar xf ControllableTalkNet/files.tar\n",
        "!rm -rf ControllableTalkNet\n",
        "\n",
        "os.chdir(\"/content/model_lists\")\n",
        "for c in custom_lists:\n",
        "    !wget \"{c}\"\n",
        "os.chdir(\"/content\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F111-Ah_KdnM",
        "cellView": "form"
      },
      "source": [
        "#@markdown Write denoiser for CPU inference\n",
        "%%writefile /content/hifi-gan/denoiser.py\n",
        "import sys\n",
        "import torch\n",
        "from stft import STFT\n",
        "\n",
        "\n",
        "class Denoiser(torch.nn.Module):\n",
        "    \"\"\" WaveGlow denoiser, adapted for HiFi-GAN \"\"\"\n",
        "\n",
        "    def __init__(\n",
        "        self, hifigan, filter_length=1024, n_overlap=4, win_length=1024, mode=\"zeros\"\n",
        "    ):\n",
        "        super(Denoiser, self).__init__()\n",
        "        self.stft = STFT(\n",
        "            filter_length=filter_length,\n",
        "            hop_length=int(filter_length / n_overlap),\n",
        "            win_length=win_length,\n",
        "        )\n",
        "        if mode == \"zeros\":\n",
        "            mel_input = torch.zeros((1, 80, 88))\n",
        "        elif mode == \"normal\":\n",
        "            mel_input = torch.randn((1, 80, 88))\n",
        "        else:\n",
        "            raise Exception(\"Mode {} if not supported\".format(mode))\n",
        "\n",
        "        with torch.no_grad():\n",
        "            bias_audio = hifigan(mel_input).view(1, -1).float()\n",
        "            bias_spec, _ = self.stft.transform(bias_audio)\n",
        "\n",
        "        self.register_buffer(\"bias_spec\", bias_spec[:, :, 0][:, :, None])\n",
        "\n",
        "    def forward(self, audio, strength=0.1):\n",
        "        audio_spec, audio_angles = self.stft.transform(audio.float())\n",
        "        audio_spec_denoised = audio_spec - self.bias_spec * strength\n",
        "        audio_spec_denoised = torch.clamp(audio_spec_denoised, 0.0)\n",
        "        audio_denoised = self.stft.inverse(audio_spec_denoised, audio_angles)\n",
        "        return audio_denoised\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "L0XOS-8yKmIL",
        "cellView": "form"
      },
      "source": [
        "#@markdown Load Controllable TalkNET\n",
        "%%writefile /content/controllable_talknet.py\n",
        "import sys\n",
        "import os\n",
        "import base64\n",
        "import dash\n",
        "from jupyter_dash import JupyterDash\n",
        "import dash_core_components as dcc\n",
        "import dash_html_components as html\n",
        "import dash_bootstrap_components as dbc\n",
        "from dash.exceptions import PreventUpdate\n",
        "import torch\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import crepe\n",
        "import scipy\n",
        "from scipy.io import wavfile\n",
        "import psola\n",
        "import io\n",
        "import nemo\n",
        "from nemo.collections.asr.models import EncDecCTCModel\n",
        "from nemo.collections.tts.models import TalkNetSpectModel\n",
        "from nemo.collections.tts.models import TalkNetPitchModel\n",
        "from nemo.collections.tts.models import TalkNetDursModel\n",
        "from talknet_singer import TalkNetSingerModel\n",
        "import json\n",
        "from tqdm import tqdm\n",
        "import gdown\n",
        "import zipfile\n",
        "import resampy\n",
        "import traceback\n",
        "import ffmpeg\n",
        "import time\n",
        "import uuid\n",
        "\n",
        "sys.path.append(\"hifi-gan\")\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",
        "\n",
        "app = JupyterDash(__name__)\n",
        "DEVICE = \"cpu\"\n",
        "CPU_PITCH = True\n",
        "RUN_PATH = os.path.dirname(os.path.realpath(__file__))\n",
        "if RUN_PATH == \"/content\":\n",
        "    UI_MODE = \"colab\"\n",
        "elif os.path.exists(\"/talknet/is_docker\"):\n",
        "    UI_MODE = \"docker\"\n",
        "else:\n",
        "    UI_MODE = \"offline\"\n",
        "torch.set_grad_enabled(False)\n",
        "if CPU_PITCH:\n",
        "    tf.config.set_visible_devices([], \"GPU\")\n",
        "\n",
        "app.title = \"Controllable TalkNet\"\n",
        "app.layout = html.Div(\n",
        "    children=[\n",
        "        html.H1(\n",
        "            id=\"header\",\n",
        "            children=\"Controllable TalkNet\",\n",
        "            style={\n",
        "                \"font-family\": \"EquestriaWebfont\",\n",
        "                \"color\": \"#ed3c96\",\n",
        "                \"font-size\": \"4em\",\n",
        "                \"text-align\": \"center\",\n",
        "                \"margin-top\": \"0em\",\n",
        "                \"margin-bottom\": \"0em\",\n",
        "            },\n",
        "        ),\n",
        "        html.Label(\"Character selection\", htmlFor=\"model-dropdown\"),\n",
        "        dbc.Select(\n",
        "            id=\"model-dropdown\",\n",
        "            options=[\n",
        "                {\n",
        "                    \"label\": \"Custom model\",\n",
        "                    \"value\": \"Custom\",\n",
        "                },\n",
        "                {\n",
        "                    \"label\": \"--- ERROR LOADING MODEL LISTS ---\",\n",
        "                    \"value\": \"\",\n",
        "                    \"disabled\": True,\n",
        "                },\n",
        "            ],\n",
        "            value=None,\n",
        "            style={\n",
        "                \"max-width\": \"90vw\",\n",
        "                \"width\": \"35em\",\n",
        "                \"margin-bottom\": \"0.7em\",\n",
        "            },\n",
        "        ),\n",
        "        html.Div(\n",
        "            children=[\n",
        "                dcc.Input(\n",
        "                    id=\"drive-id\",\n",
        "                    type=\"text\",\n",
        "                    placeholder=\"Drive ID for custom model\",\n",
        "                    style={\"width\": \"22em\"},\n",
        "                ),\n",
        "            ],\n",
        "            id=\"custom-model\",\n",
        "            style={\n",
        "                \"display\": \"none\",\n",
        "            },\n",
        "        ),\n",
        "        html.Label(\n",
        "            \"Upload reference audio to \" + RUN_PATH,\n",
        "            htmlFor=\"reference-dropdown\",\n",
        "            id=\"upload-label\",\n",
        "        ),\n",
        "        dcc.Upload(\n",
        "            id=\"upload-audio\",\n",
        "            children=html.Div([\"Drag and drop or click to select a file to upload.\"]),\n",
        "            style={\n",
        "                \"display\": \"none\",\n",
        "            },\n",
        "            multiple=True,\n",
        "        ),\n",
        "        dcc.Store(id=\"current-f0s\"),\n",
        "        dcc.Store(id=\"current-f0s-nosilence\"),\n",
        "        dcc.Store(id=\"current-filename\"),\n",
        "        dcc.Loading(\n",
        "            id=\"audio-loading\",\n",
        "            children=[\n",
        "                html.Div(\n",
        "                    [\n",
        "                        html.Button(\n",
        "                            \"Update file list\",\n",
        "                            id=\"update-button\",\n",
        "                            style={\n",
        "                                \"margin-right\": \"10px\",\n",
        "                            },\n",
        "                        ),\n",
        "                        dbc.Select(\n",
        "                            id=\"reference-dropdown\",\n",
        "                            options=[],\n",
        "                            value=None,\n",
        "                            style={\n",
        "                                \"max-width\": \"80vw\",\n",
        "                                \"width\": \"30em\",\n",
        "                            },\n",
        "                            disabled=False,\n",
        "                        ),\n",
        "                        dcc.Store(id=\"pitch-clicks\"),\n",
        "                        html.Button(\n",
        "                            \"Debug pitch\",\n",
        "                            id=\"pitch-button\",\n",
        "                            style={\n",
        "                                \"margin-left\": \"10px\",\n",
        "                            },\n",
        "                            disabled=False,\n",
        "                        ),\n",
        "                    ],\n",
        "                    style={\n",
        "                        \"width\": \"100%\",\n",
        "                        \"display\": \"flex\",\n",
        "                        \"align-items\": \"center\",\n",
        "                        \"justify-content\": \"center\",\n",
        "                        \"flex-direction\": \"row\",\n",
        "                        \"margin-left\": \"50px\",\n",
        "                        \"vertical-align\": \"middle\",\n",
        "                    },\n",
        "                ),\n",
        "                html.Audio(\n",
        "                    id=\"pitch-out\",\n",
        "                    controls=True,\n",
        "                    style={\"display\": \"none\"},\n",
        "                ),\n",
        "                html.Div(\n",
        "                    id=\"audio-loading-output\",\n",
        "                    style={\n",
        "                        \"font-style\": \"italic\",\n",
        "                        \"margin-bottom\": \"0.7em\",\n",
        "                        \"text-align\": \"center\",\n",
        "                    },\n",
        "                ),\n",
        "            ],\n",
        "            type=\"default\",\n",
        "        ),\n",
        "        html.Div(\n",
        "            [\n",
        "                dcc.Checklist(\n",
        "                    id=\"pitch-options\",\n",
        "                    options=[\n",
        "                        {\"label\": \"Change input pitch\", \"value\": \"pf\"},\n",
        "                        {\"label\": \"Auto-tune output\", \"value\": \"pc\"},\n",
        "                        {\"label\": \"Disable reference audio\", \"value\": \"dra\"},\n",
        "                    ],\n",
        "                    value=[],\n",
        "                ),\n",
        "                html.Div(\n",
        "                    [\n",
        "                        html.Label(\"Semitones\", htmlFor=\"pitch-factor\"),\n",
        "                        dcc.Input(\n",
        "                            id=\"pitch-factor\",\n",
        "                            type=\"number\",\n",
        "                            value=\"0\",\n",
        "                            style={\"width\": \"7em\"},\n",
        "                            min=-11,\n",
        "                            max=11,\n",
        "                            step=1,\n",
        "                            disabled=True,\n",
        "                        ),\n",
        "                    ],\n",
        "                    style={\n",
        "                        \"flex-direction\": \"column\",\n",
        "                        \"margin-left\": \"10px\",\n",
        "                        \"margin-bottom\": \"0.7em\",\n",
        "                    },\n",
        "                ),\n",
        "            ],\n",
        "            style={\n",
        "                \"width\": \"100%\",\n",
        "                \"display\": \"flex\",\n",
        "                \"align-items\": \"center\",\n",
        "                \"justify-content\": \"center\",\n",
        "                \"flex-direction\": \"row\",\n",
        "                \"margin-left\": \"50px\",\n",
        "                \"margin-bottom\": \"0.7em\",\n",
        "            },\n",
        "        ),\n",
        "        html.Label(\"Transcript\", htmlFor=\"transcript-input\"),\n",
        "        dcc.Textarea(\n",
        "            id=\"transcript-input\",\n",
        "            value=\"\",\n",
        "            style={\n",
        "                \"max-width\": \"90vw\",\n",
        "                \"width\": \"50em\",\n",
        "                \"height\": \"8em\",\n",
        "                \"margin-bottom\": \"0.7em\",\n",
        "            },\n",
        "        ),\n",
        "        dcc.Loading(\n",
        "            html.Div(\n",
        "                [\n",
        "                    html.Button(\n",
        "                        \"Generate\",\n",
        "                        id=\"gen-button\",\n",
        "                    ),\n",
        "                    html.Audio(\n",
        "                        id=\"audio-out\",\n",
        "                        controls=True,\n",
        "                        style={\n",
        "                            \"display\": \"none\",\n",
        "                        },\n",
        "                    ),\n",
        "                    html.Div(\n",
        "                        id=\"generated-info\",\n",
        "                        style={\n",
        "                            \"font-style\": \"italic\",\n",
        "                        },\n",
        "                    ),\n",
        "                ],\n",
        "                style={\n",
        "                    \"width\": \"100%\",\n",
        "                    \"display\": \"flex\",\n",
        "                    \"align-items\": \"center\",\n",
        "                    \"justify-content\": \"center\",\n",
        "                    \"flex-direction\": \"column\",\n",
        "                },\n",
        "            )\n",
        "        ),\n",
        "        html.Footer(\n",
        "            children=\"\"\"\n",
        "                Presented by the Pony Preservation Project.\n",
        "            \"\"\",\n",
        "            style={\"margin-top\": \"2em\", \"font-size\": \"0.7em\"},\n",
        "        ),\n",
        "    ],\n",
        "    style={\n",
        "        \"width\": \"100%\",\n",
        "        \"display\": \"flex\",\n",
        "        \"align-items\": \"center\",\n",
        "        \"justify-content\": \"center\",\n",
        "        \"flex-direction\": \"column\",\n",
        "        \"background-color\": \"#FFF\",\n",
        "    },\n",
        ")\n",
        "\n",
        "upload_display = {\n",
        "    \"width\": \"100%\",\n",
        "    \"height\": \"60px\",\n",
        "    \"lineHeight\": \"60px\",\n",
        "    \"borderWidth\": \"1px\",\n",
        "    \"borderStyle\": \"dashed\",\n",
        "    \"borderRadius\": \"5px\",\n",
        "    \"textAlign\": \"center\",\n",
        "    \"margin\": \"10px\",\n",
        "}\n",
        "\n",
        "playback_style = {\n",
        "    \"margin-top\": \"0.3em\",\n",
        "    \"margin-bottom\": \"0.3em\",\n",
        "    \"display\": \"block\",\n",
        "    \"width\": \"600px\",\n",
        "    \"max-width\": \"90vw\",\n",
        "}\n",
        "\n",
        "playback_hide = {\n",
        "    \"display\": \"none\",\n",
        "}\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    [\n",
        "        dash.dependencies.Output(\"model-dropdown\", \"options\"),\n",
        "        dash.dependencies.Output(\"upload-audio\", \"style\"),\n",
        "    ],\n",
        "    dash.dependencies.Input(\"header\", \"children\"),\n",
        ")\n",
        "def init_dropdown(value):\n",
        "    if UI_MODE == \"docker\":\n",
        "        upload_style = upload_display\n",
        "    else:\n",
        "        upload_style = playback_hide\n",
        "\n",
        "    dropdown = [\n",
        "        {\n",
        "            \"label\": \"Custom model\",\n",
        "            \"value\": \"Custom|default\",\n",
        "        }\n",
        "    ]\n",
        "    prev_values = [\"Custom|default\"]\n",
        "\n",
        "    def add_to_dropdown(entry):\n",
        "        if entry[\"value\"] in prev_values:\n",
        "            return\n",
        "        dropdown.append(entry)\n",
        "        prev_values.append(entry[\"value\"])\n",
        "\n",
        "    all_dict = {}\n",
        "    for filename in os.listdir(\"model_lists\"):\n",
        "        if len(filename) < 5 or filename[-5:].lower() != \".json\":\n",
        "            continue\n",
        "        with open(os.path.join(\"model_lists\", filename)) as f:\n",
        "            j = json.load(f)\n",
        "            for s in j:\n",
        "                for c in s[\"characters\"]:\n",
        "                    c[\"source_file\"] = filename[:-5]\n",
        "                if s[\"source\"] not in all_dict:\n",
        "                    all_dict[s[\"source\"]] = s[\"characters\"]\n",
        "                else:\n",
        "                    all_dict[s[\"source\"]].extend(s[\"characters\"])\n",
        "    for k in sorted(all_dict):\n",
        "        seen_chars = []\n",
        "        seen_ids = []\n",
        "        characters = {}\n",
        "        characters_sing = {}\n",
        "        has_singers = False\n",
        "        for c in all_dict[k]:\n",
        "            if c[\"drive_id\"] in seen_ids:\n",
        "                continue\n",
        "            seen_ids.append(c[\"drive_id\"])\n",
        "            # Handle duplicate names\n",
        "            if c[\"name\"] in seen_chars:\n",
        "                if c[\"name\"] in characters:\n",
        "                    rename = (\n",
        "                        c[\"name\"] + \" [\" + characters[c[\"name\"]][\"source_file\"] + \"]\"\n",
        "                    )\n",
        "                    characters[rename] = characters[c[\"name\"]]\n",
        "                    del characters[c[\"name\"]]\n",
        "                c[\"name\"] = c[\"name\"] + \" [\" + c[\"source_file\"] + \"]\"\n",
        "            else:\n",
        "                seen_chars.append(c[\"name\"])\n",
        "\n",
        "            characters[c[\"name\"]] = {\n",
        "                \"drive_id\": c[\"drive_id\"],\n",
        "                \"is_singing\": c[\"is_singing\"],\n",
        "                \"source_file\": c[\"source_file\"],\n",
        "            }\n",
        "            if c[\"is_singing\"]:\n",
        "                has_singers = True\n",
        "        if has_singers:\n",
        "            for ck in sorted(characters):\n",
        "                if characters[ck][\"is_singing\"]:\n",
        "                    characters_sing[ck] = characters[ck]\n",
        "                    del characters[ck]\n",
        "            separator = \"--- \" + k.strip().upper() + \" MODELS (TALKING) ---\"\n",
        "        else:\n",
        "            separator = \"--- \" + k.strip().upper() + \" MODELS ---\"\n",
        "        if len(characters) > 0:\n",
        "            add_to_dropdown(\n",
        "                {\n",
        "                    \"label\": separator,\n",
        "                    \"value\": str(uuid.uuid4()) + \"|default\",\n",
        "                    \"disabled\": True,\n",
        "                }\n",
        "            )\n",
        "            for ck in sorted(characters):\n",
        "                add_to_dropdown(\n",
        "                    {\n",
        "                        \"label\": ck,\n",
        "                        \"value\": characters[ck][\"drive_id\"] + \"|default\",\n",
        "                    }\n",
        "                )\n",
        "        if has_singers:\n",
        "            separator = \"--- \" + k.strip().upper() + \" MODELS (SINGING) ---\"\n",
        "            add_to_dropdown(\n",
        "                {\n",
        "                    \"label\": separator,\n",
        "                    \"value\": str(uuid.uuid4()) + \"|default\",\n",
        "                    \"disabled\": True,\n",
        "                }\n",
        "            )\n",
        "            for ck in sorted(characters_sing):\n",
        "                add_to_dropdown(\n",
        "                    {\n",
        "                        \"label\": ck,\n",
        "                        \"value\": characters_sing[ck][\"drive_id\"] + \"|singing\",\n",
        "                    }\n",
        "                )\n",
        "    if len(all_dict) == 0:\n",
        "        add_to_dropdown(\n",
        "            {\n",
        "                \"label\": \"--- NO MODEL LISTS FOUND ---\",\n",
        "                \"value\": str(uuid.uuid4()) + \"|default\",\n",
        "                \"disabled\": True,\n",
        "            }\n",
        "        )\n",
        "    return [dropdown, upload_style]\n",
        "\n",
        "\n",
        "def load_hifigan(model_name, conf_name):\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(DEVICE))\n",
        "    state_dict_g = torch.load(model_name, map_location=torch.device(DEVICE))\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",
        "\n",
        "def generate_json(input, outpath):\n",
        "    output = \"\"\n",
        "    sample_rate = 22050\n",
        "    lpath = input.split(\"|\")[0].strip()\n",
        "    size = os.stat(lpath).st_size\n",
        "    x = {\n",
        "        \"audio_filepath\": lpath,\n",
        "        \"duration\": size / (sample_rate * 2),\n",
        "        \"text\": input.split(\"|\")[1].strip(),\n",
        "    }\n",
        "    output += json.dumps(x) + \"\\n\"\n",
        "    with open(outpath, \"w\", encoding=\"utf8\") as w:\n",
        "        w.write(output)\n",
        "\n",
        "\n",
        "asr_model = (\n",
        "    EncDecCTCModel.from_pretrained(model_name=\"asr_talknet_aligner\").cpu().eval()\n",
        ")\n",
        "\n",
        "\n",
        "def forward_extractor(tokens, log_probs, blank):\n",
        "    \"\"\"Computes states f and p.\"\"\"\n",
        "    n, m = len(tokens), log_probs.shape[0]\n",
        "    # `f[s, t]` -- max sum of log probs for `s` first codes\n",
        "    # with `t` first timesteps with ending in `tokens[s]`.\n",
        "    f = np.empty((n + 1, m + 1), dtype=float)\n",
        "    f.fill(-(10 ** 9))\n",
        "    p = np.empty((n + 1, m + 1), dtype=int)\n",
        "    f[0, 0] = 0.0  # Start\n",
        "    for s in range(1, n + 1):\n",
        "        c = tokens[s - 1]\n",
        "        for t in range((s + 1) // 2, m + 1):\n",
        "            f[s, t] = log_probs[t - 1, c]\n",
        "            # Option #1: prev char is equal to current one.\n",
        "            if s == 1 or c == blank or c == tokens[s - 3]:\n",
        "                options = f[s : (s - 2 if s > 1 else None) : -1, t - 1]\n",
        "            else:  # Is not equal to current one.\n",
        "                options = f[s : (s - 3 if s > 2 else None) : -1, t - 1]\n",
        "            f[s, t] += np.max(options)\n",
        "            p[s, t] = np.argmax(options)\n",
        "    return f, p\n",
        "\n",
        "\n",
        "def backward_extractor(f, p):\n",
        "    \"\"\"Computes durs from f and p.\"\"\"\n",
        "    n, m = f.shape\n",
        "    n -= 1\n",
        "    m -= 1\n",
        "    durs = np.zeros(n, dtype=int)\n",
        "    if f[-1, -1] >= f[-2, -1]:\n",
        "        s, t = n, m\n",
        "    else:\n",
        "        s, t = n - 1, m\n",
        "    while s > 0:\n",
        "        durs[s - 1] += 1\n",
        "        s -= p[s, t]\n",
        "        t -= 1\n",
        "    assert durs.shape[0] == n\n",
        "    assert np.sum(durs) == m\n",
        "    assert np.all(durs[1::2] > 0)\n",
        "    return durs\n",
        "\n",
        "\n",
        "def preprocess_tokens(tokens, blank):\n",
        "    new_tokens = [blank]\n",
        "    for c in tokens:\n",
        "        new_tokens.extend([c, blank])\n",
        "    tokens = new_tokens\n",
        "    return tokens\n",
        "\n",
        "\n",
        "parser = (\n",
        "    nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset.make_vocab(\n",
        "        notation=\"phonemes\",\n",
        "        punct=True,\n",
        "        spaces=True,\n",
        "        stresses=False,\n",
        "        add_blank_at=\"last\",\n",
        "    )\n",
        ")\n",
        "\n",
        "\n",
        "def arpa_parse(input, model):\n",
        "    z = []\n",
        "    space = parser.labels.index(\" \")\n",
        "    while \"{\" in input:\n",
        "        if \"}\" not in input:\n",
        "            input.replace(\"{\", \"\")\n",
        "        else:\n",
        "            pre = input[: input.find(\"{\")]\n",
        "            if pre.strip() != \"\":\n",
        "                x = model.parse(text=pre.strip())\n",
        "                seq_ids = x.squeeze(0).cpu().detach().numpy()\n",
        "                z.extend(seq_ids)\n",
        "            z.append(space)\n",
        "\n",
        "            arpaword = input[input.find(\"{\") + 1 : input.find(\"}\")]\n",
        "            arpaword = (\n",
        "                arpaword.replace(\"0\", \"\")\n",
        "                .replace(\"1\", \"\")\n",
        "                .replace(\"2\", \"\")\n",
        "                .strip()\n",
        "                .split(\" \")\n",
        "            )\n",
        "\n",
        "            seq_ids = []\n",
        "            for x in arpaword:\n",
        "                if x == \"\":\n",
        "                    continue\n",
        "                if x.replace(\"_\", \" \") not in parser.labels:\n",
        "                    continue\n",
        "                seq_ids.append(parser.labels.index(x.replace(\"_\", \" \")))\n",
        "            seq_ids.append(space)\n",
        "            z.extend(seq_ids)\n",
        "            input = input[input.find(\"}\") + 1 :]\n",
        "    if input != \"\":\n",
        "        x = model.parse(text=input.strip())\n",
        "        seq_ids = x.squeeze(0).cpu().detach().numpy()\n",
        "        z.extend(seq_ids)\n",
        "    if z[-1] == space:\n",
        "        z = z[:-1]\n",
        "    if z[0] == space:\n",
        "        z = z[1:]\n",
        "    return [\n",
        "        z[i] for i in range(len(z)) if (i == 0) or (z[i] != z[i - 1]) or (z[i] != space)\n",
        "    ]\n",
        "\n",
        "\n",
        "def to_arpa(input):\n",
        "    arpa = \"\"\n",
        "    z = []\n",
        "    space = parser.labels.index(\" \")\n",
        "    while space in input:\n",
        "        z.append(input[: input.index(space)])\n",
        "        input = input[input.index(space) + 1 :]\n",
        "    z.append(input)\n",
        "    for y in z:\n",
        "        if len(y) == 0:\n",
        "            continue\n",
        "\n",
        "        arpaword = \" {\"\n",
        "        for s in y:\n",
        "            if parser.labels[s] == \" \":\n",
        "                arpaword += \"_ \"\n",
        "            else:\n",
        "                arpaword += parser.labels[s] + \" \"\n",
        "        arpaword += \"} \"\n",
        "        if not arpaword.replace(\"{\", \"\").replace(\"}\", \"\").replace(\" \", \"\").isalnum():\n",
        "            arpaword = arpaword.replace(\"{\", \"\").replace(\" }\", \"\")\n",
        "        arpa += arpaword\n",
        "    return arpa.replace(\"  \", \" \").replace(\" }\", \"}\").strip()\n",
        "\n",
        "\n",
        "def get_duration(wav_name, transcript, tokens):\n",
        "    if not os.path.exists(os.path.join(RUN_PATH, \"temp\")):\n",
        "        os.mkdir(os.path.join(RUN_PATH, \"temp\"))\n",
        "    if \"_\" not in transcript:\n",
        "        generate_json(\n",
        "            os.path.join(RUN_PATH, \"temp\", wav_name + \"_conv.wav\")\n",
        "            + \"|\"\n",
        "            + transcript.strip(),\n",
        "            os.path.join(RUN_PATH, \"temp\", wav_name + \".json\"),\n",
        "        )\n",
        "    else:\n",
        "        generate_json(\n",
        "            os.path.join(RUN_PATH, \"temp\", wav_name + \"_conv.wav\") + \"|\" + \"dummy\",\n",
        "            os.path.join(RUN_PATH, \"temp\", wav_name + \".json\"),\n",
        "        )\n",
        "\n",
        "    data_config = {\n",
        "        \"manifest_filepath\": os.path.join(RUN_PATH, \"temp\", wav_name + \".json\"),\n",
        "        \"sample_rate\": 22050,\n",
        "        \"batch_size\": 1,\n",
        "    }\n",
        "\n",
        "    dataset = nemo.collections.asr.data.audio_to_text._AudioTextDataset(\n",
        "        manifest_filepath=data_config[\"manifest_filepath\"],\n",
        "        sample_rate=data_config[\"sample_rate\"],\n",
        "        parser=parser,\n",
        "    )\n",
        "\n",
        "    dl = torch.utils.data.DataLoader(\n",
        "        dataset=dataset,\n",
        "        batch_size=data_config[\"batch_size\"],\n",
        "        collate_fn=dataset.collate_fn,\n",
        "        shuffle=False,\n",
        "    )\n",
        "\n",
        "    blank_id = asr_model.decoder.num_classes_with_blank - 1\n",
        "\n",
        "    for sample_idx, test_sample in tqdm(enumerate(dl), total=len(dl)):\n",
        "        log_probs, _, greedy_predictions = asr_model(\n",
        "            input_signal=test_sample[0], input_signal_length=test_sample[1]\n",
        "        )\n",
        "\n",
        "        log_probs = log_probs[0].cpu().detach().numpy()\n",
        "        target_tokens = preprocess_tokens(tokens, blank_id)\n",
        "\n",
        "        f, p = forward_extractor(target_tokens, log_probs, blank_id)\n",
        "        durs = backward_extractor(f, p)\n",
        "\n",
        "        del test_sample\n",
        "        return durs\n",
        "    return None\n",
        "\n",
        "\n",
        "def crepe_f0(wav_path, hop_length=256):\n",
        "    # sr, audio = wavfile.read(io.BytesIO(wav_data))\n",
        "    sr, audio = wavfile.read(wav_path)\n",
        "    audio_x = np.arange(0, len(audio)) / 22050.0\n",
        "    f0time, frequency, confidence, activation = crepe.predict(audio, sr, viterbi=True)\n",
        "\n",
        "    x = np.arange(0, len(audio), hop_length) / 22050.0\n",
        "    freq_interp = np.interp(x, f0time, frequency)\n",
        "    conf_interp = np.interp(x, f0time, confidence)\n",
        "    audio_interp = np.interp(x, audio_x, np.absolute(audio)) / 32768.0\n",
        "    weights = [0.5, 0.25, 0.25]\n",
        "    audio_smooth = np.convolve(audio_interp, np.array(weights)[::-1], \"same\")\n",
        "\n",
        "    conf_threshold = 0.25\n",
        "    audio_threshold = 0.0005\n",
        "    for i in range(len(freq_interp)):\n",
        "        if conf_interp[i] < conf_threshold:\n",
        "            freq_interp[i] = 0.0\n",
        "        if audio_smooth[i] < audio_threshold:\n",
        "            freq_interp[i] = 0.0\n",
        "\n",
        "    # Hack to make f0 and mel lengths equal\n",
        "    if len(audio) % hop_length == 0:\n",
        "        freq_interp = np.pad(freq_interp, pad_width=[0, 1])\n",
        "    return (\n",
        "        torch.from_numpy(freq_interp.astype(np.float32)),\n",
        "        torch.from_numpy(frequency.astype(np.float32)),\n",
        "    )\n",
        "\n",
        "\n",
        "def f0_to_audio(f0s):\n",
        "    volume = 0.2\n",
        "    sr = 22050\n",
        "    freq = 440.0\n",
        "    base_audio = (\n",
        "        np.sin(2 * np.pi * np.arange(256.0 * len(f0s)) * freq / sr) * volume\n",
        "    ).astype(np.float32)\n",
        "    shifted_audio = psola.vocode(base_audio, sr, target_pitch=f0s)\n",
        "    for i in range(len(f0s)):\n",
        "        if f0s[i] == 0.0:\n",
        "            shifted_audio[i * 256 : (i + 1) * 256] = 0.0\n",
        "    print(type(shifted_audio[0]))\n",
        "    buffer = io.BytesIO()\n",
        "    wavfile.write(buffer, sr, shifted_audio.astype(np.float32))\n",
        "    b64 = base64.b64encode(buffer.getvalue())\n",
        "    sound = \"data:audio/x-wav;base64,\" + b64.decode(\"ascii\")\n",
        "    return sound\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    dash.dependencies.Output(\"custom-model\", \"style\"),\n",
        "    dash.dependencies.Input(\"model-dropdown\", \"value\"),\n",
        ")\n",
        "def update_model(model):\n",
        "    if model is not None and model.split(\"|\")[0] == \"Custom\":\n",
        "        style = {\"margin-bottom\": \"0.7em\", \"display\": \"block\"}\n",
        "    else:\n",
        "        style = {\"display\": \"none\"}\n",
        "    return style\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    [\n",
        "        dash.dependencies.Output(\"pitch-factor\", \"disabled\"),\n",
        "        dash.dependencies.Output(\"reference-dropdown\", \"disabled\"),\n",
        "        dash.dependencies.Output(\"pitch-button\", \"disabled\"),\n",
        "    ],\n",
        "    [\n",
        "        dash.dependencies.Input(\"pitch-options\", \"value\"),\n",
        "    ],\n",
        ")\n",
        "def update_pitch_options(value):\n",
        "    return [\"pf\" not in value, \"dra\" in value, \"dra\" in value]\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    dash.dependencies.Output(\"upload-label\", \"children\"),\n",
        "    [\n",
        "        dash.dependencies.Input(\"upload-audio\", \"filename\"),\n",
        "        dash.dependencies.Input(\"upload-audio\", \"contents\"),\n",
        "    ],\n",
        ")\n",
        "def save_upload(uploaded_filenames, uploaded_file_contents):\n",
        "    try:\n",
        "        if uploaded_filenames is not None and uploaded_file_contents is not None:\n",
        "            for name, content in zip(uploaded_filenames, uploaded_file_contents):\n",
        "                if name.strip() == \"\":\n",
        "                    continue\n",
        "                data = content.encode(\"utf8\").split(b\";base64,\")[1]\n",
        "                with open(os.path.join(RUN_PATH, name), \"wb\") as fp:\n",
        "                    fp.write(base64.decodebytes(data))\n",
        "    except Exception as e:\n",
        "        return str(e)\n",
        "    return \"Uploaded \" + str(len(uploaded_filenames)) + \" file(s)\"\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    dash.dependencies.Output(\"reference-dropdown\", \"options\"),\n",
        "    [\n",
        "        dash.dependencies.Input(\"update-button\", \"n_clicks\"),\n",
        "    ],\n",
        ")\n",
        "def update_filelist(n_clicks):\n",
        "    filelist = []\n",
        "    supported_formats = [\".wav\", \".ogg\", \".mp3\", \"flac\", \".aac\"]\n",
        "    for x in sorted(os.listdir(RUN_PATH)):\n",
        "        if x[-4:].lower() in supported_formats:\n",
        "            filelist.append({\"label\": x, \"value\": x})\n",
        "    return filelist\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    [\n",
        "        dash.dependencies.Output(\"audio-loading-output\", \"children\"),\n",
        "        dash.dependencies.Output(\"current-f0s\", \"data\"),\n",
        "        dash.dependencies.Output(\"current-f0s-nosilence\", \"data\"),\n",
        "        dash.dependencies.Output(\"current-filename\", \"data\"),\n",
        "    ],\n",
        "    [\n",
        "        dash.dependencies.Input(\"reference-dropdown\", \"value\"),\n",
        "    ],\n",
        ")\n",
        "def select_file(dropdown_value):\n",
        "    if dropdown_value is not None:\n",
        "        if not os.path.exists(os.path.join(RUN_PATH, \"temp\")):\n",
        "            os.mkdir(os.path.join(RUN_PATH, \"temp\"))\n",
        "        ffmpeg.input(os.path.join(RUN_PATH, dropdown_value)).output(\n",
        "            os.path.join(RUN_PATH, \"temp\", dropdown_value + \"_conv.wav\"),\n",
        "            ar=\"22050\",\n",
        "            ac=\"1\",\n",
        "            acodec=\"pcm_s16le\",\n",
        "            map_metadata=\"-1\",\n",
        "            fflags=\"+bitexact\",\n",
        "        ).overwrite_output().run(quiet=True)\n",
        "        fo_with_silence, f0_wo_silence = crepe_f0(\n",
        "            os.path.join(RUN_PATH, \"temp\", dropdown_value + \"_conv.wav\")\n",
        "        )\n",
        "        return [\n",
        "            \"Analyzed \" + dropdown_value,\n",
        "            fo_with_silence,\n",
        "            f0_wo_silence,\n",
        "            dropdown_value,\n",
        "        ]\n",
        "    else:\n",
        "        return [\"No audio analyzed\", None, None]\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    [\n",
        "        dash.dependencies.Output(\"pitch-out\", \"src\"),\n",
        "        dash.dependencies.Output(\"pitch-out\", \"style\"),\n",
        "        dash.dependencies.Output(\"pitch-clicks\", \"data\"),\n",
        "    ],\n",
        "    [\n",
        "        dash.dependencies.Input(\"pitch-button\", \"n_clicks\"),\n",
        "        dash.dependencies.Input(\"pitch-clicks\", \"data\"),\n",
        "        dash.dependencies.Input(\"current-f0s\", \"data\"),\n",
        "    ],\n",
        ")\n",
        "def debug_pitch(n_clicks, pitch_clicks, current_f0s):\n",
        "    if not n_clicks or current_f0s is None or n_clicks <= pitch_clicks:\n",
        "        if n_clicks is not None:\n",
        "            pitch_clicks = n_clicks\n",
        "        else:\n",
        "            pitch_clicks = 0\n",
        "        return [\n",
        "            None,\n",
        "            playback_hide,\n",
        "            pitch_clicks,\n",
        "        ]\n",
        "    pitch_clicks = n_clicks\n",
        "    return [f0_to_audio(current_f0s), playback_style, pitch_clicks]\n",
        "\n",
        "\n",
        "hifigan_sr = None\n",
        "\n",
        "\n",
        "def download_model(model, custom_model):\n",
        "    try:\n",
        "        global hifigan_sr, h2, denoiser_sr\n",
        "        d = \"https://drive.google.com/uc?id=\"\n",
        "        if model == \"Custom\":\n",
        "            drive_id = custom_model\n",
        "        else:\n",
        "            drive_id = model\n",
        "        if drive_id == \"\" or drive_id is None:\n",
        "            return (\"Missing Drive ID\", None, None)\n",
        "        if not os.path.exists(os.path.join(RUN_PATH, \"models\")):\n",
        "            os.mkdir(os.path.join(RUN_PATH, \"models\"))\n",
        "        if not os.path.exists(os.path.join(RUN_PATH, \"models\", drive_id)):\n",
        "            os.mkdir(os.path.join(RUN_PATH, \"models\", drive_id))\n",
        "            zip_path = os.path.join(RUN_PATH, \"models\", drive_id, \"model.zip\")\n",
        "            gdown.download(\n",
        "                d + drive_id,\n",
        "                zip_path,\n",
        "                quiet=False,\n",
        "            )\n",
        "            if not os.path.exists(zip_path):\n",
        "                os.rmdir(os.path.join(RUN_PATH, \"models\", drive_id))\n",
        "                return (\"Model download failed\", None, None)\n",
        "            if os.stat(zip_path).st_size < 16:\n",
        "                os.remove(zip_path)\n",
        "                os.rmdir(os.path.join(RUN_PATH, \"models\", drive_id))\n",
        "                return (\"Model zip is empty\", None, None)\n",
        "            with zipfile.ZipFile(zip_path, \"r\") as zip_ref:\n",
        "                zip_ref.extractall(os.path.join(RUN_PATH, \"models\", drive_id))\n",
        "            os.remove(zip_path)\n",
        "\n",
        "        # Download super-resolution HiFi-GAN\n",
        "        sr_path = \"hifi-gan/hifisr\"\n",
        "        if not os.path.exists(sr_path):\n",
        "            gdown.download(\n",
        "                d + \"14fOprFAIlCQkVRxsfInhEPG0n-xN4QOa\", sr_path, quiet=False\n",
        "            )\n",
        "        if not os.path.exists(sr_path):\n",
        "            raise Exception(\"HiFI-GAN model failed to download!\")\n",
        "        if hifigan_sr is None:\n",
        "            hifigan_sr, h2, denoiser_sr = load_hifigan(sr_path, \"config_32k\")\n",
        "\n",
        "        return (\n",
        "            None,\n",
        "            os.path.join(RUN_PATH, \"models\", drive_id, \"TalkNetSpect.nemo\"),\n",
        "            os.path.join(RUN_PATH, \"models\", drive_id, \"hifiganmodel\"),\n",
        "        )\n",
        "    except Exception as e:\n",
        "        return (str(e), None, None)\n",
        "\n",
        "\n",
        "tnmodel, tnpath, tndurs, tnpitch = None, None, None, None\n",
        "hifigan, h, denoiser, hifipath = None, None, None, None\n",
        "\n",
        "\n",
        "@app.callback(\n",
        "    [\n",
        "        dash.dependencies.Output(\"audio-out\", \"src\"),\n",
        "        dash.dependencies.Output(\"generated-info\", \"children\"),\n",
        "        dash.dependencies.Output(\"audio-out\", \"style\"),\n",
        "        dash.dependencies.Output(\"audio-out\", \"title\"),\n",
        "    ],\n",
        "    [dash.dependencies.Input(\"gen-button\", \"n_clicks\")],\n",
        "    [\n",
        "        dash.dependencies.State(\"model-dropdown\", \"value\"),\n",
        "        dash.dependencies.State(\"drive-id\", \"value\"),\n",
        "        dash.dependencies.State(\"transcript-input\", \"value\"),\n",
        "        dash.dependencies.State(\"pitch-options\", \"value\"),\n",
        "        dash.dependencies.State(\"pitch-factor\", \"value\"),\n",
        "        dash.dependencies.State(\"current-filename\", \"data\"),\n",
        "        dash.dependencies.State(\"current-f0s\", \"data\"),\n",
        "        dash.dependencies.State(\"current-f0s-nosilence\", \"data\"),\n",
        "    ],\n",
        ")\n",
        "def generate_audio(\n",
        "    n_clicks,\n",
        "    model,\n",
        "    custom_model,\n",
        "    transcript,\n",
        "    pitch_options,\n",
        "    pitch_factor,\n",
        "    wav_name,\n",
        "    f0s,\n",
        "    f0s_wo_silence,\n",
        "):\n",
        "    global tnmodel, tnpath, tndurs, tnpitch, hifigan, h, denoiser, hifipath\n",
        "\n",
        "    if n_clicks is None:\n",
        "        raise PreventUpdate\n",
        "    if model is None:\n",
        "        return [None, \"No character selected\", playback_hide, None]\n",
        "    if transcript is None or transcript.strip() == \"\":\n",
        "        return [\n",
        "            None,\n",
        "            \"No transcript entered\",\n",
        "            playback_hide,\n",
        "            None,\n",
        "        ]\n",
        "    if wav_name is None and \"dra\" not in pitch_options:\n",
        "        return [\n",
        "            None,\n",
        "            \"No reference audio selected\",\n",
        "            playback_hide,\n",
        "            None,\n",
        "        ]\n",
        "    load_error, talknet_path, hifigan_path = download_model(\n",
        "        model.split(\"|\")[0], custom_model\n",
        "    )\n",
        "    if load_error is not None:\n",
        "        return [\n",
        "            None,\n",
        "            load_error,\n",
        "            playback_hide,\n",
        "            None,\n",
        "        ]\n",
        "\n",
        "    try:\n",
        "        with torch.no_grad():\n",
        "            if tnpath != talknet_path:\n",
        "                singer_path = os.path.join(\n",
        "                    os.path.dirname(talknet_path), \"TalkNetSinger.nemo\"\n",
        "                )\n",
        "                if os.path.exists(singer_path):\n",
        "                    tnmodel = TalkNetSingerModel.restore_from(singer_path)\n",
        "                else:\n",
        "                    tnmodel = TalkNetSpectModel.restore_from(talknet_path)\n",
        "                durs_path = os.path.join(\n",
        "                    os.path.dirname(talknet_path), \"TalkNetDurs.nemo\"\n",
        "                )\n",
        "                pitch_path = os.path.join(\n",
        "                    os.path.dirname(talknet_path), \"TalkNetPitch.nemo\"\n",
        "                )\n",
        "                if os.path.exists(durs_path):\n",
        "                    tndurs = TalkNetDursModel.restore_from(durs_path)\n",
        "                    tnmodel.add_module(\"_durs_model\", tndurs)\n",
        "                    tnpitch = TalkNetPitchModel.restore_from(pitch_path)\n",
        "                    tnmodel.add_module(\"_pitch_model\", tnpitch)\n",
        "                else:\n",
        "                    tndurs = None\n",
        "                    tnpitch = None\n",
        "                tnmodel.eval()\n",
        "                tnpath = talknet_path\n",
        "\n",
        "            token_list = arpa_parse(transcript, tnmodel)\n",
        "            tokens = torch.IntTensor(token_list).view(1, -1).to(DEVICE)\n",
        "            arpa = to_arpa(token_list)\n",
        "\n",
        "            if \"dra\" in pitch_options:\n",
        "                if tndurs is None or tnpitch is None:\n",
        "                    return [\n",
        "                        None,\n",
        "                        \"Model doesn't support pitch prediction\",\n",
        "                        playback_hide,\n",
        "                        None,\n",
        "                    ]\n",
        "                spect = tnmodel.generate_spectrogram(tokens=tokens)\n",
        "            else:\n",
        "                durs = get_duration(wav_name, transcript, token_list)\n",
        "\n",
        "                # Change pitch\n",
        "                if \"pf\" in pitch_options:\n",
        "                    f0_factor = np.power(np.e, (0.0577623 * float(pitch_factor)))\n",
        "                    f0s = [x * f0_factor for x in f0s]\n",
        "                    f0s_wo_silence = [x * f0_factor for x in f0s_wo_silence]\n",
        "\n",
        "                spect = tnmodel.force_spectrogram(\n",
        "                    tokens=tokens,\n",
        "                    durs=torch.from_numpy(durs)\n",
        "                    .view(1, -1)\n",
        "                    .type(torch.LongTensor)\n",
        "                    .to(DEVICE),\n",
        "                    f0=torch.FloatTensor(f0s).view(1, -1).to(DEVICE),\n",
        "                )\n",
        "\n",
        "            if hifipath != hifigan_path:\n",
        "                hifigan, h, denoiser = load_hifigan(hifigan_path, \"config_v1\")\n",
        "                hifipath = hifigan_path\n",
        "\n",
        "            y_g_hat = hifigan(spect.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",
        "            audio_np = (\n",
        "                audio_denoised.detach().cpu().numpy().reshape(-1).astype(np.int16)\n",
        "            )\n",
        "\n",
        "            # Auto-tuning\n",
        "            if \"pc\" in pitch_options and \"dra\" not in pitch_options:\n",
        "                _, output_freq, _, _ = crepe.predict(audio_np, 22050, viterbi=True)\n",
        "                output_pitch = torch.from_numpy(output_freq.astype(np.float32))\n",
        "                target_pitch = torch.FloatTensor(f0s_wo_silence)\n",
        "                factor = torch.mean(output_pitch) / torch.mean(target_pitch)\n",
        "\n",
        "                octaves = [0.125, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0]\n",
        "                nearest_octave = min(octaves, key=lambda x: abs(x - factor))\n",
        "                target_pitch *= nearest_octave\n",
        "                if len(target_pitch) < len(output_pitch):\n",
        "                    target_pitch = torch.nn.functional.pad(\n",
        "                        target_pitch,\n",
        "                        (0, list(output_pitch.shape)[0] - list(target_pitch.shape)[0]),\n",
        "                        \"constant\",\n",
        "                        0,\n",
        "                    )\n",
        "                if len(target_pitch) > len(output_pitch):\n",
        "                    target_pitch = target_pitch[0 : list(output_pitch.shape)[0]]\n",
        "\n",
        "                audio_np = psola.vocode(\n",
        "                    audio_np, 22050, target_pitch=target_pitch\n",
        "                ).astype(np.float32)\n",
        "                normalize = (1.0 / np.max(np.abs(audio_np))) ** 0.9\n",
        "                audio_np = audio_np * normalize * MAX_WAV_VALUE\n",
        "                audio_np = audio_np.astype(np.int16)\n",
        "\n",
        "            # Resample to 32k\n",
        "            wave = resampy.resample(\n",
        "                audio_np,\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(DEVICE))\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.detach().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 *= 4.0  # 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",
        "\n",
        "            buffer = io.BytesIO()\n",
        "            wavfile.write(buffer, 32000, sr_mix.astype(np.int16))\n",
        "            b64 = base64.b64encode(buffer.getvalue())\n",
        "            sound = \"data:audio/x-wav;base64,\" + b64.decode(\"ascii\")\n",
        "\n",
        "            output_name = \"TalkNet_\" + str(int(time.time()))\n",
        "            return [sound, arpa, playback_style, output_name]\n",
        "    except Exception:\n",
        "        return [\n",
        "            None,\n",
        "            str(traceback.format_exc()),\n",
        "            playback_hide,\n",
        "            None,\n",
        "        ]\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    app.run_server(\n",
        "        mode=\"external\",\n",
        "        debug=True,\n",
        "        dev_tools_ui=True,\n",
        "        dev_tools_hot_reload=True,\n",
        "        threaded=True,\n",
        "    )\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4wQAHcBpLKZb"
      },
      "source": [
        "# Start interface"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tOXejargIPTq",
        "cellView": "form"
      },
      "source": [
        "# @markdown Run the interface. \n",
        "\n",
        "## @markdown If you get a VersionConflict error,\n",
        "## @markdown click on Runtime -> Restart runtime, and then run this cell again.\n",
        "using_inline = True\n",
        "import pkg_resources\n",
        "from pkg_resources import DistributionNotFound, VersionConflict\n",
        "\"\"\"dependencies = [\n",
        "\"tensorflow==2.4.1\", \n",
        "\"dash\", \n",
        "\"jupyter-dash\", \n",
        "\"psola\", \n",
        "\"wget\", \n",
        "\"unidecode\", \n",
        "\"pysptk\", \n",
        "\"frozendict\", \n",
        "\"torchvision==0.9.1\", \n",
        "\"torchaudio==0.8.1\", \n",
        "\"torchtext==0.9.1\", \n",
        "\"torch_stft\", \n",
        "\"kaldiio\", \n",
        "\"pydub\", \n",
        "\"pyannote.audio\", \n",
        "\"g2p_en\", \n",
        "\"pesq\", \n",
        "\"pystoi\", \n",
        "\"crepe\", \n",
        "\"resampy\", \n",
        "\"ffmpeg-python\",\n",
        "\"numpy\",\n",
        "\"scipy\",\n",
        "\"nemo_toolkit\",\n",
        "\"tqdm\",\n",
        "\"gdown\",\n",
        "]\n",
        "pkg_resources.require(dependencies)\"\"\"\n",
        "\n",
        "from controllable_talknet import *\n",
        "app.run_server(\n",
        "    mode=\"inline\",\n",
        "    #dev_tools_ui=True,\n",
        "    #dev_tools_hot_reload=True,\n",
        "    threaded=True,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F4WDzfRIgb5Z",
        "cellView": "form"
      },
      "source": [
        "# @markdown If the above fails with a 403 error, do the following:\n",
        "# @markdown * Go to Runtime -> Restart runtime\n",
        "# @markdown * Run this cell (click the play button)\n",
        "# @markdown * Click on the googleusercontent.com link to use TalkNet in a separate tab\n",
        "try:\n",
        "    using_inline\n",
        "except:\n",
        "    using_inline = False\n",
        "if not using_inline:\n",
        "    from controllable_talknet import *\n",
        "    from google.colab.output import eval_js\n",
        "\n",
        "    print(eval_js(\"google.colab.kernel.proxyPort(8050)\"))\n",
        "    app.run_server(\n",
        "        mode=\"external\",\n",
        "        debug=False,\n",
        "        #dev_tools_ui=True,\n",
        "        #dev_tools_hot_reload=True,\n",
        "        threaded=True,\n",
        "    )"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}