public_Extract_Jukebox_Embeddings.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "A notebook for extracting embeddings from [OpenAI's Jukebox model](https://openai.com/index/jukebox/), following the approach described in [Castellon et al. (2021)](https://arxiv.org/abs/2107.05677) with some modifications followed in [Spotify's Llark paper](https://arxiv.org/pdf/2310.07160):\n",
        "\n",
        "- Source: Output of the 36th layer of the Jukebox encoder\n",
        "- Original Jukebox encoding: 4800-dimensional vectors at 345Hz\n",
        "- Audio/embeddings are chunked into 25 seconds clips as that is the max Jukebox can take in as input, any clips shorter than 25 seconds are padded before passed through Jukebox\n",
        "- Approach: Mean-pooling within 100ms frames, resulting in:\n",
        "    - Downsampled frequency: 10Hz\n",
        "    - Embedding size: 1.2 × 10^6 for a 25s audio clip.\n",
        "    - For a 25s audio clip the 2D array shape will be [250, 4800]\n",
        "- This method retains temporal information while reducing the embedding size\n",
        "\n",
        "Having a Colab notebook for this gives us an easily reproducible environment and allows us to take advantage of the cheap T4 GPU's Colab offers."
      ],
      "metadata": {
        "id": "yRq_LRQI9AMQ"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rAE4EmJ0d9yZ"
      },
      "source": [
        "Join the Jukebox community at https://discord.gg/aEqXFN9amV\n",
        "\n",
        "Speed upsampling supported. Switch to upsample mode will happen automatically if data file is detected within the folder provided.\n",
        "\n",
        "**How to handle memory problems:** In theory, this notebook is crafted to avoid Out of memory errors, but here's some tricks if you still encounter one:\n",
        "* Restart runtime: At the top of the notebook, click \"Runtime\" and then \"Restart runtime\". Then run everything again. You should do this everytime you start a second run within the same session or after you've interrupted one.\n",
        "* Decrease sample count: Choose a lower number for 'hps.n_samples'.\n",
        "\n",
        "**Please note that the core Jukebox architecture hasn't been updated since 2020. This means Jukebox is very slow and inefficient at generating samples. It will take a couple of hours to get any results.**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8qEqdj8u0gdN"
      },
      "outputs": [],
      "source": [
        "#@title Check which GPU you were assigned by running this cell.\n",
        "!nvidia-smi -i 0 -e 0\n",
        "!nvidia-smi -L\n",
        "your_lyrics = \"\"\"\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zXGQ5Lm-bUyn"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sAdFGF-bqVMY"
      },
      "outputs": [],
      "source": [
        "#@title Select your model and settings\n",
        "\n",
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "!pip install --upgrade git+https://github.com/craftmine1000/jukebox-saveopt.git\n",
        "\n",
        "import jukebox\n",
        "import torch as t\n",
        "import librosa\n",
        "import os\n",
        "from IPython.display import Audio\n",
        "from jukebox.make_models import make_vqvae, make_prior, MODELS, make_model\n",
        "from jukebox.hparams import Hyperparams, setup_hparams\n",
        "from jukebox.sample import sample_single_window, _sample, \\\n",
        "                           sample_partial_window, upsample, \\\n",
        "                           load_prompts\n",
        "from jukebox.utils.dist_utils import setup_dist_from_mpi\n",
        "from jukebox.utils.torch_utils import empty_cache\n",
        "# MPI Connect. MPI doesn't like being initialized twice, hence the following\n",
        "try:\n",
        "    if device is not None:\n",
        "        pass\n",
        "except NameError:\n",
        "    rank, local_rank, device = setup_dist_from_mpi()\n",
        "\n",
        "model = \"5b\" # @param [\"5b_lyrics\",\"5b\",\"1b_lyrics\"]\n",
        "if model == '5b':\n",
        "  your_lyrics = \"\"\"\n",
        "  \"\"\"\n",
        "\n",
        "save_and_load_models_from_drive = True\n",
        "\n",
        "#START GDRIVE MODEL LOADER\n",
        "if save_and_load_models_from_drive == True:\n",
        "  import os.path\n",
        "  !apt install pv\n",
        "  !mkdir /root/.cache ; mkdir /root/.cache/jukebox ; mkdir /root/.cache/jukebox/models ; mkdir /root/.cache/jukebox/models/1b_lyrics ; mkdir /root/.cache/jukebox/models/5b_lyrics ; mkdir /root/.cache/jukebox/models/5b\n",
        "  !mkdir /content/gdrive/MyDrive/jukebox\n",
        "  !mkdir /content/gdrive/MyDrive/jukebox/models\n",
        "  !mkdir /content/gdrive/MyDrive/jukebox/models/5b\n",
        "  !mkdir /content/gdrive/MyDrive/jukebox/models/5b_lyrics\n",
        "  !mkdir /content/gdrive/MyDrive/jukebox/models/1b_lyrics\n",
        "\n",
        "\n",
        "def load_5b_vqvae():\n",
        "    if os.path.exists(\"/root/.cache/jukebox/models/5b/vqvae.pth.tar\") == False:\n",
        "      if os.path.exists(\"/content/gdrive/MyDrive/jukebox/models/5b/vqvae.pth.tar\") == False:\n",
        "        print(\"5b_vqvae not stored in Google Drive. Downloading for the first time.\")\n",
        "        !wget https://openaipublic.azureedge.net/jukebox/models/5b/vqvae.pth.tar -O /content/gdrive/MyDrive/jukebox/models/5b/vqvae.pth.tar\n",
        "      else:\n",
        "        print(\"5b_vqvae stored in Google Drive.\")\n",
        "      print('Copying 5b VQVAE')\n",
        "      !pv /content/gdrive/MyDrive/jukebox/models/5b/vqvae.pth.tar > /root/.cache/jukebox/models/5b/vqvae.pth.tar\n",
        "\n",
        "def load_1b_lyrics_level2():\n",
        "  if os.path.exists(\"/root/.cache/jukebox/models/1b_lyrics/prior_level_2.pth.tar\") == False:\n",
        "    if os.path.exists(\"/content/gdrive/MyDrive/jukebox/models/1b_lyrics/prior_level_2.pth.tar\") == False:\n",
        "      print(\"1b_lyrics_level_2 not stored in Google Drive. Downloading for the first time. This will take a few more minutes.\")\n",
        "      !wget https://openaipublic.azureedge.net/jukebox/models/1b_lyrics/prior_level_2.pth.tar -O /content/gdrive/MyDrive/jukebox/models/1b_lyrics/prior_level_2.pth.tar\n",
        "    else:\n",
        "      print(\"1b_lyrics_level_2 stored in Google Drive.\")\n",
        "    print(\"Copying 1B_Lyrics Level 2\")\n",
        "    !pv /content/gdrive/MyDrive/jukebox/models/1b_lyrics/prior_level_2.pth.tar > /root/.cache/jukebox/models/1b_lyrics/prior_level_2.pth.tar\n",
        "\n",
        "def load_5b_lyrics_level2():\n",
        "  if os.path.exists(\"/root/.cache/jukebox/models/5b_lyrics/prior_level_2.pth.tar\") == False:\n",
        "    if os.path.exists(\"/content/gdrive/MyDrive/jukebox/models/5b_lyrics/prior_level_2.pth.tar\") == False:\n",
        "      print(\"5b_lyrics_level_2 not stored in Google Drive. Downloading for the first time. This will take up to 10-15 minutes.\")\n",
        "      !wget https://openaipublic.azureedge.net/jukebox/models/5b_lyrics/prior_level_2.pth.tar -O /content/gdrive/MyDrive/jukebox/models/5b_lyrics/prior_level_2.pth.tar\n",
        "    else:\n",
        "      print(\"5b_lyrics_level_2 stored in Google Drive.\")\n",
        "    print(\"Copying 5B_Lyrics Level 2\")\n",
        "    !pv /content/gdrive/MyDrive/jukebox/models/5b_lyrics/prior_level_2.pth.tar > /root/.cache/jukebox/models/5b_lyrics/prior_level_2.pth.tar\n",
        "\n",
        "def load_5b_level1():\n",
        "  if os.path.exists('/root/.cache/jukebox/models/5b/prior_level_1.pth.tar') == False:\n",
        "    if os.path.exists(\"/content/gdrive/MyDrive/jukebox/models/5b/prior_level_1.pth.tar\") == False:\n",
        "      print(\"5b_level_1 not stored in Google Drive. Downloading for the first time. This may take a few more minutes.\")\n",
        "      !wget https://openaipublic.azureedge.net/jukebox/models/5b/prior_level_1.pth.tar -O /content/gdrive/MyDrive/jukebox/models/5b/prior_level_1.pth.tar\n",
        "    else:\n",
        "      print(\"5b_level_1 stored in Google Drive.\")\n",
        "    print(\"Copying 5B Level 1\")\n",
        "    !pv /content/gdrive/MyDrive/jukebox/models/5b/prior_level_1.pth.tar > /root/.cache/jukebox/models/5b/prior_level_1.pth.tar\n",
        "\n",
        "def load_5b_level0():\n",
        "  if os.path.exists('/root/.cache/jukebox/models/5b/prior_level_0.pth.tar') == False:\n",
        "    if os.path.exists(\"/content/gdrive/MyDrive/jukebox/models/5b/prior_level_0.pth.tar\") == False:\n",
        "      print(\"5b_level_0 not stored in Google Drive. Downloading for the first time. This may take a few minutes.\")\n",
        "      !wget https://openaipublic.azureedge.net/jukebox/models/5b/prior_level_0.pth.tar -O /content/gdrive/MyDrive/jukebox/models/5b/prior_level_0.pth.tar\n",
        "    else:\n",
        "      print(\"5b_level_0 stored in Google Drive.\")\n",
        "    print(\"Copying 5B Level 0\")\n",
        "    !pv /content/gdrive/MyDrive/jukebox/models/5b/prior_level_0.pth.tar > /root/.cache/jukebox/models/5b/prior_level_0.pth.tar\n",
        "\n",
        "def load_5b_level2():\n",
        "  if os.path.exists('/root/.cache/jukebox/models/5b/prior_level_2.pth.tar') == False:\n",
        "    if os.path.exists(\"/content/gdrive/MyDrive/jukebox/models/5b/prior_level_2.pth.tar\") == False:\n",
        "      print(\"5b_level_2 not stored in Google Drive. Downloading for the first time. This will take up to 10-15 minutes.\")\n",
        "      !wget https://openaipublic.azureedge.net/jukebox/models/5b/prior_level_2.pth.tar -O /content/gdrive/MyDrive/jukebox/models/5b/prior_level_2.pth.tar\n",
        "    else:\n",
        "      print(\"5b_level_2 stored in Google Drive.\")\n",
        "  print(\"Copying 5B Level 2\")\n",
        "  !pv /content/gdrive/MyDrive/jukebox/models/5b/prior_level_2.pth.tar > /root/.cache/jukebox/models/5b/prior_level_2.pth.tar\n",
        "\n",
        "if save_and_load_models_from_drive == True:\n",
        "  if model == '5b_lyrics':\n",
        "    load_5b_vqvae()\n",
        "    load_5b_lyrics_level2()\n",
        "    load_5b_level1()\n",
        "    load_5b_level0()\n",
        "  if model == '5b':\n",
        "    load_5b_vqvae()\n",
        "    load_5b_level2()\n",
        "    load_5b_level1()\n",
        "    load_5b_level0()\n",
        "  elif model == '1b_lyrics':\n",
        "    load_5b_vqvae()\n",
        "    load_1b_lyrics_level2()\n",
        "    load_5b_level1()\n",
        "    load_5b_level0()\n",
        "#END GDRIVE MODEL LOADER\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gke_-z5aJ4QO"
      },
      "outputs": [],
      "source": [
        "import jukebox\n",
        "import torch as t\n",
        "import librosa\n",
        "import os\n",
        "from IPython.display import Audio\n",
        "from jukebox.make_models import make_vqvae, make_prior, MODELS, make_model\n",
        "from jukebox.hparams import Hyperparams, setup_hparams\n",
        "from jukebox.sample import sample_single_window, _sample, sample_partial_window, upsample, load_prompts\n",
        "from jukebox.utils.dist_utils import setup_dist_from_mpi\n",
        "from jukebox.utils.torch_utils import empty_cache\n",
        "\n",
        "# MPI Connect\n",
        "try:\n",
        "    if device is not None:\n",
        "        pass\n",
        "except NameError:\n",
        "    rank, local_rank, device = setup_dist_from_mpi()\n",
        "\n",
        "# Model selection and configuration\n",
        "model = \"5b\" #@param [\"5b_lyrics\", \"5b\", \"1b_lyrics\"]\n",
        "if model == '5b':\n",
        "    your_lyrics = \"\"\"\"\"\"\n",
        "\n",
        "hps = Hyperparams()\n",
        "hps.sr = 44100\n",
        "hps.n_samples = 2 #@param {type:\"integer\"}\n",
        "hps.name = '/content/gdrive/MyDrive/Project_1' #@param {type:\"string\"}\n",
        "chunk_size = 64 if model in ('5b', '5b_lyrics') else 128\n",
        "\n",
        "# GPU detection and configuration\n",
        "gpu_info = !nvidia-smi -L\n",
        "if gpu_info[0].find('Tesla T4') >= 0:\n",
        "    max_batch_size = 2\n",
        "elif gpu_info[0].find('Tesla K80') >= 0:\n",
        "    max_batch_size = 8\n",
        "elif gpu_info[0].find('Tesla P100') >= 0:\n",
        "    max_batch_size = 3\n",
        "elif gpu_info[0].find('Tesla V100') >= 0:\n",
        "    max_batch_size = 3\n",
        "elif gpu_info[0].find('L4') >= 0:\n",
        "    max_batch_size = 3\n",
        "elif gpu_info[0].find('A100') >= 0:\n",
        "    max_batch_size = 6\n",
        "else:\n",
        "    max_batch_size = 3\n",
        "\n",
        "print(f'{gpu_info[0]} detected, max_batch_size set to {max_batch_size}')\n",
        "\n",
        "hps.levels = 3\n",
        "hps.hop_fraction = [0.5, 0.5, 0.125]\n",
        "\n",
        "import torch\n",
        "\n",
        "# Load model\n",
        "vqvae, *priors = MODELS[model]\n",
        "vqvae = make_vqvae(setup_hparams(vqvae, dict(sample_length = 1048576)), device)\n",
        "top_prior = make_prior(setup_hparams(priors[-1], dict()), vqvae, device)\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "vqvae = vqvae.to(device)\n",
        "top_prior = top_prior.to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Download musicnet"
      ],
      "metadata": {
        "id": "wAEc6yep49-H"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This notebook generates embeddings for the MusicNet dataset as an example, the full dataset can be found here on HuggingFace: https://huggingface.co/datasets/jonflynn/musicnet_jukebox_embeddings"
      ],
      "metadata": {
        "id": "gIf8zVZ46L90"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iPZaI0RuOfum"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "!mkdir -p ~/.kaggle\n",
        "!cp kaggle.json ~/.kaggle/\n",
        "!chmod 600 ~/.kaggle/kaggle.json  # upload your kaggle.json key file\n",
        "\n",
        "!mkdir -p /content/musicnet\n",
        "!kaggle datasets download -d imsparsh/musicnet-dataset\n",
        "!unzip musicnet-dataset.zip -d /content/musicnet"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Auth with GCP to store embeddings"
      ],
      "metadata": {
        "id": "ejnuEjbt5G0w"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T94CoidImaAB"
      },
      "outputs": [],
      "source": [
        "from google.colab import auth\n",
        "auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O5e-gvp8C0ki"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Set the environment variable to point to your service account key file, upload your service account json key\n",
        "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = '/content/admin.json'\n",
        "\n",
        "# Authenticate with gcloud\n",
        "!gcloud auth activate-service-account --key-file $GOOGLE_APPLICATION_CREDENTIALS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vsz9UFgWR0Zk"
      },
      "outputs": [],
      "source": [
        "!pip install librosa"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set your GCP env up"
      ],
      "metadata": {
        "id": "iE7k_XTl-JHR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "GOOGLE_CLOUD_PROJECT = \"your-gcp-project\"\n",
        "GCS_BUCKET_NAME = \"jukebox-embeddings\"\n",
        "GCP_REGION = \"us-central1\"\n",
        "\n",
        "os.environ[\"GOOGLE_CLOUD_PROJECT\"] = GOOGLE_CLOUD_PROJECT"
      ],
      "metadata": {
        "id": "po15yqBu-LtK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Extract embeddings"
      ],
      "metadata": {
        "id": "DFaonUjR-TS0"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sShZQyNZOfsN"
      },
      "outputs": [],
      "source": [
        "JUKEBOX_SAMPLE_RATE = 44100\n",
        "T = 8192\n",
        "JUKEBOX_EXPECTED_SAMPLES_LEN = 1048576\n",
        "JUKEBOX_SAMPLE_SECONDS = JUKEBOX_EXPECTED_SAMPLES_LEN / JUKEBOX_SAMPLE_RATE\n",
        "ACTS_SAMPLE_RATE = T / JUKEBOX_SAMPLE_SECONDS\n",
        "\n",
        "import argparse\n",
        "import io\n",
        "import logging\n",
        "import pathlib\n",
        "import tempfile\n",
        "from functools import lru_cache\n",
        "from math import floor\n",
        "from pathlib import Path\n",
        "import librosa as lr\n",
        "from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, List, Callable\n",
        "import numpy as np\n",
        "import torch\n",
        "from google.cloud import storage\n",
        "import concurrent.futures\n",
        "import re\n",
        "\n",
        "class EmptyFileError(ValueError):\n",
        "    pass\n",
        "\n",
        "@lru_cache(None)\n",
        "def gcs_client():\n",
        "    return storage.Client()\n",
        "\n",
        "@lru_cache(None)\n",
        "def gcs_bucket(bucket_name: str):\n",
        "    return gcs_client().get_bucket(bucket_name)\n",
        "\n",
        "def to_device(tensor, device):\n",
        "    if isinstance(tensor, (list, tuple)):\n",
        "        return [to_device(t, device) for t in tensor]\n",
        "    elif isinstance(tensor, dict):\n",
        "        return {k: to_device(v, device) for k, v in tensor.items()}\n",
        "    elif isinstance(tensor, torch.Tensor):\n",
        "        return tensor.to(device)\n",
        "    return tensor\n",
        "\n",
        "def split_gcs_bucket_and_filepath(filepath: str) -> Tuple[str, str]:\n",
        "    return filepath.replace(\"gs://\", \"\").split(\"/\", maxsplit=1)\n",
        "\n",
        "def read_wav_bytes(filepath: str) -> bytes:\n",
        "    assert filepath.startswith(\"gs://\"), f\"Expected a file path on GCS, but got: {filepath!r}\"\n",
        "    bucket_name, file_name = split_gcs_bucket_and_filepath(filepath)\n",
        "    return gcs_bucket(bucket_name).blob(file_name).download_as_string()\n",
        "\n",
        "def list_gcs_folders(input_dir: str) -> Sequence[str]:\n",
        "    bucket, prefix = split_gcs_bucket_and_filepath(input_dir)\n",
        "    blobs = gcs_client().list_blobs(bucket, prefix=prefix, delimiter='/')\n",
        "\n",
        "    folder_paths = []\n",
        "    for page in blobs.pages:\n",
        "        for prefix in page.prefixes:\n",
        "            folder_paths.append(f\"gs://{bucket}/{prefix}\")\n",
        "\n",
        "    return folder_paths\n",
        "\n",
        "def load_audio_from_file(fpath: str) -> np.ndarray:\n",
        "    try:\n",
        "        audio, _ = lr.load(fpath, sr=JUKEBOX_SAMPLE_RATE)\n",
        "    except ValueError as ve:\n",
        "        raise EmptyFileError(f\"file {fpath} failed to read with exception {ve!r}; it is probably empty.\")\n",
        "    if audio.ndim == 1:\n",
        "        audio = audio[np.newaxis]\n",
        "    audio = audio.mean(axis=0)\n",
        "    norm_factor = np.abs(audio).max()\n",
        "    if norm_factor > 0:\n",
        "        audio /= norm_factor\n",
        "    return audio.flatten()\n",
        "\n",
        "def download_gcs_file(gcs_path: str) -> str:\n",
        "    bucket_name, file_name = split_gcs_bucket_and_filepath(gcs_path)\n",
        "    bucket = gcs_bucket(bucket_name)\n",
        "    blob = bucket.blob(file_name)\n",
        "    _, local_path = tempfile.mkstemp()\n",
        "    blob.download_to_filename(local_path)\n",
        "    return local_path\n",
        "\n",
        "def chunk_audio_file(audio_path: str, chunk_duration: float) -> List[np.ndarray]:\n",
        "    audio, sr = lr.load(audio_path, sr=JUKEBOX_SAMPLE_RATE)\n",
        "    total_duration = lr.get_duration(y=audio, sr=sr)\n",
        "    if audio.ndim == 1:\n",
        "        audio = audio[np.newaxis]\n",
        "    audio = audio.mean(axis=0)\n",
        "    norm_factor = np.abs(audio).max()\n",
        "    if norm_factor > 0:\n",
        "        audio /= norm_factor\n",
        "    audio = audio.flatten()\n",
        "\n",
        "    chunks = []\n",
        "\n",
        "    # If the total duration is less than or equal to chunk_duration,\n",
        "    # treat the entire audio as a single chunk\n",
        "    if total_duration <= chunk_duration:\n",
        "        chunk = audio\n",
        "        if len(chunk) < int(chunk_duration * sr):\n",
        "            chunk = np.pad(chunk, (0, int(chunk_duration * sr) - len(chunk)))\n",
        "        chunks.append((chunk, 0))\n",
        "    else:\n",
        "        for start_time in np.arange(0, total_duration, chunk_duration):\n",
        "            start_sample = int(start_time * sr)\n",
        "            end_sample = min(int(start_sample + chunk_duration * sr), len(audio))\n",
        "            chunk = audio[start_sample:end_sample]\n",
        "\n",
        "            # Pad the chunk if it's shorter than chunk_duration\n",
        "            if len(chunk) < int(chunk_duration * sr):\n",
        "                chunk = np.pad(chunk, (0, int(chunk_duration * sr) - len(chunk)))\n",
        "\n",
        "            chunks.append((chunk, start_time))\n",
        "\n",
        "    return chunks\n",
        "\n",
        "def maybe_pad_audio_to_max_len(audio: np.ndarray) -> np.ndarray:\n",
        "    if len(audio) < JUKEBOX_EXPECTED_SAMPLES_LEN:\n",
        "        audio = np.pad(audio, (0, JUKEBOX_EXPECTED_SAMPLES_LEN - len(audio)))\n",
        "    return audio\n",
        "\n",
        "def get_z(audio: np.ndarray, vqvae):\n",
        "    assert len(audio) >= JUKEBOX_EXPECTED_SAMPLES_LEN, f\"expected samples with shape {JUKEBOX_EXPECTED_SAMPLES_LEN}; got shape {audio.shape}.\"\n",
        "    audio = audio[:JUKEBOX_EXPECTED_SAMPLES_LEN]\n",
        "    audio_tensor = torch.FloatTensor(audio[np.newaxis, :, np.newaxis]).to(device)\n",
        "    zs = vqvae.encode(audio_tensor)\n",
        "    z = zs[-1].flatten()[np.newaxis, :]\n",
        "    if z.shape[-1] < T:\n",
        "        raise ValueError(\"Audio file is not long enough\")\n",
        "    return z.to(device).long()  # Convert to long and ensure it's on the correct device\n",
        "\n",
        "def get_cond(hps, top_prior):\n",
        "    sample_length_in_seconds = 62\n",
        "    hps.sample_length = (int(sample_length_in_seconds * hps.sr) // top_prior.raw_to_tokens) * top_prior.raw_to_tokens\n",
        "    metas = [dict(artist=\"unknown\", genre=\"unknown\", total_length=hps.sample_length, offset=0, lyrics=\"\"\"lyrics go here!!!\"\"\")] * hps.n_samples\n",
        "    labels = [None, None, top_prior.labeller.get_batch_labels(metas, \"cuda\")]\n",
        "    x_cond, y_cond, prime = top_prior.get_cond(None, top_prior.get_y(labels[-1], 0))\n",
        "    x_cond = x_cond[0, :T][np.newaxis, ...].to(device)\n",
        "    y_cond = y_cond[0][np.newaxis, ...].to(device)\n",
        "    return x_cond, y_cond\n",
        "\n",
        "def get_final_activations(z: torch.Tensor, x_cond: torch.Tensor, y_cond: torch.Tensor, top_prior: Any) -> torch.Tensor:\n",
        "    x = z[:, :T].to(device).long()  # Ensure x is a CUDA LongTensor\n",
        "    top_prior.prior.only_encode = True\n",
        "    out = top_prior.prior.forward(x, x_cond=x_cond, y_cond=y_cond, encoder_kv=None, fp16=False)\n",
        "    return out\n",
        "\n",
        "def windowed_average(acts: torch.Tensor, frame_len: int, ceil_mode: bool = False) -> torch.Tensor:\n",
        "    assert acts.ndim == 2, \"expected 2d inputs\"\n",
        "    assert acts.shape[1] == 4800\n",
        "    acts = torch.unsqueeze(acts, 0)\n",
        "    acts = torch.transpose(acts, 1, 2)\n",
        "    pool = torch.nn.AvgPool1d(frame_len, stride=frame_len, ceil_mode=ceil_mode)\n",
        "    acts = pool(acts)\n",
        "    return torch.transpose(acts, 1, 2)\n",
        "\n",
        "def get_acts_from_file(audio: np.ndarray, hps, vqvae, top_prior, meanpool=True, pool_frames_per_second=None):\n",
        "    input_audio_len = len(audio)\n",
        "    latent_audio_len = floor(T * input_audio_len / JUKEBOX_EXPECTED_SAMPLES_LEN)\n",
        "    audio = maybe_pad_audio_to_max_len(audio)\n",
        "    z = get_z(audio, vqvae)\n",
        "    x_cond, y_cond = get_cond(hps, top_prior)\n",
        "    acts = get_final_activations(z, x_cond, y_cond, top_prior)\n",
        "    acts = acts.squeeze().type(torch.float32)\n",
        "    acts = acts[:latent_audio_len, :]\n",
        "    if meanpool:\n",
        "        logging.warning(f\"mean pooling at f={pool_frames_per_second}\")\n",
        "        if not pool_frames_per_second:\n",
        "            acts = acts.mean(dim=0)\n",
        "        else:\n",
        "            frame_len = floor(ACTS_SAMPLE_RATE / pool_frames_per_second)\n",
        "            acts = windowed_average(acts, frame_len)\n",
        "            acts = torch.squeeze(acts, 0)\n",
        "    acts = acts.cpu().numpy()\n",
        "    logging.info(f\"acts after pooling has shape {acts.shape}\")\n",
        "    return acts\n",
        "\n",
        "def upload_to_gcs(storage_client_factory: Callable[[], storage.Client], bucket_name: str, blob_path: str, representation: np.ndarray) -> None:\n",
        "    try:\n",
        "        storage_client = storage_client_factory()\n",
        "        bucket = storage_client.bucket(bucket_name)\n",
        "        blob = bucket.blob(blob_path)\n",
        "        with blob.open(\"wb\") as f:\n",
        "            np.save(f, representation)\n",
        "        print(f\"Uploaded to gs://{bucket_name}/{blob_path}\")\n",
        "    except Exception as e:\n",
        "        print(f\"Failed to upload {blob_path} to {bucket_name}: {str(e)}\")\n",
        "\n",
        "from typing import Callable\n",
        "\n",
        "def upload_to_gcs_in_batches(upload_batch: List[Tuple[Callable, str, str, torch.Tensor]], batch_size: int):\n",
        "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
        "        for i in range(0, len(upload_batch), batch_size):\n",
        "            batch = upload_batch[i:i + batch_size]\n",
        "            executor.map(lambda args: upload_to_gcs(*args), batch)\n",
        "\n",
        "def process_musicnet_dataset(input_dir: str, output_dir: str, batch_size: int = 10):\n",
        "    storage_client_factory = lambda: storage.Client()\n",
        "    chunk_duration = 25.0\n",
        "\n",
        "    train_data_dir = os.path.join(input_dir, 'musicnet', 'musicnet', 'train_data')\n",
        "    test_data_dir = os.path.join(input_dir, 'musicnet', 'musicnet', 'test_data')\n",
        "\n",
        "    # Combine the train and test directories into a single list\n",
        "    data_dirs = [train_data_dir, test_data_dir]\n",
        "\n",
        "    # Iterate through both directories\n",
        "    for data_dir in data_dirs:\n",
        "        for subdir, dirs, files in os.walk(data_dir):\n",
        "            for file in files:\n",
        "                if file.endswith('.wav'):\n",
        "                    local_audio_path = os.path.join(subdir, file)\n",
        "                    print(f\"Processing {local_audio_path}\")\n",
        "\n",
        "                    audio_chunks = chunk_audio_file(local_audio_path, chunk_duration)\n",
        "\n",
        "                    upload_batch = []\n",
        "                    for chunk, start_time in audio_chunks:\n",
        "                        # Construct output filename\n",
        "                        output_filename = f\"{os.path.basename(local_audio_path)}_{int(start_time)}.npy\"\n",
        "\n",
        "                        if output_dir.startswith(\"gs://\"):\n",
        "                            # Construct output path in GCS\n",
        "                            bucket_name, blob_path = split_gcs_bucket_and_filepath(output_dir)\n",
        "                            output_path = os.path.join(blob_path, output_filename)\n",
        "\n",
        "                            # Check if output file already exists in GCS\n",
        "                            storage_client = storage_client_factory()\n",
        "                            out_bucket = storage_client.bucket(bucket_name)\n",
        "                            out_blob = out_bucket.blob(output_path)\n",
        "                            if out_blob.exists():\n",
        "                                print(f\"Output file {output_filename} already exists. Skipping this chunk.\")\n",
        "                                continue  # Skip processing this chunk\n",
        "\n",
        "                            # Process the chunk since it hasn't been processed yet\n",
        "                            with torch.no_grad():\n",
        "                                representation = get_acts_from_file(\n",
        "                                    chunk,\n",
        "                                    hps,\n",
        "                                    vqvae,\n",
        "                                    top_prior,\n",
        "                                    meanpool=True,\n",
        "                                    pool_frames_per_second=10,\n",
        "                                )\n",
        "\n",
        "                            # Add to upload batch\n",
        "                            upload_batch.append((storage_client_factory, bucket_name, output_path, representation))\n",
        "\n",
        "                    print(f\"Number of chunks to upload: {len(upload_batch)}\")\n",
        "                    if upload_batch:\n",
        "                        print(\"Uploading representations to GCS...\")\n",
        "                        upload_to_gcs_in_batches(upload_batch, batch_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "background_save": true
        },
        "id": "p4EghmEy_7Ui"
      },
      "outputs": [],
      "source": [
        "process_musicnet_dataset(\"/content/musicnet/\", f\"gs://{GCS_BUCKET_NAME}/musicnet_embeddings/\")"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}