steerable-nafx-tensorflow.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Tensorflow implementation** of the paper [Steerable discovery of neural audio effects](https://arxiv.org/abs/2112.02926) by [Christian J. Steinmetz](https://www.christiansteinmetz.com/) and [Joshua D. Reiss](http://www.eecs.qmul.ac.uk/~josh/)\n",
    "\n",
    "______________________________________________\n",
    "\n",
    "<div align=\"center\">\n",
    "\n",
    "# Steerable discovery of neural audio effects\n",
    "\n",
    "  [Christian J. Steinmetz](https://www.christiansteinmetz.com/)  and  [Joshua D. Reiss](http://www.eecs.qmul.ac.uk/~josh/)\n",
    "\n",
    "\n",
    "[Code](https://github.com/csteinmetz1/steerable-nafx) • [Paper](https://arxiv.org/abs/2112.02926) • [Demo](https://csteinmetz1.github.io/steerable-nafx)\t• [Slides]()\n",
    "\n",
    "<img src=\"https://csteinmetz1.github.io/steerable-nafx/assets/steerable-headline.svg\">\n",
    "\n",
    "</div>\n",
    "\n",
    "## Abtract\n",
    "Applications of deep learning for audio effects often focus on modeling analog effects or learning to control effects to emulate a trained audio engineer. \n",
    "However, deep learning approaches also have the potential to expand creativity through neural audio effects that enable new sound transformations. \n",
    "While recent work demonstrated that neural networks with random weights produce compelling audio effects, control of these effects is limited and unintuitive.\n",
    "To address this, we introduce a method for the steerable discovery of neural audio effects.\n",
    "This method enables the design of effects using example recordings provided by the user. \n",
    "We demonstrate how this method produces an effect similar to the target effect, along with interesting inaccuracies, while also providing perceptually relevant controls.\n",
    "\n",
    "\n",
    "\\* *Accepted to NeurIPS 2021 Workshop on Machine Learning for Creativity and Design*\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "import tf2onnx\n",
    "import onnx\n",
    "\n",
    "import os\n",
    "import IPython\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import librosa.display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choose computation device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "physical_devices = tf.config.list_physical_devices()\n",
    "print(f\"These are the physical devices available:\\n{physical_devices}\")\n",
    "\n",
    "try:\n",
    "    # Disable all GPUS\n",
    "    tf.config.set_visible_devices([], 'GPU')\n",
    "    visible_devices = tf.config.get_visible_devices()\n",
    "    print(f\"These are the visible devices:\\n{visible_devices}\")\n",
    "except:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "name = 'model_0'\n",
    "\n",
    "\n",
    "if not os.path.exists('models/'+name):\n",
    "    os.makedirs('models/'+name)\n",
    "else:\n",
    "    print(\"A model with the same name already exists. Please choose a new name.\")\n",
    "    exit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "keras.backend.clear_session()\n",
    "\n",
    "class TCNBlock(keras.Model):\n",
    "    def __init__(self, in_channels, out_channels, kernel_size, dilation, activation=True, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.conv = keras.layers.Conv1D(filters=out_channels, kernel_size=kernel_size, dilation_rate=dilation, padding=\"valid\", use_bias=True, kernel_initializer=\"glorot_uniform\", bias_initializer=\"zeros\")\n",
    "        # self.act = keras.layers.tanh\n",
    "        self.act = keras.layers.PReLU(alpha_initializer=tf.initializers.constant(0.25), shared_axes=[0, 1, 2])\n",
    "        self.res = keras.layers.Conv1D(out_channels, 1, use_bias=False, kernel_initializer=\"glorot_uniform\")\n",
    "        self.kernel_size = kernel_size\n",
    "        self.dilation = dilation\n",
    "\n",
    "    def call(self, x):\n",
    "        x_in = x\n",
    "        x = self.conv(x)\n",
    "        x = self.act(x)\n",
    "        x_res = self.res(x_in)\n",
    "        x_res = x_res[:, (self.kernel_size-1)*self.dilation:, :]\n",
    "        x = x + x_res\n",
    "        return x\n",
    "    \n",
    "class TCN(keras.Model):\n",
    "    def __init__(self, n_inputs=1, n_outputs=1, n_blocks=10, kernel_size=13, n_channels=64, dilation_growth=4):\n",
    "        super().__init__()\n",
    "        self.kernel_size = kernel_size\n",
    "        self.n_channels = n_channels\n",
    "        self.dilation_growth = dilation_growth\n",
    "        self.n_blocks = n_blocks\n",
    "        self.stack_size = n_blocks\n",
    "        \n",
    "        self.blocks = []\n",
    "        for n in range(n_blocks):\n",
    "            if n == 0:\n",
    "                in_ch = n_inputs\n",
    "                out_ch = n_channels\n",
    "                act = True\n",
    "            elif (n+1) == n_blocks:\n",
    "                in_ch = n_channels\n",
    "                out_ch = n_outputs\n",
    "                act = True\n",
    "            else:\n",
    "                in_ch = n_channels\n",
    "                out_ch = n_channels\n",
    "                act = True\n",
    "\n",
    "            dilation = dilation_growth ** n\n",
    "            self.blocks.append(TCNBlock(in_ch, out_ch, kernel_size, dilation, activation=act))\n",
    "\n",
    "    def call(self, x):\n",
    "        for block in self.blocks:\n",
    "            x = block(x)\n",
    "        return x\n",
    "    \n",
    "    def compute_receptive_field(self):\n",
    "        \"\"\"Compute the receptive field in samples.\"\"\"\n",
    "        rf = self.kernel_size\n",
    "        for n in range(1, self.n_blocks):\n",
    "            dilation = self.dilation_growth ** (n % self.stack_size)\n",
    "            rf = rf + ((self.kernel_size - 1) * dilation)\n",
    "        return rf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Steering (training)\n",
    "Use a pair of audio examples in order to construct neural audio effects.\n",
    "\n",
    "There are two options. Either start with the pre-loaded audio examples, or upload your own clean/processed audio recordings for the steering process."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "a.) Use some of our pre-loaded audio examples. Choose from the compressor or reverb effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Use pre-loaded audio examples for steering\n",
    "effect_type = \"Amp\" #@param [\"Compressor\", \"Reverb\", \"UltraTab\", \"Amp\"]\n",
    "\n",
    "if effect_type == \"Compressor\":\n",
    "  input_file = \"audio/drum_kit_clean.wav\"\n",
    "  output_file = \"audio/drum_kit_comp_agg.wav\"\n",
    "elif effect_type == \"Reverb\":\n",
    "  input_file = \"audio/acgtr_clean.wav\"\n",
    "  output_file = \"audio/acgtr_reverb.wav\"\n",
    "elif effect_type == \"UltraTab\":\n",
    "  input_file = \"audio/acgtr_clean.wav\"\n",
    "  output_file = \"audio/acgtr_ultratab.wav\"\n",
    "elif effect_type == \"Amp\":\n",
    "  input_file = \"audio/ts9_test1_in_FP32.wav\"\n",
    "  output_file = \"audio/ts9_test1_out_FP32.wav\"\n",
    "\n",
    "# Load and Preprocess Data ###########################################\n",
    "sample_rate, x = sp.io.wavfile.read(input_file)\n",
    "sample_rate, y = sp.io.wavfile.read(output_file)\n",
    "\n",
    "x = x.astype(np.float32)\n",
    "y = y.astype(np.float32)\n",
    "\n",
    "# x = x[..., :1]/32768.0 # when wav files are 16-bit integers\n",
    "# y = y[..., :1]/32768.0\n",
    "\n",
    "x = x.flatten()\n",
    "y = y.flatten()\n",
    "\n",
    "print(f\"x shape: {x.shape}\")\n",
    "print(f\"x = {x}\")\n",
    "print(f\"y shape: {y.shape}\")\n",
    "print(f\"y = {y}\")\n",
    "\n",
    "print(\"input file\", x.shape)\n",
    "IPython.display.display(IPython.display.Audio(data=x, rate=sample_rate))\n",
    "print(\"output file\", y.shape)\n",
    "IPython.display.display(IPython.display.Audio(data=y, rate=sample_rate))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now its time to generate the neural audio effect by training the TCN to emulate the input/output function from the target audio effect. Adjusting the parameters will enable you to tweak the optimization process. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title TCN model training parameters\n",
    "kernel_size = 13 #@param {type:\"slider\", min:3, max:32, step:1}\n",
    "n_blocks = 4 #@param {type:\"slider\", min:2, max:30, step:1}\n",
    "dilation_growth = 10 #@param {type:\"slider\", min:1, max:10, step:1}\n",
    "n_channels = 32 #@param {type:\"slider\", min:1, max:128, step:1}\n",
    "n_iters = 300 #@param {type:\"slider\", min:0, max:10000, step:1}\n",
    "length = 508032 #@param {type:\"slider\", min:0, max:524288, step:1}\n",
    "lr = 0.001 #@param {type:\"number\"}\n",
    "\n",
    "# # reshape the audio\n",
    "x_batch = x.reshape(1,-1,1)\n",
    "y_batch = y.reshape(1,-1,1)\n",
    "\n",
    "print(f\"x_batch shape: {x_batch.shape}\")\n",
    "print(f\"y_batch shape: {y_batch.shape}\")\n",
    "\n",
    "# build the model\n",
    "model = TCN(\n",
    "    n_inputs=1,\n",
    "    n_outputs=1,\n",
    "    kernel_size=kernel_size, \n",
    "    n_blocks=n_blocks, \n",
    "    dilation_growth=dilation_growth, \n",
    "    n_channels=n_channels)\n",
    "rf = model.compute_receptive_field()\n",
    "\n",
    "print(f\"Receptive field: {rf} samples or {(rf/sample_rate)*1e3:0.1f} ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "    def __init__(self, initial_learning_rate, n_iters):\n",
    "        self.initial_learning_rate = initial_learning_rate\n",
    "        self.n_iters = n_iters\n",
    "\n",
    "    @tf.function\n",
    "    def __call__(self, step):\n",
    "        if step >= tf.cast((self.n_iters * 0.8), tf.int64):\n",
    "            return self.initial_learning_rate * 0.1\n",
    "        elif step >= tf.cast(self.n_iters * 0.95, tf.int64):\n",
    "            return self.initial_learning_rate * 0.01\n",
    "        else:\n",
    "            return self.initial_learning_rate\n",
    "        \n",
    "    def get_config(self):\n",
    "        config = {\n",
    "        'initial_learning_rate': self.initial_learning_rate,\n",
    "        'niters': self.n_iters,\n",
    "        }\n",
    "        return config\n",
    "\n",
    "optimizer = keras.optimizers.Adam(learning_rate=MyLRSchedule(lr, n_iters), epsilon=1e-8)\n",
    "model.compile(optimizer=optimizer, loss='mse')\n",
    "model.build(input_shape=(None, length+rf-1, 1))\n",
    "model.summary(expand_nested=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_idx = rf \n",
    "stop_idx = start_idx + length\n",
    "\n",
    "# the data is the same with every iteration\n",
    "x_crop = x_batch[:,start_idx-rf+1:stop_idx,:]\n",
    "y_crop = y_batch[:,start_idx:stop_idx,:]\n",
    "print(f\"x_crop = {x_crop.shape}\")\n",
    "print(f\"y_crop = {y_crop.shape}\")\n",
    "\n",
    "history = model.fit(x=x_crop, y=y_crop, epochs=n_iters, batch_size=1, verbose=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_weights('models/'+name+'/'+name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run predictions\n",
    "### 0. Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights('models/'+name+'/'+name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. On the test audio data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Prediction #################################################\n",
    "# Test the model on the testing data #############################\n",
    "\n",
    "x_pad = np.pad(x_batch, ((0,0),(rf-1,0),(0,0)), mode='constant')\n",
    "\n",
    "y_hat = model.predict(x_pad)\n",
    "\n",
    "input = x_batch.flatten()\n",
    "output = y_hat.flatten()\n",
    "target = y_batch.flatten()\n",
    "\n",
    "print(f\"Input shape: {input.shape}\")\n",
    "print(f\"Output shape: {output.shape}\")\n",
    "print(f\"Target shape: {target.shape}\")\n",
    "\n",
    "# apply highpass to outpu to remove DC\n",
    "sos = sp.signal.butter(8, 20.0, fs=sample_rate, output=\"sos\", btype=\"highpass\")\n",
    "output = sp.signal.sosfilt(sos, output)\n",
    "\n",
    "input /= np.max(np.abs(input))\n",
    "output /= np.max(np.abs(output))\n",
    "target /= np.max(np.abs(target))\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, sharex=True)\n",
    "librosa.display.waveshow(target, sr=sample_rate, color='b', alpha=0.5, ax=ax, label='Target')\n",
    "librosa.display.waveshow(output, sr=sample_rate, color='r', alpha=0.5, ax=ax, label='Output')\n",
    "\n",
    "print(\"Input (clean)\")\n",
    "IPython.display.display(IPython.display.Audio(data=input, rate=sample_rate))\n",
    "print(\"Target\")\n",
    "IPython.display.display(IPython.display.Audio(data=target, rate=sample_rate))\n",
    "print(\"Output\")\n",
    "IPython.display.display(IPython.display.Audio(data=output, rate=sample_rate))\n",
    "plt.legend()\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and Preprocess Data ###########################################\n",
    "sample_rate, x_whole = sp.io.wavfile.read(\"audio/piano_clean.wav\")\n",
    "x_whole = x_whole.astype(np.float32)\n",
    "x_whole = x_whole[..., :1]/32768.0 # because wav files are 16-bit integers\n",
    "x_whole = x_whole.reshape(1,-1,1)\n",
    "\n",
    "# Padding on both sides of the receptive field\n",
    "x_whole = np.pad(x_whole, ((0,0),(rf-1,rf-1),(0,0)), mode='constant')\n",
    "\n",
    "y_whole = model.predict(x_whole)\n",
    "\n",
    "x_whole = x_whole[:, -y_whole.shape[1]:, :]\n",
    "\n",
    "y_whole /= np.abs(y_whole).max()\n",
    "\n",
    "# apply high pass filter to remove DC\n",
    "sos = sp.signal.butter(8, 20.0, fs=sample_rate, output=\"sos\", btype=\"highpass\")\n",
    "y_whole = sp.signal.sosfilt(sos, y_whole.flatten())\n",
    "\n",
    "x_whole = x_whole.flatten()\n",
    "\n",
    "y_whole = (y_whole * 0.8)\n",
    "IPython.display.display(IPython.display.Audio(data=x_whole, rate=sample_rate))\n",
    "IPython.display.display(IPython.display.Audio(data=y_whole, rate=sample_rate))\n",
    "\n",
    "x_whole /= np.max(np.abs(x_whole))\n",
    "y_whole /= np.max(np.abs(y_whole))\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, sharex=True)\n",
    "librosa.display.waveshow(y_whole, sr=sample_rate, color='r', alpha=0.5, ax=ax, label='Output')\n",
    "librosa.display.waveshow(x_whole, sr=sample_rate, alpha=0.5, ax=ax, label='Input', color=\"blue\")\n",
    "plt.legend()\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. On a number sequence (to control inference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the model simple number sequence to compare with inference #\n",
    "X_testing_2 = np.array([])\n",
    "\n",
    "for i in range(0, 2048+rf-1):\n",
    "    X_testing_2 = np.append(X_testing_2, i*0.000001)\n",
    "\n",
    "X_testing_2 = X_testing_2.reshape(1, -1, 1)\n",
    "\n",
    "print(\"Running prediction..\")\n",
    "prediction_2 = model.predict(X_testing_2)\n",
    "print(f\"prediction {prediction_2}\")\n",
    "\n",
    "print(\"X_testing_2 shape: \", X_testing_2.shape)\n",
    "print(\"prediction_2 shape: \", prediction_2.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export as tflite and onnx model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rf is the receptive field and hence the kernel size of the last layer because of the dilation\n",
    "# hence the model input size needs to be at least rf\n",
    "# still we export the model with dynamic input and batch size\n",
    "input_shape = [1, None, 1]\n",
    "\n",
    "func = tf.function(model).get_concrete_function(\n",
    "    tf.TensorSpec(input_shape, dtype=tf.float32))\n",
    "converter = tf.lite.TFLiteConverter.from_concrete_functions([func], model)\n",
    "tflite_model = converter.convert()\n",
    "\n",
    "# Save the model.\n",
    "with open(\"models/\"+name+\"/\"+\"steerable-nafx-dynamic.tflite\", 'wb') as f:\n",
    "    f.write(tflite_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export the model to onnx\n",
    "input_shape = [1, None, 1]\n",
    "\n",
    "# Define the input shape\n",
    "input_signature = [tf.TensorSpec(input_shape, tf.float32, name='x')]\n",
    "\n",
    "# Convert the model\n",
    "onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature, opset=18)\n",
    "onnx.save(proto=onnx_model, f=\"models/\"+name+\"/\"+\"steerable-nafx-tflite-dynamic.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3 TCN blocks Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title TCN model training parameters\n",
    "kernel_size = 13 #@param {type:\"slider\", min:3, max:32, step:1}\n",
    "n_blocks = 3 #@param {type:\"slider\", min:2, max:30, step:1}\n",
    "dilation_growth = 10 #@param {type:\"slider\", min:1, max:10, step:1}\n",
    "n_channels = 32 #@param {type:\"slider\", min:1, max:128, step:1}\n",
    "n_iters = 10 #@param {type:\"slider\", min:0, max:10000, step:1}\n",
    "length = 508032 #@param {type:\"slider\", min:0, max:524288, step:1}\n",
    "lr = 0.001 #@param {type:\"number\"}\n",
    "\n",
    "# # reshape the audio\n",
    "x_batch = x.reshape(1,-1,1)\n",
    "y_batch = y.reshape(1,-1,1)\n",
    "\n",
    "print(f\"x_batch shape: {x_batch.shape}\")\n",
    "print(f\"y_batch shape: {y_batch.shape}\")\n",
    "\n",
    "# build the model\n",
    "model_3_blocks = TCN(\n",
    "    n_inputs=1,\n",
    "    n_outputs=1,\n",
    "    kernel_size=kernel_size, \n",
    "    n_blocks=n_blocks, \n",
    "    dilation_growth=dilation_growth, \n",
    "    n_channels=n_channels)\n",
    "rf = model_3_blocks.compute_receptive_field()\n",
    "\n",
    "print(f\"Receptive field: {rf} samples or {(rf/sample_rate)*1e3:0.1f} ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "    def __init__(self, initial_learning_rate, n_iters):\n",
    "        self.initial_learning_rate = initial_learning_rate\n",
    "        self.n_iters = n_iters\n",
    "\n",
    "    @tf.function\n",
    "    def __call__(self, step):\n",
    "        if step >= tf.cast((self.n_iters * 0.8), tf.int64):\n",
    "            return self.initial_learning_rate * 0.1\n",
    "        elif step >= tf.cast(self.n_iters * 0.95, tf.int64):\n",
    "            return self.initial_learning_rate * 0.01\n",
    "        else:\n",
    "            return self.initial_learning_rate\n",
    "        \n",
    "    def get_config(self):\n",
    "        config = {\n",
    "        'initial_learning_rate': self.initial_learning_rate,\n",
    "        'niters': self.n_iters,\n",
    "        }\n",
    "        return config\n",
    "\n",
    "optimizer = keras.optimizers.Adam(learning_rate=MyLRSchedule(lr, n_iters), epsilon=1e-8)\n",
    "model_3_blocks.compile(optimizer=optimizer, loss='mse')\n",
    "model_3_blocks.build(input_shape=(None, length+rf-1, 1))\n",
    "model_3_blocks.summary(expand_nested=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_idx = rf \n",
    "stop_idx = start_idx + length\n",
    "\n",
    "# the data is the same with every iteration\n",
    "x_crop = x_batch[:,start_idx-rf+1:stop_idx,:]\n",
    "y_crop = y_batch[:,start_idx:stop_idx,:]\n",
    "print(f\"x_crop = {x_crop.shape}\")\n",
    "print(f\"y_crop = {y_crop.shape}\")\n",
    "\n",
    "history = model_3_blocks.fit(x=x_crop, y=y_crop, epochs=n_iters, batch_size=1, verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. On the test audio data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run Prediction #################################################\n",
    "# Test the model on the testing data #############################\n",
    "\n",
    "x_pad = np.pad(x_batch, ((0,0),(rf-1,0),(0,0)), mode='constant')\n",
    "\n",
    "y_hat = model_3_blocks.predict(x_pad)\n",
    "\n",
    "input = x_batch.flatten()\n",
    "output = y_hat.flatten()\n",
    "target = y_batch.flatten()\n",
    "\n",
    "print(f\"Input shape: {input.shape}\")\n",
    "print(f\"Output shape: {output.shape}\")\n",
    "print(f\"Target shape: {target.shape}\")\n",
    "\n",
    "# apply highpass to outpu to remove DC\n",
    "sos = sp.signal.butter(8, 20.0, fs=sample_rate, output=\"sos\", btype=\"highpass\")\n",
    "output = sp.signal.sosfilt(sos, output)\n",
    "\n",
    "input /= np.max(np.abs(input))\n",
    "output /= np.max(np.abs(output))\n",
    "target /= np.max(np.abs(target))\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, sharex=True)\n",
    "librosa.display.waveshow(target, sr=sample_rate, color='b', alpha=0.5, ax=ax, label='Target')\n",
    "librosa.display.waveshow(output, sr=sample_rate, color='r', alpha=0.5, ax=ax, label='Output')\n",
    "\n",
    "print(\"Input (clean)\")\n",
    "IPython.display.display(IPython.display.Audio(data=input, rate=sample_rate))\n",
    "print(\"Target\")\n",
    "IPython.display.display(IPython.display.Audio(data=target, rate=sample_rate))\n",
    "print(\"Output\")\n",
    "IPython.display.display(IPython.display.Audio(data=output, rate=sample_rate))\n",
    "plt.legend()\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and Preprocess Data ###########################################\n",
    "sample_rate, x_whole = sp.io.wavfile.read(\"audio/piano_clean.wav\")\n",
    "x_whole = x_whole.astype(np.float32)\n",
    "x_whole = x_whole[..., :1]/32768.0 # because wav files are 16-bit integers\n",
    "x_whole = x_whole.reshape(1,-1,1)\n",
    "\n",
    "# Padding on both sides of the receptive field\n",
    "x_whole = np.pad(x_whole, ((0,0),(rf-1,rf-1),(0,0)), mode='constant')\n",
    "\n",
    "y_whole = model_3_blocks.predict(x_whole)\n",
    "\n",
    "x_whole = x_whole[:, -y_whole.shape[1]:, :]\n",
    "\n",
    "y_whole /= np.abs(y_whole).max()\n",
    "\n",
    "# apply high pass filter to remove DC\n",
    "sos = sp.signal.butter(8, 20.0, fs=sample_rate, output=\"sos\", btype=\"highpass\")\n",
    "y_whole = sp.signal.sosfilt(sos, y_whole.flatten())\n",
    "\n",
    "x_whole = x_whole.flatten()\n",
    "\n",
    "y_whole = (y_whole * 0.8)\n",
    "IPython.display.display(IPython.display.Audio(data=x_whole, rate=sample_rate))\n",
    "IPython.display.display(IPython.display.Audio(data=y_whole, rate=sample_rate))\n",
    "\n",
    "x_whole /= np.max(np.abs(x_whole))\n",
    "y_whole /= np.max(np.abs(y_whole))\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, sharex=True)\n",
    "librosa.display.waveshow(y_whole, sr=sample_rate, color='r', alpha=0.5, ax=ax, label='Output')\n",
    "librosa.display.waveshow(x_whole, sr=sample_rate, alpha=0.5, ax=ax, label='Input', color=\"blue\")\n",
    "plt.legend()\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export as tflite and onnx model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rf is the receptive field and hence the kernel size of the last layer because of the dilation\n",
    "# hence the model input size needs to be at least rf\n",
    "# still we export the model with dynamic input and batch size\n",
    "input_shape = [1, None, 1]\n",
    "\n",
    "func = tf.function(model_3_blocks).get_concrete_function(\n",
    "    tf.TensorSpec(input_shape, dtype=tf.float32))\n",
    "converter = tf.lite.TFLiteConverter.from_concrete_functions([func], model_3_blocks)\n",
    "tflite_model_3_blocks = converter.convert()\n",
    "\n",
    "# Save the model.\n",
    "with open(\"models/\"+name+\"/\"+\"steerable-nafx-3_blocks-dynamic.tflite\", 'wb') as f:\n",
    "    f.write(tflite_model_3_blocks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export the model to onnx\n",
    "input_shape = [1, None, 1]\n",
    "\n",
    "# Define the input shape\n",
    "input_signature = [tf.TensorSpec(input_shape, tf.float32, name='x')]\n",
    "\n",
    "# Convert the model\n",
    "onnx_model_3_blocks, _ = tf2onnx.convert.from_keras(model_3_blocks, input_signature, opset=18)\n",
    "onnx.save(proto=onnx_model_3_blocks, f=\"models/\"+name+\"/\"+\"steerable-nafx-3_blocks-tflite-dynamic.onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 TCN blocks Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title TCN model training parameters\n",
    "kernel_size = 13 #@param {type:\"slider\", min:3, max:32, step:1}\n",
    "n_blocks = 2 #@param {type:\"slider\", min:2, max:30, step:1}\n",
    "dilation_growth = 10 #@param {type:\"slider\", min:1, max:10, step:1}\n",
    "n_channels = 32 #@param {type:\"slider\", min:1, max:128, step:1}\n",
    "n_iters = 10 #@param {type:\"slider\", min:0, max:10000, step:1}\n",
    "length = 508032 #@param {type:\"slider\", min:0, max:524288, step:1}\n",
    "lr = 0.001 #@param {type:\"number\"}\n",
    "\n",
    "# # reshape the audio\n",
    "x_batch = x.reshape(1,-1,1)\n",
    "y_batch = y.reshape(1,-1,1)\n",
    "\n",
    "print(f\"x_batch shape: {x_batch.shape}\")\n",
    "print(f\"y_batch shape: {y_batch.shape}\")\n",
    "\n",
    "# build the model\n",
    "model_2_blocks = TCN(\n",
    "    n_inputs=1,\n",
    "    n_outputs=1,\n",
    "    kernel_size=kernel_size, \n",
    "    n_blocks=n_blocks, \n",
    "    dilation_growth=dilation_growth, \n",
    "    n_channels=n_channels)\n",
    "rf = model_2_blocks.compute_receptive_field()\n",
    "\n",
    "print(f\"Receptive field: {rf} samples or {(rf/sample_rate)*1e3:0.1f} ms\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):\n",
    "    def __init__(self, initial_learning_rate, n_iters):\n",
    "        self.initial_learning_rate = initial_learning_rate\n",
    "        self.n_iters = n_iters\n",
    "\n",
    "    @tf.function\n",
    "    def __call__(self, step):\n",
    "        if step >= tf.cast((self.n_iters * 0.8), tf.int64):\n",
    "            return self.initial_learning_rate * 0.1\n",
    "        elif step >= tf.cast(self.n_iters * 0.95, tf.int64):\n",
    "            return self.initial_learning_rate * 0.01\n",
    "        else:\n",
    "            return self.initial_learning_rate\n",
    "        \n",
    "    def get_config(self):\n",
    "        config = {\n",
    "        'initial_learning_rate': self.initial_learning_rate,\n",
    "        'niters': self.n_iters,\n",
    "        }\n",
    "        return config\n",
    "\n",
    "optimizer = keras.optimizers.Adam(learning_rate=MyLRSchedule(lr, n_iters), epsilon=1e-8)\n",
    "model_2_blocks.compile(optimizer=optimizer, loss='mse')\n",
    "model_2_blocks.build(input_shape=(None, length+rf-1, 1))\n",
    "model_2_blocks.summary(expand_nested=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_idx = rf \n",
    "stop_idx = start_idx + length\n",
    "\n",
    "# the data is the same with every iteration\n",
    "x_crop = x_batch[:,start_idx-rf+1:stop_idx,:]\n",
    "y_crop = y_batch[:,start_idx:stop_idx,:]\n",
    "print(f\"x_crop = {x_crop.shape}\")\n",
    "print(f\"y_crop = {y_crop.shape}\")\n",
    "\n",
    "history = model_2_blocks.fit(x=x_crop, y=y_crop, epochs=n_iters, batch_size=1, verbose=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run predictions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. On the test audio data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### 1. On the test audio dat# Run Prediction #################################################\n",
    "# Test the model on the testing data #############################\n",
    "\n",
    "x_pad = np.pad(x_batch, ((0,0),(rf-1,0),(0,0)), mode='constant')\n",
    "\n",
    "y_hat = model_2_blocks.predict(x_pad)\n",
    "\n",
    "input = x_batch.flatten()\n",
    "output = y_hat.flatten()\n",
    "target = y_batch.flatten()\n",
    "\n",
    "print(f\"Input shape: {input.shape}\")\n",
    "print(f\"Output shape: {output.shape}\")\n",
    "print(f\"Target shape: {target.shape}\")\n",
    "\n",
    "# apply highpass to outpu to remove DC\n",
    "sos = sp.signal.butter(8, 20.0, fs=sample_rate, output=\"sos\", btype=\"highpass\")\n",
    "output = sp.signal.sosfilt(sos, output)\n",
    "\n",
    "input /= np.max(np.abs(input))\n",
    "output /= np.max(np.abs(output))\n",
    "target /= np.max(np.abs(target))\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, sharex=True)\n",
    "librosa.display.waveshow(target, sr=sample_rate, color='b', alpha=0.5, ax=ax, label='Target')\n",
    "librosa.display.waveshow(output, sr=sample_rate, color='r', alpha=0.5, ax=ax, label='Output')\n",
    "\n",
    "print(\"Input (clean)\")\n",
    "IPython.display.display(IPython.display.Audio(data=input, rate=sample_rate))\n",
    "print(\"Target\")\n",
    "IPython.display.display(IPython.display.Audio(data=target, rate=sample_rate))\n",
    "print(\"Output\")\n",
    "IPython.display.display(IPython.display.Audio(data=output, rate=sample_rate))\n",
    "plt.legend()\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and Preprocess Data ###########################################\n",
    "sample_rate, x_whole = sp.io.wavfile.read(\"audio/piano_clean.wav\")\n",
    "x_whole = x_whole.astype(np.float32)\n",
    "x_whole = x_whole[..., :1]/32768.0 # because wav files are 16-bit integers\n",
    "x_whole = x_whole.reshape(1,-1,1)\n",
    "\n",
    "# Padding on both sides of the receptive field\n",
    "x_whole = np.pad(x_whole, ((0,0),(rf-1,rf-1),(0,0)), mode='constant')\n",
    "\n",
    "y_whole = model_2_blocks.predict(x_whole)\n",
    "\n",
    "x_whole = x_whole[:, -y_whole.shape[1]:, :]\n",
    "\n",
    "y_whole /= np.abs(y_whole).max()\n",
    "\n",
    "# apply high pass filter to remove DC\n",
    "sos = sp.signal.butter(8, 20.0, fs=sample_rate, output=\"sos\", btype=\"highpass\")\n",
    "y_whole = sp.signal.sosfilt(sos, y_whole.flatten())\n",
    "\n",
    "x_whole = x_whole.flatten()\n",
    "\n",
    "y_whole = (y_whole * 0.8)\n",
    "IPython.display.display(IPython.display.Audio(data=x_whole, rate=sample_rate))\n",
    "IPython.display.display(IPython.display.Audio(data=y_whole, rate=sample_rate))\n",
    "\n",
    "x_whole /= np.max(np.abs(x_whole))\n",
    "y_whole /= np.max(np.abs(y_whole))\n",
    "\n",
    "fig, ax = plt.subplots(nrows=1, sharex=True)\n",
    "librosa.display.waveshow(y_whole, sr=sample_rate, color='r', alpha=0.5, ax=ax, label='Output')\n",
    "librosa.display.waveshow(x_whole, sr=sample_rate, alpha=0.5, ax=ax, label='Input', color=\"blue\")\n",
    "plt.legend()\n",
    "plt.show(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Export as tflite and onnx model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# rf is the receptive field and hence the kernel size of the last layer because of the dilation\n",
    "# hence the model input size needs to be at least rf\n",
    "# still we export the model with dynamic input and batch size\n",
    "input_shape = [1, None, 1]\n",
    "\n",
    "func = tf.function(model_2_blocks).get_concrete_function(\n",
    "    tf.TensorSpec(input_shape, dtype=tf.float32))\n",
    "converter = tf.lite.TFLiteConverter.from_concrete_functions([func], model_2_blocks)\n",
    "tflite_model_2_blocks = converter.convert()\n",
    "\n",
    "# Save the model.\n",
    "with open(\"models/\"+name+\"/\"+\"steerable-nafx-2_blocks-dynamic.tflite\", 'wb') as f:\n",
    "    f.write(tflite_model_2_blocks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export the model to onnx\n",
    "input_shape = [1, None, 1]\n",
    "\n",
    "# Define the input shape\n",
    "input_signature = [tf.TensorSpec(input_shape, tf.float32, name='x')]\n",
    "\n",
    "# Convert the model\n",
    "onnx_model_2_blocks, _ = tf2onnx.convert.from_keras(model_2_blocks, input_signature, opset=18)\n",
    "onnx.save(proto=onnx_model_2_blocks, f=\"models/\"+name+\"/\"+\"steerable-nafx-2_blocks-tflite-dynamic.onnx\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.8"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}