{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "import glob\n",
    "import torch\n",
    "import faiss\n",
    "import IPython\n",
    "import numpy as np\n",
    "import torchaudio\n",
    "import itertools\n",
    "\n",
    "from tqdm import tqdm\n",
    "from dac.model.dac2 import DAC\n",
    "from dac.model.discriminator2 import Discriminator as Discriminator_import\n",
    "from dac.nn import loss as loss_import\n",
    "from dac.utils.accelerator import Accelerator\n",
    "from dac.utils import load_model\n",
    "\n",
    "from suno_utils.models.musicfm.modeling_MusicFM import MusicFM_MERTLong\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from sklearn.cluster import KMeans, MiniBatchKMeans\n",
    "\n",
    "from suno_utils.utils.s3 import read_from_s3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Load MERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# load MERT\n",
    "model_filepath = \"s3://suno-data/minz/models/musicfm_concat_epoch=51.pt\"\n",
    "centroids_filepath = \"s3://suno-data/minz/models/musicfm_concat_centroids.npy\"\n",
    "mert_model = MusicFM_MERTLong(\n",
    "    is_flash=False,\n",
    "    stat_path=\"s3://suno-data/minz/models/mertlong_stats.json\",\n",
    "    model_path=model_filepath,\n",
    ")\n",
    "mert_model.cuda()\n",
    "mert_model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Embed audio with MERT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_VAL = True\n",
    "NUM_FRAMES = \n",
    "SAMPLE_RATE = 24_000\n",
    "\n",
    "\n",
    "class AudioDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self):\n",
    "        # get audio files\n",
    "        if USE_VAL:\n",
    "            audio_subsets = glob.glob(\n",
    "                os.path.join(f\"/app/suno/data/audio_2ch_24khz_lg/val/**\")\n",
    "            )\n",
    "        else:\n",
    "            audio_subsets = glob.glob(\n",
    "                os.path.join(f\"/app/suno/data/audio_2ch_24khz_lg/train/**\")\n",
    "            )\n",
    "\n",
    "        audio_files = []\n",
    "        for audio_subset in audio_subsets:\n",
    "            # find the first MAX_FILES_PER_SUBSET files\n",
    "            with os.scandir(audio_subset) as filepaths:\n",
    "                # Use itertools.islice to limit the iterator to the first N entries\n",
    "                first_n_files = list(itertools.islice(filepaths, 40000))\n",
    "            first_n_files = [entry.path for entry in first_n_files if entry.is_file()]\n",
    "            audio_files += first_n_files\n",
    "            print(len(first_n_files), audio_subset)\n",
    "\n",
    "        print(\"Total\", len(audio_files))\n",
    "        print(np.random.choice(audio_files, 5))\n",
    "\n",
    "        self.audio_files = audio_files\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.audio_files)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        audio_file = self.audio_files[idx]\n",
    "        num_frames = torchaudio.info(audio_file).num_frames\n",
    "        if num_frames > NUM_FRAMES:\n",
    "            frame_offset = np.random.randint(\n",
    "                0, torchaudio.info(audio_file).num_frames - NUM_FRAMES - 1\n",
    "            )\n",
    "        else:\n",
    "            frame_offset = 0\n",
    "\n",
    "        audio, sr = torchaudio.load(\n",
    "            audio_file, frame_offset=frame_offset, num_frames=NUM_FRAMES\n",
    "        )\n",
    "\n",
    "        if audio.shape[-1] < NUM_FRAMES:\n",
    "            audio = torch.cat(\n",
    "                [audio, torch.zeros(2, NUM_FRAMES - audio.shape[-1])], dim=-1\n",
    "            )\n",
    "\n",
    "        assert sr == SAMPLE_RATE\n",
    "        audio = audio / audio.abs().max().clamp(1e-8)\n",
    "        return audio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Run k-means clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run clustering\n",
    "ncentroids = 32768\n",
    "niter = 100\n",
    "nredo = 3\n",
    "verbose = True\n",
    "SAMPLE_RATE = 48_000\n",
    "d = data.shape[-1] # embed dim\n",
    "\n",
    "kmeans = faiss.Kmeans(d, ncentroids, niter=niter, nredo=nredo, verbose=verbose, gpu=True)\n",
    "kmeans.train(data_subset)\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}