{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a0d83737",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "88c3355d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\")\n",
    "\n",
    "# /home/christian/code/neon/sunoDiff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cd1bb050",
   "metadata": {},
   "outputs": [],
   "source": [
    "# base (25hz, 30s, ctx)\n",
    "# dit_model_filepath = \"/app/suno/checkpoints/2024-10-11_00-33-05_s7557/last_ckpt_infer.pt\"\n",
    "\n",
    "# prefix\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-11_02-04-25_s3533/last_ckpt_infer.pt\"\n",
    "\n",
    "# prefix2\n",
    "# dit_model_filepath = \"/app/suno/checkpoints/2024-10-13_13-04-30_s6052/last_ckpt_infer.pt\"\n",
    "\n",
    "# 100hz\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-11_16-53-57_s8050/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-16_19-23-27_s9332/last_ckpt_infer.pt\"\n",
    "\n",
    "# my 100hz with context\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-21_15-43-06_s3766/best_ckpt.pt\" \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-22_17-44-56_s9217/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz prefix\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-18_17-53-25_s1400/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz quality\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-21_18-29-39_s163/last_ckpt_infer.pt\"\n",
    "\n",
    "# latest 25hz prefix\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-26_04-39-23_s121/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-26_04-39-23_s121/step_1000000_infer.pt\"\n",
    "\n",
    "# dpo 25hz\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-28_18-52-39_s5303/last_ckpt_infer.pt\"\n",
    "#it_model_filepath = \"/app/suno/checkpoints/2024-10-28_19-12-42_s4304/last_ckpt_infer.pt\"\n",
    "\n",
    "# finetunes 25hz\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-30_20-11-31_s2639/last_ckpt_infer.pt\" # first chunk only\n",
    "\n",
    "# audio loss\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-03_20-59-48_s4933/last_ckpt_infer.pt\"\n",
    "\n",
    "#100hz \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-01_13-16-18_s411/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz dpo v2\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-04_14-51-03_s9527/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-04_15-08-19_s4277/last_ckpt_infer.pt\" # reverse labels\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-13_17-32-12_s3878/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-26_04-39-23_s121/step_1000000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-14_00-55-41_s2643/last_ckpt_infer.pt\"\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-30_15-14-54_s7904/last_ckpt_infer.pt\"\n",
    "\n",
    "#dpo 2\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-16_16-40-24_s2816/last_ckpt_infer.pt\"\n",
    "\n",
    "# tony dpo\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-08_15-04-27_s9054/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-18_01-35-12_s6013/last_ckpt_infer.pt\"\n",
    "\n",
    "#dit_model_filepath = \"s3://suno-data/georg/tmp/2b_prefix_ft.pt\"\n",
    "\n",
    "# new 100hz\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-18_14-09-23_s7157/last_ckpt_infer.pt\"\n",
    "\n",
    "# musicfm model\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-30_18-23-06_s289/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz rvq\n",
    "#it_model_filepath = \"/app/suno/checkpoints/2024-11-21_18-45-45_s8398/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz dpo tony\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-23_17-51-57_s8030/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-23_17-50-00_s7995/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-24_19-22-14_s7241/step_5000_infer.pt\"\n",
    "\n",
    "# 100hz patch size = 4\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-22_19-13-22_s8393/last_ckpt_infer.pt\"\n",
    "#100hz patch size = 10\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-22_21-16-57_s9022/last_ckpt_infer.pt\"\n",
    "\n",
    "# audio production fine tune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-25_21-22-27_s8273/last_ckpt_infer.pt\"\n",
    "\n",
    "# 16n dropout on semantic embed\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-27_03-58-38_s3695/last_ckpt_infer.pt\"\n",
    "\n",
    "# 48n 100hz paqtch size = 4\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-27_16-55-18_s8509/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-28_11-36-40_s5978/last_ckpt_infer.pt\"\n",
    "\n",
    "# no semantic finetine\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-30_15-14-54_s7904/last_ckpt_infer.pt\"\n",
    "\n",
    "# dpo 25hz\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-02_02-48-55_s3447/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-02_20-24-35_s7135/step_3000_infer.pt\"\n",
    "\n",
    "# 25hz dpo tony\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-11-30_09-50-23_s5271/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-03_05-21-13_s9053/step_5000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-03_14-34-13_s1574/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-03_16-06-09_s7094/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-03_17-34-18_s8852/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-03_22-51-23_s4390/step_1000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-05_03-42-43_s4538/step_5000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-09_04-25-13_s8586/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-10_14-30-03_s4515/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-10_16-12-25_s6751/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-10_19-53-58_s7792/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-10_19-57-41_s3580/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-12_19-26-57_s8223/step_5000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-14_06-26-29_s9980/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-23_21-22-12_s3923/step_1000_infer.pt\"\n",
    "\n",
    "# new jan 16\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-16_00-51-11_s2269/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-15_22-51-31_s6672/step_3000_infer.pt\"\n",
    "\n",
    "# 25hz sac vae convnext (minz) \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-04_22-31-34_s4925/last_ckpt_infer.pt\"\n",
    "\n",
    "# my 25hz finetune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-05_21-32-36_s2056/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz cross-attn\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-09_17-08-25_s8177/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz patch size = 5\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-09_22-57-40_s3457/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-11_17-37-56_s3412/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz logit normal \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-16_19-14-16_s9515/last_ckpt_infer.pt\"\n",
    "\n",
    "# 8n 100hz patch size = 4 6b parameter\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-10_17-01-20_s3538/last_ckpt_infer.pt\"\n",
    "\n",
    "#25hz vae fixed\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-12_19-32-59_s2767/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-12-16_15-54-26_s3031/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz mlp head (vae fixed)\n",
    "#it_model_filepath = \"/app/suno/checkpoints/2025-01-07_19-01-04_s5170/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz dropout (vae fixed)\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-07_19-59-48_s8031/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz shallow scratch depth 6 vae fixed\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-08_21-35-29_s3468/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-08_19-57-36_s2032/last_ckpt_infer.pt\" # georg ft\n",
    "\n",
    "# 25hz specaugment (vae fixed)\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-09_18-35-29_s2433/last_ckpt_infer.pt\"\n",
    "\n",
    "# new codecs to compare\n",
    "# 25hz vae dac tuned\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-20_14-28-26_s256/last_ckpt_infer.pt\"\n",
    "\n",
    "# 25hz vae convnext tuned\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-20_14-28-18_s3078/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-21_21-27-47_s3072/last_ckpt_infer.pt\" # resume\n",
    "\n",
    "# 25hz dropout_ctx 0.5\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-22_19-11-26_s6167/last_ckpt_infer.pt\"\n",
    "\n",
    "# new dpo testing\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-24_17-08-51_s8691/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-25_13-17-54_s1187/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt\"\n",
    "\n",
    "# current prod best with dpo\n",
    "#dit_model_filepath = \"s3://suno-data/tony/tmp/diff/dit_v2_dpo_t2_v1_3k.pt\"\n",
    "\n",
    "# 25hz dpo tony decay fix\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-26_02-13-54_s2704/step_3000_infer.pt\"\n",
    "\n",
    "# 5s noise finetune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-01-29_22-02-21_s9876/last_ckpt_infer.pt\"\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-08_01-36-56_s5432/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-08_19-01-33_s2586/last_ckpt_infer.pt\"\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-14_16-36-27_s8793/last_ckpt_infer.pt\"\n",
    "\n",
    "# tony dpo no shimmer?\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-18_13-15-38_s5317/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-02_00-56-00_s1708/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-04_05-30-37_s7677/step_30000_infer.pt\"\n",
    "\n",
    "# 4b wide with 96 heads, seems to be working\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-24_18-41-30_s8760/last_ckpt_infer.pt\"\n",
    "# 4b deep with 64 layers?\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-18_20-50-22_s6784/last_ckpt_infer.pt\"\n",
    "\n",
    "# infilling ckpt\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-28_22-23-39_s8204/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-03_17-46-56_s3599/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-04_22-24-12_s3548/last_ckpt_infer.pt\"\n",
    "\n",
    "# new semantic free finetune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-03_21-05-49_s9618/last_ckpt_infer.pt\"\n",
    "\n",
    "\n",
    "# tony test \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-04_05-30-37_s7677/step_30000_infer.pt\" # new\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-20_21-37-57_s6971/step_5000_infer.pt\" # prod\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-18_04-22-21_s3987/step_30000_infer.pt\" # ear cut\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-20_17-24-50_s5157/step_30000_infer.pt\" # new dpo\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-24_01-33-03_s1848/step_9000_infer.pt\" # new dpo\n",
    "\n",
    "# new models \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-02-17_16-54-01_s7787/last_ckpt_infer.pt\" # base model\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-06_20-36-11_s3822/last_ckpt_infer.pt\" # finetune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-10_17-57-21_s5734/last_ckpt_infer.pt\" # finetune v2\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-17_19-55-49_s5529/last_ckpt_infer_patch.pt\" # finetune new more noise\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-18_17-24-52_s1174/last_ckpt_infer.pt\" # finetune new with infill\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-18_22-48-27_s8275/last_ckpt_infer.pt\" # syn sft\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-19_14-04-58_s4839/last_ckpt_infer.pt\" # new syn sft\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-20_10-34-55_s7844/last_ckpt_infer.pt\" # semantic noise finetune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-20_22-59-39_s6032/last_ckpt_infer.pt\" # finetune infill v2\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-21_16-41-16_s7602/last_ckpt_infer.pt\" # finetune infill only\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-03_20-59-28_s8707/last_ckpt_infer.pt\" # finetune infill only aligned 0.9\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-21_16-32-22_s9342/last_ckpt_infer.pt\" # finetune with infill and loss mask (bugged)\n",
    "\n",
    "\n",
    "# skip finetunes\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-20_17-57-48_s5949/last_ckpt_infer.pt\" # finetune skip 2\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-31_13-36-46_s7452/last_ckpt_infer.pt\" # finetune skip 4\n",
    "\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-24_18-26-36_s5614/last_ckpt_infer.pt\" # finetune infill\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-24_22-31-58_s4919/last_ckpt_infer.pt\" # finetune infill fixed\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-24_22-31-58_s4919/step_250000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-28_18-13-20_s465/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-28_20-05-51_s9115/last_ckpt_infer.pt\"\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-28_18-46-17_s894/last_ckpt_infer.pt\" # finetune infill only\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-30_00-36-24_s4502/last_ckpt_infer.pt\" # finetune infill 2.5e-5\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-30_00-36-24_s4502/last_ckpt_infer.pt\" # finetune infill 2.5e-5 more steps\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-30_12-43-00_s7315/last_ckpt_infer.pt\" # finetune infill with ctx mode\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-31_17-26-58_s1124/last_ckpt_infer.pt\" # 4b with infill\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-03_19-36-00_s7765/last_ckpt_infer.pt\" # 4b without infill\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-13_16-02-25_s7511/last_ckpt_infer.pt\" # bottleneck ctx\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-14_18-44-38_s5003/last_ckpt_infer.pt\"  # bottleneck ctx 2\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-14_22-12-28_s2013/last_ckpt_infer.pt\" # 0.05 scale factor\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-15_11-14-10_s7859/last_ckpt_infer.pt\" # tanh on ctx (broken)\n",
    "\n",
    "# distllation \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-01_02-00-26_s861/last_ckpt_infer.pt\" # 4-step \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-07_19-12-45_s8459/last_ckpt_infer.pt\"\n",
    "\n",
    "# semantic prediction\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-22_21-09-52_s9652/last_ckpt_infer.pt\"\n",
    "\n",
    "# sample model\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-03-25_22-39-40_s843/last_ckpt_infer.pt\"\n",
    "\n",
    "# audio prod finetune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-11_14-14-12_s4538/last_ckpt_infer.pt\"\n",
    "\n",
    "# shared context\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-10_01-05-22_s3784/last_ckpt_infer.pt\"\n",
    "\n",
    "# dpo v2\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-04_15-51-42_s4429/step_3000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-10_04-13-39_s12/step_6000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-10_03-27-31_s8101/step_3000_infer.pt\"\n",
    "# 1e-6\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-10_20-37-33_s7078/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-10_04-13-39_s12/last_ckpt_infer.pt\" \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-10_04-13-39_s12/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-15_05-15-36_s7464/last_ckpt_infer.pt\"# shared ctx dpo\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-15_02-55-37_s7688/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-15_18-52-18_s7734/last_ckpt_infer.pt\"\n",
    "\n",
    "# 2b retrain with shared ctx\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-15_00-52-47_s9732/last_ckpt_infer.pt\"\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-16_11-28-40_s3492/last_ckpt_infer.pt\" # new infill\n",
    "\n",
    "# 30k \n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-12_05-22-06_s1150/step_12000_ckpt.pt\"\n",
    "\n",
    "# stems\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-04_15-44-09_s4154/last_ckpt_infer.pt\"\n",
    "\n",
    "# new infill\n",
    "# dit_model_filepath = \"/app/suno/checkpoints/2025-04-16_11-28-40_s3492/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-18_15-39-58_s2609/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-18_15-48-39_s806/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-16_11-28-40_s3492/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-16_11-28-40_s3492/step_600000_ckpt.pt\"\n",
    "\n",
    "# ctx norm\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-18_17-08-07_s7798/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-19_19-54-42_s330/last_ckpt_infer.pt\"\n",
    "\n",
    "# ear finetune\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-19_11-50-35_s4128/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-19_20-59-02_s6204/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_21-18-08_s4974/last_ckpt_infer.pt\" # infill\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_22-18-31_s3601/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_21-18-08_s4974/last_ckpt_infer.pt\"\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_21-18-08_s4974/step_100000_infer.pt\"\n",
    "\n",
    "# tony test\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_14-15-15_s7325/last_ckpt_infer.pt\" # new code\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_15-17-47_s9991/last_ckpt_infer.pt\" # old code via commit hash\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-10_04-13-39_s12/last_ckpt_infer.pt\" # old code original checkpoint\n",
    "\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_18-21-37_s2810/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2025-04-21_18-24-40_s2014/last_ckpt_infer.pt\"\n",
    "#dit_model_filepath = '/app/suno/checkpoints/2025-04-22_04-22-04_s3287/step_3000_infer.pt'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13171812",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# musicfm\n",
    "#semantic_model_filepath = \"\" #\"s3://suno-data/minz/models/musicfm_v2_epoch138.pt\"\n",
    "#semantic_clusters_filepath = \"\" #\"s3://suno-data/minz/models/musicfm_v2_fixed_centroids_5s_l6_4k.npy\"\n",
    "\n",
    "# mert25\n",
    "semantic_model_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25.pt\"\n",
    "semantic_clusters_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy\"\n",
    "\n",
    "#codec_filepath=\"/app/suno/christian/models/codecs/dac_2c_25x12.pt\"\n",
    "#codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/100hz_vae_peaq_kl_0.005.pth\"\n",
    "#codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\"\n",
    "#codec_filepath=\"/app/suno/christian/models/codecs/vae_37epoch.ckpt\"\n",
    "#codec_filepath = \"s3://suno-data/minz/models/dac_vae_fixed_25hz.pth\"\n",
    "codec_filepath = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n",
    "#codec_filepath = \"s3://suno-data/minz/models/convnext_vae_tuned_25hz.pth\"\n",
    "\n",
    "if \"2025-03-14_22-12-28_s2013\" in dit_model_filepath:\n",
    "    patch_size = 1\n",
    "    codec_scale_factor = 0.05\n",
    "elif \"dac_vae_fixed\" in codec_filepath or \"dac_vae_tuned\" in codec_filepath:\n",
    "    patch_size = 1\n",
    "    codec_scale_factor = 0.4\n",
    "elif \"convnext_vae_tuned\" in codec_filepath:\n",
    "    patch_size = 1\n",
    "    codec_scale_factor = 1.0\n",
    "else:\n",
    "    patch_size = 1\n",
    "    codec_scale_factor = 2.5\n",
    "\n",
    "#codec_scale_factor = 3.4\n",
    "\n",
    "import torch\n",
    "from suno_utils.audio import Audio\n",
    "from generation import preload_models, generate, _retrieve_models\n",
    "_ = preload_models(\n",
    "    tokenizer_filepath=\"/home/georg/notebooks/gpu_nb/tmp/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=semantic_model_filepath,\n",
    "    semantic_clusters_filepath=semantic_clusters_filepath,\n",
    "    codec_filepath=codec_filepath,\n",
    "    dit_model_filepath=dit_model_filepath,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    model_type=\"prefix\",\n",
    "    #compile=True,\n",
    "    codec_scale_factor=codec_scale_factor,\n",
    "    #patch_size=patch_size,\n",
    ")\n",
    "\n",
    "models = _retrieve_models()\n",
    "model_duration_s = 30\n",
    "if models[\"dit_model\"].ctx_len is not None:\n",
    "    model_duration_s = 6 * 60\n",
    "else:\n",
    "    model_duration_s = models[\"dit_model\"].block_size // models[\"dit_model\"].io_hz\n",
    "duration_s = 2*60 if model_duration_s >= 2*60 else 30"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4454670a",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90969db9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "687dd8ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# temp fix to ensure ctx_len is not None\n",
    "if models[\"dit_model\"].ctx_len is None:\n",
    "    models[\"dit_model\"].ctx_len = 0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac0d49e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "silence_tensor = torch.zeros(2, 48000).cuda()\n",
    "\n",
    "with torch.no_grad():\n",
    "    z = models[\"codec_model\"].encode(silence_tensor.unsqueeze(0))[\"z\"].half().float()\n",
    "    pred_audio = models[\"codec_model\"].decode(z).cpu()\n",
    "    # crop pred to 120s\n",
    "    pred_audio = pred_audio[:, :48000*120]\n",
    "\n",
    "IPython.display.display(IPython.display.Audio(pred_audio.cpu().squeeze(0), rate=48000))\n",
    "IPython.display.display(IPython.display.Audio(silence_tensor.cpu().squeeze(0), rate=48000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "776b734c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from suno_utils.audio import Audio\n",
    "import numpy as np\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\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 = \"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 = \"6727880b-f5e8-4c48-a401-5831d14c71eb\" # rock instrumental/\n",
    "#gen_id = \"8e96a40b-f528-4cfd-9b25-edcdc50080d1\" # 80s\n",
    "#gen_id = \"88fe5e81-b07f-462b-8130-04fc7c964963\" # stary night\n",
    "#gen_id = \"3617dc97-2c11-4da4-8b26-773927fa9c7b\" # makeup\n",
    "#gen_id = \"7735f201-7cce-4067-9496-67f590ec01fc\" # back pain\n",
    "#gen_id = \"77260536-b692-4577-aac7-a53c13754ec4\" # love again\n",
    "#gen_id = \"707299a3-7668-484f-8479-421e516c6916\" # italian\n",
    "#gen_id = \"4ec14116-ac9b-49e3-ae62-99eb2f1ac0c2\" # circles\n",
    "#gen_id = \"8dc2c1f3-213d-4503-84e2-65216085cc9a\" # beef\n",
    "\n",
    "# v4 examples\n",
    "#gen_id = \"6daf2794-c9b5-4d2b-ab14-58652e0014d7\" # metal lyrics\n",
    "#gen_id = \"bc095772-df08-48c0-80ac-89e02dc34e40\" # metal instrumental\n",
    "#gen_id = \"1a042042-b5d5-43e3-8b50-965d1ec0e303\" # can it be real\n",
    "#gen_id = \"1a3ac8fd-b5f8-4e78-92f5-400637e0851c\" # djent\n",
    "#gen_id = \"3a66709e-3afa-446c-8a28-ea7a21c5ed72\" # hard rock\n",
    "#gen_id = \"42985b5e-b8a9-4ca5-9aa2-4dd8df9b2f42\" # yatch rock\n",
    "#gen_id = \"b8057b18-8b02-41d4-927d-6adae708f3d4\" # jazz fusion\n",
    "#gen_id = \"1a042042-b5d5-43e3-8b50-965d1ec0e303\" # female vocal\n",
    "#gen_id = \"ebb4ce1d-8f79-4a12-8aa8-858bcc65b44a\" # optimistic\n",
    "#gen_id = \"baa78a26-936b-48e9-9b28-1ae83fdb5075\" # falling\n",
    "#gen_id = \"7719f273-bbe9-406a-a537-6075c3ff2620\"\n",
    "#gen_id = \"2e0eec9e-4f86-46c5-9a43-508829fb9d0b\"\n",
    "\n",
    "\n",
    "#gen_id = \"b8057b18-8b02-41d4-927d-6adae708f3d4\"\n",
    "#gen_id = \"72a8c50d-7e16-469d-8283-56bed483374c\"\n",
    "#gen_id = \"382b1682-8503-4b3d-94b6-5fa296933975\" # orchestral ringing\n",
    "#gen_id = \"cee9773b-9514-43bf-bba4-6561264ea94a\" # strong ringning (good test case, has decay issue)\n",
    "#gen_id = \"7c62f810-f77d-4a98-b179-d52c538a996f\" # strong ringing 2\n",
    "#gen_id = \"7c62f810-f77d-4a98-b179-d52c538a996f\"\n",
    "#gen_id = \"c04a5e04-4729-4d52-8bf4-8381635cb9a2\" # shimmer shine\n",
    "#gen_id = \"a69d5f40-62fc-4eca-9610-2216a9e20de6\"\n",
    "\n",
    "\n",
    "# metal \n",
    "#gen_id = \"99f0e17b-ae5b-4533-b049-0cb2489d77ed\"\n",
    "#gen_id = \"4f93438b-7cb2-4ad9-8c80-2475eff873e8\"\n",
    "#gen_id = \"f2cf6374-96de-4307-b797-79abacfe3461\"\n",
    "#gen_id = \"9c6614e3-5cff-4a77-b7b4-8d91f30cb4b7\"\n",
    "#gen_id = \"6281de8a-3c0d-4c24-b818-70dc4250dc07\"\n",
    "#gen_id = \"6a769a66-eb94-4219-80bc-24242e8d2702\"\n",
    "# \n",
    "\n",
    "#gen_id = \"1ec51ef1-378e-4f33-94a5-05c9b1e1448f\"\n",
    "#gen_id = \"dfcbd035-2b9d-4c43-89ba-39411d0c6e43\"\n",
    "\n",
    "# new \n",
    "#gen_id = \"24ea970a-c66b-4260-a57e-f519dca7df8b\"\n",
    "#gen_id = \"1a18fd9f-5f95-4334-8e59-143c29440238\"\n",
    "#gen_id = \"ed887516-c9c0-4c0b-9e3b-898991d7a29b\"\n",
    "#gen_id = \"47bfe5c3-aec5-46a5-a766-06fa4acb6e69\"\n",
    "#gen_id = \"51612fa6-5b54-4b99-999b-bb7b005e03ab\"\n",
    "#gen_id = \"bbf2b264-24e7-4ae7-9c01-e5faec1b4d58\"\n",
    "#gen_id = \"2ca5daed-6d21-429b-acf6-6ba1e4f94f35\"\n",
    "#gen_id = \"de6f341c-29f6-4009-9e69-ee0ed9afe8ac\"\n",
    "#gen_id = \"d149a604-2091-4515-b415-e8c5ac819a9a\"\n",
    "#gen_id = \"d3fcbd11-123a-48a0-85eb-b679a1256023\"\n",
    "#gen_id = \"3314bb17-0e88-40ec-a81b-f2e774c42a27\"\n",
    "#gen_id = \"c188e7e4-56c5-42cc-bdf1-9fbb68b9e119\" # choral\n",
    "#gen_id = \"9fd68e1b-3d91-4875-b7c7-cd1ee4bcc9de\"\n",
    "#gen_id = \"26f6502a-5aa8-4ff5-aa38-8a1a50f9ba82\"\n",
    "#gen_id = \"fb60e2cd-f694-4950-b6b2-300ee6277697\"\n",
    "#gen_id = \"1c06c231-dc26-498b-87ba-e8fd93a7de83\"\n",
    "#gen_id = \"56e61043-b11d-4758-9a61-5c013db4a0a3\"\n",
    "#gen_id = \"7464d7a8-4732-4364-84ee-bf85436f0e10\"\n",
    "\n",
    "#gen_id = \"1bb8ea30-31a2-43a6-8b94-ba15fe1e4eff\"\n",
    "\n",
    "# pronouce\n",
    "#gen_id = \"d5ffe739-f6c4-4850-8ce2-db93a422aa78\"\n",
    "\n",
    "# my songs\n",
    "#gen_id = \"d22f9b97-249c-4eda-8ad2-6c6c7b4aed35\"\n",
    "#gen_id = \"ebbddbbc-f0fa-400e-9789-94978fb0bf05\"\n",
    "\n",
    "# new codec\n",
    "#gen_id = \"85be20bd-a388-4872-8a0d-77e3385e9bde\"\n",
    "#gen_id = \"549cfa66-65cf-4bee-b264-389c0c029b45\"\n",
    "#gen_id = \"72adda71-4382-4cfa-b21f-150fb1809044\"\n",
    "#gen_id = \"7be341c0-b9af-48da-856b-136f409bba01\"\n",
    "#gen_id = \"1305dd9b-1625-4859-84ed-d40ee56b4d86\"\n",
    "#gen_id = \"6c425fb1-0194-4995-82e0-dc2764be8bb3\"\n",
    "\n",
    "# speech\n",
    "gen_id = \"1f1f19dd-174d-4c6c-8d97-6130540fd1bc\"\n",
    "\n",
    "# greek\n",
    "#gen_id = \"579dafcc-dc96-4993-9572-f904d4c37f3d\"\n",
    "\n",
    "# staging clips\n",
    "#gen_id = \"e65d270a-61f1-462f-9c38-8544bfc51a33\"\n",
    "\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",
    "\n",
    "\n",
    "print(s3_filepath)\n",
    "data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "\n",
    "vae_data = read_from_s3(vae_filepath, read_f=np.load)\n",
    "print(vae_data[\"vae_latents\"].shape)\n",
    "audio = Audio.from_s3(mp3_filepath, n_channels=2)#.get_slice(0, 120.01)\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",
    "elif \"v5.0_raw\" in data:\n",
    "    codes = data[\"v5.0_raw\"]\n",
    "else:\n",
    "    raise ValueError(\"No codes found\")\n",
    "\n",
    "\n",
    "#aws s3 cp s3://suno-data-uploads/studio/uploads/2e0eec9e-4f86-46c5-9a43-508829fb9d0b_hoot.json text_data.json\n",
    "\n",
    "#text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n",
    "os.system(f\"aws s3 cp s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json text_data.json\")\n",
    "text_data = open(\"text_data.json\", \"r\", encoding=\"utf-8\").read()\n",
    "aligned_lyrics = json.loads(text_data)\n",
    "print(aligned_lyrics)\n",
    "audio.normalize_volume().play()\n",
    "\n",
    "tags = \"pop\"\n",
    "\n",
    "lyrics = \"\"\n",
    "for elem in aligned_lyrics:\n",
    "    if \"word\" in elem:\n",
    "        lyrics += elem[\"word\"]\n",
    "\n",
    "semantic_codes = torch.from_numpy(codes[:, 0]).long()#.cuda()\n",
    "#semantic_codes = semantic_codes[:3000]\n",
    "print(semantic_codes.shape)\n",
    "\n",
    "# try to load the downbeats\n",
    "if os.path.exists(\"downbeats_{}.npz\".format(gen_id)):\n",
    "    downbeats = np.load(\"downbeats_{}.npz\".format(gen_id))[\"downbeats\"]\n",
    "    print(downbeats)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e308232",
   "metadata": {},
   "outputs": [],
   "source": [
    "# manual \n",
    "\n",
    "#audio_fp = \"/home/georg/notebooks/samples/donauwalzer_sample.mp3\"\n",
    "audio_fp = \"/home/christian/audio/reference-audio-wav/05 Another Brick In The Wall, Part 2.wav\"\n",
    "tags = \"pop\"\n",
    "lyrics = \"\"\"\n",
    "[Verse 1: Roger Waters & David Gilmour]\n",
    "We don't need no education\n",
    "We don't need no thought control\n",
    "No dark sarcasm in the classroom\n",
    "Teacher, leave them kids alone\n",
    "\n",
    "[Pre-Chorus: Roger Waters & David Gilmour]\n",
    "Hey, teacher, leave them kids alone\n",
    "\n",
    "[Chorus: Roger Waters & David Gilmour]\n",
    "All in all, it's just another brick in the wall\n",
    "All in all, you're just another brick in the wall\n",
    "\n",
    "[Verse 2: Islington Green School Students]\n",
    "We don't need no education\n",
    "We don't need no thought control\n",
    "No dark sarcasm in the classroom\n",
    "Teachers, leave them kids alone\n",
    "\n",
    "[Pre-Chorus: Islington Green School Students]\n",
    "Hey, teacher, leave us kids alone!\n",
    "\n",
    "[Chorus: Islington Green School Students]\n",
    "All in all, you're just another brick in the wall\n",
    "All in all, you're just another brick in the wall\n",
    "\n",
    "[Guitar Solo]\n",
    "\n",
    "[Outro: Roger Waters]\n",
    "Wrong, do it again\n",
    "Wrong, do it again\n",
    "If you don't eat your meat, you can't have any pudding\n",
    "(Wrong, do it again)\n",
    "How can you have any pudding if you don't eat your meat?\n",
    "(Wrong, do it again)\n",
    "(If you don't eat your meat)\n",
    "You, yes, you behind the bike sheds, stand still, laddie\n",
    "(If you don't eat your meat, you can't have any pudding\n",
    "How can you have any pudding if you don't eat your meat?)\n",
    "(You, yes, you behind the bike sheds, stand still, laddie)\n",
    "\"\"\"\n",
    "\n",
    "if True:\n",
    "    audio_fp = \"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n",
    "    lyrics = \"\"\"\n",
    "    I waited 'til I saw the sun\n",
    "    I don't know why I didn't come\n",
    "    I left you by the house of fun\n",
    "    Don't know why I didn't come\n",
    "    Don't know why I didn't come\n",
    "    When I saw the break of day\n",
    "    I wished that I could fly away\n",
    "    Instead of kneeling in the sand\n",
    "    Catching tear-drops in my hand\n",
    "    My heart is drenched in wine\n",
    "    But you'll be on my mind forever\n",
    "    Out across the endless sea\n",
    "    I would die in ecstasy\n",
    "    But I'll be a bag of bones\n",
    "    Driving down the road alone\n",
    "    My heart is drenched in wine\n",
    "    But you'll be on my mind forever\n",
    "    Something has to make you run\n",
    "    I don't know why I didn't come\n",
    "    I feel as empty as a drum\n",
    "    I don't know why I didn't come\n",
    "    Don't know why I didn't come\n",
    "    I don't know why I didn't come\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "#audio_fp = \"/home/georg/notebooks/samples/eminem_doctor.wav\"\n",
    "#tags = \"pop, rap\"\n",
    "#lyrics = \"[Chorus]\\nI'm about to lose my mind\\nYou've been gone for so long\\nI'm runnin' out of time\\nI need a doctor, call me a doctor\\nI need a doctor, doctor\\nTo bring me back to life\\n\\n[Verse 1]\\nI told the world, one day, I would pay it back\\nSay it on tape, and lay it, record it\\nSo that one day I could play it back\\nBut I don't even know if I believe it when I'm sayin' that\\nDoubts startin' to creep in, everyday, it's just so gray and black\\nHope, I just need a ray of that, 'cause no one sees my vision\\nWhen I play it for 'em, they just say it's wack\\nBut they don't know what dope is\\nAnd I don't know if I was awake or asleep when I wrote this\\nAll I know is you came to me when I was at my lowest\\nYou picked me up, breathed new life in me, I owe my life to you\\nBut for the life of me, I don't see why you don't see like I do\\nBut it just dawned on me, you lost a son\\nDemons fightin' you, it's dark\\nLet me turn on the lights and brighten me, and enlighten you\\nI don't think you realize what you mean to me\\nNot the slightest clue\\n'Cause me and you were like a crew, I was like your sidekick\\nYou gon' either wanna fight when I get off this fuckin' mic\\nOr you gon' hug me, but I'm outta options\\nThere's nothin' else I can do 'cause—\\n\\n[Chorus]\\nI'm about to lose my mind\\nYou've been gone for so long\\nI'm runnin' out of time\\nI need a doctor\\nCall me a doctor\\nI need a doctor, doctor\\nTo bring me back to life\\n\\n[Verse 2]\\nIt hurts when I see you struggle, you come to me with ideas\\nYou say they're just pieces, so I'm puzzled\\n'Cause the shit I hear is crazy\\nBut you're either gettin' lazy, or you don't believe in you no more\\nSeems like your own opinions, not one you can form\\nCan't make a decision, you keep questionin' yourself\\nSecond guessin', and it's almost like you're beggin' for my help\\nLike I'm your leader\\nYou're supposed to fuckin' be my mentor\\nI can endure no more, I demand you remember who you are\\nIt was you who believed in me\\nWhen everyone was tellin' you don't sign me\\nEveryone at the fuckin' label, let's tell the truth\\nYou risked your career for me, I know it as well as you\\nNobody wanted to fuck with the white boy\\nDre, I'm cryin' in this booth\\nYou saved my life, now, maybe it's my turn to save yours\\nBut I can never repay you, what you did for me is way more\\nBut I ain't givin' up faith, and you ain't givin' up on me\\nGet up, Dre, I'm dyin', I need you\\nCome back, for fuck's sake, 'cause—\\n\\n[Chorus]\\nI'm about to lose my mind\\nYou've been gone for so long\\nI'm runnin' out of time\\nI need a doctor, call me a doctor\\nI need a doctor, doctor\\nTo bring me back to life\\nBring me back to life\\nBring me back to life\\nI need a doctor, doctor\\nTo bring me back to life\\n\\n[Verse 3]\\nIt literally feels like a lifetime ago\\nBut I still remember the shit like it was just yesterday, though\\nYou walked in, yellow jump suit, whole room, cracked jokes\\nOnce you got inside the booth, told you, mic smoke\\nWent through friends, some of them I put on, but they just left\\nThey said they was ridin' to the death\\nBut where the fuck are they now, now that I need them?\\nI don't see none of them, all I see is Slim\\nFuck all you fair-weather friends, all I need is him\\nFuckin' backstabbers\\nWhen the chips were down, you just laughed at us\\nNow you 'bout to feel the fuckin' wrath of Aftermath, faggots\\nYou gon' see us in our lab jackets and ask where the fuck we been\\nYou can kiss my indecisive ass crack, maggots\\nAnd the cracker's ass, little cracker jack beat\\nMakin' wack-ass backwards producers, I'm back, bastards\\nOne more CD, and then I'm packin' up my bags\\nAnd as I'm leavin', I'll guarantee they'll scream\\n\\\"Dre, don't leave us like that, man,\\\" 'cause—\\n\\n[Chorus]\\nI'm about to lose my mind\\nYou've been gone for so long\\nI'm runnin' outta time\\nI need a doctor, call me a doctor\\nI need a doctor, doctor\\nTo bring me back to life\"\n",
    "\n",
    "audio = Audio.from_file(audio_fp, n_channels=2).get_slice(00.0, 240.01).convert(48000, 2, 2)\n",
    "audio.play()\n",
    "semantic_codes = audio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d82ac79d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cycle test\n",
    "import glob\n",
    "import torchaudio\n",
    "from tqdm import tqdm\n",
    "\n",
    "base_dir = \"/home/christian/audio/ear-bench\"\n",
    "output_dirname = \"codec_cycle_test\"\n",
    "output_dir = os.path.join(base_dir, output_dirname)\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# take reference audio and codec cycle it \n",
    "audio_filepaths = glob.glob(os.path.join(\"/home/christian/audio/reference-audio-wav\", \"*.wav\"))\n",
    "\n",
    "for audio_filepath in tqdm(audio_filepaths):\n",
    "    filename = os.path.basename(audio_filepath).replace(\".wav\", \"\")\n",
    "    audio_tensor, sr = torchaudio.load(audio_filepath)\n",
    "    \n",
    "    # resample to 48000\n",
    "    audio_tensor = torchaudio.functional.resample(audio_tensor, sr, 48000)\n",
    "    # crop to 120s\n",
    "    audio_tensor = audio_tensor[:, :48000*120]\n",
    "    audio_tensor = audio_tensor.to(\"cuda\")\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        z = models[\"codec_model\"].encode(audio_tensor.unsqueeze(0))[\"z\"].half().float()\n",
    "        pred_audio = models[\"codec_model\"].decode(z).cpu()\n",
    "        # crop pred to 120s\n",
    "        pred_audio = pred_audio[:, :48000*120]\n",
    "    \n",
    "    # save the original and the pred\n",
    "    output_filepath = os.path.join(output_dir, f\"{filename}_cycled_pref=False.mp3\")\n",
    "    torchaudio.save(output_filepath, pred_audio.squeeze(0).cpu(), 48000)\n",
    "    output_filepath = os.path.join(output_dir, f\"{filename}_input_pref=True.mp3\") \n",
    "    torchaudio.save(output_filepath, audio_tensor.cpu(), 48000)\n",
    "\n",
    "\n",
    "#import torchaudio\n",
    "#torchaudio.save(\"donauwalzer_sample_notebook_half_cycle.wav\", pred_audio.squeeze(0).cpu().numpy(), 48000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbdef9a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "semantic_skip_factors = [2, 3, 4, 5]\n",
    "semantic_pad_token = 4000\n",
    "semantic_codes = torch.randint(0, 4000, (100,))\n",
    "\n",
    "n = random.choice(semantic_skip_factors)\n",
    "print(n)\n",
    "\n",
    "n_phase = random.randint(0, n - 1)\n",
    "for nn in range(n - 1):\n",
    "    shifted_idx = (nn + n_phase) % n\n",
    "    semantic_codes[shifted_idx::n] = semantic_pad_token\n",
    "\n",
    "print(semantic_codes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79414f07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.mert_25 import (\n",
    "    preload_models as preload_semantic_models,\n",
    "    load_model as load_semantic_model,\n",
    "    encode as encode_semantic,\n",
    "    EMBEDDING_RATE as SEMANTIC_HZ,\n",
    ")\n",
    "import numpy as np\n",
    "\n",
    "# vae encode to get latents \n",
    "audio_tensor = audio.array_float\n",
    "audio_tensor = audio_tensor[:, :48000*120]\n",
    "audio_tensor = torch.from_numpy(audio_tensor).to(\"cuda\")\n",
    "\n",
    "with torch.no_grad():\n",
    "    z = models[\"codec_model\"].encode(audio_tensor.unsqueeze(0))[\"z\"].half().float()\n",
    "\n",
    "print(z.shape)\n",
    "\n",
    "# also encode semantic codes\n",
    "semantic_codes = encode_semantic(\n",
    "    semantic_codes.convert(sample_rate=24_000, byte_width=2, n_channels=1)\n",
    ").astype(np.int64)[:, 0]\n",
    "semantic_codes = torch.from_numpy(semantic_codes)\n",
    "print(semantic_codes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb532b47",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a8da402",
   "metadata": {},
   "outputs": [],
   "source": [
    "# also codec decode to get latents\n",
    "with torch.no_grad():\n",
    "    z_decoded = models[\"codec_model\"].decode(z)\n",
    "print(z_decoded.shape)\n",
    "\n",
    "import torchaudio\n",
    "torchaudio.save(\"audio/original-sun.mp3\", z_decoded.squeeze(0).cpu(), 48000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "5f4087d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "lyrics = \"\"\"\n",
    "[Chorus]\n",
    "It's weird\n",
    "A little uncanny\n",
    "But I can still tell\n",
    "I guess fake knows fake\n",
    "I guess fake knows fake\n",
    "\n",
    "[Verse 2]\n",
    "Felt you\n",
    "I think\n",
    "In a past life\n",
    "But that's an idea\n",
    "All you ever are is an idea\n",
    "All I ever am is an idea\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "#lyrics = \"[Instrumental]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "231d6241",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(lyrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "07b6f35c",
   "metadata": {},
   "outputs": [],
   "source": [
    "lyrics = \"\"\"\n",
    "Things could be different\n",
    "Things could change\n",
    "Things must be different\n",
    "Things won't change\n",
    "\n",
    "But... difference breeds hate\n",
    "and hate gives way to brutality.\n",
    "\n",
    "Brutality \n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a508920",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98fb6f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import IPython\n",
    "import torchaudio\n",
    "import torch\n",
    "\n",
    "steps = 16\n",
    "text_cfg_coef = 1.5\n",
    "ctx_cfg_coef = 1.0\n",
    "\n",
    "seeds = [np.random.randint(0, 2**32-1) for _ in range(3)]\n",
    "#tags = \"pop, rap\"\n",
    "\n",
    "#lyrics = \"\"\n",
    "tags = \"pop\"\n",
    "#tags = \"metal\"\n",
    "#tags = \"vocal\"\n",
    "#tags = \"classical, piano, sonata, classical piano\"\n",
    "#tags = \"Pop Rap, Trap, Emo Rap, Alt-Pop, Emo-Pop, Beat, 808, Heavy\"\n",
    "#tags = \"grammy winner, contemporary r&b, 2020, beat, trap\"\n",
    "#tags = \"poetry, spoken word, speech, podcast\"\n",
    "#tags = \"alternative, indie, 2010s, indie rock, indie pop\"\n",
    "#tags = \"midwest emo, 1990s, alternative rock, chicago school, indie rock\"\n",
    "#tags = \"folk, protest song, americana, 1960s, raspy male vocalist\"\n",
    "#tags = \"lap steel guitar, pedal steel guitar, country, instrumental, guitar, drums\"\n",
    "#tags = \"bluegrass, banjo, fiddle, acoustic, country, americana\"\n",
    "#tags = \"math pop, math rock, nuevo flamenco, complex, electric guitar, guitar solo\"\n",
    "#tags = \"contemporary folk, chamber pop, 2020, indie folk, acoustic guitar, violin, drums, introspective\"\n",
    "\n",
    "param_configs = [\n",
    "    #{\"text_cfg_coef\": 3.5, \"ctx_cfg_coef\": 1.0, \"sem_cfg_coef\": 1.0, \"steps\": 32, \"tags\": \"mono, dull, soft\"}, # prod\n",
    "    {\"text_cfg_coef\": 3.0, \"ctx_cfg_coef\": 1.0, \"sem_cfg_coef\": 1.0, \"steps\":16, \"tags\": tags}\n",
    "]\n",
    "\n",
    "#semantic_crop = semantic_codes.clone()\n",
    "#semantic_crop = semantic_crop[:750]\n",
    "#semantic_crop[375:] = 4000 # pad after 15s\n",
    "#semantic_crop[125:] = 4000 # pad after 5s\n",
    "#semantic_crop[:] = 4000\n",
    "\n",
    "semantic_pad = torch.ones(750).long()\n",
    "semantic_pad[:] = 4000#\n",
    "\n",
    "# construct prev_vae_ctx_vector\n",
    "if False:\n",
    "    vae_latents = torch.from_numpy(vae_data[\"vae_latents\"]).unsqueeze(0) * 0.4\n",
    "    #vae_latents = z.permute(0, 2, 1) * 0.4\n",
    "\n",
    "    infill_ctx_vector = vae_latents[:, 0:750, :] # next 30s of audio\n",
    "    infill_ctx_vector = infill_ctx_vector.cuda()\n",
    "    #infill_ctx_vector = torch.randn_like(z[:, :, :750])\n",
    "    #infill_ctx_vector = infill_ctx_vector.permute(0, 2, 1)\n",
    "    print(\"infill_ctx_vector.shape\", infill_ctx_vector.shape)\n",
    "    #noise = torch.randn_like(infill_ctx_vector) * 0.1\n",
    "    #infill_ctx_vector = infill_ctx_vector + noise\n",
    "\n",
    "    # zeros (False) means we will not use the infill ctx here\n",
    "    infill_ctx_mask = torch.ones(1, 750).bool().cuda()\n",
    "    infill_ctx_mask[:, 100:500] = False\n",
    "    #infill_ctx_mask[:,:] = True\n",
    "    print(\"infill_ctx_mask.shape\", infill_ctx_mask.shape)\n",
    "\n",
    "    # apply mask to infill_ctx_vector\n",
    "    infill_ctx_vector = infill_ctx_vector * infill_ctx_mask.unsqueeze(2)\n",
    "\n",
    "    semantic_codes_cropped = semantic_codes.clone()[0:750]\n",
    "    semantic_codes_cropped[100:500] = 4000\n",
    "    #semantic_codes_encoded = semantic_codes_encoded[:750]\n",
    "    #semantic_codes_encoded[100:500] = 4000\n",
    "    #semantic_codes_encoded = torch.from_numpy(semantic_codes_encoded).cuda()\n",
    "\n",
    "    #print(infill_ctx_vector.shape)\n",
    "    #print(infill_ctx_mask.shape)\n",
    "    # creat time indices in seconds\n",
    "    time_indices = torch.arange(0, 750, 1) / 25\n",
    "    import matplotlib.pyplot as plt\n",
    "    fig, axs = plt.subplots(figsize=(5, 2))\n",
    "    plt.plot(time_indices.cpu().numpy(), infill_ctx_mask.squeeze(0).cpu().numpy())\n",
    "    plt.show()\n",
    "else:\n",
    "    infill_ctx_vector = None\n",
    "    infill_ctx_mask = None\n",
    "\n",
    "semantic_noise = torch.randn(750, 768).cuda() * 0.2\n",
    "\n",
    "pred_audios = []\n",
    "pred_semantics =[]\n",
    "\n",
    "# now, take the semantic_codes and padd everywhere EXCEPT the downbeats\n",
    "#padded_semantic_codes = semantic_codes.clone()\n",
    "# downbeats are in the array downbeats\n",
    "# Create a mask where all positions are padded (set to 4000)\n",
    "#padded_semantic_codes[:] = 4000\n",
    "\n",
    "\n",
    "for param_config in param_configs:\n",
    "    print(param_config)\n",
    "    for seed in seeds:\n",
    "        print(seed)\n",
    "        pred_audio = generate(\n",
    "            semantic_codes,\n",
    "            lyrics=lyrics,\n",
    "            tags=\"poerty, spoken word\", \n",
    "            text_cfg_coef=param_config[\"text_cfg_coef\"],\n",
    "            ctx_cfg_coef=param_config[\"ctx_cfg_coef\"],\n",
    "            sem_cfg_coef=param_config[\"sem_cfg_coef\"],\n",
    "            steps=param_config[\"steps\"],\n",
    "            seed=seed,\n",
    "            sigma_min=0.5,\n",
    "            sigma_max=80.0,\n",
    "            rho=1.0,\n",
    "            semantic_skip_factor=2,\n",
    "            downscale_ctx_vector=False,\n",
    "            noise_ctx_vector=0.75,\n",
    "            noise_ctx_vector_pad_size=0,\n",
    "            semantic_noise=None,\n",
    "            #ctx_vector_scale=1.00,\n",
    "            #codec_cycle_ctx_vector=False,\n",
    "            #history_ctx_vae=history_ctx_vector,\n",
    "            #infill_ctx_vae=infill_ctx_vector,\n",
    "            #infill_ctx_mask=infill_ctx_mask,\n",
    "            #infill_variation=True, # use the entire from the diffusion\n",
    "            #chunk_size=750,\n",
    "            sampler_eta=1.0,\n",
    "            sampler_s_noise=1.0,\n",
    "\n",
    "            # for infill upsample\n",
    "            #ctx_len_used=250, # how much to pass to infill_ctx\n",
    "            #chunk_size=750-250, # how much to predict in new window\n",
    "            \n",
    "            #use_ctx_vector=True,\n",
    "            #ctx_vector_clip=None,\n",
    "            # oracle_first_chunk=False,|\n",
    "            #se_codec=False\n",
    "            use_ctx_vector=True,\n",
    "            distilled=False,\n",
    "            perturbed_semantic_skip_factor=1,\n",
    "        )\n",
    "\n",
    "        pred_audio.play()\n",
    "        #pred_semantics.append(pred_audio)\n",
    "        #torchaudio.save(f\"audio/noise_ctx2/metal_{noise_ctx}.mp3\", torch.tensor(pred_audio.array_float), 48000)\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd27b21f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(lyrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5edf90d",
   "metadata": {},
   "outputs": [],
   "source": [
    "models[\"info\"][\"weights_precision\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "08bbb036",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save pred semantics to file npz\n",
    "np.savez(\"/home/christian/code/christian/notebooks/semantics/semantic_preds.npz\", *pred_semantics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cb85d57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "semantic_pred_0 = np.load(\"/home/christian/code/christian/notebooks/semantics/semantic_pred_1.npy\")\n",
    "print(semantic_pred_0.shape)\n",
    "semantic_pred_0 = torch.from_numpy(semantic_pred_0).long()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63764c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.mert_25 import (\n",
    "    preload_models as preload_semantic_models,\n",
    "    load_model as load_semantic_model,\n",
    "    encode as encode_semantic,\n",
    "    EMBEDDING_RATE as SEMANTIC_HZ,\n",
    ")\n",
    "import numpy as np\n",
    "\n",
    "semantic_model = load_semantic_model(\n",
    "    semantic_model_filepath,\n",
    "    semantic_clusters_filepath,\n",
    "    device=\"cuda\"\n",
    ")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "898bdab0",
   "metadata": {},
   "outputs": [],
   "source": [
    "semantic_pred = torch.from_numpy(semantic_preds[\"arr_2\"]).unsqueeze(0).to(\"cuda\")\n",
    "semantic_codes = semantic_model[\"cluster_model\"].encode(semantic_pred)\n",
    "\n",
    "semantic_codes = semantic_codes.squeeze(0).cpu()[:,0]\n",
    "print(semantic_codes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9179b59",
   "metadata": {},
   "outputs": [],
   "source": [
    "# histogram of vae pad embed\n",
    "vae_pad_embed = models[\"dit_model\"].vae_pad_embed.detach().cpu().float().numpy()\n",
    "plt.hist(vae_pad_embed, bins=10, alpha=0.5)\n",
    "plt.show()\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0e7df4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sr = 48000\n",
    "n_fft = int(1 * sr)\n",
    "hop_size = n_fft // 2\n",
    "\n",
    "print(audio.array_float.shape)\n",
    "print(pred_audio.array_float.shape)\n",
    "\n",
    "input_centroid = compute_spectral_centroid(torch.tensor(audio.array_float).mean(dim=0, keepdim=True), sr, hop_size=hop_size, n_fft=n_fft)\n",
    "\n",
    "start_frame = 0\n",
    "end_frame = input_centroid.shape[0]\n",
    "\n",
    "frame_indices = torch.arange(start_frame, end_frame, 1)\n",
    "# convert frame indices to seconds\n",
    "frame_indices = frame_indices * (hop_size / sr)\n",
    "\n",
    "diff_centroids = []\n",
    "for idx, pred_audio_array in enumerate(pred_audios):\n",
    "    pred_centroid = compute_spectral_centroid(torch.tensor(pred_audio_array).mean(dim=0, keepdim=True), sr, hop_size=hop_size, n_fft=n_fft)\n",
    "    plt.plot(frame_indices, pred_centroid[start_frame:end_frame].squeeze(0).cpu().numpy(), label=f\"diff {idx}\")\n",
    "\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "#plt.plot(frame_indices, input_centroid[start_frame:end_frame].squeeze(0).cpu().numpy(), label=\"input\")\n",
    "#plt.plot(frame_indices, pred_centroid[start_frame:end_frame].squeeze(0).cpu().numpy(), label=\"pred\")\n",
    "\n",
    "plt.legend(loc=\"upper right\")\n",
    "\n",
    "ticks = np.arange(0, frame_indices[-1], 30)\n",
    "plt.grid(c=\"lightgray\")\n",
    "plt.xticks(ticks, ticks)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d9c4e0cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_spectral_centroid(audio_tensor, sample_rate, n_fft=4096, hop_size=2048):\n",
    "    # Split into frames\n",
    "    frame_length = n_fft\n",
    "    hop_length = hop_size\n",
    "    frames = audio_tensor.unfold(1, frame_length, hop_length)\n",
    "    print(frames.shape)\n",
    "    # Compute FFT for each frame\n",
    "    spectrum = torch.fft.rfft(frames.squeeze(0))  # [num_frames, n_fft//2 + 1]\n",
    "    freqs = torch.fft.rfftfreq(n_fft, d=1/sample_rate)  # [n_fft//2 + 1]\n",
    "    \n",
    "    # Compute magnitudes for each frame\n",
    "    magnitudes =torch.abs(spectrum)  # [num_frames, n_fft//2 + 1]\n",
    "    \n",
    "    # Compute centroid for each frame\n",
    "    numerator = torch.sum(freqs.view(1, -1) * magnitudes, dim=1)  # Sum over frequencies for each frame\n",
    "    denominator = torch.sum(magnitudes, dim=1)\n",
    "\n",
    "\n",
    "    # Compute mean centroid across all frames\n",
    "    centroid = numerator / (denominator + 1e-8)\n",
    "\n",
    "    return centroid\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "177ba920",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "audio = Audio.from_file(audio_filepath, n_channels=2).get_slice(00.0, 60.01).convert(48000, 2, 2)\n",
    "\n",
    "audio_tensor = torch.from_numpy(audio.array_float).unsqueeze(0).float().to(\"cuda\")\n",
    "with torch.no_grad():\n",
    "    z = models[\"codec_model\"].encode(audio_tensor)[\"z\"].half().float()\n",
    "    print(z.shape)\n",
    "    pred_audio = models[\"codec_model\"].decode(z)\n",
    "    IPython.display.display(IPython.display.Audio(pred_audio.squeeze(0).cpu().numpy(), rate=48000))\n",
    "\n",
    "\n",
    "lyrics = \"\"\"\n",
    "Now here you go again\n",
    "You say you want your freedom\n",
    "Well, who am I to keep you down?\n",
    "It's only right that you should\n",
    "Play the way you feel it\n",
    "But listen carefully\n",
    "To the sound of your loneliness\n",
    "Like a heartbeat drives you mad\n",
    "In the stillness of remembering what you had\n",
    "And what you lost\n",
    "And what you had\n",
    "And what you lost\n",
    "Oh, thunder only happens when it's rainin'\n",
    "Players only love you when they're playin'\n",
    "Say women, they will come and they will go\n",
    "When the rain washes you clean, you'll know\n",
    "You'll know\n",
    "Now here I go again\n",
    "I see the crystal visions\n",
    "I keep my visions to myself\n",
    "It's only me who wants to wrap around your dreams\n",
    "And have you any dreams you'd like to sell?\n",
    "Dreams of loneliness\n",
    "Like a heartbeat drives you mad\n",
    "In the stillness of remembering what you had\n",
    "And what you lost\n",
    "And what you had\n",
    "Ooh, what you lost\n",
    "Thunder only happens when it's rainin'\n",
    "Players only love you when they're playin'\n",
    "Women, they will come and they will go\n",
    "When the rain washes you clean, you'll know\n",
    "Oh, thunder only happens when it's rainin'\n",
    "Players only love you when they're playin'\n",
    "Say women, they will come and they will go\n",
    "When the rain washes you clean, you'll know\n",
    "You'll know\n",
    "You will know\n",
    "Oh, you'll know\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "205d77fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# try to cycle pred audio again\n",
    "with torch.no_grad():\n",
    "    z = model_25hz.encode(pred_audio.to(\"cuda\"))[\"z\"].half().float()\n",
    "    print(z.shape)\n",
    "    pred_audio_2 = model_25hz.decode(z)\n",
    "    IPython.display.display(IPython.display.Audio(pred_audio_2.squeeze(0).cpu().numpy(), rate=48000))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "656adb9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dac.model.dac4 import DAC\n",
    "import funcy\n",
    "\n",
    "device = \"cuda\"\n",
    "checkpoint_filepath = \"s3://suno-data/christian/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",
    "model_25hz = DAC(**sd[\"metadata\"][\"kwargs\"])\n",
    "model_25hz.load_state_dict(sd[\"state_dict\"])\n",
    "model_25hz.eval()\n",
    "model_25hz.to(device)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5455d18",
   "metadata": {},
   "outputs": [],
   "source": [
    "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "    #mp3_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "    #mp3_filepath = \"/home/christian/audio/50_genre_songs/BTS (방탄소년단) 'Dynamite' Official MV.mp3\"\n",
    "\n",
    "if \"s3://\" in mp3_filepath:\n",
    "    audio = Audio.from_s3(mp3_filepath, n_channels=2).get_slice(00.0, 60.01).convert(48000, 2, 2)\n",
    "else:\n",
    "    audio = Audio.from_file(mp3_filepath, n_channels=2).get_slice(00.0, 60.01).convert(48000, 2, 2)\n",
    "audio_tensor = torch.from_numpy(audio.array_float).unsqueeze(0).float().to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    z = model_25hz.encode(audio_tensor)[\"z\"].half().float()\n",
    "    print(z.shape)\n",
    "    pred_audio = model_25hz.decode(z)\n",
    "    IPython.display.display(IPython.display.Audio(pred_audio.squeeze(0).cpu().numpy(), rate=48000))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a7641d74",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def reconstruct_audio(\n",
    "    frames: torch.Tensor, original_length: int, overlap: float = 0.0\n",
    ") -> torch.Tensor:\n",
    "    \"\"\"\n",
    "    Reconstruct audio signal from frames. Uses overlap-add only when overlap > 0.\n",
    "\n",
    "    Args:\n",
    "        frames: Input tensor of shape (batch_size, channels, num_frames, frame_size)\n",
    "        original_length: Length of the original sequence\n",
    "        overlap: Overlap used in framing\n",
    "    \"\"\"\n",
    "    batch_size, channels, num_frames, frame_size = frames.shape\n",
    "    hop_size = int(frame_size * (1 - overlap))\n",
    "\n",
    "    # For no overlap, we can just reshape\n",
    "    if overlap == 0:\n",
    "        # Check if the frames can be directly reshaped\n",
    "        expected_length = num_frames * frame_size\n",
    "        if expected_length == original_length:\n",
    "            return frames.reshape(batch_size, channels, -1)\n",
    "        else:\n",
    "            # If not exact match, still do frame-by-frame to handle partial frames\n",
    "            output = torch.zeros(\n",
    "                batch_size,\n",
    "                channels,\n",
    "                original_length,\n",
    "                dtype=frames.dtype,\n",
    "                device=frames.device,\n",
    "            )\n",
    "            for i in range(num_frames):\n",
    "                start_idx = i * frame_size\n",
    "                end_idx = min(start_idx + frame_size, original_length)\n",
    "                output[:, :, start_idx:end_idx] = frames[\n",
    "                    :, :, i, : (end_idx - start_idx)\n",
    "                ]\n",
    "            return output\n",
    "\n",
    "    # For overlap > 0, use overlap-add (Hann window sum to 1)\n",
    "    output = torch.zeros(\n",
    "        batch_size, channels, original_length, dtype=frames.dtype, device=frames.device\n",
    "    )\n",
    "\n",
    "    for i in range(num_frames):\n",
    "        start_idx = i * hop_size\n",
    "        end_idx = start_idx + frame_size\n",
    "        output[:, :, start_idx:end_idx] += frames[:, :, i]\n",
    "\n",
    "    return output[:, :, :original_length]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09da74ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import torchaudio\n",
    "from suno_utils.audio import Audio\n",
    "import numpy as np\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "#output_dir = \"outputs/prefix-25hz-prod-dpo-s4538-step=5000-12052024\"\n",
    "#output_dir = \"outputs/prefix-25hz-prod-ft-s2056-step=16-12062024\"\n",
    "\n",
    "original_output_dir = \"outputs/original-03062025\"\n",
    "#mut_output_dir = \"outputs/prefix-25hz-dpo-s2269\"\n",
    "mut_output_dir = \"outputs/v4-25hz-s1848-dpo\"\n",
    "#mut_output_dir = \"outputs/v4-25hz-s5157-dpo-ctx-noise_0_5\"\n",
    "\n",
    "#mut_output_dir = \"outputs/v4-25hz-base\"\n",
    "#mut_output_dir = \"outputs/v4-25hz-s6971-prod\"\n",
    "\n",
    "#mut_output_dir = \"outputs/dit_v2_dpo_2025-01-26_02-13-54_s2704\"\n",
    "os.makedirs(original_output_dir, exist_ok=True)\n",
    "os.makedirs(mut_output_dir, exist_ok=True)\n",
    "\n",
    "if True:\n",
    "    gen_list = [\n",
    "        (\"bc095772-df08-48c0-80ac-89e02dc34e40\", \"\"),\n",
    "        (\"1a042042-b5d5-43e3-8b50-965d1ec0e303\", \"\"),\n",
    "        (\"1a3ac8fd-b5f8-4e78-92f5-400637e0851c\", \"\"),\n",
    "        (\"42985b5e-b8a9-4ca5-9aa2-4dd8df9b2f42\", \"\"),\n",
    "        (\"382b1682-8503-4b3d-94b6-5fa296933975\", \"\"),\n",
    "        (\"c04a5e04-4729-4d52-8bf4-8381635cb9a2\", \"\"),\n",
    "        (\"7c62f810-f77d-4a98-b179-d52c538a996f\", \"\"),\n",
    "        (\"b8057b18-8b02-41d4-927d-6adae708f3d4\", \"\"),\n",
    "        (\"ebb4ce1d-8f79-4a12-8aa8-858bcc65b44a\", \"\"),\n",
    "        (\"a5e2198a-f352-4abb-9a24-7f81b143ded3\", \"\"),\n",
    "        (\"f3b8bced-ff94-4e86-946b-28218cbc81ce\", \"\"),\n",
    "        (\"baa78a26-936b-48e9-9b28-1ae83fdb5075\", \"\"),\n",
    "    ]\n",
    "\n",
    "    gen_list2 = [\n",
    "        (\"cff54b12-5e94-4b0b-bc68-d5b6ec2f5957\", \"Ancient Mysterious Sanskrit Lama Mali Ninja Beats, Female Ancient Mystic Japanese Red-Skinned Ninja Dakini Queen, Produced and Mixed by the Female-Ancient-Ninja-Wisdom-Artist-Owl edited by Pythia\"),\n",
    "        (\"aa7d1f7c-854f-498c-ba5d-921120b9b559\", \"A raw, intense rock piece with a relentless beat, fragmented melody, emotional turmoil, complex drum-percussion interplay, and vocal fry adding an abrasive edge to the passionate vocals\"),\n",
    "        (\"f876800d-53b9-4862-8166-65ffe17d6c24\", \"afro beat, acid jazz, groove, latin house, 120bmp\"),\n",
    "        (\"f961cb32-4c7f-4c98-b391-c40dfda43669\", \"female vocalist pop pop rock rock adult contemporary passionate Spanish rock\"),\n",
    "        (\"64e89e7d-1134-47be-9ceb-6609e4536d73\", \"[GENRES: Traditional Japanese, Ritualistic, noh theatre, Edo Period Soundtrack] [STYLE: Minimalist, Dramatic, Sacred, Atmospheric]\"),\n",
    "        (\"3df01248-c250-4bb8-81e5-b1bcd215e53b\", \"Calypso Recipe\"),\n",
    "        (\"9082c473-6a53-44be-9d94-689042ea06c7\", \"trap, drill, rap, cloudrap\"),\n",
    "        (\"2da26679-4d97-4a35-be97-a8b63b1633aa\", \"high-energy, electric, electronic, techno, bass-driven\"),\n",
    "        (\"90ef738f-1b3e-4b82-a472-4be42f8b0979\", \"Powerful female vocals,Modern metalcore, Alternative Metal, Emotional, Melancholic, Melodic,dramatic tension,emotional depth,cinematic atmosphere,vulnerability,defiance,building to a sweeping climax\"),\n",
    "        (\"ee6ed078-ffe2-4363-820a-144f2b53332c\", \"Rock & Roll ballad, blues, 1960, 60's, soft instrumentation, melancholic, longing, male vocal, clean vocals\"),\n",
    "        (\"5bec9ed1-c252-4d29-aa40-b92a8f069be4\", \"Catchy intro melody, Pop, hip hop, K-POP\"),\n",
    "        (\"1a042042-b5d5-43e3-8b50-965d1ec0e303\", \"nostalgic acoustic folk, female vocals\"),\n",
    "        (\"05353e4a-4be3-42b9-bf7b-0f5e0ca2a0a2\", \"bluegrass, banjo\"),\n",
    "        (\"a69d5f40-62fc-4eca-9610-2216a9e20de6\", \"japanese jazz, vaporwave, sound art\"),\n",
    "        (\"d48055c6-7de3-4264-9e06-fb1f0914a3b3\", \"\"),\n",
    "        (\"9b9bab11-a399-4650-9a36-d7c912ef038f\", \"funky Africa tribal, new wave, Miami,Miami,Miami,Miami,Miami,Miami,Miami,Miami,\"),\n",
    "        (\"ebb4ce1d-8f79-4a12-8aa8-858bcc65b44a\", \"pop, female vocal\")\n",
    "    ]\n",
    "else:\n",
    "    gen_list = [\n",
    "        (\"18a5e173-cc9f-44d0-9c0b-af58b4a53fa3\", \"bossa nova, vocal harmonies, jazz voices, ambient, jazz chords, acoustic sounds, live sounds, live voices, full choir, multiple voices\"),\n",
    "        (\"e4f2ee1d-8a48-47a8-804c-5e0b7a331f30\", \"pop\"),\n",
    "        (\"9af31091-b6da-45a2-a567-5b0d70276048\", \"pop, female vocal\"),\n",
    "        (\"55f062c1-a6fc-45c4-bd4e-12d02d91ca43\", \"pop, male vocal\"),\n",
    "        (\"45451fbf-800e-41ab-8aec-b985df877840\", \"rap\"),\n",
    "        (\"e2bcfba9-967f-41c7-b2ab-c20fb3bff8bb\", \"piano\"),\n",
    "        (\"6ab41abb-af83-4188-a4db-4fe59b9f07d8\", \"djent\"),\n",
    "    ]\n",
    "\n",
    "# set numpy seed\n",
    "np.random.seed(42)\n",
    "\n",
    "pred_audios = []\n",
    "full_gen_list = gen_list2 + gen_list\n",
    "\n",
    "for seed, (gen_id, gen_tags) in enumerate(full_gen_list):\n",
    "    s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
    "    data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "\n",
    "    seed += 1000\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",
    "\n",
    "    text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n",
    "    aligned_lyrics = json.loads(text_data)\n",
    "    print(aligned_lyrics)\n",
    "    #audio.normalize_volume().play()\n",
    "\n",
    "\n",
    "    mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "    audio = Audio.from_s3(mp3_filepath, n_channels=2)#.get_slice(00.0, 60.01).convert(48000, 2, 2)\n",
    "\n",
    "    lyrics = \"\"\n",
    "    for elem in aligned_lyrics:\n",
    "        if \"word\" in elem:\n",
    "            lyrics += elem[\"word\"]\n",
    "\n",
    "    print(codes.shape)\n",
    "    semantic_codes = torch.from_numpy(codes[:, 0]).long()\n",
    "    #semantic_codes = semantic_codes[:1500]\n",
    "    print(semantic_codes.shape)    \n",
    "\n",
    "    cfg_scales = [2.0]\n",
    "    steps = 12\n",
    "    for cfg_scale in cfg_scales:\n",
    "        print(f\"Generating with cfg_scale={cfg_scale}\")\n",
    "        pred_audio = generate(\n",
    "            semantic_codes, \n",
    "            lyrics=lyrics,\n",
    "            #aligned_lyrics=aligned_lyrics,\n",
    "            tags=gen_tags, \n",
    "            text_cfg_coef=cfg_scale,\n",
    "            ctx_cfg_coef=1.0,\n",
    "            sem_cfg_coef=1.0,\n",
    "            downscale_ctx_vector=True,\n",
    "            noise_ctx_vector=0.0,\n",
    "            steps=steps,\n",
    "            seed=seed,\n",
    "            sigma_max=50.0,\n",
    "            #sampling_method=sampling_method\n",
    "        )\n",
    "        pred_audios.append(pred_audio.array_float)\n",
    "\n",
    "        # save the entire audio\n",
    "        #output_filepath = os.path.join(mut_output_dir, f\"{gen_id}_pred.mp3\")\n",
    "        #pred_audio.write_hq_mp3(output_filepath)\n",
    "\n",
    "        # save out audio in chunks of 10 seconds \n",
    "        # from beginning, middle and end\n",
    "        if True:\n",
    "            chunks = [\"beginning\", \"middle\", \"end\"]\n",
    "            duration_s = pred_audio.array_float.shape[-1] / 48000\n",
    "            for chunk in chunks:\n",
    "                if chunk == \"beginning\":\n",
    "                    start_s = 0\n",
    "                    end_s = 10\n",
    "                elif chunk == \"middle\":\n",
    "                    start_s = duration_s // 2\n",
    "                    end_s = start_s + 10\n",
    "                elif chunk == \"end\":\n",
    "                    start_s = duration_s - 10\n",
    "                    end_s = duration_s\n",
    "                \n",
    "                # convert s to samples\n",
    "                start_samp = int(start_s * 48000)\n",
    "                end_samp = int(end_s * 48000)   \n",
    "                audio_tensor = torch.tensor(pred_audio.array_float[:, start_samp:end_samp]).float()\n",
    "                input_audio_tensor = torch.tensor(audio.array_float[:, start_samp:end_samp]).float()\n",
    "                torchaudio.save(f\"{mut_output_dir}/{gen_id}_{chunk}.mp3\", audio_tensor, 48000)\n",
    "                torchaudio.save(f\"{original_output_dir}/{gen_id}_{chunk}.mp3\", input_audio_tensor, 48000)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943836bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save out audios \n",
    "os.makedirs(\"outputs/dpo-25012025\", exist_ok=True)\n",
    "for i, pred_audio in enumerate(pred_audios):\n",
    "    audio_tensor = torch.tensor(pred_audio).float()\n",
    "    print(audio_tensor.shape)\n",
    "    torchaudio.save(f\"outputs/dpo-25012025/{i}.mp3\", audio_tensor, 48000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "908498d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "sr = 48000\n",
    "hop_size = n_fft // 2\n",
    "n_fft = int(sr * 10)\n",
    "\n",
    "\n",
    "pred_centroid_list = []\n",
    "delta_centroid_list = []\n",
    "for pred_audio_array in pred_audios:\n",
    "    pred_centroid = compute_spectral_centroid(torch.tensor(pred_audio_array).mean(dim=0, keepdim=True), sr, hop_size=hop_size, n_fft=n_fft)\n",
    "    #pred_centroid = pred_centroid[:-1] # drop first frame\n",
    "\n",
    "    if pred_centroid.shape[0] < 30:\n",
    "        continue\n",
    "\n",
    "    pred_centroid_list.append(pred_centroid)\n",
    "    frame_indices = torch.arange(0, pred_centroid.shape[0], 1)\n",
    "    frame_seconds = frame_indices * (hop_size / sr)\n",
    "\n",
    "    print(pred_centroid.shape, frame_seconds.shape)\n",
    "\n",
    "    delta_centroid = (pred_centroid[1:] - pred_centroid[:-1]).mean()\n",
    "    print(delta_centroid)\n",
    "    delta_centroid_list.append(delta_centroid)\n",
    "\n",
    "    if torch.mean(pred_centroid) > 0:\n",
    "        color = \"red\"\n",
    "    else:\n",
    "        color = \"blue\"\n",
    "    plt.plot(frame_seconds, pred_centroid.squeeze(0).cpu().numpy(), linewidth=0.8, color=color, alpha=0.5)\n",
    "\n",
    "# plot the mean of the centroinds\n",
    "# crop to the smallest length\n",
    "min_length = min([centroid.shape[0] for centroid in pred_centroid_list])\n",
    "pred_centroid_list = [centroid[:min_length] for centroid in pred_centroid_list]\n",
    "pred_centroid_mean = torch.mean(torch.stack(pred_centroid_list), dim=0)\n",
    "frame_indices = torch.arange(0, pred_centroid_mean.shape[0], 1)\n",
    "frame_seconds = frame_indices * (hop_size / sr)\n",
    "\n",
    "plt.plot(frame_seconds, pred_centroid_mean.squeeze(0).cpu().numpy(), label=\"mean\", color=\"black\")\n",
    "\n",
    "#plt.plot(frame_indices, input_centroid[start_frame:end_frame].squeeze(0).cpu().numpy(), label=\"input\")\n",
    "#plt.plot(frame_indices, pred_centroid[start_frame:end_frame].squeeze(0).cpu().numpy(), label=\"pred\")\n",
    "\n",
    "#plt.legend(loc=\"upper right\")\n",
    "\n",
    "ticks = np.arange(0, frame_seconds[-1], 30)\n",
    "plt.grid(c=\"lightgray\")\n",
    "plt.xticks(ticks, ticks)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60a5b247",
   "metadata": {},
   "outputs": [],
   "source": [
    "# histogra of delta_centroid\n",
    "plt.hist(delta_centroid_list, bins=10)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "566058c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# diffusion cycled the audio (no lyrics)\n",
    "\n",
    "base_dir = \"/home/christian/audio/ear-bench\"\n",
    "output_dirname = \"diffusion_cycle_test\"\n",
    "output_dir = os.path.join(base_dir, output_dirname)\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "# take reference audio and codec cycle it \n",
    "audio_filepaths = glob.glob(os.path.join(\"/home/christian/audio/reference-audio-wav\", \"*.wav\"))\n",
    "\n",
    "for audio_filepath in tqdm(audio_filepaths):\n",
    "    filename = os.path.basename(audio_filepath).replace(\".wav\", \"\")\n",
    "    clean_filename = filename.replace(\" \", \"_\")\n",
    "    audio = Audio.from_file(audio_filepath, n_channels=2).get_slice(00.0, 120.01).convert(48000, 2, 2)\n",
    "    \n",
    "    pred_audio = generate(\n",
    "        audio, \n",
    "        lyrics=\"\",\n",
    "        tags=\"\", \n",
    "        steps=16,\n",
    "        text_cfg_coef=1.0,\n",
    "        ctx_cfg_coef=1.0,\n",
    "        downscale_ctx_vector=False,\n",
    "    )\n",
    "    output_filepath = os.path.join(output_dir, f\"{clean_filename}_cycled_pref=False.mp3\")\n",
    "    pred_audio.write_hq_mp3(output_filepath)\n",
    "    output_filepath = os.path.join(output_dir, f\"{clean_filename}_input_pref=True.mp3\")\n",
    "    audio.write_hq_mp3(output_filepath)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb494e8a",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9843759c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # align audio/lyrics\n",
    "# from suno_utils.tasks.hoot import get_word_timing_from_audio_and_lyrics, preload_models\n",
    "# _ = preload_models(\n",
    "#     checkpoint_filepath=\"/app/suno/models/hoot_v3.pt\",\n",
    "#     tokenizer_filepath=\"/app/suno/models/hoot_v3_tokenizer.model\",\n",
    "# )\n",
    "# aligned_lyrics = get_word_timing_from_audio_and_lyrics(audio, lyrics)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c946bc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import torch\n",
    "# models = _retrieve_models()\n",
    "# extra_args = torch.load(\"extra_args.pt\", weights_only=True)\n",
    "# torch.manual_seed(0)\n",
    "# x = torch.randn([1, 128, 100*30], device=\"cuda\")*50\n",
    "# with torch.no_grad():\n",
    "#     out = models[\"dit_model\"].forward_inference(\n",
    "#         x,\n",
    "#         x.new_ones([x.shape[0]])*50,\n",
    "#         cross_attn_cond=extra_args[\"cross_attn_cond\"],\n",
    "#         empty_cross_attn_cond=extra_args[\"empty_cross_attn_cond\"],\n",
    "#         cfg_scale=extra_args[\"cfg_scale\"]\n",
    "#     )\n",
    "# print(out.shape)\n",
    "# out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a93b889",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aa732de",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e5084b6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c4bd6b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acda5fe8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
