{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.dac_vae_100hz_peaq import (  # NOTE: works for 25hz as well\n",
    "    preload_models as preload_codec_models,\n",
    "    load_model as load_codec_model,\n",
    "    decode as codec_decode,\n",
    ")\n",
    "ckpt_path = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n",
    "preload_codec_models(checkpoint_filepath=ckpt_path)\n",
    "codec_model = load_codec_model()\n",
    "\n",
    "for param in codec_model.parameters():\n",
    "    param.requires_grad = False\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import auraloss\n",
    "\n",
    "mrstft_loss = auraloss.freq.MultiResolutionSTFTLoss(\n",
    "    fft_sizes=[1024, 2048, 512],\n",
    "    hop_sizes=[120, 240, 50],\n",
    "    win_lengths=[600, 1200, 240],\n",
    ").cuda()\n",
    "\n",
    "target_latents = torch.randn(4, 128, 750).cuda()\n",
    "pred_latents = torch.randn(4, 128, 750).cuda()\n",
    "\n",
    "target_audio = codec_model.decode(target_latents).detach()\n",
    "pred_audio = codec_model.decode(pred_latents)\n",
    "\n",
    "#loss = torch.nn.functional.mse_loss(y, target)\n",
    "#print(loss)\n",
    "loss = mrstft_loss(pred_audio, target_audio)\n",
    "print(loss)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "semantic_codes = torch.arange(750)\n",
    "\n",
    "n = 3\n",
    "phase = rand\n",
    "for nn in range(n - 1):\n",
    "    # Shift starting position by phase\n",
    "    shifted_idx = ((nn + phase) % n)\n",
    "    semantic_codes[shifted_idx::n] = 4000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(semantic_codes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env2",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}