MelGAN_VC_MLAIforSoundWorkshop.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "MelGAN_VC_MLAIforSoundWorkshop.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "SJjEEQOhlmIw",
        "cellView": "form"
      },
      "source": [
        "#@title Clone MelGAN-VC Repository\n",
        "! git clone https://github.com/moiseshorta/MelGAN-VC.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NXZxS_YtUDrq",
        "cellView": "form"
      },
      "source": [
        "#@title Mount your Google Drive\n",
        "#Mount your Google Drive account\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-n9J8OiI2Ftq"
      },
      "source": [
        "#Get Example Datasets\n",
        "%cd /content/\n",
        "\n",
        "#Target Audio = Antonio Zepeda - Templo Mayor \n",
        "!wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1-1pKhCyccFc6cDIHe_vOZFSRak6eh5Dy' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=1-1pKhCyccFc6cDIHe_vOZFSRak6eh5Dy\" -O az.zip && rm -rf /tmp/cookies.txt\n",
        "\n",
        "#Source Audio = Native American flutes \n",
        "!wget --load-cookies /tmp/cookies.txt \"https://docs.google.com/uc?export=download&confirm=$(wget --quiet --save-cookies /tmp/cookies.txt --keep-session-cookies --no-check-certificate 'https://docs.google.com/uc?export=download&id=1-BR2SlBmxskjBKHKqazUiyIL0GGdyeg7' -O- | sed -rn 's/.*confirm=([0-9A-Za-z_]+).*/\\1\\n/p')&id=1-BR2SlBmxskjBKHKqazUiyIL0GGdyeg7\" -O flautaindigena.zip && rm -rf /tmp/cookies.txt\n",
        "\n",
        "!unzip az.zip\n",
        "!unzip flautaindigena.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V00rptcdKSbq"
      },
      "source": [
        "#@title Import Tensorflow and torchaudio\n",
        "#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 torch==1.8.0\n",
        "!pip install --no-deps torchaudio==0.8\n",
        "!pip uninstall tensorflow\n",
        "!pip install tensorflow-gpu==2.1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LEvqwT96l_Yq",
        "cellView": "form"
      },
      "source": [
        "#@title 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",
        "cellView": "form"
      },
      "source": [
        "#@title Hyperparameters\n",
        "#Hyperparameters\n",
        "\n",
        "hop=512               #hop size (window size = 6*hop)\n",
        "sr=44100              #sampling rate\n",
        "min_level_db=-100     #reference values to normalize data\n",
        "ref_level_db=20\n",
        "\n",
        "shape=64              #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",
        "cellView": "form"
      },
      "source": [
        "#@title Waveform to Spectrogram converter\n",
        "#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",
        "cellView": "form"
      },
      "source": [
        "#@title Helper Functions\n",
        "#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",
        "#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",
        "cellView": "form"
      },
      "source": [
        "#@title Generating Mel-Spectrogram dataset. Audio files must be 44.1Khz/16bit .wav \n",
        "source_audio_directory = \"/content/flautaindigena\" #@param {type:\"string\"}\n",
        "target_audio_directory = \"/content/AntonioZepeda\" #@param {type:\"string\"}\n",
        "#Generating Mel-Spectrogram dataset (Uncomment where needed)\n",
        "#adata: source spectrograms\n",
        "#bdata: target spectrograms\n",
        "\n",
        "#SOURCE\n",
        "awv = audio_array(source_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 length\n",
        "#TARGET\n",
        "bwv = audio_array(target_audio_directory)\n",
        "bspec = tospec(bwv)\n",
        "bdata = splitcut(bspec)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qSesIbwr_GyO",
        "cellView": "form"
      },
      "source": [
        "#@title Creates Tensorflow Datasets\n",
        "#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",
        "cellView": "form"
      },
      "source": [
        "#@title Adding Spectral Normalization to convolutional layers\n",
        "#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",
        "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",
        "class DenseSN(Dense):\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",
        "cellView": "form"
      },
      "source": [
        "#@title Networks Architecture\n",
        "#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",
        "cellView": "form"
      },
      "source": [
        "#@title Save in training loop\n",
        "checkpoint_save_directory = \"/content/MelGAN-VC/checkpoint\" #@param {type:\"string\"}\n",
        "#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=checkpoint_save_directory):                 #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",
        "cellView": "form"
      },
      "source": [
        "#@title Losses\n",
        "#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",
        "cellView": "form"
      },
      "source": [
        "#@title Get models and optimizers & Set up learning rate\n",
        "#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",
        "cellView": "form"
      },
      "source": [
        "#@title Training Functions set to Voice Modality\n",
        "#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",
        "    #loss_id = 0.\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": {
        "cellView": "form",
        "id": "ukAs1Fo3sE5X"
      },
      "source": [
        "#@title Training Functions set to Music Modality\n",
        "#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",
        "    loss_id = 0.\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",
        "cellView": "form"
      },
      "source": [
        "#@title Training Loop\n",
        "#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",
        "cellView": "form"
      },
      "source": [
        "#@title Build models and initialize optimizers.  If resume_model=True, specify the path where the models are saved\n",
        "resume_model = True #@param {type:\"boolean\"}\n",
        "model_checkpoints_directory = \"/content/MelGAN-VC/checkpoint\" #@param {type:\"string\"}\n",
        "#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=resume_model, path=model_checkpoints_directory)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BknKCA-8yqap",
        "cellView": "form"
      },
      "source": [
        "#@title Training. epoch_save = how many epochs between each saving and displaying of results. n_epochs = how many epochs the model will train\n",
        "epoch_save = 10#@param {type:\"integer\"}\n",
        "n_epochs = 3000 #@param {type:\"integer\"}\n",
        "#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(n_epochs, batch_size=bs, lr=0.0002, n_save=epoch_save, gupt=3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A-f6nSiF95H-",
        "cellView": "form"
      },
      "source": [
        "#@title After Training, use these functions to convert data with the generator and save the results\n",
        "#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": "9PgZZwGznvZV",
        "cellView": "form"
      },
      "source": [
        "#@title Generator and wav to wav conversion\n",
        "input_wav = \"/content/flautaindigena/23.wav\" #@param {type:\"string\"}\n",
        "output_name = \"flauta_to_AntonioZepeda\" #@param {type:\"string\"}\n",
        "output_directory = \"/content/MelGAN-VC\" #@param {type:\"string\"}\n",
        "#Wav to wav conversion\n",
        "\n",
        "wv, sr = librosa.load(input_wav, sr=44100)  #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=output_name, path=output_directory)          #Convert and save wav"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}