MELGAN_VC.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sg83HLkK3UD2",
        "outputId": "de2d5d8e-1568-4a54-a091-b2072f8df4ba"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/drive/MyDrive/\"Voice Swapper\"/cmu_us_jmk_arctic"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wrE-0prX3YBL",
        "outputId": "0b1e37cb-5bd0-49ea-c0d3-9a8fabe2e025"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/MyDrive/Voice Swapper/cmu_us_jmk_arctic\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install wav2mel"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dnOPJXl6KLjL",
        "outputId": "48053648-e58d-44ff-e585-2080bb32e704"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: wav2mel in /usr/local/lib/python3.8/dist-packages (0.3.0)\n",
            "Requirement already satisfied: numpy==1.19.0 in /usr/local/lib/python3.8/dist-packages (from wav2mel) (1.19.0)\n",
            "Requirement already satisfied: librosa==0.7.2 in /usr/local/lib/python3.8/dist-packages (from wav2mel) (0.7.2)\n",
            "Requirement already satisfied: scipy==1.5.1 in /usr/local/lib/python3.8/dist-packages (from wav2mel) (1.5.1)\n",
            "Requirement already satisfied: numba==0.48 in /usr/local/lib/python3.8/dist-packages (from wav2mel) (0.48.0)\n",
            "Requirement already satisfied: joblib>=0.12 in /usr/local/lib/python3.8/dist-packages (from librosa==0.7.2->wav2mel) (1.2.0)\n",
            "Requirement already satisfied: six>=1.3 in /usr/local/lib/python3.8/dist-packages (from librosa==0.7.2->wav2mel) (1.15.0)\n",
            "Requirement already satisfied: audioread>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from librosa==0.7.2->wav2mel) (3.0.0)\n",
            "Requirement already satisfied: scikit-learn!=0.19.0,>=0.14.0 in /usr/local/lib/python3.8/dist-packages (from librosa==0.7.2->wav2mel) (1.0.2)\n",
            "Requirement already satisfied: decorator>=3.0.0 in /usr/local/lib/python3.8/dist-packages (from librosa==0.7.2->wav2mel) (4.4.2)\n",
            "Requirement already satisfied: resampy>=0.2.2 in /usr/local/lib/python3.8/dist-packages (from librosa==0.7.2->wav2mel) (0.3.1)\n",
            "Requirement already satisfied: soundfile>=0.9.0 in /usr/local/lib/python3.8/dist-packages (from librosa==0.7.2->wav2mel) (0.11.0)\n",
            "Requirement already satisfied: llvmlite<0.32.0,>=0.31.0dev0 in /usr/local/lib/python3.8/dist-packages (from numba==0.48->wav2mel) (0.31.0)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.8/dist-packages (from numba==0.48->wav2mel) (57.4.0)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn!=0.19.0,>=0.14.0->librosa==0.7.2->wav2mel) (3.1.0)\n",
            "Requirement already satisfied: cffi>=1.0 in /usr/local/lib/python3.8/dist-packages (from soundfile>=0.9.0->librosa==0.7.2->wav2mel) (1.15.1)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.8/dist-packages (from cffi>=1.0->soundfile>=0.9.0->librosa==0.7.2->wav2mel) (2.21)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wav2mel /wav/arctic_a0001.wav"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "81GhbfGlXdkA",
        "outputId": "9db68711-1f39-4ed1-8c29-8a8af35569f6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/bin/bash: wav2mel: command not found\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2Fq6Mg0BlbYo"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torchvision\n",
        "import torch.nn as nn\n",
        "import math"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import math\n",
        "\n",
        "import torch\n",
        "import torchvision\n",
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "class Discriminator(nn.Module):\n",
        "  def __init__(self,input_nc=3, filters=64, n_layers=0):\n",
        "    netD = []\n",
        "        \n",
        "    # input is (nc) x 256 x 256\n",
        "    netD.append(nn.Conv2d(input_nc, filters, kernel_size=4,stride=2))\n",
        "    netD.append(nn.LeakyReLU(0.2, True))\n",
        "    \n",
        "    nf_mult = 1\n",
        "    nf_mult_prev = 1\n",
        "    for n in range(1,n_layers):\n",
        "      nf_mult_prev = nf_mult\n",
        "      nf_mult = math.min(2^n,8)\n",
        "      netD.append(nn.Conv2d(filters * nf_mult_prev, filters * nf_mult, kernel_size=4, stride=2))\n",
        "      netD.append(nn.BatchNorm2d(filters * nf_mult))\n",
        "      netD.append(nn.LeakyReLU(0.2, True))\n",
        "    \n",
        "    # state size: (filters*M) x N x N\n",
        "    nf_mult_prev = nf_mult\n",
        "    nf_mult = math.min(2^n_layers,8)\n",
        "    netD.append(nn.Conv2d(filters * nf_mult_prev, filters * nf_mult, kernel_size=4,stride=1))\n",
        "    netD.append(nn.BatchNorm2d(filters * nf_mult))\n",
        "    netD.append(nn.LeakyReLU(0.2, True))\n",
        "    \n",
        "    # state size: (filters*M*2) x (N-1) x (N-1)\n",
        "    netD.append(nn.Conv2d(filters * nf_mult, 1, kernel_size=4,stride=1))\n",
        "    # state size: 1 x (N-2) x (N-2)\n",
        "    \n",
        "    netD.append(nn.Sigmoid())\n",
        "    # state size: 1 x (N-2) x (N-2)\n",
        "    \n",
        "    self.model = nn.Sequential(*netD)\n",
        "\n",
        "  def forward(self,x):\n",
        "    return self.model(x)"
      ],
      "metadata": {
        "id": "9_kPTOB7l1Fn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\" Parts of the U-Net model \"\"\"\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "class DoubleConv(nn.Module):\n",
        "    \"\"\"(convolution => [BN] => ReLU) * 2\"\"\"\n",
        "\n",
        "    def __init__(self, in_channels, out_channels, mid_channels=None):\n",
        "        super().__init__()\n",
        "        if not mid_channels:\n",
        "            mid_channels = out_channels\n",
        "        self.double_conv = nn.Sequential(\n",
        "            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(mid_channels),\n",
        "            nn.ReLU(inplace=True),\n",
        "            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),\n",
        "            nn.BatchNorm2d(out_channels),\n",
        "            nn.ReLU(inplace=True)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.double_conv(x)\n",
        "\n",
        "\n",
        "class Encode(nn.Module):\n",
        "    \"\"\"Downscaling with maxpool then double conv\"\"\"\n",
        "\n",
        "    def __init__(self, in_channels, out_channels):\n",
        "        super().__init__()\n",
        "        self.maxpool_conv = nn.Sequential(\n",
        "            nn.MaxPool2d(2),\n",
        "            DoubleConv(in_channels, out_channels)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.maxpool_conv(x)\n",
        "\n",
        "\n",
        "class Decode(nn.Module):\n",
        "    \"\"\"Upscaling then double conv\"\"\"\n",
        "\n",
        "    def __init__(self, in_channels, out_channels, bilinear=True):\n",
        "        super().__init__()\n",
        "\n",
        "        # if bilinear, use the normal convolutions to reduce the number of channels\n",
        "        if bilinear:\n",
        "            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
        "            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)\n",
        "        else:\n",
        "            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)\n",
        "            self.conv = DoubleConv(in_channels, out_channels)\n",
        "\n",
        "    def forward(self, x1, x2):\n",
        "        x1 = self.up(x1)\n",
        "        # input is CHW\n",
        "        diffY = x2.size()[2] - x1.size()[2]\n",
        "        diffX = x2.size()[3] - x1.size()[3]\n",
        "\n",
        "        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,\n",
        "                        diffY // 2, diffY - diffY // 2])\n",
        "        # if you have padding issues, see\n",
        "        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a\n",
        "        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd\n",
        "        x = torch.cat([x2, x1], dim=1)\n",
        "        return self.conv(x)\n",
        "\n",
        "\n",
        "class OutConv(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels):\n",
        "        super(OutConv, self).__init__()\n",
        "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.conv(x)"
      ],
      "metadata": {
        "id": "8SK5B11VCTPr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from .unet_parts import *\n",
        "\n",
        "class Generator(nn.Module):\n",
        "  def __init__(self, in_channels, out_channels, bilinear=False):\n",
        "    super().__init__()\n",
        "\n",
        "    self.in_channels = in_channels\n",
        "    self.out_channels = out_channels\n",
        "    self.bilinear = bilinear\n",
        "\n",
        "    # self.inc = (DoubleConv(in_channels, 64))\n",
        "    self.encode1 = (Encode(in_channels, 64))\n",
        "    self.encode2 = (Encode(64, 128))\n",
        "    self.encode3 = (Encode(128, 256))\n",
        "    factor = 2 if bilinear else 1\n",
        "    self.encode4 = (Encode(256, 512 // factor))\n",
        "    self.decode1 = (Decode(512, 256 // factor, bilinear))\n",
        "    self.decode2 = (Decode(256, 128 // factor, bilinear))\n",
        "    self.decode3 = (Decode(128, 64 // factor, bilinear))\n",
        "    self.decode4 = (Decode(64, out_channels, bilinear))\n",
        "    self.out = nn.Tanh()\n",
        "    # self.outc = (OutConv(64, out_channels))\n",
        "  \n",
        "  def forward(self, x): #x is noise\n",
        "    x_l1,x_l2 = np.split(x,2,axis=1)\n",
        "\n",
        "    x2 = self.encode1(x_l1)\n",
        "    x3 = self.encode2(x2)\n",
        "    x4 = self.encode3(x3)\n",
        "    x5 = self.encode4(x4)\n",
        "    x = self.decode1(x5, x4)\n",
        "    x = self.decode2(x, x3)\n",
        "    x = self.decode3(x, x2)\n",
        "    x = self.decode4(x, x1)\n",
        "    output_l1 = self.outc(x)\n",
        "\n",
        "    x2 = self.encode1(x_l2)\n",
        "    x3 = self.encode2(x2)\n",
        "    x4 = self.encode3(x3)\n",
        "    x5 = self.encode4(x4)\n",
        "    x = self.decode1(x5, x4)\n",
        "    x = self.decode2(x, x3)\n",
        "    x = self.decode3(x, x2)\n",
        "    x = self.decode4(x, x1)\n",
        "    output_l2 = self.outc(x)\n",
        "\n",
        "    return np.concat((output_l1,output_l2),axis=1)"
      ],
      "metadata": {
        "id": "mIp2c9X_nuv5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "in_channels = MEL_CHANNELS\n",
        "out_channels = MEL_CHANNELS\n",
        "\n",
        "netG = Generator(in_channels,out_channels)\n",
        "netD = Discriminator()\n",
        "fft = Audio2Mel()"
      ],
      "metadata": {
        "id": "fn7cXfn8Qy_a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.optim as optim\n",
        "\n",
        "criterion = nn.BCELoss()\n",
        "optimizerG = optim.Adam(netG.parameters(), lr=1e-4)\n",
        "optimizerD = optim.Adam(netD.parameters(), lr=4e-4)"
      ],
      "metadata": {
        "id": "kfD9HaZqnwOl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "Isbt37OXuQ_2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "for epoch in range(NUM_EPOCHS):\n",
        "  for i,real_audio in enumerate(dataloader):\n",
        "    # convert to spec\n",
        "    x = fft(real_audio)\n",
        "    # get generated spec\n",
        "    x_pred = netG(x)\n",
        "    # get output of discriminator on real spectrogram\n",
        "    netD.zero_grad()\n",
        "    source_spec,target_spec = data[0],data[1]\n",
        "    output = netD(target_spec)\n",
        "    errD_real = criterion(output,1) #assume 1 is real and 0 is fake\n",
        "    errD_real.backward()\n",
        "    D_x = output.mean().item()\n",
        "\n",
        "    # train generator\n",
        "    fake = netG(source_spec)\n",
        "    output = netD(fake.detach()).view(-1)\n",
        "    errD_fake = criterion(output,0)\n",
        "    errD_fake.backward()\n",
        "    D_G_z1 = output.mean().item()\n",
        "\n",
        "    # Compute error of D as sum over the fake and the real batches\n",
        "    errD = errD_real + errD_fake\n",
        "\n",
        "    optimizerD.step()\n",
        "\n",
        "    ############################\n",
        "    # (2) Update G network: maximize log(D(G(z)))\n",
        "    ###########################\n",
        "    netG.zero_grad()\n",
        "    label.fill_(real_label)  # fake labels are real for generator cost\n",
        "    # Since we just updated D, perform another forward pass of all-fake batch through D\n",
        "    output = netD(fake).view(-1)\n",
        "    # Calculate G's loss based on this output\n",
        "    errG = criterion(output, label)\n",
        "    # Calculate gradients for G\n",
        "    errG.backward()\n",
        "    D_G_z2 = output.mean().item()\n",
        "    # Update G\n",
        "    optimizerG.step()"
      ],
      "metadata": {
        "id": "niNP3MdQnmOC"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}