Latent_Soundings_workshop_RAVE.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/hdparmar/AI-Music/blob/main/Latent_Soundings_workshop_RAVE.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pMrgS5Yn7Us7",
        "outputId": "54a62606-2e78-4db6-c0ee-0daf8b8043c4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n",
            "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n",
            "yes: stdout: Broken pipe\n",
            "      ___           ___           ___           ___     \n",
            "     /\\  \\         /\\  \\         /\\__\\         /\\  \\    \n",
            "    /::\\  \\       /::\\  \\       /:/  /        /::\\  \\   \n",
            "   /:/\\:\\  \\     /:/\\:\\  \\     /:/  /        /:/\\:\\  \\  \n",
            "  /::\\~\\:\\  \\   /::\\~\\:\\  \\   /:/__/  ___   /::\\~\\:\\  \\ \n",
            " /:/\\:\\ \\:\\__\\ /:/\\:\\ \\:\\__\\  |:|  | /\\__\\ /:/\\:\\ \\:\\__\\\n",
            " \\/_|::\\/:/  / \\/__\\:\\/:/  /  |:|  |/:/  / \\:\\~\\:\\ \\/__/\n",
            "    |:|::/  /       \\::/  /   |:|__/:/  /   \\:\\ \\:\\__\\  \n",
            "    |:|\\/__/        /:/  /     \\::::/__/     \\:\\ \\/__/  \n",
            "    |:|  |         /:/  /       ~~~~          \\:\\__\\    \n",
            "     \\|__|         \\/__/                       \\/__/    \n",
            "\n",
            "\n",
            "\n",
            "\n",
            "\n",
            "RAVE: Realtime Audio Variational autoEncoder | Authors: Antoine Caillon, Acids-IRCAM | https://github.com/acids-ircam/RAVE | https://arxiv.org/abs/2111.05011\n",
            "Custom collab implementation by Moisés Horta | http://twitter.com/hexorcismos | http://github.com/moiseshorta\n",
            "\n",
            "\n",
            "Last Modified: 4.October.2022\n"
          ]
        }
      ],
      "source": [
        "#@title README\n",
        "!yes | pip install --quiet pyfiglet\n",
        "\n",
        "import pyfiglet\n",
        "  \n",
        "resulta = pyfiglet.figlet_format(f\"RAVE\", font = \"isometric1\", )\n",
        "# resultb = pyfiglet.figlet_format(f\"TRAIN\", font = \"isometric2\", )\n",
        "print(resulta)\n",
        "# print(resultb)\n",
        "print(\"\\n\")\n",
        "print(\"\\n\")\n",
        "print(\"RAVE: Realtime Audio Variational autoEncoder | Authors: Antoine Caillon, Acids-IRCAM | https://github.com/acids-ircam/RAVE | https://arxiv.org/abs/2111.05011\")\n",
        "print(\"Custom collab implementation by Moisés Horta | http://twitter.com/hexorcismos | http://github.com/moiseshorta\")\n",
        "print(\"\\n\")\n",
        "print(\"Last Modified: 4.October.2022\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HQapDcKr7YV-",
        "outputId": "a7fb359c-439a-4ab9-bf98-f77e4e129af3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "zsh:1: command not found: nvidia-smi\n"
          ]
        }
      ],
      "source": [
        "#@title Check GPU\n",
        "\n",
        "import subprocess\n",
        "!nvidia-smi\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_6GvzMiT7jGo"
      },
      "outputs": [],
      "source": [
        "#@title Anti-disconnect from Colab\n",
        "\n",
        "import IPython\n",
        "js_code = '''\n",
        "function ClickConnect(){\n",
        "console.log(\"Working\");\n",
        "document.querySelector(\"colab-toolbar-button#connect\").click()\n",
        "}\n",
        "setInterval(ClickConnect,60000)\n",
        "'''\n",
        "display(IPython.display.Javascript(js_code))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Yjj8Yj-PY0Q9"
      },
      "outputs": [],
      "source": [
        "#@title Mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "V4wWTVf57mv8"
      },
      "outputs": [],
      "source": [
        "#@title Clone repository & install dependencies\n",
        "\n",
        "%cd /content/drive/MyDrive\n",
        "!git clone --quiet --depth 1 https://github.com/acids-ircam/RAVE.git RAVE_COLLAB\n",
        "print(\"Cloned RAVE repository\")\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB\n",
        "!curl -L https://repo.anaconda.com/miniconda/Miniconda3-py39_4.12.0-Linux-x86_64.sh -o miniconda.sh\n",
        "!chmod +x miniconda.sh\n",
        "!sh miniconda.sh -b -p /content/miniconda\n",
        "!/content/miniconda/bin/pip install --quiet -r requirements.txt\n",
        "!/content/miniconda/bin/pip install --quiet --upgrade ipython ipykernel\n",
        "!pip install youtube-dl\n",
        "print(\"Succesfully installed requirements\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kLEt6PLoPwhi"
      },
      "outputs": [],
      "source": [
        "#@title [Optional] Download audio dataset from Youtube\n",
        "\n",
        "!mkdir dataset \n",
        "%cd dataset\n",
        "\n",
        "youtube_link = \"https://www.youtube.com/watch?v=h2OuSwK7yn0\" #@param {type:\"string\"}\n",
        "\n",
        "!youtube-dl -x --audio-format wav --audio-quality 0 $youtube_link "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wh3qXCKKTsKV"
      },
      "outputs": [],
      "source": [
        "#@title [Optional] Apply audio augmentation to your dataset and check total duration\n",
        "\n",
        "dataset = \"/path/to/your/dataset\" #@param {type:\"string\"}\n",
        "sample_rate = 48000 #@param {type:\"integer\"}\n",
        "\n",
        "%cd $dataset\n",
        "\n",
        "!/content/miniconda/bin/resample --input $dataset --sr $sample_rate --output $dataset/resampled/\n",
        "\n",
        "%cd $dataset/resampled/\n",
        "\n",
        "!/content/miniconda/bin/duration\n",
        "\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "An090S938lha"
      },
      "source": [
        "## *Training*"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "pAWZGYYfb0i3"
      },
      "source": [
        "## Configure your training\n",
        "\n",
        "Here you can define the parameters for your training, also called *run*. This is a collab version of the `cli_helper.py` script.\n",
        "\n",
        "Input dataset is the path to your audio dataset. \n",
        "\n",
        "**PRO TIP**: using the left panel you can navigate to your Drive, find your dataset folder, right click and copy its path ! "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3uv_XWr3bsL7"
      },
      "outputs": [],
      "source": [
        "training_name = \"test\"  #@param {type:\"string\"}\n",
        "input_dataset = \"/path/to/audio/dataset.wav\"  #@param {type:\"string\"}\n",
        "sampling_rate = 48000  #@param {type:\"integer\"}\n",
        "multiband_number = 16  #@param {type:\"integer\"}\n",
        "n_signal = 65536  #@param {type:\"integer\"}\n",
        "size = \"default\"  #@param [\"small\", \"default\", \"large\"]\n",
        "prior = 32  #@param {type:\"integer\"}\n",
        "fidelity = 99  #@param {type:\"slider\", min:80, max:99, step:1}\n",
        "no_latency_mode = False  #@param {type:\"boolean\"}\n",
        "regularization_strength = 0.1  #@param {type:\"slider\", min:0.001, max:1, step:0.001}\n",
        "validation_every = 2500  #@param {type:\"integer\"}\n",
        "\n",
        "training_name = training_name.replace(\" \", \"_\").lower()\n",
        "no_latency_mode = \"true\" if no_latency_mode else \"false\"\n",
        "# sampling_rate = str(sampling_rate)\n",
        "# multiband_number = str(multiband_number)\n",
        "# n_signal = str(n_signal)\n",
        "# prior = str(prior)\n",
        "# fidelity = str(fidelity)\n",
        "# regularization_strength = str(regularization_strength)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SKBWN8f-ByrJ"
      },
      "outputs": [],
      "source": [
        "#@title Setup Tensorboard to monitor training\n",
        "\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir . --bind_all "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KuORQP0N7q4z"
      },
      "outputs": [],
      "source": [
        "#@title Begin training!\n",
        "\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB\n",
        "\n",
        "!/content/miniconda/bin/python train_rave.py \\\n",
        "  -c $size \\\n",
        "  --data-size $multiband_number \\\n",
        "  --no-latency $no_latency_mode \\\n",
        "  --min-kl $regularization_strength \\\n",
        "  --max-kl $regularization_strength \\\n",
        "  --ckpt /content/drive/MyDrive/RAVE_COLLAB/runs/$training_name/rave/ \\\n",
        "  --preprocessed /content/preprocessed/rave \\\n",
        "  --wav $input_dataset \\\n",
        "  --sr $sampling_rate \\\n",
        "  --n-signal $n_signal \\\n",
        "  --val-every $validation_every \\\n",
        "  --name $training_name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PN2X9WkjBKFD"
      },
      "outputs": [],
      "source": [
        "#@title Resume training\n",
        "\n",
        "\n",
        "checkpoint = \"/path/to/rave.ckpt\" #@param {type:\"string\"}\n",
        "\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB\n",
        "!/content/miniconda/bin/python train_rave.py \\\n",
        "  -c $size --data-size $multiband_number \\\n",
        "  --no-latency $no_latency_mode \\\n",
        "  --min-kl $regularization_strength \\\n",
        "  --max-kl $regularization_strength \\\n",
        "  --ckpt $checkpoint \\\n",
        "  --preprocessed /content/preprocessed/rave \\\n",
        "  --wav $input_dataset \\\n",
        "  --sr $sampling_rate \\\n",
        "  --n-signal $n_signal \\\n",
        "  --val-every $validation_every \\\n",
        "  --name $training_name"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "B9iz90MYHe6i"
      },
      "outputs": [],
      "source": [
        "#@title Export trained RAVE as torchscript (.ts)\n",
        "\n",
        "name_of_run = \"test\" #@param {type:\"string\"}\n",
        "# checkpoint = \"/path/to/checkpoint.ckpt\" #@param {type:\"string\"}\n",
        "\n",
        "# !python /content/RAVE/export_rave.py \\\n",
        "#   --name $name_of_run \\\n",
        "#   --run $checkpoint \\\n",
        "#   --cached true \\\n",
        "#   --fidelity 0.95 \\\n",
        "#   --stereo true\n",
        "\n",
        "checkpoint = \"/path/to/rave.ckpt\"  #@param {type:\"string\"}\n",
        "\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB\n",
        "\n",
        "!/content/miniconda/bin/python export_rave.py --run $checkpoint --cached false --name $name_of_run\n",
        "!/content/miniconda/bin/python export_rave.py --run $checkpoint --cached true --name {name_of_run}_rt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "E3BFoe4sJS_R"
      },
      "outputs": [],
      "source": [
        "#@title Begin training Prior\n",
        "\n",
        "name_of_run = \"test\" #@param {type:\"string\"}\n",
        "audio_dataset = \"/path/to/audio/dataset\" #@param {type:\"string\"}\n",
        "pretrained_rave = \"/path/to/pretrained-rave.ts\" #@param {type:\"string\"}\n",
        "# preprocessed_audio_location = \"/content/temp4\"\n",
        "\n",
        "# !python /content/drive/MyDrive/RAVE_repo/update/RAVE/train_prior.py \\\n",
        "#   --name $name_of_run \\\n",
        "#   --wav $audio_dataset \\\n",
        "#   --preprocessed $preprocessed_audio_location \\\n",
        "#   --pretrained-vae $pretrained_rave\n",
        "\n",
        "\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB/\n",
        "\n",
        "!/content/miniconda/bin/python train_prior.py \\\n",
        "  --pretrained-vae rave_{name_of_run}.ts \\\n",
        "  --preprocessed /content/preprocessed/prior \\\n",
        "  --wav $audio_dataset \\\n",
        "  --name $name_of_run \\\n",
        "  --val-every $validation_every"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "DRG7-MKc86Ou"
      },
      "source": [
        "## *Synthesize audio*"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3W8onpK19MSI"
      },
      "outputs": [],
      "source": [
        "#@title Perform Timbre Transfer on a folder of wavs\n",
        "\n",
        "checkpoint = \"/path/to/checkpoint.ckpt\" #@param {type:\"string\"}\n",
        "wav_folder = \"/path/to/wav/folder\" #@param {type:\"string\"}\n",
        "\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB/\n",
        "\n",
        "!/content/miniconda/bin/python reconstruct.py \\\n",
        "  --ckpt $checkpoint \\\n",
        "  --wav-folder $wav_folder"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RZ7P0lDDBkhA"
      },
      "outputs": [],
      "source": [
        "#@title Modify the latent dimensions for timbre transfer exploration\n",
        "\n",
        "exported_rave_ts = \"/path/to/pretrained-rave.ts\" #@param {type:\"string\"}\n",
        "input_audio = \"/path/to/audio.wav\" #@param {type:\"string\"}\n",
        "output_audio = \"/content/rave_out.wav\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown > Latent Dimension Parameters\n",
        "d0 = 1.09 #@param {type:\"slider\", min:-3, max:3, step:0.01}\n",
        "d1 = 0 #@param {type:\"slider\", min:-3, max:3, step:0.01}\n",
        "d2 = 0 #@param {type:\"slider\", min:-3, max:3, step:0.01}\n",
        "d3 = 0 #@param {type:\"slider\", min:-3, max:3, step:0.01}\n",
        "\n",
        "import torch\n",
        "import librosa as li\n",
        "import soundfile as sf\n",
        "import IPython\n",
        "\n",
        "torch.set_grad_enabled(False)\n",
        "\n",
        "model = torch.jit.load(exported_rave_ts).eval()\n",
        "print(\"Loaded\",exported_rave_ts)\n",
        "x = li.load(input_audio)[0]\n",
        "\n",
        "#ENCODE AUDIO TO LATENT REPRESENTATION\n",
        "x = torch.from_numpy(x).reshape(1, 1, -1)\n",
        "z = model.encode(x)\n",
        "print(\"Encoded\",input_audio)\n",
        "\n",
        "#BIAS LATENT DIMENSIONS\n",
        "z[:, 0] += torch.linspace(d0,d0, z.shape[-1])\n",
        "z[:, 1] += torch.linspace(d1,d1, z.shape[-1])\n",
        "z[:, 2] += torch.linspace(d2,d2, z.shape[-1])\n",
        "z[:, 3] += torch.linspace(d3,d3, z.shape[-1])\n",
        "\n",
        "#SYNTHESIZE AUDIO\n",
        "y = model.decode(z).numpy().reshape(-1)\n",
        "sf.write(output_audio, y, 48000)\n",
        "\n",
        "print(\"RAVE output:\")\n",
        "IPython.display.Audio(output_audio)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "2ItYyvzFP0_W"
      },
      "outputs": [],
      "source": [
        "#@title Unconditional neural audio synthesis using RAVE + Prior\n",
        "\n",
        "prior_checkpoint = \"/path/to/prior.ckpt\" #@param {type:\"string\"}\n",
        "exported_rave = \"/path/to/pretrained-rave.ts\" #@param {type:\"string\"}\n",
        "seconds_to_generate = 15 #@param {type:\"integer\"}\n",
        "output_audio = \"/content/rave+prior_output.wav\" #@param {type:\"string\"}\n",
        "\n",
        "%cd /content/drive/MyDrive/RAVE_COLLAB/\n",
        "\n",
        "\n",
        "!/content/miniconda/bin/python generation.py \\\n",
        "  --prior-ckpt $prior_checkpoint \\\n",
        "  --exported-rave $exported_rave \\\n",
        "  --length $seconds_to_generate \\\n",
        "  --out-path $output_audio\n",
        "\n",
        "print(\"RAVE+Prior output:\")\n",
        "IPython.display.Audio(output_audio)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "include_colab_link": true,
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}