{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import funcy\n",
    "import numpy as np\n",
    "import IPython\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "from dac.model.dac4 import DAC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda:0\"\n",
    "checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\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",
    "vae_model = DAC(**sd[\"metadata\"][\"kwargs\"])\n",
    "vae_model.load_state_dict(sd[\"state_dict\"])\n",
    "vae_model.eval()\n",
    "vae_model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "VAE_T_MEMMAP = 1000\n",
    "VAE_DIM = 128\n",
    "\n",
    "data_dir = \"/app/suno/christian/data/suno_diffusion_tiktok_covers_lyrics/\"\n",
    "vae_memmap_filepath = os.path.join(data_dir, \"vae_val.bin\")\n",
    "# load output memmap\n",
    "vae_data = np.memmap(os.path.join(vae_memmap_filepath), dtype=np.float32, mode=\"r\")\n",
    "vae_data = vae_data.reshape(-1, VAE_DIM, VAE_T_MEMMAP)\n",
    "\n",
    "metas = read_jsonl(\"/app/suno/christian/data/suno_diffusion_tiktok_covers_lyrics/val_metas.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idx = np.random.randint(0, len(metas))\n",
    "print(rand_idx)\n",
    "print(metas[rand_idx])\n",
    "vae_seq = vae_data[rand_idx]\n",
    "vae_seq = torch.from_numpy(vae_seq.copy()).to(device).unsqueeze(0).float()\n",
    "\n",
    "audio = vae_model.decode(vae_seq)[0].detach().cpu()         \n",
    "audio /= audio.abs().max().clamp(1e-8)\n",
    "print(audio.mean())\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(audio.numpy(), rate=44100))"
   ]
  },
  {
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}