{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n",
      "Loading G2P model... Done!\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "import yaml\n",
    "import numpy\n",
    "import torch\n",
    "!mkdir -p ../wavs\n",
    "import onnxruntime\n",
    "from sys import path\n",
    "from tqdm import tqdm\n",
    "SAMPLING_RATE = 22050\n",
    "os.chdir('../Fastspeech2_HS')\n",
    "path.append(\"hifigan\")\n",
    "from env import AttrDict\n",
    "from models import Generator\n",
    "from IPython.display import Audio\n",
    "from scipy.io.wavfile import write\n",
    "from meldataset import MAX_WAV_VALUE\n",
    "from espnet_onnx.export import TTSModelExport\n",
    "from espnet2.bin.tts_inference import Text2Speech\n",
    "from espnet_onnx import Text2Speech as Text2SpeechInference\n",
    "from text_preprocess_for_inference import TTSDurAlignPreprocessor, CharTextPreprocessor, TTSPreprocessor"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Original Inference\n",
    "* uses the environment defined in Fastspeech2_HS repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_hifigan_vocoder(language, gender, family, device):\n",
    "    vocoder_config = f\"vocoder/{gender}/{family}/hifigan/config.json\"\n",
    "    vocoder_generator = f\"vocoder/{gender}/{family}/hifigan/generator\"\n",
    "    with open(vocoder_config, 'r') as f: json_config = json.load(f)\n",
    "    h = AttrDict(json_config)\n",
    "    torch.manual_seed(h.seed)\n",
    "    device = torch.device(device)\n",
    "    generator = Generator(h).to(device)\n",
    "    state_dict_g = torch.load(vocoder_generator, device)\n",
    "    generator.load_state_dict(state_dict_g['generator'])\n",
    "    generator.eval()\n",
    "    generator.remove_weight_norm()\n",
    "    return generator\n",
    "def load_fastspeech2_model(language, gender, device):\n",
    "    with open(f\"{language}/{gender}/model/config.yaml\", \"r\") as file: config = yaml.safe_load(file)\n",
    "    current_working_directory = os.getcwd()\n",
    "    feat = \"model/feats_stats.npz\"\n",
    "    pitch = \"model/pitch_stats.npz\"\n",
    "    energy = \"model/energy_stats.npz\"\n",
    "    feat_path = os.path.join(current_working_directory, language, gender, feat)\n",
    "    pitch_path = os.path.join(current_working_directory, language, gender, pitch)\n",
    "    energy_path = os.path.join(current_working_directory, language, gender, energy)\n",
    "    config[\"normalize_conf\"][\"stats_file\"] = feat_path\n",
    "    config[\"pitch_normalize_conf\"][\"stats_file\"] = pitch_path\n",
    "    config[\"energy_normalize_conf\"][\"stats_file\"] = energy_path\n",
    "    with open(f\"{language}/{gender}/model/config.yaml\", \"w\") as file: yaml.dump(config, file)\n",
    "    tts_model = f\"{language}/{gender}/model/model.pth\"\n",
    "    tts_config = f\"{language}/{gender}/model/config.yaml\"\n",
    "    return Text2Speech(train_config=tts_config, model_file=tts_model, device=device)\n",
    "def text_synthesis(language, gender, sample_text, vocoder, MAX_WAV_VALUE, device):\n",
    "    with torch.no_grad():\n",
    "        model = load_fastspeech2_model(language, gender, device)\n",
    "        out = model(sample_text, decode_conf={\"alpha\": 1})\n",
    "        x = out[\"feat_gen_denorm\"].T.unsqueeze(0) * 2.3262\n",
    "        x = x.to(device)\n",
    "        y_g_hat = vocoder(x)\n",
    "        audio = y_g_hat.squeeze()\n",
    "        audio = audio * MAX_WAV_VALUE\n",
    "        audio = audio.cpu().numpy().astype('int16')\n",
    "        return audio\n",
    "def text2speech(language, gender, family, sample_text, device):\n",
    "    vocoder = load_hifigan_vocoder(language, gender, family, device)\n",
    "    if language == \"urdu\" or language == \"punjabi\": preprocessor = CharTextPreprocessor()\n",
    "    elif language == \"english\": preprocessor = TTSPreprocessor()\n",
    "    else: preprocessor = TTSDurAlignPreprocessor()\n",
    "    preprocessed_text, phrases = preprocessor.preprocess(sample_text, language, gender)\n",
    "    preprocessed_text = \" \".join(preprocessed_text)\n",
    "    audio = text_synthesis(language, gender, preprocessed_text, vocoder, MAX_WAV_VALUE, device)\n",
    "    output_file = f\"../wavs/{language}_{gender}-{family}_orig_output.wav\"\n",
    "    write(output_file, SAMPLING_RATE, audio)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Original Inference Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removing weight norm...\n",
      "length: 26624 array: [945 894 605 ...  10  12  30]\n",
      "abssum: 42271788 min: -16131 max: 10878\n"
     ]
    }
   ],
   "source": [
    "audio_orig = text_synthesis('english', 'male', 'this is a sentence', load_hifigan_vocoder('english', 'male', 'aryan', 'cpu'), MAX_WAV_VALUE, 'cpu')\n",
    "print('length:', len(audio_orig), 'array:', audio_orig)\n",
    "print('abssum:', numpy.abs(audio_orig).sum(), 'min:', audio_orig.min(), 'max:', audio_orig.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removing weight norm...\n",
      "length: 26624 array: [945 894 605 ...  10  12  30]\n",
      "abssum: 42271796 min: -16131 max: 10878\n"
     ]
    }
   ],
   "source": [
    "audio_orig = text_synthesis('english', 'male', 'this is a sentence', load_hifigan_vocoder('english', 'male', 'aryan', 'cuda'), MAX_WAV_VALUE, 'cuda')\n",
    "print('length:', len(audio_orig), 'array:', audio_orig)\n",
    "print('abssum:', numpy.abs(audio_orig).sum(), 'min:', audio_orig.min(), 'max:', audio_orig.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Latest Pytorch Inference Results\n",
    "* similar environment as defined in scripts/perform_onnx_conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/nikhilesh/miniforge3/envs/iitm-tts-latest-pytorch/lib/python3.10/site-packages/torch/nn/utils/weight_norm.py:30: UserWarning: torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\n",
      "  warnings.warn(\"torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Removing weight norm...\n",
      "length: 26624 array: [945 894 605 ...  10  12  30]\n",
      "abssum: 42271783 min: -16131 max: 10878\n"
     ]
    }
   ],
   "source": [
    "audio_orig = text_synthesis('english', 'male', 'this is a sentence', load_hifigan_vocoder('english', 'male', 'aryan', 'cpu'), MAX_WAV_VALUE, 'cpu')\n",
    "print('length:', len(audio_orig), 'array:', audio_orig)\n",
    "print('abssum:', numpy.abs(audio_orig).sum(), 'min:', audio_orig.min(), 'max:', audio_orig.max())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_orig = text_synthesis('english', 'male', 'this is a sentence', load_hifigan_vocoder('english', 'male', 'aryan', 'cuda'), MAX_WAV_VALUE, 'cuda')\n",
    "print('length:', len(audio_orig), 'array:', audio_orig)\n",
    "print('abssum:', numpy.abs(audio_orig).sum(), 'min:', audio_orig.min(), 'max:', audio_orig.max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ORT Conversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p ../ort_models\n",
    "def convert_to_ort(language, gender, family):\n",
    "    vocoder = load_hifigan_vocoder(language, gender, family, 'cpu')\n",
    "    model = load_fastspeech2_model(language, gender, 'cpu')\n",
    "    if language == \"urdu\" or language == \"punjabi\": preprocessor = CharTextPreprocessor()\n",
    "    elif language == \"english\": preprocessor = TTSPreprocessor()\n",
    "    else: preprocessor = TTSDurAlignPreprocessor()\n",
    "    preprocessed_text, phrases = preprocessor.preprocess('this is a sentence', language, gender)\n",
    "    preprocessed_text = \" \".join(preprocessed_text)\n",
    "    exporter = TTSModelExport()\n",
    "    exporter.export(model, f'{language}-{gender}-ort', quantize=False)\n",
    "    out = model(preprocessed_text, decode_conf={\"alpha\": 1})\n",
    "    x = out[\"feat_gen_denorm\"].T.unsqueeze(0) * 2.3262\n",
    "    torch.onnx.export(vocoder, x, f'../ort_models/vocoders/{gender}-{family}-vocoder.onnx', input_names=['input'], output_names=['output'], dynamic_axes={'input': [0, 2], 'output': [0]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "convert_to_ort('english', 'male', 'aryan')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ORT Inference\n",
    "* environment as defined in triton_models/tts/envbuilder.sh\n",
    "* you can delete the ort_models folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_hifigan_vocoder(language, gender, family, device): return onnxruntime.InferenceSession(f\"../ort_models/vocoders/{gender}-{family}-vocoder.onnx\", providers=['CPUExecutionProvider' if device == 'cpu' else 'CUDAExecutionProvider'])\n",
    "def load_fastspeech2_model(language, gender, device): return Text2SpeechInference(f'{language}-{gender}-ort', providers=['CPUExecutionProvider' if device == 'cpu' else 'CUDAExecutionProvider'])\n",
    "def text_synthesis(language, gender, sample_text, vocoder, MAX_WAV_VALUE, device):\n",
    "    model = load_fastspeech2_model(language, gender, device)\n",
    "    x = numpy.expand_dims(model.postprocess(model.tts_model(model.preprocess.token_id_converter.tokens2ids(model.preprocess.tokenizer.text2tokens(sample_text)))['feat_gen']).T, axis=0) * 2.3262\n",
    "    y_g_hat = vocoder.run(None, {'input': x})[0]\n",
    "    audio = y_g_hat.squeeze()\n",
    "    audio = audio * MAX_WAV_VALUE\n",
    "    audio = audio.astype('int16')\n",
    "    return audio\n",
    "def text2speech(language, gender, family, sample_text, device):\n",
    "    vocoder = load_hifigan_vocoder(language, gender, family, device)\n",
    "    if language == \"urdu\" or language == \"punjabi\": preprocessor = CharTextPreprocessor()\n",
    "    elif language == \"english\": preprocessor = TTSPreprocessor()\n",
    "    else: preprocessor = TTSDurAlignPreprocessor()\n",
    "    preprocessed_text, phrases = preprocessor.preprocess(sample_text, language, gender)\n",
    "    preprocessed_text = \" \".join(preprocessed_text)\n",
    "    audio = text_synthesis(language, gender, preprocessed_text, vocoder, MAX_WAV_VALUE, device)\n",
    "    output_file = f\"../wavs/{language}_{gender}-{family}_ort_output.wav\"\n",
    "    write(output_file, SAMPLING_RATE, audio)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ORT Inference Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_orig = text_synthesis('english', 'male', 'this', load_hifigan_vocoder('english', 'male', 'aryan', 'cpu'), MAX_WAV_VALUE, 'cpu')\n",
    "print('length:', len(audio_orig), 'array:', audio_orig)\n",
    "print('abssum:', numpy.abs(audio_orig).sum(), 'min:', audio_orig.min(), 'max:', audio_orig.max())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "iitm-tts",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}