MelSpecVAE_v1.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "MelSpecVAE_v1.5.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "jW7iOo_KIOBX",
        "kWqhd1LILIxD",
        "uTuKspchQjY0"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QGsLH8ZeE-EK"
      },
      "source": [
        "# MelSpecVAE v.1.5\n",
        " Author: Moisés Horta Valenzuela, 2021\n",
        "> Website: [moiseshorta.audio](https://moiseshorta.audio)\n",
        "\n",
        "> Twitter: [@hexorcismos](https://twitter.com/hexorcismos)\n",
        "\n",
        "\n",
        "```\n",
        "MelSpecVAE is a Variational Autoencoder which synthesizes Mel-Spectrograms thay can be inverted into raw audio waveform.\n",
        "Currently you can train it with any dataset of .wav audio at 44.1khz Sample Rate and 16bit bitdepth.\n",
        "\n",
        "> Features:\n",
        "* Interpolate through 2 different points in the latent space and synthesize the 'in between' sounds.\n",
        "* Generate short one-shot audio\n",
        "* Synthesize arbitrarily long audio samples by generating seeds and sampling from the latent space.\n",
        "\n",
        "> Credits:\n",
        "* VAE neural network architecture coded following 'The Sound of AI' Youtube tutorial series by Valerio Velardo\n",
        "* Some utility functions from Marco Passini's MelGAN-VC Jupyter Notebook.\n",
        "\n",
        "> Last update:\n",
        "* 12.12.2021: Added experimental feature 'Timbre Transfer'\n",
        "```\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "dbS8aCU-fMU8"
      },
      "source": [
        "#@title Mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J8p_G3UIaanv"
      },
      "source": [
        "## Run the next cells first for training, generating or timbre transfer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "otzXYjIG_wmh",
        "cellView": "form"
      },
      "source": [
        "#@title Import Tensorflow and torchaudio \n",
        "\n",
        "!pip install tensorflow-gpu==2.0.0\n",
        "!pip install h5py==2.10.0 --force-reinstall\n",
        "!pip install soundfile                    #to save wav files\n",
        "!pip install --no-deps torchaudio==0.5\n",
        "!pip install git+https://github.com/pvigier/perlin-numpy #for generating perlin and fractal noise"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kbf9hPRG7Pjr",
        "cellView": "form"
      },
      "source": [
        "#@title Import libraries\n",
        "from glob import glob\n",
        "import librosa\n",
        "import librosa.display\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from numpy import asarray\n",
        "from numpy.random import randn\n",
        "from numpy.random import randint\n",
        "from numpy import linspace\n",
        "import soundfile as sf\n",
        "import time\n",
        "import IPython\n",
        "import tensorflow as tf\n",
        "from perlin_numpy import (\n",
        "    generate_fractal_noise_2d, generate_perlin_noise_2d, \n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JdZmiOYdBVSq",
        "cellView": "form"
      },
      "source": [
        "#@title Hyperparameters \n",
        "learning_rate = 0.001 #@param {type:\"raw\"}\n",
        "num_epochs_to_train =  10#@param {type:\"integer\"}\n",
        "batch_size =  32#@param {type:\"integer\"}\n",
        "vector_dimension = 64 #@param {type:\"integer\"}\n",
        "\n",
        "hop=256               #hop size (window size = 4*hop)\n",
        "sr=44100              #sampling rate\n",
        "min_level_db=-100     #reference values to normalize data\n",
        "ref_level_db=20\n",
        "\n",
        "LEARNING_RATE = learning_rate\n",
        "BATCH_SIZE = batch_size\n",
        "EPOCHS = num_epochs_to_train\n",
        "VECTOR_DIM=vector_dimension\n",
        "\n",
        "shape=128           #length of time axis of split specrograms         \n",
        "spec_split=1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Y2kIlpC7_lD",
        "cellView": "form"
      },
      "source": [
        "#@title Waveform to Spectrogram conversion\n",
        "\n",
        "''' Decorsière, Rémi, Peter L. Søndergaard, Ewen N. MacDonald, and Torsten Dau. \n",
        "\"Inversion of auditory spectrograms, traditional spectrograms, and other envelope representations.\" \n",
        "IEEE/ACM Transactions on Audio, Speech, and Language Processing 23, no. 1 (2014): 46-56.'''\n",
        "\n",
        "#ORIGINAL CODE FROM https://github.com/yoyololicon/spectrogram-inversion\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from tqdm import tqdm\n",
        "from functools import partial\n",
        "import math\n",
        "import heapq\n",
        "from torchaudio.transforms import MelScale, Spectrogram\n",
        "\n",
        "torch.set_default_tensor_type('torch.cuda.FloatTensor')\n",
        "\n",
        "specobj = Spectrogram(n_fft=4*hop, win_length=4*hop, hop_length=hop, pad=0, power=2, normalized=False)\n",
        "specfunc = specobj.forward\n",
        "melobj = MelScale(n_mels=hop, sample_rate=sr, f_min=0.)\n",
        "melfunc = melobj.forward\n",
        "\n",
        "def melspecfunc(waveform):\n",
        "  specgram = specfunc(waveform)\n",
        "  mel_specgram = melfunc(specgram)\n",
        "  return mel_specgram\n",
        "\n",
        "def spectral_convergence(input, target):\n",
        "    return 20 * ((input - target).norm().log10() - target.norm().log10())\n",
        "\n",
        "def GRAD(spec, transform_fn, samples=None, init_x0=None, maxiter=1000, tol=1e-6, verbose=1, evaiter=10, lr=0.002):\n",
        "\n",
        "    spec = torch.Tensor(spec)\n",
        "    samples = (spec.shape[-1]*hop)-hop\n",
        "\n",
        "    if init_x0 is None:\n",
        "        init_x0 = spec.new_empty((1,samples)).normal_(std=1e-6)\n",
        "    x = nn.Parameter(init_x0)\n",
        "    T = spec\n",
        "\n",
        "    criterion = nn.L1Loss()\n",
        "    optimizer = torch.optim.Adam([x], lr=lr)\n",
        "\n",
        "    bar_dict = {}\n",
        "    metric_func = spectral_convergence\n",
        "    bar_dict['spectral_convergence'] = 0\n",
        "    metric = 'spectral_convergence'\n",
        "\n",
        "    init_loss = None\n",
        "    with tqdm(total=maxiter, disable=not verbose) as pbar:\n",
        "        for i in range(maxiter):\n",
        "            optimizer.zero_grad()\n",
        "            V = transform_fn(x)\n",
        "            loss = criterion(V, T)\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            lr = lr*0.9999\n",
        "            for param_group in optimizer.param_groups:\n",
        "              param_group['lr'] = lr\n",
        "\n",
        "            if i % evaiter == evaiter - 1:\n",
        "                with torch.no_grad():\n",
        "                    V = transform_fn(x)\n",
        "                    bar_dict[metric] = metric_func(V, spec).item()\n",
        "                    l2_loss = criterion(V, spec).item()\n",
        "                    pbar.set_postfix(**bar_dict, loss=l2_loss)\n",
        "                    pbar.update(evaiter)\n",
        "\n",
        "    return x.detach().view(-1).cpu()\n",
        "\n",
        "def normalize(S):\n",
        "  return np.clip((((S - min_level_db) / -min_level_db)*2.)-1., -1, 1)\n",
        "\n",
        "def denormalize(S):\n",
        "  return (((np.clip(S, -1, 1)+1.)/2.) * -min_level_db) + min_level_db\n",
        "\n",
        "def prep(wv,hop=192):\n",
        "  S = np.array(torch.squeeze(melspecfunc(torch.Tensor(wv).view(1,-1))).detach().cpu())\n",
        "  S = librosa.power_to_db(S)-ref_level_db\n",
        "  return normalize(S)\n",
        "\n",
        "def deprep(S):\n",
        "  S = denormalize(S)+ref_level_db\n",
        "  S = librosa.db_to_power(S)\n",
        "  wv = GRAD(np.expand_dims(S,0), melspecfunc, maxiter=2500, evaiter=10, tol=1e-8)\n",
        "  return np.array(np.squeeze(wv))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YNRYjsCDqDjF",
        "cellView": "form"
      },
      "source": [
        "#@title Helper functions\n",
        "\n",
        "#Generate spectrograms from waveform array\n",
        "def tospec(data):\n",
        "  specs=np.empty(data.shape[0], dtype=object)\n",
        "  for i in range(data.shape[0]):\n",
        "    x = data[i]\n",
        "    S=prep(x)\n",
        "    S = np.array(S, dtype=np.float32)\n",
        "    specs[i]=np.expand_dims(S, -1)\n",
        "  print(specs.shape)\n",
        "  return specs\n",
        "\n",
        "#Generate multiple spectrograms with a determined length from single wav file\n",
        "def tospeclong(path, length=4*44100):\n",
        "  x, sr = librosa.load(path,sr=44100)\n",
        "  x,_ = librosa.effects.trim(x)\n",
        "  loudls = librosa.effects.split(x, top_db=50)\n",
        "  xls = np.array([])\n",
        "  for interv in loudls:\n",
        "    xls = np.concatenate((xls,x[interv[0]:interv[1]]))\n",
        "  x = xls\n",
        "  num = x.shape[0]//length\n",
        "  specs=np.empty(num, dtype=object)\n",
        "  for i in range(num-1):\n",
        "    a = x[i*length:(i+1)*length]\n",
        "    S = prep(a)\n",
        "    S = np.array(S, dtype=np.float32)\n",
        "    try:\n",
        "      sh = S.shape\n",
        "      specs[i]=S\n",
        "    except AttributeError:\n",
        "      print('spectrogram failed')\n",
        "  print(specs.shape)\n",
        "  return specs\n",
        "\n",
        "#Waveform array from path of folder containing wav files\n",
        "def audio_array(path):\n",
        "  ls = glob(f'{path}/*.wav')\n",
        "  adata = []\n",
        "  for i in range(len(ls)):\n",
        "    x, sr = tf.audio.decode_wav(tf.io.read_file(ls[i]), 1)\n",
        "    x = np.array(x, dtype=np.float32)\n",
        "    adata.append(x)\n",
        "  return np.array(adata)\n",
        "\n",
        "#Waveform array from path of folder containing wav files\n",
        "def single_audio_array(path):\n",
        "  ls = glob(path)\n",
        "  adata = []\n",
        "  for i in range(len(ls)):\n",
        "    x, sr = tf.audio.decode_wav(tf.io.read_file(ls[i]), 1)\n",
        "    x = np.array(x, dtype=np.float32)\n",
        "    adata.append(x)\n",
        "  return np.array(adata)\n",
        "\n",
        "\n",
        "#Concatenate spectrograms in array along the time axis\n",
        "def testass(a):\n",
        "  but=False\n",
        "  con = np.array([])\n",
        "  nim = a.shape[0]\n",
        "  for i in range(nim):\n",
        "    im = a[i]\n",
        "    im = np.squeeze(im)\n",
        "    if not but:\n",
        "      con=im\n",
        "      but=True\n",
        "    else:\n",
        "      con = np.concatenate((con,im), axis=1)\n",
        "  return np.squeeze(con)\n",
        "\n",
        "#Split spectrograms in chunks with equal size\n",
        "def splitcut(data):\n",
        "  ls = []\n",
        "  mini = 0\n",
        "  minifinal = spec_split*shape   #max spectrogram length\n",
        "  for i in range(data.shape[0]-1):\n",
        "    if data[i].shape[1]<=data[i+1].shape[1]:\n",
        "      mini = data[i].shape[1]\n",
        "    else:\n",
        "      mini = data[i+1].shape[1]\n",
        "    if mini>=3*shape and mini<minifinal:\n",
        "      minifinal = mini\n",
        "  for i in range(data.shape[0]):\n",
        "    x = data[i]\n",
        "    if x.shape[1]>=3*shape:\n",
        "      for n in range(x.shape[1]//minifinal):\n",
        "        ls.append(x[:,n*minifinal:n*minifinal+minifinal,:])\n",
        "      ls.append(x[:,-minifinal:,:])\n",
        "  return np.array(ls)\n",
        "\n",
        "# Generates timestamp string of \"day_month_year_hourMin\" \n",
        "def get_time_stamp():\n",
        "  secondsSinceEpoch = time.time()\n",
        "  timeObj = time.localtime(secondsSinceEpoch)\n",
        "  x = ('%d_%d_%d_%d%d' % (timeObj.tm_mday, timeObj.tm_mon, timeObj.tm_year, timeObj.tm_hour, timeObj.tm_min))\n",
        "  return x\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jW7iOo_KIOBX"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h7B5LHi1-jRK",
        "cellView": "form"
      },
      "source": [
        "#@title Import folder containing .wav files for training\n",
        "\n",
        "#Generating Mel-Spectrogram dataset (Uncomment where needed)\n",
        "#adata: source spectrograms\n",
        "\n",
        "audio_directory = \"/path/to/your/audio/dataset\" #@param {type:\"string\"}\n",
        "\n",
        "#AUDIO TO CONVERT\n",
        "awv = audio_array(audio_directory)         #get waveform array from folder containing wav files\n",
        "aspec = tospec(awv)                        #get spectrogram array\n",
        "adata = splitcut(aspec)                    #split spectrogams to fixed \n",
        "print(np.shape(adata))\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ddE4uNQMxNCG",
        "cellView": "form"
      },
      "source": [
        "#@title Build VAE Neural Network\n",
        "\n",
        "#VAE\n",
        "\n",
        "import os\n",
        "import pickle\n",
        "\n",
        "from tensorflow.keras import Model\n",
        "from tensorflow.keras.layers import Input, Conv2D, ReLU, BatchNormalization, Flatten, Dense, Reshape, Conv2DTranspose, Activation, Lambda\n",
        "from tensorflow.keras import backend as K\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.losses import MeanSquaredError\n",
        "from tensorflow.keras.callbacks import ModelCheckpoint\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "tf.compat.v1.disable_eager_execution()\n",
        "\n",
        "\n",
        "class VAE:\n",
        "  \"\"\"\n",
        "  VAE represents a Deep Convolutional autoencoder architecture\n",
        "  with mirrored encoder and decoder components.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               input_shape, #shape of the input data\n",
        "               conv_filters, #convolutional network filters\n",
        "               conv_kernels, #convNet kernel size\n",
        "               conv_strides, #convNet strides\n",
        "               latent_space_dim):\n",
        "    self.input_shape = input_shape # [28, 28, 1], in this case is 28 x 28 pixels on 1 channel for greyscale\n",
        "    self.conv_filters = conv_filters # is a list for each layer, i.e. [2, 4, 8]\n",
        "    self.conv_kernels = conv_kernels # list of kernels per layer, [1,2,3]\n",
        "    self.conv_strides = conv_strides # stride for each filter [1, 2, 2], note: 2 means you are downsampling the data in half\n",
        "    self.latent_space_dim = latent_space_dim # how many neurons on bottleneck\n",
        "    self.reconstruction_loss_weight = 1000000\n",
        "\n",
        "    self.encoder = None\n",
        "    self.decoder = None\n",
        "    self.model = None\n",
        "\n",
        "    self._num_conv_layers = len(conv_filters)\n",
        "    self._shape_before_bottleneck = None\n",
        "    self._model_input = None\n",
        "\n",
        "    self._build()\n",
        "\n",
        "  def summary(self):\n",
        "    self.encoder.summary()\n",
        "    print(\"\\n\")\n",
        "    self.decoder.summary()\n",
        "    print(\"\\n\")\n",
        "    self.model.summary()\n",
        "\n",
        "  def _build(self):\n",
        "    self._build_encoder()\n",
        "    self._build_decoder()\n",
        "    self._build_autoencoder()\n",
        "\n",
        "  def compile(self, learning_rate=0.0001):\n",
        "    optimizer = Adam(learning_rate=learning_rate)\n",
        "    self.model.compile(optimizer=optimizer, loss=self._calculate_combined_loss,\n",
        "                       metrics=[self._calculate_reconstruction_loss,\n",
        "                                self._calculate_kl_loss])\n",
        "\n",
        "  def train(self, x_train, batch_size, num_epochs):\n",
        "   # checkpoint = ModelCheckpoint(\"best_model.hdf5\", monitor='loss', verbose=1,\n",
        "   #                              save_best_only=True, mode='auto', period=1)\n",
        "    self.model.fit(x_train,\n",
        "                   x_train,\n",
        "                   batch_size=batch_size,\n",
        "                   epochs=num_epochs,\n",
        "                   shuffle=True)\n",
        "                   #callbacks=[checkpoint])\n",
        "\n",
        "  def save(self, save_folder=\".\"):\n",
        "    self._create_folder_if_it_doesnt_exist(save_folder)\n",
        "    self._save_parameters(save_folder)\n",
        "    self._save_weights(save_folder)\n",
        "\n",
        "  def load_weights(self, weights_path):\n",
        "    self.model.load_weights(weights_path)\n",
        "\n",
        "  def reconstruct(self, spec):\n",
        "      latent_representations = self.encoder.predict(spec)\n",
        "      reconstructed_spec = self.decoder.predict(latent_representations)\n",
        "      return reconstructed_spec, latent_representations\n",
        "  \n",
        "  def encode(self, spec):\n",
        "    latent_representation = self.encoder.predict(spec)\n",
        "    return latent_representation\n",
        "  \n",
        "  def sample_from_latent_space(self, z):\n",
        "      z_vector = self.decoder.predict(z)\n",
        "      return z_vector\n",
        "\n",
        "  @classmethod\n",
        "  def load(cls, save_folder=\".\"):\n",
        "      parameters_path = os.path.join(save_folder, \"parameters.pkl\")\n",
        "      with open(parameters_path, \"rb\") as f:\n",
        "          parameters = pickle.load(f)\n",
        "      autoencoder = VAE(*parameters)\n",
        "      weights_path = os.path.join(save_folder, \"weights.h5\")\n",
        "      autoencoder.load_weights(weights_path)\n",
        "      return autoencoder\n",
        "\n",
        "  def _calculate_combined_loss(self, y_target, y_predicted):\n",
        "    reconstruction_loss = self._calculate_reconstruction_loss(y_target, y_predicted)\n",
        "    kl_loss = self._calculate_kl_loss(y_target, y_predicted)\n",
        "    combined_loss = self.reconstruction_loss_weight * reconstruction_loss + kl_loss\n",
        "    return combined_loss\n",
        "  \n",
        "  def _calculate_reconstruction_loss(self, y_target, y_predicted):\n",
        "    error = y_target - y_predicted\n",
        "    reconstruction_loss = K.mean(K.square(error), axis=[1, 2, 3])\n",
        "    return reconstruction_loss\n",
        "\n",
        "  def _calculate_kl_loss(self, y_target, y_predicted):\n",
        "    kl_loss = -0.5 * K.sum(1 + self.log_variance - K.square(self.mu) -\n",
        "                           K.exp(self.log_variance), axis =1)\n",
        "    return kl_loss\n",
        "\n",
        "  def _create_folder_if_it_doesnt_exist(self, folder):\n",
        "      if not os.path.exists(folder):\n",
        "          os.makedirs(folder)\n",
        "\n",
        "  def _save_parameters(self, save_folder):\n",
        "      parameters = [\n",
        "          self.input_shape,\n",
        "          self.conv_filters,\n",
        "          self.conv_kernels,\n",
        "          self.conv_strides,\n",
        "          self.latent_space_dim\n",
        "      ]\n",
        "      save_path = os.path.join(save_folder, \"parameters.pkl\")\n",
        "      with open(save_path, \"wb\") as f:\n",
        "          pickle.dump(parameters, f)\n",
        "\n",
        "  def _save_weights(self, save_folder):\n",
        "      save_path = os.path.join(save_folder, \"weights.h5\")\n",
        "      self.model.save_weights(save_path)\n",
        "\n",
        "#-----------AUTOENCODER----------#\n",
        "\n",
        "  def _build_autoencoder(self):\n",
        "    model_input = self._model_input\n",
        "    model_output = self.decoder(self.encoder(model_input))\n",
        "    self.model = Model(model_input, model_output, name=\"autoencoder\")\n",
        "\n",
        "#--------------DECODER------------#\n",
        "\n",
        "  def _build_decoder(self):\n",
        "    decoder_input = self._add_decoder_input()\n",
        "    dense_layer = self._add_dense_layer(decoder_input)\n",
        "    reshape_layer = self._add_reshape_layer(dense_layer)\n",
        "    conv_transpose_layers = self._add_conv_transpose_layers(reshape_layer)\n",
        "    decoder_output = self._add_decoder_output(conv_transpose_layers)\n",
        "    self.decoder = Model(decoder_input, decoder_output, name=\"decoder\")\n",
        "\n",
        "  def _add_decoder_input(self):\n",
        "    return Input(shape=self.latent_space_dim, name=\"decoder_input\")\n",
        "\n",
        "  def _add_dense_layer(self, decoder_input):\n",
        "    num_neurons = np.prod(self._shape_before_bottleneck) # [ 1, 2, 4] -> 8\n",
        "    dense_layer = Dense(num_neurons, name=\"decoder_dense\")(decoder_input)\n",
        "    return dense_layer\n",
        "\n",
        "  def _add_reshape_layer(self, dense_layer):\n",
        "    return Reshape(self._shape_before_bottleneck)(dense_layer)\n",
        "\n",
        "  def _add_conv_transpose_layers(self, x):\n",
        "    \"\"\"Add conv transpose blocks.\"\"\"\n",
        "    # Loop through all the conv layers in reverse order and\n",
        "    # stop at the first layer\n",
        "    for layer_index in reversed(range(1, self._num_conv_layers)):\n",
        "      x = self._add_conv_transpose_layer(layer_index, x)\n",
        "    return x\n",
        "\n",
        "  def _add_conv_transpose_layer(self, layer_index, x):\n",
        "    layer_num = self._num_conv_layers - layer_index\n",
        "    conv_transpose_layer = Conv2DTranspose(\n",
        "        filters=self.conv_filters[layer_index],\n",
        "        kernel_size = self.conv_kernels[layer_index],\n",
        "        strides = self.conv_strides[layer_index],\n",
        "        padding = \"same\",\n",
        "        name=f\"decoder_conv_transpose_layer_{layer_num}\"\n",
        "    )\n",
        "    x = conv_transpose_layer(x)\n",
        "    x = ReLU(name=f\"decoder_relu_{layer_num}\")(x)\n",
        "    x = BatchNormalization(name=f\"decoder_bn_{layer_num}\")(x)\n",
        "    return x\n",
        "\n",
        "  def _add_decoder_output(self, x):\n",
        "    conv_transpose_layer = Conv2DTranspose(\n",
        "        filters = 1,\n",
        "        kernel_size = self.conv_kernels[0],\n",
        "        strides = self.conv_strides[0],\n",
        "        padding = \"same\",\n",
        "        name=f\"decoder_conv_transpose_layer_{self._num_conv_layers}\"\n",
        "    )\n",
        "    x = conv_transpose_layer(x)\n",
        "    output_layer = Activation(\"sigmoid\", name=\"sigmoid_output_layer\")(x)\n",
        "    return output_layer\n",
        "\n",
        "#----------------ENCODER-----------------#\n",
        "\n",
        "  def _build_encoder(self):\n",
        "    encoder_input = self._add_encoder_input()\n",
        "    conv_layers = self._add_conv_layers(encoder_input)\n",
        "    bottleneck =  self._add_bottleneck(conv_layers)\n",
        "    self._model_input = encoder_input\n",
        "    self.encoder = Model(encoder_input, bottleneck, name=\"encoder\")\n",
        "\n",
        "  def _add_encoder_input(self):\n",
        "    return Input(shape=self.input_shape, name=\"encoder_input\")\n",
        "\n",
        "  def _add_conv_layers(self, encoder_input):\n",
        "    \"\"\"Creates all convolutional blocks in encoder\"\"\"\n",
        "    x = encoder_input\n",
        "    for layer_index in range(self._num_conv_layers):\n",
        "      x = self._add_conv_layer(layer_index, x)\n",
        "    return x\n",
        "  \n",
        "  def _add_conv_layer(self, layer_index, x):\n",
        "    \"\"\"Adds a convolutional block to a graph of layers, consisting\n",
        "    of Conv 2d + ReLu activation + batch normalization.\n",
        "    \"\"\"\n",
        "    layer_number = layer_index + 1\n",
        "    conv_layer = Conv2D(\n",
        "        filters= self.conv_filters[layer_index],\n",
        "        kernel_size = self.conv_kernels[layer_index],\n",
        "        strides = self.conv_strides[layer_index],\n",
        "        padding = \"same\",\n",
        "        name = f\"encoder_conv_layer_{layer_number}\"\n",
        "    )\n",
        "    x = conv_layer(x)\n",
        "    x = ReLU(name=f\"encoder_relu_{layer_number}\")(x)\n",
        "    x = BatchNormalization(name=f\"encoder_bn_{layer_number}\")(x)\n",
        "    return x\n",
        "\n",
        "#-------------Bottleneck (Latent Space)-------------#\n",
        "\n",
        "  def _add_bottleneck(self, x):\n",
        "    \"\"\"Flatten data and add bottleneck with Gaussian sampling (Dense layer)\"\"\"\n",
        "    self._shape_before_bottleneck = K.int_shape(x)[1:]\n",
        "    x = Flatten()(x)\n",
        "    self.mu = Dense(self.latent_space_dim,name=\"mu\")(x)\n",
        "    self.log_variance = Dense(self.latent_space_dim,\n",
        "                              name=\"log_variance\")(x)\n",
        "    \n",
        "    def sample_point_from_normal_distribution(args):\n",
        "      mu, log_variance = args\n",
        "      epsilon = K.random_normal(shape=K.shape(self.mu), mean=0., \n",
        "                                stddev=1.)\n",
        "      sampled_point = mu + K.exp(log_variance / 2) * epsilon\n",
        "\n",
        "      return sampled_point\n",
        "\n",
        "    x = Lambda(sample_point_from_normal_distribution, \n",
        "               name=\"encoder_output\")([self.mu, self.log_variance])\n",
        "    return x\n",
        "\n",
        "print(\"VAE successfully built\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BRJI_kQsHgXP",
        "cellView": "form"
      },
      "source": [
        "#@title Training functions\n",
        "\n",
        "def train(x_train, learning_rate, batch_size, epochs): \n",
        "  vae = VAE(\n",
        "      input_shape = (hop, shape*spec_split, 1),\n",
        "      conv_filters=(512, 256, 128, 64, 32),\n",
        "      conv_kernels=(3, 3, 3, 3, 3),\n",
        "      conv_strides=(2, 2, 2, 2, (2,1)),\n",
        "      latent_space_dim = VECTOR_DIM\n",
        "  )\n",
        "  vae.summary()\n",
        "  vae.compile(learning_rate)\n",
        "  vae.train(x_train, batch_size, epochs)\n",
        "  return vae\n",
        "\n",
        "def train_tfdata(x_train, learning_rate, epochs=10): \n",
        "  vae = VAE(\n",
        "      input_shape = (hop, 3*shape, 1),\n",
        "      conv_filters=(512, 256, 128, 64, 32),\n",
        "      conv_kernels=(3, 3, 3, 3, 3),\n",
        "      conv_strides=(2, 2, 2, 2, (2,1)),\n",
        "      latent_space_dim = VECTOR_DIM\n",
        "  )\n",
        "  vae.summary()\n",
        "  vae.compile(learning_rate)\n",
        "  vae.train(x_train, num_epochs=epochs)\n",
        "  return vae\n",
        "\n",
        "def continue_training(checkpoint):\n",
        "  vae = VAE.load(checkpoint)\n",
        "  vae.summary()\n",
        "  vae.compile(LEARNING_RATE)\n",
        "  vae.train(adata,BATCH_SIZE,EPOCHS)\n",
        "  return vae\n",
        "\n",
        "def load_model(checkpoint):\n",
        "  vae = VAE.load(checkpoint)\n",
        "  vae.summary()\n",
        "  vae.compile(LEARNING_RATE)\n",
        "  return vae\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7yzHZt-PJc65",
        "cellView": "form"
      },
      "source": [
        "#@title Start training from scratch or resume training\n",
        "\n",
        "training_run_name = \"my_melspecvae_model\" #@param {type:\"string\"}\n",
        "checkpoint_save_directory = \"/path/to/your/checkpoints/\" #@param {type:\"string\"}\n",
        "resume_training = False #@param {type:\"boolean\"}\n",
        "resume_training_checkpoint_path = \"/path/to/your/checkpoints/\" #@param {type:\"string\"}\n",
        "\n",
        "current_time = get_time_stamp()\n",
        "\n",
        "if not resume_training:\n",
        "  vae = train(adata, LEARNING_RATE, BATCH_SIZE, EPOCHS)\n",
        " #vae = train_tfdata(dsa, LEARNING_RATE, EPOCHS)\n",
        "  vae.save(f\"{checkpoint_save_directory}{training_run_name}_{current_time}_h{hop}_w{shape}_z{VECTOR_DIM}\")\n",
        "else:\n",
        "  vae = continue_training(resume_training_checkpoint_path)\n",
        "  vae.save(f\"{checkpoint_save_directory}{training_run_name}_{current_time}_h{hop}_w{shape}_z{VECTOR_DIM}\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kWqhd1LILIxD"
      },
      "source": [
        "## Generation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ceWvu0rMPCGx"
      },
      "source": [
        "#@title Build VAE Neural Network\n",
        "\n",
        "#VAE\n",
        "\n",
        "import os\n",
        "import pickle\n",
        "\n",
        "from tensorflow.keras import Model\n",
        "from tensorflow.keras.layers import Input, Conv2D, ReLU, BatchNormalization, Flatten, Dense, Reshape, Conv2DTranspose, Activation, Lambda\n",
        "from tensorflow.keras import backend as K\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.losses import MeanSquaredError\n",
        "from tensorflow.keras.callbacks import ModelCheckpoint\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "tf.compat.v1.disable_eager_execution()\n",
        "\n",
        "\n",
        "class VAE:\n",
        "  \"\"\"\n",
        "  VAE represents a Deep Convolutional autoencoder architecture\n",
        "  with mirrored encoder and decoder components.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               input_shape, #shape of the input data\n",
        "               conv_filters, #convolutional network filters\n",
        "               conv_kernels, #convNet kernel size\n",
        "               conv_strides, #convNet strides\n",
        "               latent_space_dim):\n",
        "    self.input_shape = input_shape # [28, 28, 1], in this case is 28 x 28 pixels on 1 channel for greyscale\n",
        "    self.conv_filters = conv_filters # is a list for each layer, i.e. [2, 4, 8]\n",
        "    self.conv_kernels = conv_kernels # list of kernels per layer, [1,2,3]\n",
        "    self.conv_strides = conv_strides # stride for each filter [1, 2, 2], note: 2 means you are downsampling the data in half\n",
        "    self.latent_space_dim = latent_space_dim # how many neurons on bottleneck\n",
        "    self.reconstruction_loss_weight = 1000000\n",
        "\n",
        "    self.encoder = None\n",
        "    self.decoder = None\n",
        "    self.model = None\n",
        "\n",
        "    self._num_conv_layers = len(conv_filters)\n",
        "    self._shape_before_bottleneck = None\n",
        "    self._model_input = None\n",
        "\n",
        "    self._build()\n",
        "\n",
        "  def summary(self):\n",
        "    self.encoder.summary()\n",
        "    print(\"\\n\")\n",
        "    self.decoder.summary()\n",
        "    print(\"\\n\")\n",
        "    self.model.summary()\n",
        "\n",
        "  def _build(self):\n",
        "    self._build_encoder()\n",
        "    self._build_decoder()\n",
        "    self._build_autoencoder()\n",
        "\n",
        "  def compile(self, learning_rate=0.0001):\n",
        "    optimizer = Adam(learning_rate=learning_rate)\n",
        "    self.model.compile(optimizer=optimizer, loss=self._calculate_combined_loss,\n",
        "                       metrics=[self._calculate_reconstruction_loss,\n",
        "                                self._calculate_kl_loss])\n",
        "\n",
        "  def train(self, x_train, batch_size, num_epochs):\n",
        "   # checkpoint = ModelCheckpoint(\"best_model.hdf5\", monitor='loss', verbose=1,\n",
        "   #                              save_best_only=True, mode='auto', period=1)\n",
        "    self.model.fit(x_train,\n",
        "                   x_train,\n",
        "                   batch_size=batch_size,\n",
        "                   epochs=num_epochs,\n",
        "                   shuffle=True)\n",
        "                   #callbacks=[checkpoint])\n",
        "\n",
        "  def save(self, save_folder=\".\"):\n",
        "    self._create_folder_if_it_doesnt_exist(save_folder)\n",
        "    self._save_parameters(save_folder)\n",
        "    self._save_weights(save_folder)\n",
        "\n",
        "  def load_weights(self, weights_path):\n",
        "    self.model.load_weights(weights_path)\n",
        "\n",
        "  def reconstruct(self, spec):\n",
        "      latent_representations = self.encoder.predict(spec)\n",
        "      reconstructed_spec = self.decoder.predict(latent_representations)\n",
        "      return reconstructed_spec, latent_representations\n",
        "  \n",
        "  def encode(self, spec):\n",
        "    latent_representation = self.encoder.predict(spec)\n",
        "    return latent_representation\n",
        "  \n",
        "  def sample_from_latent_space(self, z):\n",
        "      z_vector = self.decoder.predict(z)\n",
        "      return z_vector\n",
        "\n",
        "  @classmethod\n",
        "  def load(cls, save_folder=\".\"):\n",
        "      parameters_path = os.path.join(save_folder, \"parameters.pkl\")\n",
        "      with open(parameters_path, \"rb\") as f:\n",
        "          parameters = pickle.load(f)\n",
        "      autoencoder = VAE(*parameters)\n",
        "      weights_path = os.path.join(save_folder, \"weights.h5\")\n",
        "      autoencoder.load_weights(weights_path)\n",
        "      return autoencoder\n",
        "\n",
        "  def _calculate_combined_loss(self, y_target, y_predicted):\n",
        "    reconstruction_loss = self._calculate_reconstruction_loss(y_target, y_predicted)\n",
        "    kl_loss = self._calculate_kl_loss(y_target, y_predicted)\n",
        "    combined_loss = self.reconstruction_loss_weight * reconstruction_loss + kl_loss\n",
        "    return combined_loss\n",
        "  \n",
        "  def _calculate_reconstruction_loss(self, y_target, y_predicted):\n",
        "    error = y_target - y_predicted\n",
        "    reconstruction_loss = K.mean(K.square(error), axis=[1, 2, 3])\n",
        "    return reconstruction_loss\n",
        "\n",
        "  def _calculate_kl_loss(self, y_target, y_predicted):\n",
        "    kl_loss = -0.5 * K.sum(1 + self.log_variance - K.square(self.mu) -\n",
        "                           K.exp(self.log_variance), axis =1)\n",
        "    return kl_loss\n",
        "\n",
        "  def _create_folder_if_it_doesnt_exist(self, folder):\n",
        "      if not os.path.exists(folder):\n",
        "          os.makedirs(folder)\n",
        "\n",
        "  def _save_parameters(self, save_folder):\n",
        "      parameters = [\n",
        "          self.input_shape,\n",
        "          self.conv_filters,\n",
        "          self.conv_kernels,\n",
        "          self.conv_strides,\n",
        "          self.latent_space_dim\n",
        "      ]\n",
        "      save_path = os.path.join(save_folder, \"parameters.pkl\")\n",
        "      with open(save_path, \"wb\") as f:\n",
        "          pickle.dump(parameters, f)\n",
        "\n",
        "  def _save_weights(self, save_folder):\n",
        "      save_path = os.path.join(save_folder, \"weights.h5\")\n",
        "      self.model.save_weights(save_path)\n",
        "\n",
        "#-----------AUTOENCODER----------#\n",
        "\n",
        "  def _build_autoencoder(self):\n",
        "    model_input = self._model_input\n",
        "    model_output = self.decoder(self.encoder(model_input))\n",
        "    self.model = Model(model_input, model_output, name=\"autoencoder\")\n",
        "\n",
        "#--------------DECODER------------#\n",
        "\n",
        "  def _build_decoder(self):\n",
        "    decoder_input = self._add_decoder_input()\n",
        "    dense_layer = self._add_dense_layer(decoder_input)\n",
        "    reshape_layer = self._add_reshape_layer(dense_layer)\n",
        "    conv_transpose_layers = self._add_conv_transpose_layers(reshape_layer)\n",
        "    decoder_output = self._add_decoder_output(conv_transpose_layers)\n",
        "    self.decoder = Model(decoder_input, decoder_output, name=\"decoder\")\n",
        "\n",
        "  def _add_decoder_input(self):\n",
        "    return Input(shape=self.latent_space_dim, name=\"decoder_input\")\n",
        "\n",
        "  def _add_dense_layer(self, decoder_input):\n",
        "    num_neurons = np.prod(self._shape_before_bottleneck) # [ 1, 2, 4] -> 8\n",
        "    dense_layer = Dense(num_neurons, name=\"decoder_dense\")(decoder_input)\n",
        "    return dense_layer\n",
        "\n",
        "  def _add_reshape_layer(self, dense_layer):\n",
        "    return Reshape(self._shape_before_bottleneck)(dense_layer)\n",
        "\n",
        "  def _add_conv_transpose_layers(self, x):\n",
        "    \"\"\"Add conv transpose blocks.\"\"\"\n",
        "    # Loop through all the conv layers in reverse order and\n",
        "    # stop at the first layer\n",
        "    for layer_index in reversed(range(1, self._num_conv_layers)):\n",
        "      x = self._add_conv_transpose_layer(layer_index, x)\n",
        "    return x\n",
        "\n",
        "  def _add_conv_transpose_layer(self, layer_index, x):\n",
        "    layer_num = self._num_conv_layers - layer_index\n",
        "    conv_transpose_layer = Conv2DTranspose(\n",
        "        filters=self.conv_filters[layer_index],\n",
        "        kernel_size = self.conv_kernels[layer_index],\n",
        "        strides = self.conv_strides[layer_index],\n",
        "        padding = \"same\",\n",
        "        name=f\"decoder_conv_transpose_layer_{layer_num}\"\n",
        "    )\n",
        "    x = conv_transpose_layer(x)\n",
        "    x = ReLU(name=f\"decoder_relu_{layer_num}\")(x)\n",
        "    x = BatchNormalization(name=f\"decoder_bn_{layer_num}\")(x)\n",
        "    return x\n",
        "\n",
        "  def _add_decoder_output(self, x):\n",
        "    conv_transpose_layer = Conv2DTranspose(\n",
        "        filters = 1,\n",
        "        kernel_size = self.conv_kernels[0],\n",
        "        strides = self.conv_strides[0],\n",
        "        padding = \"same\",\n",
        "        name=f\"decoder_conv_transpose_layer_{self._num_conv_layers}\"\n",
        "    )\n",
        "    x = conv_transpose_layer(x)\n",
        "    output_layer = Activation(\"sigmoid\", name=\"sigmoid_output_layer\")(x)\n",
        "    return output_layer\n",
        "\n",
        "#----------------ENCODER-----------------#\n",
        "\n",
        "  def _build_encoder(self):\n",
        "    encoder_input = self._add_encoder_input()\n",
        "    conv_layers = self._add_conv_layers(encoder_input)\n",
        "    bottleneck =  self._add_bottleneck(conv_layers)\n",
        "    self._model_input = encoder_input\n",
        "    self.encoder = Model(encoder_input, bottleneck, name=\"encoder\")\n",
        "\n",
        "  def _add_encoder_input(self):\n",
        "    return Input(shape=self.input_shape, name=\"encoder_input\")\n",
        "\n",
        "  def _add_conv_layers(self, encoder_input):\n",
        "    \"\"\"Creates all convolutional blocks in encoder\"\"\"\n",
        "    x = encoder_input\n",
        "    for layer_index in range(self._num_conv_layers):\n",
        "      x = self._add_conv_layer(layer_index, x)\n",
        "    return x\n",
        "  \n",
        "  def _add_conv_layer(self, layer_index, x):\n",
        "    \"\"\"Adds a convolutional block to a graph of layers, consisting\n",
        "    of Conv 2d + ReLu activation + batch normalization.\n",
        "    \"\"\"\n",
        "    layer_number = layer_index + 1\n",
        "    conv_layer = Conv2D(\n",
        "        filters= self.conv_filters[layer_index],\n",
        "        kernel_size = self.conv_kernels[layer_index],\n",
        "        strides = self.conv_strides[layer_index],\n",
        "        padding = \"same\",\n",
        "        name = f\"encoder_conv_layer_{layer_number}\"\n",
        "    )\n",
        "    x = conv_layer(x)\n",
        "    x = ReLU(name=f\"encoder_relu_{layer_number}\")(x)\n",
        "    x = BatchNormalization(name=f\"encoder_bn_{layer_number}\")(x)\n",
        "    return x\n",
        "\n",
        "#-------------Bottleneck (Latent Space)-------------#\n",
        "\n",
        "  def _add_bottleneck(self, x):\n",
        "    \"\"\"Flatten data and add bottleneck with Gaussian sampling (Dense layer)\"\"\"\n",
        "    self._shape_before_bottleneck = K.int_shape(x)[1:]\n",
        "    x = Flatten()(x)\n",
        "    self.mu = Dense(self.latent_space_dim,name=\"mu\")(x)\n",
        "    self.log_variance = Dense(self.latent_space_dim,\n",
        "                              name=\"log_variance\")(x)\n",
        "    \n",
        "    def sample_point_from_normal_distribution(args):\n",
        "      mu, log_variance = args\n",
        "      epsilon = K.random_normal(shape=K.shape(self.mu), mean=0., \n",
        "                                stddev=1.)\n",
        "      sampled_point = mu + K.exp(log_variance / 2) * epsilon\n",
        "\n",
        "      return sampled_point\n",
        "\n",
        "    x = Lambda(sample_point_from_normal_distribution, \n",
        "               name=\"encoder_output\")([self.mu, self.log_variance])\n",
        "    return x\n",
        "\n",
        "print(\"VAE successfully built\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qe62172WPEYD",
        "cellView": "form"
      },
      "source": [
        "#@title Load Checkpoint for Generating\n",
        "\n",
        "checkpoint_load_directory = \"/path/to/your/checkpoint\" #@param {type:\"string\"}\n",
        "\n",
        "#-------LOAD MODEL FOR GENERATING-------------#\n",
        "\n",
        "vae = VAE.load(checkpoint_load_directory)\n",
        "print(\"Loaded checkpoint\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K-pa9kR6Pe8B",
        "cellView": "form"
      },
      "source": [
        "#@title Import synthesis utility functions\n",
        "\n",
        "#-----TESTING FUNCTIONS ----------- #\n",
        "\n",
        "def select_spec(spec, labels, num_spec=10):\n",
        "    sample_spec_index = np.random.choice(range(len(spec)), num_spec)\n",
        "    sample_spec = spec[sample_spec_index]\n",
        "    sample_labels = labels[sample_spec_index]\n",
        "    return sample_spec, sample_labels\n",
        "\n",
        "\n",
        "def plot_reconstructed_spec(spec, reconstructed_spec):\n",
        "    fig = plt.figure(figsize=(15, 3))\n",
        "    num_spec = len(spec)\n",
        "    for i, (image, reconstructed_image) in enumerate(zip(spec, reconstructed_spec)):\n",
        "        image = image.squeeze()\n",
        "        ax = fig.add_subplot(2, num_spec, i + 1)\n",
        "        ax.axis(\"off\")\n",
        "        ax.imshow(image, cmap=\"gray_r\")\n",
        "        reconstructed_image = reconstructed_image.squeeze()\n",
        "        ax = fig.add_subplot(2, num_spec, i + num_spec + 1)\n",
        "        ax.axis(\"off\")\n",
        "        ax.imshow(reconstructed_image, cmap=\"gray_r\")\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def plot_spec_encoded_in_latent_space(latent_representations, sample_labels):\n",
        "    plt.figure(figsize=(10, 10))\n",
        "    plt.scatter(latent_representations[:, 0],\n",
        "                latent_representations[:, 1],\n",
        "                cmap=\"rainbow\",\n",
        "                c=sample_labels,\n",
        "                alpha=0.5,\n",
        "                s=2)\n",
        "    plt.colorbar()\n",
        "    plt.show()\n",
        "\n",
        "#---------------NOISE GENERATOR FUNCTIONS ------------#\n",
        "\n",
        "def generate_random_z_vect(seed=1001,size_z=1,scale=1.0):\n",
        "    np.random.seed(seed)\n",
        "    x = np.random.uniform(low=(scale * -1.0), high=scale, size=(size_z,VECTOR_DIM))\n",
        "    return x\n",
        "\n",
        "def generate_z_vect_from_perlin_noise(seed=1001, size_z=1, scale=1.0):\n",
        "    np.random.seed(seed)\n",
        "    x = generate_perlin_noise_2d((size_z, VECTOR_DIM), (1,1))\n",
        "    x = x*scale\n",
        "    return x\n",
        "\n",
        "def generate_z_vect_from_fractal_noise(seed=1001, size_z=1, scale=1.0,):\n",
        "    np.random.seed(seed)\n",
        "    x = generate_fractal_noise_2d((size_z, VECTOR_DIM), (1,1),)\n",
        "    x = x*scale\n",
        "    return x\n",
        "\n",
        "\n",
        "#-------SPECTROGRAM AND SOUND SYNTHESIS UTILITY FUNCTIONS -------- #\n",
        "\n",
        "#Assembling generated Spectrogram chunks into final Spectrogram\n",
        "def specass(a,spec):\n",
        "  but=False\n",
        "  con = np.array([])\n",
        "  nim = a.shape[0]\n",
        "  for i in range(nim-1):\n",
        "    im = a[i]\n",
        "    im = np.squeeze(im)\n",
        "    if not but:\n",
        "      con=im\n",
        "      but=True\n",
        "    else:\n",
        "      con = np.concatenate((con,im), axis=1)\n",
        "  diff = spec.shape[1]-(nim*shape)\n",
        "  a = np.squeeze(a)\n",
        "  con = np.concatenate((con,a[-1,:,-diff:]), axis=1)\n",
        "  return np.squeeze(con)\n",
        "\n",
        "#Splitting input spectrogram into different chunks to feed to the generator\n",
        "def chopspec(spec):\n",
        "  dsa=[]\n",
        "  for i in range(spec.shape[1]//shape):\n",
        "    im = spec[:,i*shape:i*shape+shape]\n",
        "    im = np.reshape(im, (im.shape[0],im.shape[1],1))\n",
        "    dsa.append(im)\n",
        "  imlast = spec[:,-shape:]\n",
        "  imlast = np.reshape(imlast, (imlast.shape[0],imlast.shape[1],1))\n",
        "  dsa.append(imlast)\n",
        "  return np.array(dsa, dtype=np.float32)\n",
        "\n",
        "#Converting from source Spectrogram to target Spectrogram\n",
        "def towave_reconstruct(spec, spec1, name, path='../content/', show=False, save=False):\n",
        "  specarr = chopspec(spec)\n",
        "  specarr1 = chopspec(spec1)\n",
        "  print(specarr.shape)\n",
        "  a = specarr\n",
        "  print('Generating...')\n",
        "  ab = specarr1\n",
        "  print('Assembling and Converting...')\n",
        "  a = specass(a,spec)\n",
        "  ab = specass(ab,spec1)\n",
        "  awv = deprep(a)\n",
        "  abwv = deprep(ab)\n",
        "  if save:\n",
        "    print('Saving...')\n",
        "    pathfin = f'{path}/{name}'\n",
        "    sf.write(f'{pathfin}.wav', awv, sr)\n",
        "    print('Saved WAV!')\n",
        "  IPython.display.display(IPython.display.Audio(np.squeeze(abwv), rate=sr))\n",
        "  IPython.display.display(IPython.display.Audio(np.squeeze(awv), rate=sr))\n",
        "  if show:\n",
        "    fig, axs = plt.subplots(ncols=2)\n",
        "    axs[0].imshow(np.flip(a, -2), cmap=None)\n",
        "    axs[0].axis('off')\n",
        "    axs[0].set_title('Reconstructed')\n",
        "    axs[1].imshow(np.flip(ab, -2), cmap=None)\n",
        "    axs[1].axis('off')\n",
        "    axs[1].set_title('Input')\n",
        "    plt.show()\n",
        "  return abwv\n",
        "\n",
        "#Converting from Z vector generated spectrogram to waveform\n",
        "def towave_from_z(spec, name, path='../content/', show=False, save=False):\n",
        "  specarr = chopspec(spec)\n",
        "  print(specarr.shape)\n",
        "  a = specarr\n",
        "  print('Generating...')\n",
        "  print('Assembling and Converting...')\n",
        "  a = specass(a,spec)\n",
        "  awv = deprep(a)\n",
        "  if save:\n",
        "    print('Saving...')\n",
        "    pathfin = f'{path}/{name}'\n",
        "    sf.write(f'{pathfin}.wav', awv, sr)\n",
        "    print('Saved WAV!')\n",
        "  IPython.display.display(IPython.display.Audio(np.squeeze(awv), rate=sr))\n",
        "  if show:\n",
        "    fig, axs = plt.subplots(ncols=1)\n",
        "    axs.imshow(np.flip(a, -2), cmap=None)\n",
        "    axs.axis('off')\n",
        "    axs.set_title('Decoder Synthesis')\n",
        "    plt.show()\n",
        "  return awv\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5HfNl_oQQDQ-",
        "cellView": "form"
      },
      "source": [
        "#@title Compare resynthesized MelSpec with ground truth MelSpec. Note: Only works if you imported an audio dataset\n",
        "\n",
        "num_spec_to_resynthesize =  5 #@param {type:\"integer\"}\n",
        "\n",
        "num_sample_spec_to_show = num_spec_to_resynthesize\n",
        "sample_spec, _ = select_spec(adata, adata, num_sample_spec_to_show)\n",
        "reconstructed_spec, _ = vae.reconstruct(sample_spec)\n",
        "plot_reconstructed_spec(sample_spec, reconstructed_spec)\n",
        "\n",
        "reconst = num_sample_spec_to_show\n",
        "\n",
        "for i in range(reconst):\n",
        "  y = towave_reconstruct(reconstructed_spec[i],sample_spec[i],name='reconstructions',show=True, save=False)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cwHuwGDb0zvD",
        "cellView": "form"
      },
      "source": [
        "#@title Generate one-shot samples from latent space with random or manual seed\n",
        "num_samples_to_generate =   10#@param {type:\"integer\"}\n",
        "use_seed = False #@param {type:\"boolean\"}\n",
        "seed =  0 #@param {type:\"slider\", min:0, max:4294967295, step:1}\n",
        "scale_z_vectors =  1.5 #@param {type:\"slider\", min:-5.0, max:5.0, step:0.1}\n",
        "save_audio = False #@param {type:\"boolean\"}\n",
        "audio_name = \"one_shot\" #@param {type:\"string\"}\n",
        "audio_save_directory = \"/content/\" #@param {type:\"string\"}\n",
        "\n",
        "y = np.random.randint(0, 2**32-1)                          # generated random int to pass and convert into vector\n",
        "i=0\n",
        "while i < num_samples_to_generate:\n",
        "  if not use_seed:\n",
        "    z = generate_random_z_vect(y, num_samples_to_generate,scale=scale_z_vectors)  \n",
        "  else:\n",
        "    z = generate_random_z_vect(seed, num_samples_to_generate,scale=scale_z_vectors)\n",
        "  z_sample = np.array(vae.sample_from_latent_space(z))\n",
        "  towave_from_z(z_sample[i], name=f'{audio_name}_{i}',path=audio_save_directory,show=True, save=save_audio)\n",
        "  i+=1\n",
        "\n",
        "if not use_seed:\n",
        "  print(\"Generated from seed:\", y)\n",
        "else:\n",
        "  print(\"Generated from seed:\", seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8akqOemGJTIG",
        "cellView": "form"
      },
      "source": [
        "#@title Generate arbitrary long audio from latent space with random or custom seed using uniform, Perlin or fractal noise\n",
        "num_seeds_to_generate = 32#@param {type:\"integer\"}\n",
        "noise_type = \"uniform\" #@param [\"uniform\", \"perlin\", \"fractal\"]\n",
        "use_seed = False #@param {type:\"boolean\"}\n",
        "seed =  0 #@param {type:\"slider\", min:0, max:4294967295, step:1}\n",
        "scale_z_vectors =  1.5 #@param {type:\"slider\", min:-5.0, max:5.0, step:0.1}\n",
        "save_audio = False #@param {type:\"boolean\"}\n",
        "audio_name = \"VAE_synthesis2\" #@param {type:\"string\"}\n",
        "audio_save_directory = \"/content\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "y = np.random.randint(0, 2**32-1)                         # generated random int to pass and convert into vector\n",
        "if not use_seed:\n",
        "  if noise_type == \"uniform\":\n",
        "    z = generate_random_z_vect(y, num_seeds_to_generate,scale_z_vectors)            # vectors to input into latent space\n",
        "  if noise_type == \"perlin\":\n",
        "    z = generate_z_vect_from_perlin_noise(y, num_seeds_to_generate,scale_z_vectors)            # vectors to input into latent space\n",
        "  if noise_type == \"fractal\":\n",
        "    z = generate_z_vect_from_fractal_noise(y, num_seeds_to_generate,scale_z_vectors)            # vectors to input into latent space\n",
        "if use_seed:\n",
        "  if noise_type == \"uniform\":\n",
        "    z = generate_random_z_vect(seed, num_seeds_to_generate,scale_z_vectors)            # vectors to input into latent space\n",
        "  if noise_type == \"perlin\":\n",
        "    z = generate_z_vect_from_perlin_noise(seed, num_seeds_to_generate,scale_z_vectors)            # vectors to input into latent space\n",
        "  if noise_type == \"fractal\":\n",
        "    z = generate_z_vect_from_fractal_noise(seed, num_seeds_to_generate,scale_z_vectors)            # vectors to input into latent space\n",
        "z_sample = np.array(vae.sample_from_latent_space(z))\n",
        "assembled_spec = testass(z_sample)\n",
        "towave_from_z(assembled_spec,audio_name,audio_save_directory,show=True,save=save_audio)\n",
        "\n",
        "if not use_seed:\n",
        "  print(\"Generated from seed:\", y)\n",
        "else:\n",
        "  print(\"Generated from seed:\", seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "dcnq9i8UcC4g"
      },
      "source": [
        "#@title Interpolate between two seeds for n-amount of steps\n",
        "\n",
        "use_seed = False #@param {type:\"boolean\"}\n",
        "seed_a =  0 #@param {type:\"slider\", min:0, max:4294967295, step:1}\n",
        "seed_b =  4294967295 #@param {type:\"slider\", min:0, max:4294967295, step:1}\n",
        "num_interpolation_steps =   8#@param {type:\"integer\"}\n",
        "scale_z_vectors =  1.5 #@param {type:\"slider\", min:-5.0, max:5.0, step:0.1}\n",
        "scale_interpolation_ratio =  1 #@param {type:\"slider\", min:-5.0, max:5.0, step:0.1}\n",
        "save_audio = False #@param {type:\"boolean\"}\n",
        "audio_name = \"random_seeds_interpolation\" #@param {type:\"string\"}\n",
        "audio_save_directory = \"/content/\" #@param {type:\"string\"}\n",
        "\n",
        "# generate points in latent space as input for the generator\n",
        "def generate_latent_points(latent_dim, n_samples, n_classes=10):\n",
        "\t# generate points in the latent space\n",
        "\tx_input = randn(latent_dim * n_samples)\n",
        "\t# reshape into a batch of inputs for the network\n",
        "\tz_input = x_input.reshape(n_samples, latent_dim)\n",
        "\treturn z_input\n",
        " \n",
        "# uniform interpolation between two points in latent space\n",
        "def interpolate_points(p1, p2,scale, n_steps=10):\n",
        "\t# interpolate ratios between the points\n",
        "\tratios = linspace(-scale, scale, num=n_steps)\n",
        "\t# linear interpolate vectors\n",
        "\tvectors = list()\n",
        "\tfor ratio in ratios:\n",
        "\t\tv = (1.0 - ratio) * p1 + ratio * p2\n",
        "\t\tvectors.append(v)\n",
        "\treturn asarray(vectors)\n",
        "\n",
        "y = np.random.randint(0, 2**32-1)\n",
        "\n",
        "if not use_seed:\n",
        "  pts = generate_random_z_vect(y,10,scale_z_vectors)\n",
        "  interpolated = interpolate_points(pts[0], pts[1], scale_interpolation_ratio, num_interpolation_steps) # interpolate points in latent space\n",
        "else:\n",
        "  pts_a = generate_random_z_vect(seed_a,10,scale_z_vectors)\n",
        "  pts_b = generate_random_z_vect(seed_b,10,scale_z_vectors)\n",
        "  interpolated = interpolate_points(pts_a[0], pts_b[0], scale_interpolation_ratio, num_interpolation_steps) # interpolate points in latent space\n",
        "\n",
        "interp = np.array(vae.sample_from_latent_space(interpolated))\n",
        "assembled_spec = testass(interp)\n",
        "towave_from_z(assembled_spec,audio_name,audio_save_directory,show=True, save=save_audio)\n",
        "\n",
        "\n",
        "if not use_seed:\n",
        "  print(\"Generated from seed:\", y)\n",
        "else:\n",
        "  print(\"Generated from seed:\", seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## EXPERIMENTAL: Timbre Transfer\n",
        "### Note: You need to Restart Runtime (Runtime > Restart Runtime) per each Timbre Transfer"
      ],
      "metadata": {
        "id": "uTuKspchQjY0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Import wav\n",
        "\n",
        "input_audio = \"your_audio_to_timbre_transfer.wav\" #@param {type:\"string\"}\n",
        "\n",
        "audio_in = single_audio_array(input_audio)\n",
        "aspec = tospec(audio_in)\n",
        "out_data = splitcut(aspec)"
      ],
      "metadata": {
        "id": "Z4LwV0anOG6S",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "rp7lQDQFQuct"
      },
      "source": [
        "#@title Build VAE Neural Network\n",
        "\n",
        "#VAE\n",
        "\n",
        "import os\n",
        "import pickle\n",
        "\n",
        "from tensorflow.keras import Model\n",
        "from tensorflow.keras.layers import Input, Conv2D, ReLU, BatchNormalization, Flatten, Dense, Reshape, Conv2DTranspose, Activation, Lambda\n",
        "from tensorflow.keras import backend as K\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.losses import MeanSquaredError\n",
        "from tensorflow.keras.callbacks import ModelCheckpoint\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "tf.compat.v1.disable_eager_execution()\n",
        "\n",
        "\n",
        "class VAE:\n",
        "  \"\"\"\n",
        "  VAE represents a Deep Convolutional autoencoder architecture\n",
        "  with mirrored encoder and decoder components.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self,\n",
        "               input_shape, #shape of the input data\n",
        "               conv_filters, #convolutional network filters\n",
        "               conv_kernels, #convNet kernel size\n",
        "               conv_strides, #convNet strides\n",
        "               latent_space_dim):\n",
        "    self.input_shape = input_shape # [28, 28, 1], in this case is 28 x 28 pixels on 1 channel for greyscale\n",
        "    self.conv_filters = conv_filters # is a list for each layer, i.e. [2, 4, 8]\n",
        "    self.conv_kernels = conv_kernels # list of kernels per layer, [1,2,3]\n",
        "    self.conv_strides = conv_strides # stride for each filter [1, 2, 2], note: 2 means you are downsampling the data in half\n",
        "    self.latent_space_dim = latent_space_dim # how many neurons on bottleneck\n",
        "    self.reconstruction_loss_weight = 1000000\n",
        "\n",
        "    self.encoder = None\n",
        "    self.decoder = None\n",
        "    self.model = None\n",
        "\n",
        "    self._num_conv_layers = len(conv_filters)\n",
        "    self._shape_before_bottleneck = None\n",
        "    self._model_input = None\n",
        "\n",
        "    self._build()\n",
        "\n",
        "  def summary(self):\n",
        "    self.encoder.summary()\n",
        "    print(\"\\n\")\n",
        "    self.decoder.summary()\n",
        "    print(\"\\n\")\n",
        "    self.model.summary()\n",
        "\n",
        "  def _build(self):\n",
        "    self._build_encoder()\n",
        "    self._build_decoder()\n",
        "    self._build_autoencoder()\n",
        "\n",
        "  def compile(self, learning_rate=0.0001):\n",
        "    optimizer = Adam(learning_rate=learning_rate)\n",
        "    self.model.compile(optimizer=optimizer, loss=self._calculate_combined_loss,\n",
        "                       metrics=[self._calculate_reconstruction_loss,\n",
        "                                self._calculate_kl_loss])\n",
        "\n",
        "  def train(self, x_train, batch_size, num_epochs):\n",
        "   # checkpoint = ModelCheckpoint(\"best_model.hdf5\", monitor='loss', verbose=1,\n",
        "   #                              save_best_only=True, mode='auto', period=1)\n",
        "    self.model.fit(x_train,\n",
        "                   x_train,\n",
        "                   batch_size=batch_size,\n",
        "                   epochs=num_epochs,\n",
        "                   shuffle=True)\n",
        "                   #callbacks=[checkpoint])\n",
        "\n",
        "  def save(self, save_folder=\".\"):\n",
        "    self._create_folder_if_it_doesnt_exist(save_folder)\n",
        "    self._save_parameters(save_folder)\n",
        "    self._save_weights(save_folder)\n",
        "\n",
        "  def load_weights(self, weights_path):\n",
        "    self.model.load_weights(weights_path)\n",
        "\n",
        "  def reconstruct(self, spec):\n",
        "      latent_representations = self.encoder.predict(spec)\n",
        "      reconstructed_spec = self.decoder.predict(latent_representations)\n",
        "      return reconstructed_spec, latent_representations\n",
        "  \n",
        "  def encode(self, spec):\n",
        "    latent_representation = self.encoder.predict(spec)\n",
        "    return latent_representation\n",
        "  \n",
        "  def sample_from_latent_space(self, z):\n",
        "      z_vector = self.decoder.predict(z)\n",
        "      return z_vector\n",
        "\n",
        "  @classmethod\n",
        "  def load(cls, save_folder=\".\"):\n",
        "      parameters_path = os.path.join(save_folder, \"parameters.pkl\")\n",
        "      with open(parameters_path, \"rb\") as f:\n",
        "          parameters = pickle.load(f)\n",
        "      autoencoder = VAE(*parameters)\n",
        "      weights_path = os.path.join(save_folder, \"weights.h5\")\n",
        "      autoencoder.load_weights(weights_path)\n",
        "      return autoencoder\n",
        "\n",
        "  def _calculate_combined_loss(self, y_target, y_predicted):\n",
        "    reconstruction_loss = self._calculate_reconstruction_loss(y_target, y_predicted)\n",
        "    kl_loss = self._calculate_kl_loss(y_target, y_predicted)\n",
        "    combined_loss = self.reconstruction_loss_weight * reconstruction_loss + kl_loss\n",
        "    return combined_loss\n",
        "  \n",
        "  def _calculate_reconstruction_loss(self, y_target, y_predicted):\n",
        "    error = y_target - y_predicted\n",
        "    reconstruction_loss = K.mean(K.square(error), axis=[1, 2, 3])\n",
        "    return reconstruction_loss\n",
        "\n",
        "  def _calculate_kl_loss(self, y_target, y_predicted):\n",
        "    kl_loss = -0.5 * K.sum(1 + self.log_variance - K.square(self.mu) -\n",
        "                           K.exp(self.log_variance), axis =1)\n",
        "    return kl_loss\n",
        "\n",
        "  def _create_folder_if_it_doesnt_exist(self, folder):\n",
        "      if not os.path.exists(folder):\n",
        "          os.makedirs(folder)\n",
        "\n",
        "  def _save_parameters(self, save_folder):\n",
        "      parameters = [\n",
        "          self.input_shape,\n",
        "          self.conv_filters,\n",
        "          self.conv_kernels,\n",
        "          self.conv_strides,\n",
        "          self.latent_space_dim\n",
        "      ]\n",
        "      save_path = os.path.join(save_folder, \"parameters.pkl\")\n",
        "      with open(save_path, \"wb\") as f:\n",
        "          pickle.dump(parameters, f)\n",
        "\n",
        "  def _save_weights(self, save_folder):\n",
        "      save_path = os.path.join(save_folder, \"weights.h5\")\n",
        "      self.model.save_weights(save_path)\n",
        "\n",
        "#-----------AUTOENCODER----------#\n",
        "\n",
        "  def _build_autoencoder(self):\n",
        "    model_input = self._model_input\n",
        "    model_output = self.decoder(self.encoder(model_input))\n",
        "    self.model = Model(model_input, model_output, name=\"autoencoder\")\n",
        "\n",
        "#--------------DECODER------------#\n",
        "\n",
        "  def _build_decoder(self):\n",
        "    decoder_input = self._add_decoder_input()\n",
        "    dense_layer = self._add_dense_layer(decoder_input)\n",
        "    reshape_layer = self._add_reshape_layer(dense_layer)\n",
        "    conv_transpose_layers = self._add_conv_transpose_layers(reshape_layer)\n",
        "    decoder_output = self._add_decoder_output(conv_transpose_layers)\n",
        "    self.decoder = Model(decoder_input, decoder_output, name=\"decoder\")\n",
        "\n",
        "  def _add_decoder_input(self):\n",
        "    return Input(shape=self.latent_space_dim, name=\"decoder_input\")\n",
        "\n",
        "  def _add_dense_layer(self, decoder_input):\n",
        "    num_neurons = np.prod(self._shape_before_bottleneck) # [ 1, 2, 4] -> 8\n",
        "    dense_layer = Dense(num_neurons, name=\"decoder_dense\")(decoder_input)\n",
        "    return dense_layer\n",
        "\n",
        "  def _add_reshape_layer(self, dense_layer):\n",
        "    return Reshape(self._shape_before_bottleneck)(dense_layer)\n",
        "\n",
        "  def _add_conv_transpose_layers(self, x):\n",
        "    \"\"\"Add conv transpose blocks.\"\"\"\n",
        "    # Loop through all the conv layers in reverse order and\n",
        "    # stop at the first layer\n",
        "    for layer_index in reversed(range(1, self._num_conv_layers)):\n",
        "      x = self._add_conv_transpose_layer(layer_index, x)\n",
        "    return x\n",
        "\n",
        "  def _add_conv_transpose_layer(self, layer_index, x):\n",
        "    layer_num = self._num_conv_layers - layer_index\n",
        "    conv_transpose_layer = Conv2DTranspose(\n",
        "        filters=self.conv_filters[layer_index],\n",
        "        kernel_size = self.conv_kernels[layer_index],\n",
        "        strides = self.conv_strides[layer_index],\n",
        "        padding = \"same\",\n",
        "        name=f\"decoder_conv_transpose_layer_{layer_num}\"\n",
        "    )\n",
        "    x = conv_transpose_layer(x)\n",
        "    x = ReLU(name=f\"decoder_relu_{layer_num}\")(x)\n",
        "    x = BatchNormalization(name=f\"decoder_bn_{layer_num}\")(x)\n",
        "    return x\n",
        "\n",
        "  def _add_decoder_output(self, x):\n",
        "    conv_transpose_layer = Conv2DTranspose(\n",
        "        filters = 1,\n",
        "        kernel_size = self.conv_kernels[0],\n",
        "        strides = self.conv_strides[0],\n",
        "        padding = \"same\",\n",
        "        name=f\"decoder_conv_transpose_layer_{self._num_conv_layers}\"\n",
        "    )\n",
        "    x = conv_transpose_layer(x)\n",
        "    output_layer = Activation(\"sigmoid\", name=\"sigmoid_output_layer\")(x)\n",
        "    return output_layer\n",
        "\n",
        "#----------------ENCODER-----------------#\n",
        "\n",
        "  def _build_encoder(self):\n",
        "    encoder_input = self._add_encoder_input()\n",
        "    conv_layers = self._add_conv_layers(encoder_input)\n",
        "    bottleneck =  self._add_bottleneck(conv_layers)\n",
        "    self._model_input = encoder_input\n",
        "    self.encoder = Model(encoder_input, bottleneck, name=\"encoder\")\n",
        "\n",
        "  def _add_encoder_input(self):\n",
        "    return Input(shape=self.input_shape, name=\"encoder_input\")\n",
        "\n",
        "  def _add_conv_layers(self, encoder_input):\n",
        "    \"\"\"Creates all convolutional blocks in encoder\"\"\"\n",
        "    x = encoder_input\n",
        "    for layer_index in range(self._num_conv_layers):\n",
        "      x = self._add_conv_layer(layer_index, x)\n",
        "    return x\n",
        "  \n",
        "  def _add_conv_layer(self, layer_index, x):\n",
        "    \"\"\"Adds a convolutional block to a graph of layers, consisting\n",
        "    of Conv 2d + ReLu activation + batch normalization.\n",
        "    \"\"\"\n",
        "    layer_number = layer_index + 1\n",
        "    conv_layer = Conv2D(\n",
        "        filters= self.conv_filters[layer_index],\n",
        "        kernel_size = self.conv_kernels[layer_index],\n",
        "        strides = self.conv_strides[layer_index],\n",
        "        padding = \"same\",\n",
        "        name = f\"encoder_conv_layer_{layer_number}\"\n",
        "    )\n",
        "    x = conv_layer(x)\n",
        "    x = ReLU(name=f\"encoder_relu_{layer_number}\")(x)\n",
        "    x = BatchNormalization(name=f\"encoder_bn_{layer_number}\")(x)\n",
        "    return x\n",
        "\n",
        "#-------------Bottleneck (Latent Space)-------------#\n",
        "\n",
        "  def _add_bottleneck(self, x):\n",
        "    \"\"\"Flatten data and add bottleneck with Gaussian sampling (Dense layer)\"\"\"\n",
        "    self._shape_before_bottleneck = K.int_shape(x)[1:]\n",
        "    x = Flatten()(x)\n",
        "    self.mu = Dense(self.latent_space_dim,name=\"mu\")(x)\n",
        "    self.log_variance = Dense(self.latent_space_dim,\n",
        "                              name=\"log_variance\")(x)\n",
        "    \n",
        "    def sample_point_from_normal_distribution(args):\n",
        "      mu, log_variance = args\n",
        "      epsilon = K.random_normal(shape=K.shape(self.mu), mean=0., \n",
        "                                stddev=1.)\n",
        "      sampled_point = mu + K.exp(log_variance / 2) * epsilon\n",
        "\n",
        "      return sampled_point\n",
        "\n",
        "    x = Lambda(sample_point_from_normal_distribution, \n",
        "               name=\"encoder_output\")([self.mu, self.log_variance])\n",
        "    return x\n",
        "\n",
        "print(\"VAE successfully built\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Jf74jyfyQ2SF"
      },
      "source": [
        "#@title Load Checkpoint for Generating\n",
        "\n",
        "checkpoint_load_directory = \"/path/to/your/checkpoint\" #@param {type:\"string\"}\n",
        "\n",
        "#-------LOAD MODEL FOR GENERATING-------------#\n",
        "\n",
        "vae = VAE.load(checkpoint_load_directory)\n",
        "print(\"Loaded checkpoint\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Z2JISyMiQ6zt"
      },
      "source": [
        "#@title Import synthesis utility functions\n",
        "\n",
        "#-----TESTING FUNCTIONS ----------- #\n",
        "\n",
        "def select_spec(spec, labels, num_spec=10):\n",
        "    sample_spec_index = np.random.choice(range(len(spec)), num_spec)\n",
        "    sample_spec = spec[sample_spec_index]\n",
        "    sample_labels = labels[sample_spec_index]\n",
        "    return sample_spec, sample_labels\n",
        "\n",
        "\n",
        "def plot_reconstructed_spec(spec, reconstructed_spec):\n",
        "    fig = plt.figure(figsize=(15, 3))\n",
        "    num_spec = len(spec)\n",
        "    for i, (image, reconstructed_image) in enumerate(zip(spec, reconstructed_spec)):\n",
        "        image = image.squeeze()\n",
        "        ax = fig.add_subplot(2, num_spec, i + 1)\n",
        "        ax.axis(\"off\")\n",
        "        ax.imshow(image, cmap=\"gray_r\")\n",
        "        reconstructed_image = reconstructed_image.squeeze()\n",
        "        ax = fig.add_subplot(2, num_spec, i + num_spec + 1)\n",
        "        ax.axis(\"off\")\n",
        "        ax.imshow(reconstructed_image, cmap=\"gray_r\")\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def plot_spec_encoded_in_latent_space(latent_representations, sample_labels):\n",
        "    plt.figure(figsize=(10, 10))\n",
        "    plt.scatter(latent_representations[:, 0],\n",
        "                latent_representations[:, 1],\n",
        "                cmap=\"rainbow\",\n",
        "                c=sample_labels,\n",
        "                alpha=0.5,\n",
        "                s=2)\n",
        "    plt.colorbar()\n",
        "    plt.show()\n",
        "\n",
        "#---------------NOISE GENERATOR FUNCTIONS ------------#\n",
        "\n",
        "def generate_random_z_vect(seed=1001,size_z=1,scale=1.0):\n",
        "    np.random.seed(seed)\n",
        "    x = np.random.uniform(low=(scale * -1.0), high=scale, size=(size_z,VECTOR_DIM))\n",
        "    return x\n",
        "\n",
        "def generate_z_vect_from_perlin_noise(seed=1001, size_z=1, scale=1.0):\n",
        "    np.random.seed(seed)\n",
        "    x = generate_perlin_noise_2d((size_z, VECTOR_DIM), (1,1))\n",
        "    x = x*scale\n",
        "    return x\n",
        "\n",
        "def generate_z_vect_from_fractal_noise(seed=1001, size_z=1, scale=1.0,):\n",
        "    np.random.seed(seed)\n",
        "    x = generate_fractal_noise_2d((size_z, VECTOR_DIM), (1,1),)\n",
        "    x = x*scale\n",
        "    return x\n",
        "\n",
        "\n",
        "#-------SPECTROGRAM AND SOUND SYNTHESIS UTILITY FUNCTIONS -------- #\n",
        "\n",
        "#Assembling generated Spectrogram chunks into final Spectrogram\n",
        "def specass(a,spec):\n",
        "  but=False\n",
        "  con = np.array([])\n",
        "  nim = a.shape[0]\n",
        "  for i in range(nim-1):\n",
        "    im = a[i]\n",
        "    im = np.squeeze(im)\n",
        "    if not but:\n",
        "      con=im\n",
        "      but=True\n",
        "    else:\n",
        "      con = np.concatenate((con,im), axis=1)\n",
        "  diff = spec.shape[1]-(nim*shape)\n",
        "  a = np.squeeze(a)\n",
        "  con = np.concatenate((con,a[-1,:,-diff:]), axis=1)\n",
        "  return np.squeeze(con)\n",
        "\n",
        "#Splitting input spectrogram into different chunks to feed to the generator\n",
        "def chopspec(spec):\n",
        "  dsa=[]\n",
        "  for i in range(spec.shape[1]//shape):\n",
        "    im = spec[:,i*shape:i*shape+shape]\n",
        "    im = np.reshape(im, (im.shape[0],im.shape[1],1))\n",
        "    dsa.append(im)\n",
        "  imlast = spec[:,-shape:]\n",
        "  imlast = np.reshape(imlast, (imlast.shape[0],imlast.shape[1],1))\n",
        "  dsa.append(imlast)\n",
        "  return np.array(dsa, dtype=np.float32)\n",
        "\n",
        "#Converting from source Spectrogram to target Spectrogram\n",
        "def towave_reconstruct(spec, spec1, name, path='../content/', show=False, save=False):\n",
        "  specarr = chopspec(spec)\n",
        "  specarr1 = chopspec(spec1)\n",
        "  print(specarr.shape)\n",
        "  a = specarr\n",
        "  print('Generating...')\n",
        "  ab = specarr1\n",
        "  print('Assembling and Converting...')\n",
        "  a = specass(a,spec)\n",
        "  ab = specass(ab,spec1)\n",
        "  awv = deprep(a)\n",
        "  abwv = deprep(ab)\n",
        "  if save:\n",
        "    print('Saving...')\n",
        "    pathfin = f'{path}/{name}'\n",
        "    sf.write(f'{pathfin}.wav', awv, sr)\n",
        "    print('Saved WAV!')\n",
        "  IPython.display.display(IPython.display.Audio(np.squeeze(abwv), rate=sr))\n",
        "  IPython.display.display(IPython.display.Audio(np.squeeze(awv), rate=sr))\n",
        "  if show:\n",
        "    fig, axs = plt.subplots(ncols=2)\n",
        "    axs[0].imshow(np.flip(a, -2), cmap=None)\n",
        "    axs[0].axis('off')\n",
        "    axs[0].set_title('Reconstructed')\n",
        "    axs[1].imshow(np.flip(ab, -2), cmap=None)\n",
        "    axs[1].axis('off')\n",
        "    axs[1].set_title('Input')\n",
        "    plt.show()\n",
        "  return abwv\n",
        "\n",
        "#Converting from Z vector generated spectrogram to waveform\n",
        "def towave_from_z(spec, name, path='../content/', show=False, save=False):\n",
        "  specarr = chopspec(spec)\n",
        "  print(specarr.shape)\n",
        "  a = specarr\n",
        "  print('Generating...')\n",
        "  print('Assembling and Converting...')\n",
        "  a = specass(a,spec)\n",
        "  awv = deprep(a)\n",
        "  if save:\n",
        "    print('Saving...')\n",
        "    pathfin = f'{path}/{name}'\n",
        "    sf.write(f'{pathfin}.wav', awv, sr)\n",
        "    print('Saved WAV!')\n",
        "  IPython.display.display(IPython.display.Audio(np.squeeze(awv), rate=sr))\n",
        "  if show:\n",
        "    fig, axs = plt.subplots(ncols=1)\n",
        "    axs.imshow(np.flip(a, -2), cmap=None)\n",
        "    axs.axis('off')\n",
        "    axs.set_title('Decoder Synthesis')\n",
        "    plt.show()\n",
        "  return awv\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Perform Timbre Transfer\n",
        "\n",
        "save_audio = True #@param {type:\"boolean\"}\n",
        "audio_name = \"timbre_transfered_audio\" #@param {type:\"string\"}\n",
        "audio_save_directory = \"/content/\" #@param {type:\"string\"}\n",
        "\n",
        "encoded_spec = vae.encode(out_data)\n",
        "print(np.shape(encoded_spec))\n",
        "\n",
        "z_sample = np.array(vae.sample_from_latent_space(encoded_spec))\n",
        "\n",
        "assembled_spec = testass(z_sample)\n",
        "towave_from_z(assembled_spec, name=audio_name,path=audio_save_directory,show=True, save=save_audio)"
      ],
      "metadata": {
        "id": "nt4lLm7Dxgfi",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}