{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7af49ce5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME\n",
    "\n",
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1b4ee5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "input_vaes = np.load(\"/home/tony/test.npz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38aefd2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_vae_latents = input_vaes[\"vae_latents\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42ba8c68",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(input_vae_latents.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "923c06e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# input_vae_latents = input_vae_latents[:200]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab6cfc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# traditional codec\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    decode_stream_to_full_audio,\n",
    "    preload_models,\n",
    "    load_model,\n",
    "    decode,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873a518e",
   "metadata": {},
   "outputs": [],
   "source": [
    "codec_model = load_model(\"/app/suno/data/dpo/models/dac_vae_tuned_25hz.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b2654d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_audio = decode_stream_to_full_audio(input_vae_latents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1db72c56",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9cc56ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "def plot_decoded_audio_waveforms(\n",
    "    audio_a,\n",
    "    audio_b,\n",
    "    label_a=\"default_audio\",\n",
    "    label_b=\"default_audio_2\",\n",
    "    sample_rate=48000,\n",
    "):\n",
    "    \"\"\"\n",
    "    Plots two decoded audio waveforms for comparison.\n",
    "\n",
    "    Parameters:\n",
    "        audio_a: Audio object or np.ndarray\n",
    "        audio_b: Audio object or np.ndarray\n",
    "        label_a: str, label for the first audio\n",
    "        label_b: str, label for the second audio\n",
    "        sample_rate: int, sample rate in Hz (default: 44100)\n",
    "    \"\"\"\n",
    "    # Ensure both audio outputs are numpy arrays\n",
    "    arr1 = audio_a.array_float[0]\n",
    "    arr2 = audio_b.array_float[0]\n",
    "    if arr1.shape != arr2.shape:\n",
    "        print(f\"Audio shapes do not match: {arr1.shape} != {arr2.shape}\")\n",
    "        arr2 = arr2[: arr1.shape[0]]\n",
    "\n",
    "    # Create time axis in seconds\n",
    "    time_axis = np.arange(arr1.shape[0]) / sample_rate\n",
    "\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    plt.plot(time_axis, arr1 - arr2, label=label_a)\n",
    "    plt.title(\"Comparison of Decoded Audio Waveforms\")\n",
    "    plt.xlabel(\"Time (seconds)\")\n",
    "    plt.ylabel(\"Amplitude\")\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# Example usage:\n",
    "# plot_decoded_audio_waveforms(default_audio, default_audio_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b29d887d",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_audio_2 = decode(input_vae_latents, chunksize_s=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa102fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_decoded_audio_waveforms(default_audio, default_audio_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebafce5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.codec_engine import CodecEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87279954",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_engine = CodecEngine(\n",
    "    model=codec_model,\n",
    "    n_stride_tokens=20,\n",
    "    n_overlap_tokens=5,\n",
    "    is_vae=True,\n",
    "    compile=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9fbe4a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.codec_engine import Request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7386e667",
   "metadata": {},
   "outputs": [],
   "source": [
    "request = Request(\n",
    "    id=\"test\",\n",
    "    tokens=input_vae_latents,\n",
    "    input_tokens_finished=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d2caf37",
   "metadata": {},
   "outputs": [],
   "source": [
    "request_output = test_engine.run_request(request)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c25838a",
   "metadata": {},
   "outputs": [],
   "source": [
    "engine_audio = Audio.concatenate(request_output.generated_audios)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e774bcc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_decoded_audio_waveforms(default_audio, engine_audio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1379729",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env_dev",
   "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": 5
}