{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0e09c52f-da2a-4a21-b659-b58383cba6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME\n",
    "\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"\n",
    "\n",
    "import sys\n",
    "#ys.path.insert(0, \"/home/tony/Work/neon/sunoDiff/\")\n",
    "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\")\n",
    "sys.path.insert(0, \"/home/m4burns/glockenspiel/suno_utils/notebooks/shimmerscore/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c046f7b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.audio import Audio\n",
    "from generation import preload_models, generate, _retrieve_models\n",
    "CKPT_DIR = \"/home/christian/code/neon/stable-audio-tools/harmonai_train/\"\n",
    "_ = preload_models(\n",
    "    tokenizer_filepath=\"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=\"/app/suno/data/dpo/models/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"/app/suno/data/dpo/models/mert_25_2x4k.npy\",\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    "    #codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/dac_vae_fixed_25hz_2.pth\",\n",
    "    codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",\n",
    "    codec_scale_factor=0.4,\n",
    "\n",
    "    # dit_model_filepath=\"/app/suno/data/dpo/models/diff_v2_2b_2mil_ft_v0.pt\", # reference diff v2\n",
    "    dit_model_filepath=\"/app/suno/checkpoints/2025-04-10_04-13-39_s12/step_9000_infer.pt\", # t6_cs (v2 dpo without filter)\n",
    "\n",
    "    #dit_model_filepath=\"/app/suno/checkpoints/2025-02-17_16-54-01_s7787/last_ckpt_infer.pt\", # cs\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": "code",
   "execution_count": 3,
   "id": "bac651fe-cfda-49f4-a4f2-b8878b89b128",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.dac_2c_12cb import preload_models as preload_codec_models\n",
    "from suno_utils.tasks.dac_2c_12cb import (\n",
    "    encode as dac_codec_encode,\n",
    "    decode_stream_to_full_audio as dac_codec_decode,\n",
    "    decode as dac_decode\n",
    ")\n",
    "_ = preload_codec_models(\"/app/suno/data/dpo/models/dac_2c_25x12.pt\", device=\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf696991",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "\n",
    "# gen_id = \"afc12c52-4cad-4f0c-b57c-1449637dcb15\" # The Machine 2.0\n",
    "# gen_tags = \"Industrial, dark EDM with fast beats, varied, deep basslines, dynamic drops, and haunting synths, Build tension with breaks, powerful bass drops, and an intense, energetic flow\"\n",
    "\n",
    "# gen_id = \"09704b3b-dd3e-4a31-bae0-18817bd5e12d\" # Therapy\n",
    "# gen_tags = \"nightmare on wax, robot rock, bassment jack, vegas in death, break beat, chemical dust, liquid swing, afro beat, remixed, sample heavy, old school dj turntablism, big breakbeat, future funk\"\n",
    "\n",
    "gen_id = \"1ec51ef1-378e-4f33-94a5-05c9b1e1448f\"\n",
    "gen_tags = \"Heavy, distorted electric guitars with chugging riffs, aggressive double-kick drumming, prominent bass providing low-end drive, electronic synth layers adding atmospheric texture, screamed and clean vocal interplay, dynamic shifts between intense verses and melodic choruses, breakdown sections featuring syncopation, layered vocal harmonies in chorus, use of vocal effects such as reverb and distortion, overall dark and cinematic production typical of modern metalcore; not a duet; no major charting or awards known\"\n",
    "\n",
    "# gen_id_audio = \"cb105e06-6e20-4acc-9d1c-e7c9b36b24e2\"\n",
    "gen_id_audio = gen_id\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_audio}.mp3\"\n",
    "print(s3_filepath)\n",
    "data = read_from_s3(s3_filepath, read_f=np.load)\n",
    "\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",
    "elif \"v5.0_raw\" in data:\n",
    "    codes = data[\"v5.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",
    "lyrics = \"\".join(w[\"word\"] for w in aligned_lyrics if \"word\" in w)\n",
    "audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6e8953",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from shimmerscore_new import shimmer_score\n",
    "from tempfile import NamedTemporaryFile\n",
    "\n",
    "score_means_first30 = []\n",
    "score_means_next30 = []\n",
    "\n",
    "for init_seed in [42, 43, 44, 45]:\n",
    "    np.random.seed(init_seed)\n",
    "    seeds = np.random.randint(0, 2**32 - 1, 2)\n",
    "    sem_codes = torch.tensor(codes[:, 0]).to(torch.long)\n",
    "\n",
    "    steps = 10\n",
    "    text_cfg_coef = 2.0\n",
    "    ctx_cfg_coef = 1.0\n",
    "\n",
    "    audios_first30 = []\n",
    "    audios_next30 = []\n",
    "    for n in seeds:\n",
    "        pred_audio = generate(\n",
    "            sem_codes[:25*60],\n",
    "            lyrics=lyrics,\n",
    "            tags=gen_tags,\n",
    "            text_cfg_coef=text_cfg_coef, \n",
    "            ctx_cfg_coef=ctx_cfg_coef,\n",
    "            steps=steps,\n",
    "            seed=n,\n",
    "            downscale_ctx_vector=False,\n",
    "        ).normalize_volume()\n",
    "        pred_audio.play()\n",
    "        audios_first30.append(pred_audio.get_segment(from_s=0, to_s=30))\n",
    "        audios_next30.append(pred_audio.get_segment(from_s=30))\n",
    "\n",
    "    scores_first30 = []\n",
    "    scores_next30 = []\n",
    "\n",
    "    for in_arr, out_arr, means_arr in [(audios_first30, scores_first30, score_means_first30), (audios_next30, scores_next30, score_means_next30)]:\n",
    "        for audio in in_arr:\n",
    "            with NamedTemporaryFile(suffix=\".wav\") as f:\n",
    "                audio.write_wav(f.name)\n",
    "                out_arr.append(shimmer_score(f.name))\n",
    "        means_arr.append(np.mean(out_arr))\n",
    "\n",
    "print(\"First 30s: \", [round(score, 2) for score in score_means_first30])\n",
    "print(\"Next 30s: \", [round(score, 2) for score in score_means_next30])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6999f34",
   "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
}
