{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "import glob\n",
    "import torch\n",
    "import IPython\n",
    "import numpy as np\n",
    "import funcy\n",
    "import torchaudio\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.utils.s3 import read_from_s3\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filter args\n",
    "device = \"cuda:0\"\n",
    "checkpoint_filepath = \"s3://suno-data/georg/models/codec/dac_2c_25x12.pt\"\n",
    "\n",
    "load_f = funcy.partial(torch.load, map_location=\"cpu\")\n",
    "\n",
    "if checkpoint_filepath.startswith(\"s3://\"):\n",
    "    sd = read_from_s3(checkpoint_filepath, read_f=load_f)\n",
    "else:\n",
    "    sd = load_f(checkpoint_filepath)\n",
    "\n",
    "sd[\"metadata\"][\"kwargs\"] = {\n",
    "    k: v\n",
    "    for k, v in sd[\"metadata\"][\"kwargs\"].items()\n",
    "    if k in DAC.__init__.__code__.co_varnames\n",
    "}\n",
    "model = DAC(**sd[\"metadata\"][\"kwargs\"])\n",
    "model.load_state_dict(sd[\"state_dict\"])\n",
    "model.eval()\n",
    "model.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load some audio\n",
    "audio_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "audio, sr = torchaudio.load(audio_filepath)\n",
    "audio = torchaudio.functional.resample(audio, sr, 48000)\n",
    "start_frame = audio.shape[-1] // 2\n",
    "end_frame = start_frame + int(480000 * 2)\n",
    "audio = audio[:,start_frame:end_frame]\n",
    "print(audio.shape)\n",
    "audio = audio.to(device)\n",
    "vocab_size = 2048\n",
    "\n",
    "codebook_corrupt_probs = [0.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]\n",
    "\n",
    "# encode\n",
    "with torch.no_grad():\n",
    "    res = model.encode(audio.unsqueeze(0))\n",
    "    codes = res[\"codes\"]\n",
    "    print(codes)\n",
    "\n",
    "    corrupt_codes = codes.clone()\n",
    "    print(corrupt_codes.shape)\n",
    "    n_codes = codes.shape[-1]\n",
    "    n_codebooks = codes.shape[1]\n",
    "    for codebook_idx in range(n_codebooks):\n",
    "        if np.random.rand() < codebook_corrupt_probs[codebook_idx]:\n",
    "            print(\"corrupting codebook\", codebook_idx)\n",
    "            for code_idx in range(n_codes):\n",
    "                if np.random.rand() < 0.5:\n",
    "                    corrupt_codes[0, codebook_idx, code_idx] = torch.randint(0, vocab_size-1, (1,))\n",
    "\n",
    "    # decode \n",
    "    z_q, _, _ = model.quantizer.from_codes(codes)\n",
    "    z_q_corrupt, _, _ = model.quantizer.from_codes(corrupt_codes)\n",
    "\n",
    "    print(z_q.shape)\n",
    "\n",
    "    audio_cycled = model.decode(z_q)\n",
    "    audio_cycled_corrupt = model.decode(z_q_corrupt)\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(audio_cycled.squeeze().cpu().numpy(), rate=48000))\n",
    "IPython.display.display(IPython.display.Audio(audio_cycled_corrupt.squeeze().cpu().numpy(), rate=48000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "language": "python",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}