{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"5\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c046f7b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job\n",
    "#from suno_utils.tasks.dac_vae_100hz_peaq import decode_stream_to_full_audio\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import decode_stream_to_full_audio, load_model\n",
    "\n",
    "import torch\n",
    "from suno_utils.audio import Audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c8b9fef",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion_gen.preload_tokenizer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caf93fd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_model(\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a718fb2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "diffusion_gen.preload_dit_model(\n",
    "    #dit_model_filepath=\"/app/suno/victor/2b_prefix_ft.pt\",\n",
    "    #codec_filepath=\"/app/suno/victor/25hz_vae_peaq_kl_0.005.pth\",\n",
    "    dit_model_filepath=\"/app/suno/checkpoints/2025-02-17_16-54-01_s7787/last_ckpt_infer.pt\",\n",
    "    #codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf696991",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "# gen_id = \"99bffa17-7e59-47b4-a048-5528cbda05d5\" # sister\n",
    "# gen_id = \"081d73c4-7805-4212-9c80-8db1137ca3c4\" # friends\n",
    "# gen_id = \"562f762d-6ced-4080-9af1-910ee3d0a5dc\" # something real\n",
    "# gen_id = \"23c15c62-494d-422d-8a60-8b0454044322\" # rubber duck\n",
    "# gen_id = \"4b140a9e-964b-422c-85b5-5861ad1a9d38\" # once\n",
    "# gen_id = \"7b214347-fa38-4e9b-96f4-f7ec65adea45\" # rock n roll\n",
    "gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"  # stone\n",
    "# gen_id = \"05ffc918-0bde-4418-8828-5105cd5717e1\"  # sara bad click\n",
    "# gen_id = \"7f5e7819-b7a1-449d-a76c-9c38950c4274\" # canada\n",
    "# gen_id = \"27135940-957f-4152-9e40-72e5c10a46bd\" # drunk\n",
    "# gen_id = \"caf1d7ae-ce81-4c75-b737-fd41e47a38ba\" # linger\n",
    "# gen_id = \"ac5aa29a-f18f-406f-a904-c47493f1beec\" # linger vox\n",
    "# gen_id = \"b04719c3-012a-4a4a-9b02-3cb2875d1744\" # courtney\n",
    "# gen_id = \"7407c7b3-8aa2-4f0a-82e8-fa19933f6e94\" # opera\n",
    "# gen_id = \"5b9624ef-b416-4095-a5f3-4945f09bba86\" # butterflies\n",
    "# gen_id = \"e757aa19-5c33-4945-aaa1-1519b7890f0b\" # bubblegum\n",
    "# gen_id = \"7f774078-1672-4858-a37f-acad373c5a84\" # golden\n",
    "# gen_id = \"1cee79db-a0ee-44e4-9646-d312cd620997\"  # chinese\n",
    "# gen_id = \"e404a4f0-64e0-423a-93d9-486b64c606ce\"  # mikey click\n",
    "\n",
    "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "vae_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}_vae.npz\"\n",
    "print(s3_filepath)\n",
    "data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "vae_data = read_from_s3(vae_filepath, read_f=np.load)\n",
    "audio = Audio.from_s3(mp3_filepath, n_channels=2)\n",
    "\n",
    "if \"v3.0_raw\" in data:\n",
    "    codes = data[\"v3.0_raw\"]\n",
    "elif \"v3.5_raw\" in data:\n",
    "    codes = data[\"v3.5_raw\"]\n",
    "elif \"v4.0_raw\" in data:\n",
    "    codes = data[\"v4.0_raw\"]\n",
    "else:\n",
    "    raise ValueError(\"No codes found\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ce61455",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(vae_data.keys())\n",
    "vae_latents = vae_data[\"vae_latents\"]\n",
    "vae_latents.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c0c34cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "engine = UpsampleEngine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b3fcef",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "    audio=audio,\n",
    "    steps=16,\n",
    ")\n",
    "request = Request(\n",
    "    id=\"dummy\",\n",
    "    generation_config=gen_cfg,\n",
    "    tokens=codes[:750, 0],\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "result = engine.run_request(request)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0aeca0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vae_latents = torch.concat(result.vae_latents)\n",
    "decode_stream_to_full_audio(vae_latents / 0.4).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8628f293",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_gpt45",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}