audiolm_pytorch_demo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "n337KoD2om3L",
        "outputId": "97ada0c6-f21c-483e-d63d-08abddd49004"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mon Jan 30 20:47:47 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   73C    P0    32W /  70W |  10692MiB / 15360MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|    0   N/A  N/A      5896      C                                   10689MiB |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "!nvidia-smi\n",
        "\n",
        "# If this doesn't work, there's no GPU available or detected"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install audiolm-pytorch"
      ],
      "metadata": {
        "id": "TLJAcUHpvmp4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "95bcda95-a484-40c6-e5a7-47f4378759a8"
      },
      "execution_count": 13,
      "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: audiolm-pytorch in /usr/local/lib/python3.8/dist-packages (0.7.5)\n",
            "Requirement already satisfied: ema-pytorch in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.1.4)\n",
            "Requirement already satisfied: sentencepiece in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.1.97)\n",
            "Requirement already satisfied: beartype in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.12.0)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (1.0.2)\n",
            "Requirement already satisfied: torchaudio in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.13.1+cu116)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (1.2.0)\n",
            "Requirement already satisfied: torch>=1.6 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (1.13.1+cu116)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (4.26.0)\n",
            "Requirement already satisfied: Mega-pytorch in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.0.12)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (4.64.1)\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.15.0)\n",
            "Requirement already satisfied: vector-quantize-pytorch>=0.10.15 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.10.15)\n",
            "Requirement already satisfied: einops>=0.6 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.6.0)\n",
            "Requirement already satisfied: local-attention>=1.5.7 in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (1.5.8)\n",
            "Requirement already satisfied: fairseq in /usr/local/lib/python3.8/dist-packages (from audiolm-pytorch) (0.12.2)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch>=1.6->audiolm-pytorch) (4.4.0)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch) (21.3)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch) (6.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch) (1.21.6)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.8/dist-packages (from accelerate->audiolm-pytorch) (5.4.8)\n",
            "Requirement already satisfied: bitarray in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch) (2.6.2)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch) (2022.6.2)\n",
            "Requirement already satisfied: hydra-core<1.1,>=1.0.7 in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch) (1.0.7)\n",
            "Requirement already satisfied: cffi in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch) (1.15.1)\n",
            "Requirement already satisfied: cython in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch) (0.29.33)\n",
            "Requirement already satisfied: omegaconf<2.1 in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch) (2.0.6)\n",
            "Requirement already satisfied: sacrebleu>=1.4.12 in /usr/local/lib/python3.8/dist-packages (from fairseq->audiolm-pytorch) (2.3.1)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.8/dist-packages (from Mega-pytorch->audiolm-pytorch) (1.7.3)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from scikit-learn->audiolm-pytorch) (3.1.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch) (2.25.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch) (0.12.0)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch) (0.13.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.8/dist-packages (from transformers->audiolm-pytorch) (3.9.0)\n",
            "Requirement already satisfied: antlr4-python3-runtime==4.8 in /usr/local/lib/python3.8/dist-packages (from hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch) (4.8)\n",
            "Requirement already satisfied: importlib-resources in /usr/local/lib/python3.8/dist-packages (from hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch) (5.10.2)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.8/dist-packages (from packaging>=20.0->accelerate->audiolm-pytorch) (3.0.9)\n",
            "Requirement already satisfied: portalocker in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch) (2.7.0)\n",
            "Requirement already satisfied: tabulate>=0.8.9 in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch) (0.8.10)\n",
            "Requirement already satisfied: colorama in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch) (0.4.6)\n",
            "Requirement already satisfied: lxml in /usr/local/lib/python3.8/dist-packages (from sacrebleu>=1.4.12->fairseq->audiolm-pytorch) (4.9.2)\n",
            "Requirement already satisfied: pycparser in /usr/local/lib/python3.8/dist-packages (from cffi->fairseq->audiolm-pytorch) (2.21)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch) (2022.12.7)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch) (1.24.3)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->transformers->audiolm-pytorch) (4.0.0)\n",
            "Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.8/dist-packages (from importlib-resources->hydra-core<1.1,>=1.0.7->fairseq->audiolm-pytorch) (3.11.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup\n",
        "\n",
        "Includes:\n",
        "\n",
        "- How to generate a placeholder dataset if you haven't already, just the basics to run \"training\" e2e on a tiny dataset\n",
        "- How to download a dataset from OpenSLR"
      ],
      "metadata": {
        "id": "xuNcsDJsvQwh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Imports & paths"
      ],
      "metadata": {
        "id": "jBxNK5cKW--_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# imports\n",
        "import math\n",
        "import wave\n",
        "import struct\n",
        "import os\n",
        "import urllib.request\n",
        "import tarfile\n",
        "from audiolm_pytorch import SoundStream, SoundStreamTrainer, HubertWithKmeans, SemanticTransformer, SemanticTransformerTrainer, HubertWithKmeans, CoarseTransformer, CoarseTransformerWrapper, CoarseTransformerTrainer, FineTransformer, FineTransformerWrapper, FineTransformerTrainer, AudioLM\n",
        "from torch import nn\n",
        "import torch\n",
        "import torchaudio\n",
        "\n",
        "\n",
        "# define all dataset paths, checkpoints, etc\n",
        "dataset_folder = \"placeholder_dataset\"\n",
        "soundstream_ckpt = \"results/soundstream.8.pt\" # this can change depending on number of steps\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"
      ],
      "metadata": {
        "id": "OrNeKngVVM0L"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Data"
      ],
      "metadata": {
        "id": "pA56YODZXBtf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Placeholder data generation\n",
        "def get_sinewave(freq=440.0, duration_ms=200, volume=1.0, sample_rate=44100.0):\n",
        "  # code adapted from https://stackoverflow.com/a/33913403\n",
        "  audio = []\n",
        "  num_samples = duration_ms * (sample_rate / 1000.0)\n",
        "  for x in range(int(num_samples)):\n",
        "    audio.append(volume * math.sin(2 * math.pi * freq * (x / sample_rate)))\n",
        "  return audio\n",
        "\n",
        "def save_wav(file_name, audio, sample_rate=44100.0):\n",
        "  # Open up a wav file\n",
        "  wav_file=wave.open(file_name,\"w\")\n",
        "  # wav params\n",
        "  nchannels = 1\n",
        "  sampwidth = 2\n",
        "  # 44100 is the industry standard sample rate - CD quality.  If you need to\n",
        "  # save on file size you can adjust it downwards. The stanard for low quality\n",
        "  # is 8000 or 8kHz.\n",
        "  nframes = len(audio)\n",
        "  comptype = \"NONE\"\n",
        "  compname = \"not compressed\"\n",
        "  wav_file.setparams((nchannels, sampwidth, sample_rate, nframes, comptype, compname))\n",
        "  # WAV files here are using short, 16 bit, signed integers for the \n",
        "  # sample size.  So we multiply the floating point data we have by 32767, the\n",
        "  # maximum value for a short integer.  NOTE: It is theortically possible to\n",
        "  # use the floating point -1.0 to 1.0 data directly in a WAV file but not\n",
        "  # obvious how to do that using the wave module in python.\n",
        "  for sample in audio:\n",
        "      wav_file.writeframes(struct.pack('h', int( sample * 32767.0 )))\n",
        "  wav_file.close()\n",
        "  return\n",
        "\n",
        "def make_placeholder_dataset():\n",
        "  # Make a placeholder dataset with a few .wav files that you can \"train\" on, just to verify things work e2e\n",
        "  if os.path.isdir(dataset_folder):\n",
        "    return\n",
        "  os.makedirs(dataset_folder)\n",
        "  save_wav(f\"{dataset_folder}/example.wav\", get_sinewave())\n",
        "  save_wav(f\"{dataset_folder}/example2.wav\", get_sinewave(duration_ms=500))\n",
        "  os.makedirs(f\"{dataset_folder}/subdirectory\")\n",
        "  save_wav(f\"{dataset_folder}/subdirectory/example.wav\", get_sinewave(freq=330.0))\n",
        "\n",
        "make_placeholder_dataset()"
      ],
      "metadata": {
        "id": "6nnPceFWwedh"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Get actual dataset. Uncomment this if you want to try training on real data\n",
        "\n",
        "# full dataset: https://www.openslr.org/12\n",
        "# We'll use https://us.openslr.org/resources/12/dev-clean.tar.gz development set, \"clean\" speech.\n",
        "# We *should* train on, well, training, but this is just to demo running things end-to-end at all so I just picked a small clean set.\n",
        "\n",
        "# url = \"https://us.openslr.org/resources/12/dev-clean.tar.gz\"\n",
        "# filename = \"dev-clean\"\n",
        "# filename_targz = filename + \".tar.gz\"\n",
        "# if not os.path.isfile(filename_targz):\n",
        "#   urllib.request.urlretrieve(url, filename_targz)\n",
        "# if not os.path.isdir(filename):\n",
        "#   # open file\n",
        "#   with tarfile.open(filename_targz) as t:\n",
        "#     t.extractall(filename)"
      ],
      "metadata": {
        "id": "jwYCbFpHvmRI"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training\n",
        "\n",
        "Now that we have a dataset, we can train AudioLM.\n",
        "\n",
        "**Note**: do NOT type \"y\" to overwrite previous experiments/ checkpoints when running through the cells here unless you're ready to the entire results folder! Otherwise you will end up erasing things (e.g. you train SoundStream first, and if you choose \"overwrite\" then you lose the SoundStream checkpoint when you then train SemanticTransformer)."
      ],
      "metadata": {
        "id": "PYcI0aXEwuxR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### SoundStream"
      ],
      "metadata": {
        "id": "T7GiyBcBWiZV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "soundstream = SoundStream(\n",
        "    codebook_size = 1024,\n",
        "    rq_num_quantizers = 8,\n",
        ")\n",
        "\n",
        "trainer = SoundStreamTrainer(\n",
        "    soundstream,\n",
        "    folder = dataset_folder,\n",
        "    batch_size = 4,\n",
        "    grad_accum_every = 8,         # effective batch size of 32\n",
        "    data_max_length = 320 * 32,\n",
        "    save_results_every = 2,\n",
        "    save_model_every = 4,\n",
        "    num_train_steps = 9\n",
        ").cuda()\n",
        "# NOTE: I changed num_train_steps to 9 (aka 8 + 1) from 10000 to make things go faster for demo purposes\n",
        "# adjusting save_*_every variables for the same reason\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "nGU0OZiOwPEO",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "21dd959c-6458-4477-8403-cf810166f38d"
      },
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "training with dataset of 2 samples and validating with randomly splitted 1 samples\n",
            "0: soundstream total loss: 167.262, soundstream recon loss: 1.123 | discr (scale 1) loss: 2.003 | discr (scale 0.5) loss: 1.999 | discr (scale 0.25) loss: 1.999\n",
            "0: saving to results\n",
            "0: saving model to results\n",
            "1: soundstream total loss: 182.282, soundstream recon loss: 1.389 | discr (scale 1) loss: 1.938 | discr (scale 0.5) loss: 1.928 | discr (scale 0.25) loss: 1.928\n",
            "2: soundstream total loss: 196.668, soundstream recon loss: 1.450 | discr (scale 1) loss: 1.845 | discr (scale 0.5) loss: 1.842 | discr (scale 0.25) loss: 1.843\n",
            "2: saving to results\n",
            "3: soundstream total loss: 216.329, soundstream recon loss: 1.451 | discr (scale 1) loss: 1.751 | discr (scale 0.5) loss: 1.750 | discr (scale 0.25) loss: 1.757\n",
            "4: soundstream total loss: 206.804, soundstream recon loss: 1.167 | discr (scale 1) loss: 1.671 | discr (scale 0.5) loss: 1.706 | discr (scale 0.25) loss: 1.724\n",
            "4: saving to results\n",
            "4: saving model to results\n",
            "5: soundstream total loss: 195.325, soundstream recon loss: 0.929 | discr (scale 1) loss: 1.348 | discr (scale 0.5) loss: 1.372 | discr (scale 0.25) loss: 1.482\n",
            "6: soundstream total loss: 245.195, soundstream recon loss: 1.054 | discr (scale 1) loss: 1.060 | discr (scale 0.5) loss: 1.244 | discr (scale 0.25) loss: 1.288\n",
            "6: saving to results\n",
            "7: soundstream total loss: 245.724, soundstream recon loss: 0.970 | discr (scale 1) loss: 1.092 | discr (scale 0.5) loss: 1.358 | discr (scale 0.25) loss: 1.079\n",
            "8: soundstream total loss: 202.707, soundstream recon loss: 0.786 | discr (scale 1) loss: 0.733 | discr (scale 0.5) loss: 0.687 | discr (scale 0.25) loss: 0.790\n",
            "8: saving to results\n",
            "8: saving model to results\n",
            "training complete\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### SemanticTransformer"
      ],
      "metadata": {
        "id": "lqjN28L4Wc5Q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# hubert checkpoints can be downloaded at\n",
        "# https://github.com/facebookresearch/fairseq/tree/main/examples/hubert\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 = f'./{hubert_ckpt}',\n",
        "    kmeans_path = f'./{hubert_quantizer}'\n",
        ")\n",
        "\n",
        "semantic_transformer = SemanticTransformer(\n",
        "    num_semantic_tokens = wav2vec.codebook_size,\n",
        "    dim = 1024,\n",
        "    depth = 6\n",
        ").cuda()\n",
        "\n",
        "\n",
        "trainer = SemanticTransformerTrainer(\n",
        "    transformer = semantic_transformer,\n",
        "    wav2vec = wav2vec,\n",
        "    folder = dataset_folder,\n",
        "    batch_size = 1,\n",
        "    data_max_length = 320 * 32,\n",
        "    num_train_steps = 1\n",
        ")\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "qgd962eSvDzS",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b0550cde-0c8b-4a39-f896-f6f813f50f8c"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.8/dist-packages/sklearn/base.py:329: UserWarning: Trying to unpickle estimator MiniBatchKMeans from version 0.24.0 when using version 1.0.2. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
            "https://scikit-learn.org/stable/modules/model_persistence.html#security-maintainability-limitations\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "training with dataset of 2 samples and validating with randomly splitted 1 samples\n",
            "do you want to clear previous experiment checkpoints and results? (y/n) n\n",
            "0: loss: 6.648584365844727\n",
            "0: valid loss 5.763116359710693\n",
            "0: saving model to results\n",
            "training complete\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### CoarseTransformer"
      ],
      "metadata": {
        "id": "4eEvIzhEWwRz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "wav2vec = HubertWithKmeans(\n",
        "    checkpoint_path = f'./{hubert_ckpt}',\n",
        "    kmeans_path = f'./{hubert_quantizer}'\n",
        ")\n",
        "\n",
        "soundstream = SoundStream(\n",
        "    codebook_size = 1024,\n",
        "    rq_num_quantizers = 8,\n",
        ")\n",
        "\n",
        "soundstream.load(f\"./{soundstream_ckpt}\")\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",
        ")\n",
        "\n",
        "trainer = CoarseTransformerTrainer(\n",
        "    transformer = coarse_transformer,\n",
        "    soundstream = soundstream,\n",
        "    wav2vec = wav2vec,\n",
        "    folder = dataset_folder,\n",
        "    batch_size = 1,\n",
        "    data_max_length = 320 * 32,\n",
        "    save_results_every = 2,\n",
        "    save_model_every = 4,\n",
        "    num_train_steps = 9\n",
        ")\n",
        "# NOTE: I changed num_train_steps to 9 (aka 8 + 1) from 10000 to make things go faster for demo purposes\n",
        "# adjusting save_*_every variables for the same reason\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "1LeWmaNHzzY9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7e7ecb3b-f59e-4d18-c8c9-64762e9b43fc"
      },
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.8/dist-packages/sklearn/base.py:329: UserWarning: Trying to unpickle estimator MiniBatchKMeans from version 0.24.0 when using version 1.0.2. This might lead to breaking code or invalid results. Use at your own risk. For more info please refer to:\n",
            "https://scikit-learn.org/stable/modules/model_persistence.html#security-maintainability-limitations\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "training with dataset of 2 samples and validating with randomly splitted 1 samples\n",
            "do you want to clear previous experiment checkpoints and results? (y/n) n\n",
            "0: loss: 63.983970642089844\n",
            "0: valid loss 63.398582458496094\n",
            "0: saving model to results\n",
            "1: loss: 65.85967254638672\n",
            "2: loss: 62.4722900390625\n",
            "2: valid loss 50.01605987548828\n",
            "3: loss: 11.735434532165527\n",
            "4: loss: 3.976104497909546\n",
            "4: valid loss 46.094608306884766\n",
            "4: saving model to results\n",
            "5: loss: 58.27140426635742\n",
            "6: loss: 41.68347930908203\n",
            "6: valid loss 45.54595184326172\n",
            "7: loss: 2.2387890815734863\n",
            "8: loss: 0.4718627631664276\n",
            "8: valid loss 39.10848617553711\n",
            "8: saving model to results\n",
            "training complete\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### FineTransformer"
      ],
      "metadata": {
        "id": "fRvj7qOJWzmw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "soundstream = SoundStream(\n",
        "    codebook_size = 1024,\n",
        "    rq_num_quantizers = 8,\n",
        ")\n",
        "\n",
        "soundstream.load(f\"./{soundstream_ckpt}\")\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",
        ")\n",
        "\n",
        "trainer = FineTransformerTrainer(\n",
        "    transformer = fine_transformer,\n",
        "    soundstream = soundstream,\n",
        "    folder = dataset_folder,\n",
        "    batch_size = 1,\n",
        "    data_max_length = 320 * 32,\n",
        "    num_train_steps = 9\n",
        ")\n",
        "# NOTE: I changed num_train_steps to 9 (aka 8 + 1) from 10000 to make things go faster for demo purposes\n",
        "# adjusting save_*_every variables for the same reason\n",
        "\n",
        "trainer.train()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZRaEhRRKWg8F",
        "outputId": "7cc166c4-c8e9-45ef-8293-8f5381c2d3af"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "training with dataset of 2 samples and validating with randomly splitted 1 samples\n",
            "do you want to clear previous experiment checkpoints and results? (y/n) n\n",
            "0: loss: 70.90608215332031\n",
            "0: valid loss 65.99951171875\n",
            "0: saving model to results\n",
            "1: loss: 43.6014289855957\n",
            "2: loss: 8.300681114196777\n",
            "3: loss: 61.23375701904297\n",
            "4: loss: 63.34052276611328\n",
            "5: loss: 2.010118246078491\n",
            "6: loss: 56.52588653564453\n",
            "7: loss: 0.5423888564109802\n",
            "8: loss: 0.005095238331705332\n",
            "training complete\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Inference"
      ],
      "metadata": {
        "id": "QoHgkgA3XKXH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Everything together\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",
        ")\n",
        "\n",
        "generated_wav = audiolm(batch_size = 1)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rzghrux5WinW",
        "outputId": "9dd39f7f-0046-4a5f-826e-a442345987af"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "generating semantic:   0%|          | 10/2048 [00:00<00:25, 78.55it/s]\n",
            "generating coarse: 100%|██████████| 512/512 [00:14<00:00, 34.83it/s]\n",
            "generating fine: 100%|██████████| 512/512 [02:56<00:00,  2.91it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "output_path = \"out.wav\"\n",
        "sample_rate = 44100\n",
        "torchaudio.save(output_path, generated_wav.cpu(), sample_rate)"
      ],
      "metadata": {
        "id": "4rQPHTSRngEr"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "is9wLY_ncDYK"
      },
      "execution_count": 22,
      "outputs": []
    }
  ]
}