{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import funcy\n",
    "import IPython\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"5\"\n",
    "\n",
    "from dac.model.dac4 import DAC\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "from stable_audio_tools.inference.sampling import get_alphas_sigmas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_frame_rate = 25\n",
    "device = \"cuda:0\"\n",
    "\n",
    "if vae_frame_rate == 25:\n",
    "    # load VAE model\n",
    "    checkpoint_filepath = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n",
    "elif vae_frame_rate == 100:\n",
    "    checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n",
    "else:\n",
    "    raise ValueError(f\"VAE frame rate not supported: {vae_frame_rate}\")\n",
    "\n",
    "#checkpoint_filepath = \"/home/christian/code/christian/checkpoints/25hz_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)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reconstruction from generation (real audio)\n",
    "from suno_utils.tasks.data_loader import load_audio_mp\n",
    "\n",
    "audio_path = \"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/04 Fuckwithmeyouknowigotit.wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/Speak For Me [omeNvD8IddM].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/09 Sounds Like Hallelujah.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/halo_24.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/Beyoncé - Halo (Lyrics) [wekDNXDWGjM].wav\"\n",
    "#audio_path = \"/home/christian/audio/bad-audio/bill-evans-intro.wav\"'\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/Crazy [CKTOvHw8qFM].wav\"\n",
    "#audio_path = \"/home/christian/audio/reference-audio-wav/03 Your New Aesthetic.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/linger-30s.wav\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/turkey.mp3\"\n",
    "#audio_path = \"/home/christian/code/neon/stable-audio-tools/outputs/Cash Cobain & J. Cole - Grippy (AUDIO) [9wn_ARfYMw4].mp3\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/Lost in Tokyo.m4a\"\n",
    "#audio_path = \"/home/christian/code/christian/outputs/i don't trust you (ben camp x suno).mp3\"\n",
    "\n",
    "audio_arrays_48khz = load_audio_mp(\n",
    "    [audio_path],\n",
    "    target_sample_rate=48000,\n",
    "    normalize_volume=True,\n",
    "    num_workers=12,\n",
    "    n_channels=2,\n",
    ")\n",
    "\n",
    "\n",
    "# load some audio and encode with vae to get latents\n",
    "# encode audio with vae\n",
    "with torch.no_grad():\n",
    "    latents = vae_model.encode(audio_arrays_48khz[0].unsqueeze(0).cuda())[\"z\"].detach()    \n",
    "\n",
    "print(latents.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_factor = 2.5\n",
    "t = torch.ones(1).cuda() * 0.0\n",
    "alphas, sigmas = get_alphas_sigmas(t)\n",
    "print(alphas, sigmas)\n",
    "\n",
    "diffusion_input = latents * scale_factor\n",
    "\n",
    "# Combine the ground truth data and the noise\n",
    "alphas = alphas[:, None, None]\n",
    "sigmas = sigmas[:, None, None]\n",
    "noise = torch.randn_like(diffusion_input)\n",
    "\n",
    "noised_inputs = diffusion_input * alphas + noise * sigmas\n",
    "targets = noise * alphas - diffusion_input * sigmas\n",
    "\n",
    "# decode the targets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    print(\"noised_inputs\")\n",
    "    pred_audio = vae_model.decode(noised_inputs / scale_factor)[0].detach().cpu()         \n",
    "    pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "    IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n",
    "\n",
    "\n",
    "    print(\"noised_inputs\")\n",
    "    pred_audio = vae_model.decode(noised_inputs / scale_factor)[0].detach().cpu()         \n",
    "    pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "    IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n",
    "\n",
    "    print(\"targets\")\n",
    "    pred_audio = vae_model.decode(targets / scale_factor)[0].detach().cpu()         \n",
    "    pred_audio /= pred_audio.abs().max().clamp(1e-8)\n",
    "    IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "def get_alphas_sigmas(t, sequence_length):\n",
    "    \"\"\"\n",
    "    Returns the scaling factors for the clean latent (alpha) and for the\n",
    "    noise (sigma), given a timestep and latent sequence length.\n",
    "    \n",
    "    Args:\n",
    "    t (torch.Tensor): The timestep(s)\n",
    "    sequence_length (int): The length of the latent sequence\n",
    "    \n",
    "    Returns:\n",
    "    tuple: (alpha, sigma) scaling factors\n",
    "    \"\"\"\n",
    "    # Define a base sequence length\n",
    "    base_sequence_length = 128  # This could be adjusted based on your specific use case\n",
    "    \n",
    "    # Calculate the sequence length factor using a more flexible approach\n",
    "    length_factor = math.log(sequence_length / base_sequence_length) / math.log(2)\n",
    "    \n",
    "    # Adjust the schedule based on sequence length\n",
    "    adjusted_t = t * (1 + 0.05 * length_factor)\n",
    "    \n",
    "    # Clip adjusted_t to ensure it stays within [0, 1]\n",
    "    adjusted_t = torch.clamp(adjusted_t, 0, 1)\n",
    "    \n",
    "    return torch.cos(adjusted_t * math.pi / 2), torch.sin(adjusted_t * math.pi / 2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = torch.linspace(0, 1, 100)\n",
    "alphas, sigmas = get_alphas_sigmas(t, 128)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig, axs = plt.subplots(2, 1)\n",
    "\n",
    "axs[0].plot(t, alphas)\n",
    "axs[0].plot(t, sigmas)\n",
    "\n",
    "alphas, sigmas = get_alphas_sigmas(t, 12000)\n",
    "axs[1].plot(t, alphas)\n",
    "axs[1].plot(t, sigmas)\n",
    "plt.show()"
   ]
  },
  {
   "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
}