{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import json\n",
    "import torch\n",
    "import funcy\n",
    "import IPython\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "\n",
    "from suno_utils.utils.s3 import read_from_s3, download_s3_files\n",
    "from dac.model.dac4 import DAC\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ----------------  load vae model for embedding ----------------\n",
    "device = \"cuda\"\n",
    "print(\"Loading VAE model...\")\n",
    "#checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n",
    "checkpoint_filepath = (\n",
    "        \"/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.1/best/dac/weights.pth\"\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",
    "vae_model = DAC(**sd[\"metadata\"][\"kwargs\"])\n",
    "vae_model.load_state_dict(sd[\"state_dict\"])\n",
    "vae_model.eval()\n",
    "vae_model.to(device)\n"
   ]
  },
  {
   "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\"\n",
    "vae_memmap_filepath = os.path.join(data_dir, \"vae_100hz_kl_0.1_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/val_metas.jsonl\")\n",
    "\n",
    "print(vae_data.shape[0], len(metas))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idx = np.random.randint(0, vae_data.shape[0])\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": [
    "root_dir = \"/app/suno/christian/data/genius_hq\"\n",
    "vae_npz = glob.glob(os.path.join(root_dir, \"vae\", \"*.npz\"))\n",
    "print(f\"Found {len(vae_npz)} npz files\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_npz_idx = np.random.randint(0, len(vae_npz))\n",
    "vae_npz_filepath = vae_npz[vae_npz_idx]\n",
    "\n",
    "print(f\"Loading {vae_npz_filepath}\")\n",
    "vae_data = np.load(vae_npz_filepath)\n",
    "vae_latents = torch.from_numpy(vae_data[\"vae_latents\"]).to(device)\n",
    "\n",
    "meta_filepath = vae_npz_filepath.replace(\"vae\", \"meta\").replace(\".npz\", \".json\")\n",
    "with open(meta_filepath, \"r\") as f:\n",
    "    meta = json.load(f)\n",
    "\n",
    "for key, val in meta.items():\n",
    "    print(key, val)\n",
    "\n",
    "vae_latents = vae_latents[:, :, :6000]\n",
    "\n",
    "with torch.no_grad():\n",
    "    pred_audio = vae_model.decode(vae_latents)[0].detach().cpu()         \n",
    "    pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "    print(pred_audio.mean())\n",
    "\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(pred_audio.numpy(), rate=48000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "s3_filepath = \"s3://suno-data/datasets/harvest/youtube_music/audio/Ys37DilVPBU.webm\"\n",
    "download_s3_files([s3_filepath], [f\"/home/christian/code/christian/notebooks/outputs/s3/{os.path.basename(s3_filepath)}\"])\n",
    "audio, sr = torchaudio.load(f\"/home/christian/code/christian/notebooks/outputs/s3/{os.path.basename(s3_filepath)}\")\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(audio.numpy(), rate=sr))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metas = read_jsonl(\"/home/christian/code/christian/metadata/genius_hq_metas.jsonl\", progress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_idx = np.random.randint(0, len(metas))\n",
    "for key, val in metas[rand_idx].items():\n",
    "    print(key, val)"
   ]
  },
  {
   "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
}