{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "889dc7bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "sys.path.append(\"/home/minz/dev_old/neon_codec/sunoCodec\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ec08365d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from models.codec_dac_vae import DACVAE\n",
    "from suno_utils.audio import Audio\n",
    "\n",
    "# S = torch.load(\"/app2/suno/checkpoints/2025-11-12_05-14-45/best_checkpoint.pt\")  # dac-large, float32\n",
    "# S = torch.load(\"/app2/suno/checkpoints/2025-11-12_05-03-52/best_checkpoint.pt\")  # dac-small, bfloat16\n",
    "# S = torch.load(\"/app2/suno/checkpoints/2025-11-12_05-04-09/best_checkpoint.pt\")  # dac finetuned, bfloat16\n",
    "# S = torch.load(\"/app2/suno/checkpoints/2025-11-12_14-48-06/best_checkpoint.pt\")  # dac finetuned decoder, bfloat16\n",
    "S = torch.load(\"/app2/suno/checkpoints/2025-11-13_05-36-14/checkpoint_iter_20000.pt\")  # dac finetuned, adv only\n",
    "model = DACVAE(**S[\"codec_args\"])\n",
    "model.load_state_dict(S[\"codec_model\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ca99cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = Audio.from_file(\"/home/minz/temp/country_road.mp3\", sample_rate=48000, n_channels=2).get_segment(from_s=0, to_s=30.0)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1db4ff5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "wav = audio.array_float[:, :48000*30]\n",
    "wav = torch.from_numpy(wav).unsqueeze(0)\n",
    "out = model(wav)[\"audio\"]\n",
    "Audio.from_array_float(out.detach().cpu().numpy()[0], sample_rate=48000).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4667ad36",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02415fe4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nenv",
   "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
}