audiolm_e2e_replace_hubert_and_soundstream.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h9muioqT4g7a"
      },
      "outputs": [],
      "source": [
        "# run audiolm with replaced soundstream etc-- see changes starting with\n",
        "# https://github.com/LWprogramming/audiolm-pytorch/commit/37be3c512cadeecab2184def3dd9bb12171b1bca\n",
        "\n",
        "!rm -rf audiolm_pytorch/ setup.py audiolm-pytorch.zip # clean out any old stuff floating around\n",
        "\n",
        "!pip install torch datasets boto3\n",
        "\n",
        "# download audiolm_pytorch manually so i can inject print statements\n",
        "# !pip uninstall -y audiolm_pytorch\n",
        "\n",
        "raise AssertionError(\"don't forget to upload the customized version of audiolm_pytorch with print statements\")\n",
        "\n",
        "# !zip -r audiolm_pytorch.zip audiolm_pytorch/\n",
        "\n",
        "import urllib.request\n",
        "import os\n",
        "import zipfile\n",
        "if not os.path.isfile(\"audiolm-pytorch.zip\"):\n",
        "  urllib.request.urlretrieve(\"https://github.com/LWProgramming/audiolm-pytorch/archive/refs/heads/personal_hacks.zip\", \"audiolm-pytorch.zip\")\n",
        "if not os.path.isdir(\"audiolm-pytorch\"):\n",
        "  with zipfile.ZipFile(\"audiolm-pytorch.zip\", 'r') as zip_ref:\n",
        "    zip_ref.extractall(\"audiolm-pytorch\")\n",
        "!mv audiolm-pytorch/audiolm-pytorch-personal_hacks/audiolm_pytorch .\n",
        "\n",
        "# install necessary files for patched audiolm-pytorch\n",
        "!mv audiolm-pytorch/audiolm-pytorch-personal_hacks/setup.py .\n",
        "!pip install . # install requirements from the patched audiolm-pytorch dir\n",
        "!rm -rf audiolm-pytorch # not the one with underscore which is the actual library"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Semantic (MERT)"
      ],
      "metadata": {
        "id": "eM_xUgNP5Coo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# original semantic transformer\n",
        "import torch\n",
        "from audiolm_pytorch import HubertWithKmeans, SemanticTransformer, SemanticTransformerWrapper, SemanticTransformerTrainer\n",
        "import os\n",
        "import urllib\n",
        "\n",
        "# hubert checkpoints can be downloaded at\n",
        "# https://github.com/facebookresearch/fairseq/tree/main/examples/hubert\n",
        "\n",
        "hubert_ckpt = 'hubert/hubert_base_ls960.pt'\n",
        "hubert_quantizer = f'hubert/hubert_base_ls960_L9_km500.bin' # listed in row \"HuBERT Base (~95M params)\", column Quantizer\n",
        "if not os.path.isdir(\"hubert\"):\n",
        "  os.makedirs(\"hubert\")\n",
        "if not os.path.isfile(hubert_ckpt):\n",
        "  hubert_ckpt_download = f\"https://dl.fbaipublicfiles.com/{hubert_ckpt}\"\n",
        "  urllib.request.urlretrieve(hubert_ckpt_download, f\"./{hubert_ckpt}\")\n",
        "if not os.path.isfile(hubert_quantizer):\n",
        "  hubert_quantizer_download = f\"https://dl.fbaipublicfiles.com/{hubert_quantizer}\"\n",
        "  urllib.request.urlretrieve(hubert_quantizer_download, f\"./{hubert_quantizer}\")\n",
        "\n",
        "wav2vec = HubertWithKmeans(\n",
        "    # checkpoint_path = './hubert/hubert_base_ls960.pt',\n",
        "    checkpoint_path = None,\n",
        "    kmeans_path = './hubert/hubert_base_ls960_L9_km500.bin',\n",
        "    use_mert = True\n",
        ")\n",
        "\n",
        "semantic_transformer = SemanticTransformer(\n",
        "    num_semantic_tokens = wav2vec.codebook_size,\n",
        "    dim = 1024,\n",
        "    depth = 6\n",
        ")\n",
        "\n"
      ],
      "metadata": {
        "id": "Z8-mZuqO43pP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Coarse (encodec)"
      ],
      "metadata": {
        "id": "24swIU7X5KII"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from audiolm_pytorch import CoarseTransformer\n",
        "\n",
        "assert 'wav2vec' in locals() # expect that we have a wav2vec from semantic part\n",
        "\n",
        "coarse_transformer = CoarseTransformer(\n",
        "    num_semantic_tokens = wav2vec.codebook_size,\n",
        "    codebook_size = 1024,\n",
        "    num_coarse_quantizers = 3,\n",
        "    dim = 512,\n",
        "    depth = 6\n",
        ")"
      ],
      "metadata": {
        "id": "Y16xJG9k5L5k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Fine (encodec)"
      ],
      "metadata": {
        "id": "G5C1PTMg5NHg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from audiolm_pytorch import FineTransformer\n",
        "\n",
        "fine_transformer = FineTransformer(\n",
        "    num_coarse_quantizers = 3,\n",
        "    num_fine_quantizers = 5,\n",
        "    codebook_size = 1024,\n",
        "    dim = 512,\n",
        "    depth = 6\n",
        ")"
      ],
      "metadata": {
        "id": "3QlS_wo65WFP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate"
      ],
      "metadata": {
        "id": "gom0cmwu51CK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from audiolm_pytorch import AudioLM\n",
        "\n",
        "audiolm = AudioLM(\n",
        "    wav2vec = wav2vec,\n",
        "    soundstream = soundstream,\n",
        "    semantic_transformer = semantic_transformer,\n",
        "    coarse_transformer = coarse_transformer,\n",
        "    fine_transformer = fine_transformer\n",
        ").cuda() # TODO: possibly things are already cuda but I'm not sure\n",
        "\n",
        "generated_wav = audiolm(batch_size = 1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 171
        },
        "id": "PERm5jim5oVz",
        "outputId": "501927e5-2edf-4569-b198-9046334d40c6"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "error",
          "ename": "AssertionError",
          "evalue": "ignored",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-2-a871fdc9ebee>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0;32massert\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
            "\u001b[0;31mAssertionError\u001b[0m: "
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "x = True\n",
        "y = True\n",
        "not(x and y)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XeJ90I-A5o3Z",
        "outputId": "88eaac62-373f-4170-a39e-ae368d1756ba"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "uQ0ubwK5D6rj"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}