MelGAN_VC.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "MelGAN_VC.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/marcoppasini/MelGAN-VC/blob/master/MelGAN_VC.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V00rptcdKSbq"
      },
      "source": [
        "#We'll be using TF 2.1 and torchaudio\n",
        "\n",
        "try:\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "  pass\n",
        "import tensorflow as tf\n",
        "!pip install soundfile                    #to save wav files\n",
        "!pip install --no-deps torchaudio==0.5.0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CAmiyxtl2J5s"
      },
      "source": [
        "#Connecting Drive to save model checkpoints during training and to use custom data, uncomment if needed\n",
        "\n",
        "# import os\n",
        "# from google.colab import drive\n",
        "# drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TNXtq27kl_Ym"
      },
      "source": [
        "#Dataset download (Uncomment where needed)\n",
        "\n",
        "#Arctic dataset for voice conversion\n",
        "\n",
        "# !wget --header=\"Host: festvox.org\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3\" --header=\"Accept-Language: it-IT,it;q=0.9,en-US;q=0.8,en;q=0.7\" --header=\"Referer: http://festvox.org/cmu_arctic/cmu_arctic/packed/\" \"http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_bdl_arctic-0.95-release.zip\" -O \"cmu_us_bdl_arctic-0.95-release.zip\" -c\n",
        "# !unzip -qq cmu_us_bdl_arctic-0.95-release.zip #MALE1\n",
        "# !wget --header=\"Host: festvox.org\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3\" --header=\"Accept-Language: it-IT,it;q=0.9,en-US;q=0.8,en;q=0.7\" --header=\"Referer: http://festvox.org/cmu_arctic/cmu_arctic/packed/\" \"http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_clb_arctic-0.95-release.zip\" -O \"cmu_us_clb_arctic-0.95-release.zip\" -c\n",
        "# !unzip -qq cmu_us_clb_arctic-0.95-release.zip #FEMALE1\n",
        "# !wget --header=\"Host: festvox.org\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3\" --header=\"Accept-Language: it-IT,it;q=0.9,en-US;q=0.8,en;q=0.7\" --header=\"Referer: http://festvox.org/cmu_arctic/cmu_arctic/packed/\" \"http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_rms_arctic-0.95-release.zip\" -O \"cmu_us_rms_arctic-0.95-release.zip\" -c\n",
        "# !unzip -qq cmu_us_rms_arctic-0.95-release.zip #MALE2\n",
        "# !wget --header=\"Host: festvox.org\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/76.0.3809.100 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3\" --header=\"Accept-Language: it-IT,it;q=0.9,en-US;q=0.8,en;q=0.7\" --header=\"Referer: http://festvox.org/cmu_arctic/cmu_arctic/packed/\" \"http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_slt_arctic-0.95-release.zip\" -O \"cmu_us_slt_arctic-0.95-release.zip\" -c\n",
        "# !unzip -qq cmu_us_slt_arctic-0.95-release.zip #FEMALE2\n",
        "\n",
        "!wget http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_bdl_arctic-0.95-release.tar.bz2 #MALE1\n",
        "!tar -xf cmu_us_bdl_arctic-0.95-release.tar.bz2\n",
        "!wget http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_clb_arctic-0.95-release.tar.bz2 #FEMALE1\n",
        "!tar -xf cmu_us_clb_arctic-0.95-release.tar.bz2\n",
        "!wget http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_rms_arctic-0.95-release.tar.bz2 #MALE2\n",
        "!tar -xf cmu_us_rms_arctic-0.95-release.tar.bz2\n",
        "!wget http://festvox.org/cmu_arctic/cmu_arctic/packed/cmu_us_slt_arctic-0.95-release.tar.bz2 #FEMALE2\n",
        "!tar -xf cmu_us_slt_arctic-0.95-release.tar.bz2\n",
        "\n",
        "#GTZAN dataset for music genre transfer\n",
        "# !wget --header=\"Host: opihi.cs.uvic.ca\" --header=\"User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/77.0.3865.90 Safari/537.36\" --header=\"Accept: text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3\" --header=\"Accept-Language: it-IT,it;q=0.9,en-US;q=0.8,en;q=0.7\" --header=\"Referer: http://marsyas.info/downloads/datasets.html\" \"http://opihi.cs.uvic.ca/sound/genres.tar.gz\" -O \"genres.tar.gz\" -c\n",
        "# !tar -xzf genres.tar.gz\n",
        "\n",
        "!ls"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LEvqwT96l_Yq"
      },
      "source": [
        "#Imports\n",
        "\n",
        "from __future__ import print_function, division\n",
        "from glob import glob\n",
        "import scipy\n",
        "import soundfile as sf\n",
        "import matplotlib.pyplot as plt\n",
        "from IPython.display import clear_output\n",
        "from tensorflow.keras.layers import Input, Dense, Reshape, Flatten, Concatenate, Conv2D, Conv2DTranspose, GlobalAveragePooling2D, UpSampling2D, LeakyReLU, ReLU, Add, Multiply, Lambda, Dot, BatchNormalization, Activation, ZeroPadding2D, Cropping2D, Cropping1D\n",
        "from tensorflow.keras.models import Sequential, Model, load_model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "from tensorflow.keras.initializers import TruncatedNormal, he_normal\n",
        "import tensorflow.keras.backend as K\n",
        "import datetime\n",
        "import numpy as np\n",
        "import random\n",
        "import matplotlib.pyplot as plt\n",
        "import collections\n",
        "from PIL import Image\n",
        "from skimage.transform import resize\n",
        "import imageio\n",
        "import librosa\n",
        "import librosa.display\n",
        "from librosa.feature import melspectrogram\n",
        "import os\n",
        "import time\n",
        "import IPython"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KbaM4WKrvO7r"
      },
      "source": [
        "#Hyperparameters\n",
        "\n",
        "hop=192               #hop size (window size = 6*hop)\n",
        "sr=16000              #sampling rate\n",
        "min_level_db=-100     #reference values to normalize data\n",
        "ref_level_db=20\n",
        "\n",
        "shape=24              #length of time axis of split specrograms to feed to generator            \n",
        "vec_len=128           #length of vector generated by siamese vector\n",
        "bs = 16               #batch size\n",
        "delta = 2.            #constant for siamese loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K9pIPj9hnyJ0"
      },
      "source": [
        "#There seems to be a problem with Tensorflow STFT, so we'll be using pytorch to handle offline mel-spectrogram generation and waveform reconstruction\n",
        "#For waveform reconstruction, a gradient-based method is used:\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=6*hop, win_length=6*hop, hop_length=hop, pad=0, power=2, normalized=True)\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.003):\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=2000, evaiter=10, tol=1e-8)\n",
        "  return np.array(np.squeeze(wv))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YNRYjsCDqDjF"
      },
      "source": [
        "#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*16000):\n",
        "  x, sr = librosa.load(path,sr=16000)\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",
        "#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 = 10*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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tK_UnhfMELHD"
      },
      "source": [
        "#Generating Mel-Spectrogram dataset (Uncomment where needed)\n",
        "#adata: source spectrograms\n",
        "#bdata: target spectrograms\n",
        "\n",
        "#MALE1\n",
        "awv = audio_array('../content/cmu_us_clb_arctic/wav')                               #get waveform array from folder containing wav files\n",
        "aspec = tospec(awv)                                                                 #get spectrogram array\n",
        "adata = splitcut(aspec)                                                             #split spectrogams to fixed length\n",
        "#FEMALE1\n",
        "bwv = audio_array('../content/cmu_us_bdl_arctic/wav')\n",
        "bspec = tospec(bwv)\n",
        "bdata = splitcut(bspec)\n",
        "# #MALE2\n",
        "# awv = audio_array('../content/cmu_us_rms_arctic/wav')\n",
        "# aspec = tospec(awv)\n",
        "# adata = splitcut(aspec)\n",
        "# #FEMALE2\n",
        "# bwv = audio_array('../content/cmu_us_slt_arctic/wav')\n",
        "# bspec = tospec(bwv)\n",
        "# bdata = splitcut(bspec)\n",
        "\n",
        "#JAZZ MUSIC\n",
        "# awv = audio_array('../content/genres/jazz')\n",
        "# aspec = tospec(awv)\n",
        "# adata = splitcut(aspec)\n",
        "#CLASSICAL MUSIC\n",
        "# bwv = audio_array('../content/genres/classical')\n",
        "# bspec = tospec(bwv)\n",
        "# bdata = splitcut(bspec)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qSesIbwr_GyO"
      },
      "source": [
        "#Creating Tensorflow Datasets\n",
        "\n",
        "@tf.function\n",
        "def proc(x):\n",
        "  return tf.image.random_crop(x, size=[hop, 3*shape, 1])\n",
        "\n",
        "dsa = tf.data.Dataset.from_tensor_slices(adata).repeat(50).map(proc, num_parallel_calls=tf.data.experimental.AUTOTUNE).shuffle(10000).batch(bs, drop_remainder=True)\n",
        "dsb = tf.data.Dataset.from_tensor_slices(bdata).repeat(50).map(proc, num_parallel_calls=tf.data.experimental.AUTOTUNE).shuffle(10000).batch(bs, drop_remainder=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AHnP2zr7Ypgi"
      },
      "source": [
        "#Adding Spectral Normalization to convolutional layers\n",
        "\n",
        "from tensorflow.python.keras.utils import conv_utils\n",
        "from tensorflow.python.ops import array_ops\n",
        "from tensorflow.python.ops import math_ops\n",
        "from tensorflow.python.ops import sparse_ops\n",
        "from tensorflow.python.ops import gen_math_ops\n",
        "from tensorflow.python.ops import standard_ops\n",
        "from tensorflow.python.eager import context\n",
        "from tensorflow.python.framework import tensor_shape\n",
        "\n",
        "def l2normalize(v, eps=1e-12):\n",
        "    return v / (tf.norm(v) + eps)\n",
        "\n",
        "\n",
        "class ConvSN2D(tf.keras.layers.Conv2D):\n",
        "\n",
        "    def __init__(self, filters, kernel_size, power_iterations=1, **kwargs):\n",
        "        super(ConvSN2D, self).__init__(filters, kernel_size, **kwargs)\n",
        "        self.power_iterations = power_iterations\n",
        "\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        super(ConvSN2D, self).build(input_shape)\n",
        "\n",
        "        if self.data_format == 'channels_first':\n",
        "            channel_axis = 1\n",
        "        else:\n",
        "            channel_axis = -1\n",
        "\n",
        "        self.u = self.add_weight(self.name + '_u',\n",
        "            shape=tuple([1, self.kernel.shape.as_list()[-1]]), \n",
        "            initializer=tf.initializers.RandomNormal(0, 1),\n",
        "            trainable=False\n",
        "        )\n",
        "\n",
        "    def compute_spectral_norm(self, W, new_u, W_shape):\n",
        "        for _ in range(self.power_iterations):\n",
        "\n",
        "            new_v = l2normalize(tf.matmul(new_u, tf.transpose(W)))\n",
        "            new_u = l2normalize(tf.matmul(new_v, W))\n",
        "            \n",
        "        sigma = tf.matmul(tf.matmul(new_v, W), tf.transpose(new_u))\n",
        "        W_bar = W/sigma\n",
        "\n",
        "        with tf.control_dependencies([self.u.assign(new_u)]):\n",
        "          W_bar = tf.reshape(W_bar, W_shape)\n",
        "\n",
        "        return W_bar\n",
        "\n",
        "\n",
        "    def call(self, inputs):\n",
        "        W_shape = self.kernel.shape.as_list()\n",
        "        W_reshaped = tf.reshape(self.kernel, (-1, W_shape[-1]))\n",
        "        new_kernel = self.compute_spectral_norm(W_reshaped, self.u, W_shape)\n",
        "        outputs = self.convolution_op(inputs, new_kernel)\n",
        "\n",
        "        if self.use_bias:\n",
        "            if self.data_format == 'channels_first':\n",
        "                    outputs = tf.nn.bias_add(outputs, self.bias, data_format='NCHW')\n",
        "            else:\n",
        "                outputs = tf.nn.bias_add(outputs, self.bias, data_format='NHWC')\n",
        "        if self.activation is not None:\n",
        "            return self.activation(outputs)\n",
        "\n",
        "        return outputs\n",
        "\n",
        "\n",
        "class ConvSN2DTranspose(tf.keras.layers.Conv2DTranspose):\n",
        "\n",
        "    def __init__(self, filters, kernel_size, power_iterations=1, **kwargs):\n",
        "        super(ConvSN2DTranspose, self).__init__(filters, kernel_size, **kwargs)\n",
        "        self.power_iterations = power_iterations\n",
        "\n",
        "\n",
        "    def build(self, input_shape):\n",
        "        super(ConvSN2DTranspose, self).build(input_shape)\n",
        "\n",
        "        if self.data_format == 'channels_first':\n",
        "            channel_axis = 1\n",
        "        else:\n",
        "            channel_axis = -1\n",
        "\n",
        "        self.u = self.add_weight(self.name + '_u',\n",
        "            shape=tuple([1, self.kernel.shape.as_list()[-1]]), \n",
        "            initializer=tf.initializers.RandomNormal(0, 1),\n",
        "            trainable=False\n",
        "        )\n",
        "\n",
        "    def compute_spectral_norm(self, W, new_u, W_shape):\n",
        "        for _ in range(self.power_iterations):\n",
        "\n",
        "            new_v = l2normalize(tf.matmul(new_u, tf.transpose(W)))\n",
        "            new_u = l2normalize(tf.matmul(new_v, W))\n",
        "            \n",
        "        sigma = tf.matmul(tf.matmul(new_v, W), tf.transpose(new_u))\n",
        "        W_bar = W/sigma\n",
        "\n",
        "        with tf.control_dependencies([self.u.assign(new_u)]):\n",
        "          W_bar = tf.reshape(W_bar, W_shape)\n",
        "\n",
        "        return W_bar\n",
        "\n",
        "    def call(self, inputs):\n",
        "        W_shape = self.kernel.shape.as_list()\n",
        "        W_reshaped = tf.reshape(self.kernel, (-1, W_shape[-1]))\n",
        "        new_kernel = self.compute_spectral_norm(W_reshaped, self.u, W_shape)\n",
        "\n",
        "        inputs_shape = array_ops.shape(inputs)\n",
        "        batch_size = inputs_shape[0]\n",
        "        if self.data_format == 'channels_first':\n",
        "          h_axis, w_axis = 2, 3\n",
        "        else:\n",
        "          h_axis, w_axis = 1, 2\n",
        "\n",
        "        height, width = inputs_shape[h_axis], inputs_shape[w_axis]\n",
        "        kernel_h, kernel_w = self.kernel_size\n",
        "        stride_h, stride_w = self.strides\n",
        "\n",
        "        if self.output_padding is None:\n",
        "          out_pad_h = out_pad_w = None\n",
        "        else:\n",
        "          out_pad_h, out_pad_w = self.output_padding\n",
        "\n",
        "        out_height = conv_utils.deconv_output_length(height,\n",
        "                                                    kernel_h,\n",
        "                                                    padding=self.padding,\n",
        "                                                    output_padding=out_pad_h,\n",
        "                                                    stride=stride_h,\n",
        "                                                    dilation=self.dilation_rate[0])\n",
        "        out_width = conv_utils.deconv_output_length(width,\n",
        "                                                    kernel_w,\n",
        "                                                    padding=self.padding,\n",
        "                                                    output_padding=out_pad_w,\n",
        "                                                    stride=stride_w,\n",
        "                                                    dilation=self.dilation_rate[1])\n",
        "        if self.data_format == 'channels_first':\n",
        "          output_shape = (batch_size, self.filters, out_height, out_width)\n",
        "        else:\n",
        "          output_shape = (batch_size, out_height, out_width, self.filters)\n",
        "\n",
        "        output_shape_tensor = array_ops.stack(output_shape)\n",
        "        outputs = K.conv2d_transpose(\n",
        "            inputs,\n",
        "            new_kernel,\n",
        "            output_shape_tensor,\n",
        "            strides=self.strides,\n",
        "            padding=self.padding,\n",
        "            data_format=self.data_format,\n",
        "            dilation_rate=self.dilation_rate)\n",
        "\n",
        "        if not context.executing_eagerly():\n",
        "          out_shape = self.compute_output_shape(inputs.shape)\n",
        "          outputs.set_shape(out_shape)\n",
        "\n",
        "        if self.use_bias:\n",
        "          outputs = tf.nn.bias_add(\n",
        "              outputs,\n",
        "              self.bias,\n",
        "              data_format=conv_utils.convert_data_format(self.data_format, ndim=4))\n",
        "\n",
        "        if self.activation is not None:\n",
        "          return self.activation(outputs)\n",
        "        return outputs  \n",
        "    \n",
        "    \n",
        "class DenseSN(Dense):\n",
        "    \n",
        "    def build(self, input_shape):\n",
        "        super(DenseSN, self).build(input_shape)\n",
        "\n",
        "        self.u = self.add_weight(self.name + '_u',\n",
        "            shape=tuple([1, self.kernel.shape.as_list()[-1]]), \n",
        "            initializer=tf.initializers.RandomNormal(0, 1),\n",
        "            trainable=False)\n",
        "        \n",
        "    def compute_spectral_norm(self, W, new_u, W_shape):\n",
        "        new_v = l2normalize(tf.matmul(new_u, tf.transpose(W)))\n",
        "        new_u = l2normalize(tf.matmul(new_v, W))\n",
        "        sigma = tf.matmul(tf.matmul(new_v, W), tf.transpose(new_u))\n",
        "        W_bar = W/sigma\n",
        "        with tf.control_dependencies([self.u.assign(new_u)]):\n",
        "          W_bar = tf.reshape(W_bar, W_shape)\n",
        "        return W_bar\n",
        "        \n",
        "    def call(self, inputs):\n",
        "        W_shape = self.kernel.shape.as_list()\n",
        "        W_reshaped = tf.reshape(self.kernel, (-1, W_shape[-1]))\n",
        "        new_kernel = self.compute_spectral_norm(W_reshaped, self.u, W_shape)\n",
        "        rank = len(inputs.shape)\n",
        "        if rank > 2:\n",
        "          outputs = standard_ops.tensordot(inputs, new_kernel, [[rank - 1], [0]])\n",
        "          if not context.executing_eagerly():\n",
        "            shape = inputs.shape.as_list()\n",
        "            output_shape = shape[:-1] + [self.units]\n",
        "            outputs.set_shape(output_shape)\n",
        "        else:\n",
        "          inputs = math_ops.cast(inputs, self._compute_dtype)\n",
        "          if K.is_sparse(inputs):\n",
        "            outputs = sparse_ops.sparse_tensor_dense_matmul(inputs, new_kernel)\n",
        "          else:\n",
        "            outputs = gen_math_ops.mat_mul(inputs, new_kernel)\n",
        "        if self.use_bias:\n",
        "          outputs = tf.nn.bias_add(outputs, self.bias)\n",
        "        if self.activation is not None:\n",
        "          return self.activation(outputs)\n",
        "        return outputs\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eX41awYeHE1N"
      },
      "source": [
        "#Networks Architecture\n",
        "\n",
        "init = tf.keras.initializers.he_uniform()\n",
        "\n",
        "def conv2d(layer_input, filters, kernel_size=4, strides=2, padding='same', leaky=True, bnorm=True, sn=True):\n",
        "  if leaky:\n",
        "    Activ = LeakyReLU(alpha=0.2)\n",
        "  else:\n",
        "    Activ = ReLU()\n",
        "  if sn:\n",
        "    d = ConvSN2D(filters, kernel_size=kernel_size, strides=strides, padding=padding, kernel_initializer=init, use_bias=False)(layer_input)\n",
        "  else:\n",
        "    d = Conv2D(filters, kernel_size=kernel_size, strides=strides, padding=padding, kernel_initializer=init, use_bias=False)(layer_input)\n",
        "  if bnorm:\n",
        "    d = BatchNormalization()(d)\n",
        "  d = Activ(d)\n",
        "  return d\n",
        "\n",
        "def deconv2d(layer_input, layer_res, filters, kernel_size=4, conc=True, scalev=False, bnorm=True, up=True, padding='same', strides=2):\n",
        "  if up:\n",
        "    u = UpSampling2D((1,2))(layer_input)\n",
        "    u = ConvSN2D(filters, kernel_size, strides=(1,1), kernel_initializer=init, use_bias=False, padding=padding)(u)\n",
        "  else:\n",
        "    u = ConvSN2DTranspose(filters, kernel_size, strides=strides, kernel_initializer=init, use_bias=False, padding=padding)(layer_input)\n",
        "  if bnorm:\n",
        "    u = BatchNormalization()(u)\n",
        "  u = LeakyReLU(alpha=0.2)(u)\n",
        "  if conc:\n",
        "    u = Concatenate()([u,layer_res])\n",
        "  return u\n",
        "\n",
        "#Extract function: splitting spectrograms\n",
        "def extract_image(im):\n",
        "  im1 = Cropping2D(((0,0), (0, 2*(im.shape[2]//3))))(im)\n",
        "  im2 = Cropping2D(((0,0), (im.shape[2]//3,im.shape[2]//3)))(im)\n",
        "  im3 = Cropping2D(((0,0), (2*(im.shape[2]//3), 0)))(im)\n",
        "  return im1,im2,im3\n",
        "\n",
        "#Assemble function: concatenating spectrograms\n",
        "def assemble_image(lsim):\n",
        "  im1,im2,im3 = lsim\n",
        "  imh = Concatenate(2)([im1,im2,im3])\n",
        "  return imh\n",
        "\n",
        "#U-NET style architecture\n",
        "def build_generator(input_shape):\n",
        "  h,w,c = input_shape\n",
        "  inp = Input(shape=input_shape)\n",
        "  #downscaling\n",
        "  g0 = tf.keras.layers.ZeroPadding2D((0,1))(inp)\n",
        "  g1 = conv2d(g0, 256, kernel_size=(h,3), strides=1, padding='valid')\n",
        "  g2 = conv2d(g1, 256, kernel_size=(1,9), strides=(1,2))\n",
        "  g3 = conv2d(g2, 256, kernel_size=(1,7), strides=(1,2))\n",
        "  #upscaling\n",
        "  g4 = deconv2d(g3,g2, 256, kernel_size=(1,7), strides=(1,2))\n",
        "  g5 = deconv2d(g4,g1, 256, kernel_size=(1,9), strides=(1,2), bnorm=False)\n",
        "  g6 = ConvSN2DTranspose(1, kernel_size=(h,1), strides=(1,1), kernel_initializer=init, padding='valid', activation='tanh')(g5)\n",
        "  return Model(inp,g6, name='G')\n",
        "\n",
        "#Siamese Network\n",
        "def build_siamese(input_shape):\n",
        "  h,w,c = input_shape\n",
        "  inp = Input(shape=input_shape)\n",
        "  g1 = conv2d(inp, 256, kernel_size=(h,3), strides=1, padding='valid', sn=False)\n",
        "  g2 = conv2d(g1, 256, kernel_size=(1,9), strides=(1,2), sn=False)\n",
        "  g3 = conv2d(g2, 256, kernel_size=(1,7), strides=(1,2), sn=False)\n",
        "  g4 = Flatten()(g3)\n",
        "  g5 = Dense(vec_len)(g4)\n",
        "  return Model(inp, g5, name='S')\n",
        "\n",
        "#Discriminator (Critic) Network\n",
        "def build_critic(input_shape):\n",
        "  h,w,c = input_shape\n",
        "  inp = Input(shape=input_shape)\n",
        "  g1 = conv2d(inp, 512, kernel_size=(h,3), strides=1, padding='valid', bnorm=False)\n",
        "  g2 = conv2d(g1, 512, kernel_size=(1,9), strides=(1,2), bnorm=False)\n",
        "  g3 = conv2d(g2, 512, kernel_size=(1,7), strides=(1,2), bnorm=False)\n",
        "  g4 = Flatten()(g3)\n",
        "  g4 = DenseSN(1, kernel_initializer=init)(g4)\n",
        "  return Model(inp, g4, name='C')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4fXJmItOzrhC"
      },
      "source": [
        "#Load past models from path to resume training or test\n",
        "def load(path):\n",
        "  gen = build_generator((hop,shape,1))\n",
        "  siam = build_siamese((hop,shape,1))\n",
        "  critic = build_critic((hop,3*shape,1))\n",
        "  gen.load_weights(path+'/gen.h5')\n",
        "  critic.load_weights(path+'/critic.h5')\n",
        "  siam.load_weights(path+'/siam.h5')\n",
        "  return gen,critic,siam\n",
        "\n",
        "#Build models\n",
        "def build():\n",
        "  gen = build_generator((hop,shape,1))\n",
        "  siam = build_siamese((hop,shape,1))\n",
        "  critic = build_critic((hop,3*shape,1))                                          #the discriminator accepts as input spectrograms of triple the width of those generated by the generator\n",
        "  return gen,critic,siam\n",
        "\n",
        "#Generate a random batch to display current training results\n",
        "def testgena():\n",
        "  sw = True\n",
        "  while sw:\n",
        "    a = np.random.choice(aspec)\n",
        "    if a.shape[1]//shape!=1:\n",
        "      sw=False\n",
        "  dsa = []\n",
        "  if a.shape[1]//shape>6:\n",
        "    num=6\n",
        "  else:\n",
        "    num=a.shape[1]//shape\n",
        "  rn = np.random.randint(a.shape[1]-(num*shape))\n",
        "  for i in range(num):\n",
        "    im = a[:,rn+(i*shape):rn+(i*shape)+shape]\n",
        "    im = np.reshape(im, (im.shape[0],im.shape[1],1))\n",
        "    dsa.append(im)\n",
        "  return np.array(dsa, dtype=np.float32)\n",
        "\n",
        "#Show results mid-training\n",
        "def save_test_image_full(path):\n",
        "  a = testgena()\n",
        "  print(a.shape)\n",
        "  ab = gen(a, training=False)\n",
        "  ab = testass(ab)\n",
        "  a = testass(a)\n",
        "  abwv = deprep(ab)\n",
        "  awv = deprep(a)\n",
        "  sf.write(path+'/new_file.wav', abwv, sr)\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",
        "  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('Source')\n",
        "  axs[1].imshow(np.flip(ab, -2), cmap=None)\n",
        "  axs[1].axis('off')\n",
        "  axs[1].set_title('Generated')\n",
        "  plt.show()\n",
        "\n",
        "#Save in training loop\n",
        "def save_end(epoch,gloss,closs,mloss,n_save=3,save_path='../content/'):                 #use custom save_path (i.e. Drive '../content/drive/My Drive/')\n",
        "  if epoch % n_save == 0:\n",
        "    print('Saving...')\n",
        "    path = f'{save_path}/MELGANVC-{str(gloss)[:9]}-{str(closs)[:9]}-{str(mloss)[:9]}'\n",
        "    os.mkdir(path)\n",
        "    gen.save_weights(path+'/gen.h5')\n",
        "    critic.save_weights(path+'/critic.h5')\n",
        "    siam.save_weights(path+'/siam.h5')\n",
        "    save_test_image_full(path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fn2s65AxjDJ8"
      },
      "source": [
        "#Losses\n",
        "\n",
        "def mae(x,y):\n",
        "  return tf.reduce_mean(tf.abs(x-y))\n",
        "\n",
        "def mse(x,y):\n",
        "  return tf.reduce_mean((x-y)**2)\n",
        "\n",
        "def loss_travel(sa,sab,sa1,sab1):\n",
        "  l1 = tf.reduce_mean(((sa-sa1) - (sab-sab1))**2)\n",
        "  l2 = tf.reduce_mean(tf.reduce_sum(-(tf.nn.l2_normalize(sa-sa1, axis=[-1]) * tf.nn.l2_normalize(sab-sab1, axis=[-1])), axis=-1))\n",
        "  return l1+l2\n",
        "\n",
        "def loss_siamese(sa,sa1):\n",
        "  logits = tf.sqrt(tf.reduce_sum((sa-sa1)**2, axis=-1, keepdims=True))\n",
        "  return tf.reduce_mean(tf.square(tf.maximum((delta - logits), 0)))\n",
        "\n",
        "def d_loss_f(fake):\n",
        "  return tf.reduce_mean(tf.maximum(1 + fake, 0))\n",
        "\n",
        "def d_loss_r(real):\n",
        "  return tf.reduce_mean(tf.maximum(1 - real, 0))\n",
        "\n",
        "def g_loss_f(fake):\n",
        "  return tf.reduce_mean(- fake)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fgjxHjyIhPwl"
      },
      "source": [
        "#Get models and optimizers\n",
        "def get_networks(shape, load_model=False, path=None):\n",
        "  if not load_model:\n",
        "    gen,critic,siam = build()\n",
        "  else:\n",
        "    gen,critic,siam = load(path)\n",
        "  print('Built networks')\n",
        "\n",
        "  opt_gen = Adam(0.0001, 0.5)\n",
        "  opt_disc = Adam(0.0001, 0.5)\n",
        "\n",
        "  return gen,critic,siam, [opt_gen,opt_disc]\n",
        "\n",
        "#Set learning rate\n",
        "def update_lr(lr):\n",
        "  opt_gen.learning_rate = lr\n",
        "  opt_disc.learning_rate = lr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WGWjgHqDWR78"
      },
      "source": [
        "#Training Functions\n",
        "\n",
        "#Train Generator, Siamese and Critic\n",
        "@tf.function\n",
        "def train_all(a,b):\n",
        "  #splitting spectrogram in 3 parts\n",
        "  aa,aa2,aa3 = extract_image(a) \n",
        "  bb,bb2,bb3 = extract_image(b)\n",
        "\n",
        "  with tf.GradientTape() as tape_gen, tf.GradientTape() as tape_disc:\n",
        "\n",
        "    #translating A to B\n",
        "    fab = gen(aa, training=True)\n",
        "    fab2 = gen(aa2, training=True)\n",
        "    fab3 = gen(aa3, training=True)\n",
        "    #identity mapping B to B                                                        COMMENT THESE 3 LINES IF THE IDENTITY LOSS TERM IS NOT NEEDED\n",
        "    fid = gen(bb, training=True) \n",
        "    fid2 = gen(bb2, training=True)\n",
        "    fid3 = gen(bb3, training=True)\n",
        "    #concatenate/assemble converted spectrograms\n",
        "    fabtot = assemble_image([fab,fab2,fab3])\n",
        "\n",
        "    #feed concatenated spectrograms to critic\n",
        "    cab = critic(fabtot, training=True)\n",
        "    cb = critic(b, training=True)\n",
        "    #feed 2 pairs (A,G(A)) extracted spectrograms to Siamese\n",
        "    sab = siam(fab, training=True)\n",
        "    sab2 = siam(fab3, training=True)\n",
        "    sa = siam(aa, training=True)\n",
        "    sa2 = siam(aa3, training=True)\n",
        "\n",
        "    #identity mapping loss\n",
        "    loss_id = (mae(bb,fid)+mae(bb2,fid2)+mae(bb3,fid3))/3.                         #loss_id = 0. IF THE IDENTITY LOSS TERM IS NOT NEEDED\n",
        "    #travel loss\n",
        "    loss_m = loss_travel(sa,sab,sa2,sab2)+loss_siamese(sa,sa2)\n",
        "    #generator and critic losses\n",
        "    loss_g = g_loss_f(cab)\n",
        "    loss_dr = d_loss_r(cb)\n",
        "    loss_df = d_loss_f(cab)\n",
        "    loss_d = (loss_dr+loss_df)/2.\n",
        "    #generator+siamese total loss\n",
        "    lossgtot = loss_g+10.*loss_m+0.5*loss_id                                       #CHANGE LOSS WEIGHTS HERE  (COMMENT OUT +w*loss_id IF THE IDENTITY LOSS TERM IS NOT NEEDED)\n",
        "  \n",
        "  #computing and applying gradients\n",
        "  grad_gen = tape_gen.gradient(lossgtot, gen.trainable_variables+siam.trainable_variables)\n",
        "  opt_gen.apply_gradients(zip(grad_gen, gen.trainable_variables+siam.trainable_variables))\n",
        "\n",
        "  grad_disc = tape_disc.gradient(loss_d, critic.trainable_variables)\n",
        "  opt_disc.apply_gradients(zip(grad_disc, critic.trainable_variables))\n",
        "  \n",
        "  return loss_dr,loss_df,loss_g,loss_id\n",
        "\n",
        "#Train Critic only\n",
        "@tf.function\n",
        "def train_d(a,b):\n",
        "  aa,aa2,aa3 = extract_image(a)\n",
        "  with tf.GradientTape() as tape_disc:\n",
        "\n",
        "    fab = gen(aa, training=True)\n",
        "    fab2 = gen(aa2, training=True)\n",
        "    fab3 = gen(aa3, training=True)\n",
        "    fabtot = assemble_image([fab,fab2,fab3])\n",
        "\n",
        "    cab = critic(fabtot, training=True)\n",
        "    cb = critic(b, training=True)\n",
        "\n",
        "    loss_dr = d_loss_r(cb)\n",
        "    loss_df = d_loss_f(cab)\n",
        "\n",
        "    loss_d = (loss_dr+loss_df)/2.\n",
        "  \n",
        "  grad_disc = tape_disc.gradient(loss_d, critic.trainable_variables)\n",
        "  opt_disc.apply_gradients(zip(grad_disc, critic.trainable_variables))\n",
        "\n",
        "  return loss_dr,loss_df"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aVwL-Ry-nNru"
      },
      "source": [
        "#Training Loop\n",
        "\n",
        "def train(epochs, batch_size=16, lr=0.0001, n_save=6, gupt=5):\n",
        "  \n",
        "  update_lr(lr)\n",
        "  df_list = []\n",
        "  dr_list = []\n",
        "  g_list = []\n",
        "  id_list = []\n",
        "  c = 0\n",
        "  g = 0\n",
        "  \n",
        "  for epoch in range(epochs):\n",
        "        bef = time.time()\n",
        "        \n",
        "        for batchi,(a,b) in enumerate(zip(dsa,dsb)):\n",
        "          \n",
        "            if batchi%gupt==0:\n",
        "              dloss_t,dloss_f,gloss,idloss = train_all(a,b)\n",
        "            else:\n",
        "              dloss_t,dloss_f = train_d(a,b)\n",
        "\n",
        "            df_list.append(dloss_f)\n",
        "            dr_list.append(dloss_t)\n",
        "            g_list.append(gloss)\n",
        "            id_list.append(idloss)\n",
        "            c += 1\n",
        "            g += 1\n",
        "\n",
        "            if batchi%600==0:\n",
        "                print(f'[Epoch {epoch}/{epochs}] [Batch {batchi}] [D loss f: {np.mean(df_list[-g:], axis=0)} ', end='')\n",
        "                print(f'r: {np.mean(dr_list[-g:], axis=0)}] ', end='')\n",
        "                print(f'[G loss: {np.mean(g_list[-g:], axis=0)}] ', end='')\n",
        "                print(f'[ID loss: {np.mean(id_list[-g:])}] ', end='')\n",
        "                print(f'[LR: {lr}]')\n",
        "                g = 0\n",
        "            nbatch=batchi\n",
        "\n",
        "        print(f'Time/Batch {(time.time()-bef)/nbatch}')\n",
        "        save_end(epoch,np.mean(g_list[-n_save*c:], axis=0),np.mean(df_list[-n_save*c:], axis=0),np.mean(id_list[-n_save*c:], axis=0),n_save=n_save)\n",
        "        print(f'Mean D loss: {np.mean(df_list[-c:], axis=0)} Mean G loss: {np.mean(g_list[-c:], axis=0)} Mean ID loss: {np.mean(id_list[-c:], axis=0)}')\n",
        "        c = 0\n",
        "                      "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JruweKNrl_ZD"
      },
      "source": [
        "#Build models and initialize optimizers\n",
        "\n",
        "#If load_model=True, specify the path where the models are saved\n",
        "\n",
        "gen,critic,siam, [opt_gen,opt_disc] = get_networks(shape, load_model=False, path='../content/drive/My Drive/')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BknKCA-8yqap",
        "cellView": "both"
      },
      "source": [
        "#Training\n",
        "\n",
        "#n_save = how many epochs between each saving and displaying of results\n",
        "#gupt = how many discriminator updates for generator+siamese update\n",
        "\n",
        "train(5000, batch_size=bs, lr=0.0002, n_save=1, gupt=3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A-f6nSiF95H-"
      },
      "source": [
        "#After Training, use these functions to convert data with the generator and save the results\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(spec, name, path='../content/', show=False):\n",
        "  specarr = chopspec(spec)\n",
        "  print(specarr.shape)\n",
        "  a = specarr\n",
        "  print('Generating...')\n",
        "  ab = gen(a, training=False)\n",
        "  print('Assembling and Converting...')\n",
        "  a = specass(a,spec)\n",
        "  ab = specass(ab,spec)\n",
        "  awv = deprep(a)\n",
        "  abwv = deprep(ab)\n",
        "  print('Saving...')\n",
        "  pathfin = f'{path}/{name}'\n",
        "  os.mkdir(pathfin)\n",
        "  sf.write(pathfin+'/AB.wav', abwv, sr)\n",
        "  sf.write(pathfin+'/A.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('Source')\n",
        "    axs[1].imshow(np.flip(ab, -2), cmap=None)\n",
        "    axs[1].axis('off')\n",
        "    axs[1].set_title('Generated')\n",
        "    plt.show()\n",
        "  return abwv"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6FZE91V1BIJX"
      },
      "source": [
        "#Wav to wav conversion\n",
        "\n",
        "wv, sr = librosa.load(librosa.util.example_audio_file(), sr=16000)  #Load waveform\n",
        "print(wv.shape)\n",
        "speca = prep(wv)                                                    #Waveform to Spectrogram\n",
        "\n",
        "plt.figure(figsize=(50,1))                                          #Show Spectrogram\n",
        "plt.imshow(np.flip(speca, axis=0), cmap=None)\n",
        "plt.axis('off')\n",
        "plt.show()\n",
        "\n",
        "abwv = towave(speca, name='FILENAME1', path='../content/')           #Convert and save wav"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}