test_onnx.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0900bc22-a3d8-4e5d-8124-d934a0228976",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import onnxruntime as ort\n",
    "import soundfile as sf\n",
    "# import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb087d04-4a0e-461b-8a5b-6a1ed28e776a",
   "metadata": {},
   "outputs": [],
   "source": [
    "index_list=['ก', 'ข', 'ค', 'ฆ', 'ง', 'จ', 'ฉ', 'ช', 'ซ', 'ฌ', 'ญ', 'ฎ', 'ฏ', 'ฐ', 'ฑ', 'ฒ', 'ณ', 'ด', 'ต', 'ถ', 'ท', 'ธ', 'น', 'บ', 'ป', 'ผ', 'ฝ', 'พ', 'ฟ', 'ภ', 'ม', 'ย', 'ร', 'ฤ', 'ล', 'ว', 'ศ', 'ษ', 'ส', 'ห', 'ฬ', 'อ', 'ฮ', 'ะ', 'ั', 'า', 'ำ', 'ิ', 'ี', 'ึ', 'ื', 'ุ', 'ู', 'เ', 'แ', 'โ', 'ใ', 'ไ', 'ๅ', '็', '่', '้', '๊', '๋', '์', ' ']\n",
    "dict_idx={k:i for i,k in enumerate(index_list)}\n",
    "\n",
    "def clean(text):\n",
    "    seq = np.array([[66]+[dict_idx[i] for i in text if i]+[67]])\n",
    "    _s=np.array([len(seq[0])])\n",
    "    return seq,_s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "49658e77-a9ae-48d5-945d-59445aa640dd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[66, 20, 17, 38, 41, 23, 67]]), array([7]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "clean(\"ทดสอบ\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45660d4b-39df-424a-8b4b-602f83aaba38",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_mel_channels = 80\n",
    "n_frames_per_step = 1\n",
    "attention_rnn_dim = 1024\n",
    "decoder_rnn_dim=1024\n",
    "encoder_embedding_dim=512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "70e0c23f-5a37-4fc5-b849-2de9ebaae984",
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_decoder_states(memory):\n",
    "    B = memory.shape[0]\n",
    "    MAX_TIME = memory.shape[1]\n",
    "\n",
    "    attention_hidden = np.zeros((B, attention_rnn_dim), dtype=np.float32)\n",
    "    attention_cell = np.zeros((B, attention_rnn_dim), dtype=np.float32)\n",
    "\n",
    "    decoder_hidden = np.zeros((B, decoder_rnn_dim), dtype=np.float32)\n",
    "    decoder_cell = np.zeros((B, decoder_rnn_dim), dtype=np.float32)\n",
    "\n",
    "    attention_weights = np.zeros((B, MAX_TIME), dtype=np.float32)\n",
    "    attention_weights_cum = np.zeros((B, MAX_TIME), dtype=np.float32)\n",
    "    attention_context = np.zeros((B, encoder_embedding_dim), dtype=np.float32)\n",
    "\n",
    "    return (\n",
    "        attention_hidden,\n",
    "        attention_cell,\n",
    "        decoder_hidden,\n",
    "        decoder_cell,\n",
    "        attention_weights,\n",
    "        attention_weights_cum,\n",
    "        attention_context,\n",
    "    )\n",
    "\n",
    "\n",
    "def get_go_frame(memory):\n",
    "    B = memory.shape[0]\n",
    "    decoder_input = np.zeros((B, n_mel_channels*n_frames_per_step), dtype=np.float32)\n",
    "    return decoder_input\n",
    "\n",
    "\n",
    "def sigmoid(x):\n",
    "    return np.exp(-np.logaddexp(0, -x))\n",
    "\n",
    "\n",
    "def parse_decoder_outputs(mel_outputs, gate_outputs, alignments):\n",
    "    # (T_out, B) -> (B, T_out)\n",
    "    alignments = np.stack(alignments).transpose((1, 0, 2, 3))\n",
    "    # (T_out, B) -> (B, T_out)\n",
    "    # Add a -1 to prevent squeezing the batch dimension in case\n",
    "    # batch is 1\n",
    "    gate_outputs = np.stack(gate_outputs).squeeze(-1).transpose((1, 0, 2))\n",
    "    # (T_out, B, n_mel_channels) -> (B, T_out, n_mel_channels)\n",
    "    mel_outputs = np.stack(mel_outputs).transpose((1, 0, 2, 3))\n",
    "    # decouple frames per step\n",
    "    mel_outputs = mel_outputs.reshape(mel_outputs.shape[0], -1, n_mel_channels)\n",
    "    # (B, T_out, n_mel_channels) -> (B, n_mel_channels, T_out)\n",
    "    mel_outputs = mel_outputs.transpose((0, 2, 1))\n",
    "\n",
    "    return mel_outputs, gate_outputs, alignments\n",
    "\n",
    "\n",
    "# only numpy operations\n",
    "def test_inference(encoder, decoder_iter, postnet):\n",
    "    sequences, sequence_lengths = clean(\"สวัสดีภาษาไทย    ทดสอบ จาก โอเอ็นเอ็นเอ๊ก\")\n",
    "\n",
    "    print(\"Running Tacotron2 Encoder\")\n",
    "    inputs = {\"seq\": sequences, \"seq_len\": sequence_lengths}\n",
    "    memory, processed_memory, _ = encoder.run(None, inputs)\n",
    "\n",
    "    print(\"Running Tacotron2 Decoder\")\n",
    "    mel_lengths = np.zeros([memory.shape[0]], dtype=np.int32)\n",
    "    not_finished = np.ones([memory.shape[0]], dtype=np.int32)\n",
    "    mel_outputs, gate_outputs, alignments = [], [], []\n",
    "    gate_threshold = 0.5\n",
    "    max_decoder_steps = 5000\n",
    "    first_iter = True\n",
    "\n",
    "    (\n",
    "        attention_hidden,\n",
    "        attention_cell,\n",
    "        decoder_hidden,\n",
    "        decoder_cell,\n",
    "        attention_weights,\n",
    "        attention_weights_cum,\n",
    "        attention_context,\n",
    "    ) = initialize_decoder_states(memory)\n",
    "\n",
    "    decoder_input = get_go_frame(memory)\n",
    "\n",
    "    while True:\n",
    "        inputs = {\n",
    "            \"decoder_input\": decoder_input,\n",
    "            \"attention_hidden\": attention_hidden,\n",
    "            \"attention_cell\": attention_cell,\n",
    "            \"decoder_hidden\": decoder_hidden,\n",
    "            \"decoder_cell\": decoder_cell,\n",
    "            \"attention_weights\": attention_weights,\n",
    "            \"attention_weights_cum\": attention_weights_cum,\n",
    "            \"attention_context\": attention_context,\n",
    "            \"memory\": memory,\n",
    "            \"processed_memory\": processed_memory,\n",
    "        }\n",
    "        (\n",
    "            mel_output,\n",
    "            gate_output,\n",
    "            attention_hidden,\n",
    "            attention_cell,\n",
    "            decoder_hidden,\n",
    "            decoder_cell,\n",
    "            attention_weights,\n",
    "            attention_weights_cum,\n",
    "            attention_context,\n",
    "        ) = decoder_iter.run(None, inputs)\n",
    "\n",
    "        if first_iter:\n",
    "            mel_outputs = [np.expand_dims(mel_output, 2)]\n",
    "            gate_outputs = [np.expand_dims(gate_output, 2)]\n",
    "            alignments = [np.expand_dims(attention_weights, 2)]\n",
    "            first_iter = False\n",
    "        else:\n",
    "            mel_outputs += [np.expand_dims(mel_output, 2)]\n",
    "            gate_outputs += [np.expand_dims(gate_output, 2)]\n",
    "            alignments += [np.expand_dims(attention_weights, 2)]\n",
    "\n",
    "        dec = np.less(sigmoid(gate_output), gate_threshold)\n",
    "        dec = np.squeeze(dec, axis=1)\n",
    "        not_finished = not_finished * dec\n",
    "        mel_lengths += not_finished\n",
    "\n",
    "        if not_finished.sum() == 0:\n",
    "            print(\"Stopping after \", len(mel_outputs), \" decoder steps\")\n",
    "            break\n",
    "        if len(mel_outputs) == max_decoder_steps:\n",
    "            print(\"Warning! Reached max decoder steps\")\n",
    "            break\n",
    "\n",
    "        decoder_input = mel_output\n",
    "\n",
    "    mel_outputs, gate_outputs, alignments = parse_decoder_outputs(\n",
    "        mel_outputs, gate_outputs, alignments\n",
    "    )\n",
    "\n",
    "    print(\"Running Tacotron2 PostNet\")\n",
    "    inputs = {\"mel_spec\": mel_outputs}\n",
    "    mel_outputs_postnet = postnet.run(None, inputs)\n",
    "\n",
    "    return mel_outputs_postnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "06dae132-e5f4-4998-9d8d-7aca7ed968f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running Tacotron2 Encoder\n",
      "Running Tacotron2 Decoder\n",
      "Stopping after  153  decoder steps\n",
      "Running Tacotron2 PostNet\n"
     ]
    }
   ],
   "source": [
    "encoder = ort.InferenceSession(\"tacotron2encoder-th.onnx\")\n",
    "decoder = ort.InferenceSession(\"tacotron2decoder-th.onnx\")\n",
    "postnet = ort.InferenceSession(\"tacotron2postnet-th.onnx\")\n",
    "\n",
    "mel = test_inference(encoder, decoder, postnet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1ce1cda6-33b0-47d8-a9ac-ceacd0bd3081",
   "metadata": {},
   "outputs": [],
   "source": [
    "hifi = ort.InferenceSession(\"vocoder.onnx\") # the best\n",
    "audio = hifi.run(None, {\"spec\": mel[0]})\n",
    "audio = audio[0][0, 0, :]\n",
    "sf.write(\"speech32-1.wav\", audio, 22050, format=\"WAV\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8bd6321-09b3-4ad4-a6de-d338860c82cd",
   "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
}