notebooks/234-encodec-audio-compression/234-encodec-audio-compression.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "285a34a4-18ff-493c-8322-71adf1f09622",
   "metadata": {},
   "source": [
    "# Audio compression with EnCodec and OpenVINO\n",
    "\n",
    "Compression is an important part of the Internet today because it enables people to easily share high-quality photos, listen to audio messages, stream their favorite shows, and so much more. Even when using today’s state-of-the-art techniques, enjoying these rich multimedia experiences requires a high speed Internet connection and plenty of storage space. AI helps to overcome these limitations: \"Imagine listening to a friend’s audio message in an area with low connectivity and not having it stall or glitch.\"\n",
    "\n",
    "This tutorial considers ways to use OpenVINO and EnCodec algorithm for hyper compression of audio.\n",
    "EnCodec is a real-time, high-fidelity audio codec that uses AI to compress audio files without losing quality. It was introduced in [High Fidelity Neural Audio Compression](https://arxiv.org/pdf/2210.13438.pdf) paper by Meta AI. The researchers claimed they achieved an approximate 10x compression rate without loss of quality and made it work for CD-quality audio. More details about this approach can be found in [Meta AI blog](https://ai.facebook.com/blog/ai-powered-audio-compression-technique/) and original [repo](https://github.com/facebookresearch/encodec).\n",
    "\n",
    "![image.png](https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/17546d66-12b9-4841-9293-cc878258a186)\n",
    "\n",
    "#### Table of content:\n",
    "- [Prerequisites](#Prerequisites-Uparrow)\n",
    "- [Instantiate audio compression pipeline](#Instantiate-audio-compression-pipeline-Uparrow)\n",
    "- [Explore EnCodec pipeline](#Explore-EnCodec-pipeline-Uparrow)\n",
    "    - [Preprocessing](#Preprocessing-Uparrow)\n",
    "    - [Encoding](#Encoding-Uparrow)\n",
    "    - [Decompression](#Decompression-Uparrow)\n",
    "- [Convert model to OpenVINO Intermediate Representation format](#Convert-model-to-OpenVINO-Intermediate-Representation-format-Uparrow)\n",
    "- [Integrate OpenVINO to EnCodec pipeline](#Integrate-OpenVINO-to-EnCodec-pipeline-Uparrow)\n",
    "    - [Select inference device](#Select-inference-device-Uparrow)\n",
    "- [Run EnCodec with OpenVINO](#Run-EnCodec-with-OpenVINO-Uparrow)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e328a76f-03b8-4972-afa9-36d3f8cc8364",
   "metadata": {},
   "source": [
    "## Prerequisites [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Install required dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2ec9aa3-00c2-4d49-954c-c1af2fe43f2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82279f60-a018-4ea8-9b0b-77f0f40903cf",
   "metadata": {},
   "source": [
    "## Instantiate audio compression pipeline [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "[Codecs](https://en.wikipedia.org/wiki/Codec), which act as encoders and decoders for streams of data, help empower most of the audio compression people currently use online. Some examples of commonly used codecs include MP3, Opus, and EVS. Classic codecs like these decompose the signal between different frequencies and encode as efficiently as possible. Most classic codecs leverage human hearing knowledge (psychoacoustics) but have a finite or given set of handcrafted ways to efficiently encode and decode the file. EnCodec, a neural network that is trained from end to end to reconstruct the input signal, was introduced as an attempt to overcome this limitation. It consists of three parts:\n",
    "\n",
    "* The **encoder**, which takes the uncompressed data in and transforms it into a higher dimensional and lower frame rate representation.\n",
    "\n",
    "* The **quantizer**, which compresses this representation to the target size. This compressed representation is what is stored on disk or will be sent through the network.\n",
    "\n",
    "* The **decoder** is the final step. It turns the compressed signal back into a waveform that is as similar as possible to the original. The key to lossless compression is to identify changes that will not be perceivable by humans, as perfect reconstruction is impossible at low bit rates.\n",
    "\n",
    "![encodec_compression](https://github.com/openvinotoolkit/openvino_notebooks/assets/29454499/5cd9a482-b42b-4dea-85a5-6d66b20ce13d))\n",
    "\n",
    "\n",
    "The authors provide two multi-bandwidth models:\n",
    "* `encodec_model_24khz` - a causal model operating at 24 kHz on monophonic audio trained on a variety of audio data.\n",
    "* `encodec_model_48khz` - a non-causal model operating at 48 kHz on stereophonic audio trained on music-only data.\n",
    "\n",
    "In this tutorial, we will use `encodec_model_24khz` as an example, but the same actions are also applicable to `encodec_model_48khz` model as well.\n",
    "To start working with this model, we need to instantiate model class using `EncodecModel.encodec_model_24khz()` and select required compression bandwidth among available: 1.5, 3, 6, 12 or 24 kbps for 24 kHz model and 3, 6, 12 and 24 kbps for 48 kHz model. We will use 6 kbps bandwidth."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9855472-5c0d-4853-b284-ed659885c76e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from encodec import EncodecModel, compress, decompress\n",
    "from encodec.utils import convert_audio, save_audio\n",
    "import torchaudio\n",
    "import torch\n",
    "import typing as tp\n",
    "\n",
    "# Instantiate a pretrained EnCodec model\n",
    "model = EncodecModel.encodec_model_24khz()\n",
    "model.set_target_bandwidth(6.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2146fbd2-6f15-4e31-b9f6-7c8e69daabd0",
   "metadata": {},
   "source": [
    "## Explore EnCodec pipeline [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Let us explore model capabilities on example audio:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84f727cf-9414-480a-b824-cae3bb97371e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import librosa\n",
    "import matplotlib.pyplot as plt\n",
    "import librosa.display\n",
    "import IPython.display as ipd\n",
    "\n",
    "sys.path.append(\"../utils\")\n",
    "\n",
    "from notebook_utils import download_file\n",
    "\n",
    "test_data_url = \"https://github.com/facebookresearch/encodec/raw/main/test_24k.wav\"\n",
    "\n",
    "sample_file = 'test_24k.wav'\n",
    "download_file(test_data_url, sample_file)\n",
    "audio, sr = librosa.load(sample_file)\n",
    "plt.figure(figsize=(14, 5))\n",
    "librosa.display.waveshow(audio, sr=sr)\n",
    "\n",
    "ipd.Audio(sample_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "970e162d-ea7e-4d32-bb35-3c3ed29f9720",
   "metadata": {},
   "source": [
    "### Preprocessing [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "To achieve the best result, audio should have the number of channels and sample rate expected by the model. If audio does not fulfill these requirements, it can be converted to the desired sample rate and the number of channels using the `convert_audio` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2704048-da45-46ca-bd76-2083194c9208",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_sr, model_channels = model.sample_rate, model.channels\n",
    "print(f\"Model expected sample rate {model_sr}\")\n",
    "print(f\"Model expected audio format {'mono' if model_channels == 1 else 'stereo'}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98235cce-41d5-44dc-b543-85f6ec074a97",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and pre-process the audio waveform\n",
    "wav, sr = torchaudio.load(sample_file)\n",
    "\n",
    "wav = convert_audio(wav, sr, model_sr, model_channels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f51bfa76-25ad-427c-851a-fb61481f02e3",
   "metadata": {},
   "source": [
    "### Encoding [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "Audio waveform should be split by chunks and then encoded by Encoder model, then compressed by quantizer for reducing memory. The result of compression is a binary file with `ecdc` extension, a special format for storing EnCodec compressed audio on disc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a137999-a1f8-4ac5-ac64-5afe38cb4b1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "\n",
    "out_file = Path(\"compressed.ecdc\")\n",
    "b = compress(model, wav)\n",
    "out_file.write_bytes(b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22b28212-9cf1-4f89-ac92-da469239e87e",
   "metadata": {},
   "source": [
    "Let us compare obtained compression result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cac0be4-3d01-4431-b56c-c4dfb67eb69d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "orig_file_stats = os.stat(sample_file)\n",
    "compressed_file_stats = os.stat(\"compressed.ecdc\")\n",
    "print(f\"size before compression in Bytes: {orig_file_stats.st_size}\")\n",
    "print(f\"size after compression in Bytes: {compressed_file_stats.st_size}\")\n",
    "print(f\"Compression file size ratio: {orig_file_stats.st_size / compressed_file_stats.st_size:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94c1a9dd-09bc-4369-8b55-935db09c729b",
   "metadata": {},
   "source": [
    "Great! Now, we see the power of hyper compression. Binary size of a file becomes 60 times smaller and more suitable for sending via network."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41301b36-2757-48c7-91dd-5c2d38d5c467",
   "metadata": {},
   "source": [
    "### Decompression [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "After successful sending of the compressed audio, it should be decompressed on the recipient's side. The decoder model is responsible for restoring the compressed signal back into a waveform that is as similar as possible to the original."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cc5d08d-1a65-493e-b3fb-fd70dd2efbf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "out, out_sr = decompress(out_file.read_bytes())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "196ca5a7-b9b0-4993-b1f0-3d81be1ce425",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_file = \"decopressed.wav\"\n",
    "save_audio(out, output_file, out_sr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40e394db-512b-4b87-bfd0-fd5afe8ef1a3",
   "metadata": {},
   "source": [
    "The decompressed audio will be saved to the `decompressed.wav` file when decompression is finished. We can compare result with the original audio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ecafd48-f8db-48bb-a9d7-a41f00964c28",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio, sr = librosa.load(output_file)\n",
    "plt.figure(figsize=(14, 5))\n",
    "librosa.display.waveshow(audio, sr=sr)\n",
    "\n",
    "ipd.Audio(output_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac4cc2e5-fd24-409e-8104-77d8c24bfa7d",
   "metadata": {},
   "source": [
    "Nice! Audio sounds close to original."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44271e54-ea69-4da6-a84b-838b7896c586",
   "metadata": {},
   "source": [
    "## Convert model to OpenVINO Intermediate Representation format [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "For best results with OpenVINO, it is recommended to convert the model to OpenVINO IR format. OpenVINO supports PyTorch via ONNX conversion. We will use `torch.onnx.export` for exporting the ONNX model from PyTorch. We need to provide initialized model's instance and example of inputs for shape inference. We will use `ov.convert_model` functionality to convert the ONNX models. The `ov.convert_model` Python function returns an OpenVINO model ready to load on the device and start making predictions. We can save it on disk for the next usage with `ov.save_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f37ce32-452e-4aa4-8311-16c7439da8e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FrameEncoder(torch.nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "\n",
    "    def forward(self, x: torch.Tensor):\n",
    "        codes, scale = self.model._encode_frame(x)\n",
    "        if not self.model.normalize:\n",
    "            return codes\n",
    "        return codes, scale        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1948a427-7a5d-4520-a75d-ea4a0f11204f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FrameDecoder(torch.nn.Module):\n",
    "    def __init__(self, model):\n",
    "        super().__init__()\n",
    "        self.model = model\n",
    "\n",
    "    def forward(self, codes, scale=None):\n",
    "        return model._decode_frame((codes, scale))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a7f491-c571-464f-a2be-7eb364e78325",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder = FrameEncoder(model)\n",
    "decoder = FrameDecoder(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36df95d8-a254-4061-9a56-7aa4a03d11a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openvino as ov\n",
    "\n",
    "\n",
    "core = ov.Core()\n",
    "\n",
    "OV_ENCODER_PATH = Path(\"encodec_encoder.xml\")\n",
    "if not OV_ENCODER_PATH.exists():\n",
    "    torch.onnx.export(encoder, torch.zeros(1, 1, 480000), \"encodec_encoder.onnx\")\n",
    "    encoder_ov = ov.convert_model(\"encodec_encoder.onnx\")\n",
    "    ov.save_model(encoder_ov, OV_ENCODER_PATH)\n",
    "else:\n",
    "    encoder_ov = core.read_model(OV_ENCODER_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "480bb4e7-0958-46ec-b24d-7f4af77bc671",
   "metadata": {},
   "outputs": [],
   "source": [
    "OV_DECODER_PATH = Path(\"encodec_decoder.xml\")\n",
    "if not OV_DECODER_PATH.exists():\n",
    "    torch.onnx.export(decoder, torch.zeros([1, 8, 1500], dtype=torch.long), \"encodec_decoder.onnx\", input_names=[\"codes\", \"scale\"])\n",
    "    decoder_ov = ov.convert_model(\"encodec_decoder.onnx\")\n",
    "    ov.save_model(decoder_ov, OV_DECODER_PATH)\n",
    "else:\n",
    "    decoder_ov = core.read_model(OV_DECODER_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc8cf59c-a3fa-48ea-9f18-0979860f0aa9",
   "metadata": {},
   "source": [
    "## Integrate OpenVINO to EnCodec pipeline [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "The following steps are required for integration of OpenVINO to EnCodec pipeline:\n",
    "\n",
    "1. Load the model to a device.\n",
    "2. Define audio frame processing functions.\n",
    "3. Replace the original frame processing functions with OpenVINO based algorithms."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e32cea0b-e92c-40e7-912f-b5d47a43f921",
   "metadata": {},
   "source": [
    "### Select inference device [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "select device from dropdown list for running inference using OpenVINO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43944e3c-4734-45ac-9c1f-2a70adc047e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipywidgets as widgets\n",
    "\n",
    "device = widgets.Dropdown(\n",
    "    options=core.available_devices + [\"AUTO\"],\n",
    "    value='AUTO',\n",
    "    description='Device:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba53f8f9-18b6-4f68-9f9e-a800c2f5b1a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "compiled_encoder = core.compile_model(encoder_ov, device.value)\n",
    "encoder_out = compiled_encoder.output(0)\n",
    "\n",
    "compiled_decoder = core.compile_model(decoder_ov, device.value)\n",
    "decoder_out = compiled_decoder.output(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "967846a5-164b-4778-93b0-9b5947d93104",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode_frame(x: torch.Tensor):\n",
    "    has_scale = len(compiled_encoder.outputs) == 2\n",
    "    result = compiled_encoder(x)\n",
    "    codes = torch.from_numpy(result[encoder_out])\n",
    "    if has_scale:\n",
    "        scale = torch.from_numpy(result[compiled_encoder.output(1)])\n",
    "    else:\n",
    "        scale = None\n",
    "    return codes, scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0886546-08bb-4b01-8bb3-c8ed7c722fd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "EncodedFrame = tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]\n",
    "\n",
    "\n",
    "def decode_frame(encoded_frame: EncodedFrame):\n",
    "    codes, scale = encoded_frame\n",
    "    inputs = [codes]\n",
    "    if scale is not None:\n",
    "        inputs.append(scale)\n",
    "    return torch.from_numpy(compiled_decoder(inputs)[decoder_out])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16681a9a-cd29-4ce6-838c-09f4514eed7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model._encode_frame = encode_frame\n",
    "model._decode_frame = decode_frame"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21b2e5b4-035c-4123-8257-00cd3f4af3fc",
   "metadata": {},
   "source": [
    "## Run EnCodec with OpenVINO [$\\Uparrow$](#Table-of-content:)\n",
    "\n",
    "The process of running encodec with OpenVINO under hood will be the same like with the original PyTorch models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9213d6b6-56c1-4bbe-aa18-4bd1006513b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = compress(model, wav, use_lm=False)\n",
    "out_file = Path(\"compressed_ov.ecdc\")\n",
    "out_file.write_bytes(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce575606-fe03-4996-8fe2-f64deda0e766",
   "metadata": {},
   "outputs": [],
   "source": [
    "out, out_sr = decompress(out_file.read_bytes())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a297c13b-de70-4c96-85a0-d173667de241",
   "metadata": {},
   "outputs": [],
   "source": [
    "ov_output_file = \"decopressed_ov.wav\"\n",
    "save_audio(out, ov_output_file, out_sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd5d876a-648c-43b6-bbf3-def1f4cbb118",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio, sr = librosa.load(ov_output_file)\n",
    "plt.figure(figsize=(14, 5))\n",
    "librosa.display.waveshow(audio, sr=sr)\n",
    "\n",
    "ipd.Audio(ov_output_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "030a34e1-4b3d-4aec-aae1-633772e95992",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gradio as gr\n",
    "from typing import Tuple\n",
    "import numpy as np\n",
    "\n",
    "def preprocess(input, sample_rate, model_sr, model_channels):\n",
    "    input = torch.tensor(input, dtype=torch.float32)\n",
    "    input = input / 2**15  # adjust to int16 scale\n",
    "    input = input.unsqueeze(0)\n",
    "    input = convert_audio(input, sample_rate, model_sr, model_channels)\n",
    "    \n",
    "    return input\n",
    "\n",
    "\n",
    "def postprocess(output):\n",
    "    output = output.squeeze()\n",
    "    output = output * 2**15  # adjust to [-1, 1] scale\n",
    "    output = output.numpy(force=True)\n",
    "    output = output.astype(np.int16)\n",
    "\n",
    "    return output\n",
    "    \n",
    "\n",
    "def _compress(input: Tuple[int, np.ndarray]):\n",
    "    sample_rate, waveform = input\n",
    "    waveform = preprocess(waveform, sample_rate, model_sr, model_channels)\n",
    "    \n",
    "    b = compress(model, waveform, use_lm=False)\n",
    "    out, out_sr = decompress(b)\n",
    "    \n",
    "    out = postprocess(out)\n",
    "    return out_sr, out\n",
    "    \n",
    "demo = gr.Interface(\n",
    "    _compress,\n",
    "    'audio',\n",
    "    'audio',\n",
    "    examples=['test_24k.wav']\n",
    ")\n",
    "    \n",
    "try:\n",
    "    demo.launch(debug=True)\n",
    "except Exception:\n",
    "    demo.launch(share=True, debug=True)\n",
    "# if you are launching remotely, specify server_name and server_port\n",
    "# demo.launch(server_name='your server name', server_port='server port in int')\n",
    "# Read more in the docs: https://gradio.app/docs/"
   ]
  }
 ],
 "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.8.10"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}