{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0d83737",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "88c3355d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "import sys\n",
    "#sys.path.insert(0, \"/home/georg/code/neon/sunoDiff/\")\n",
    "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\")"
   ]
  },
  {
   "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=\"/home/georg/notebooks/gpu_nb/tmp/tokenizer_60k.json\",\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",
    "    weights_precision=torch.float16,\n",
    "    \n",
    "    ### 25hz ###\n",
    "#     codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\",\n",
    "    \n",
    "    # 30s\n",
    "#     dit_model_filepath=os.path.join(CKPT_DIR, \"tt96sbbx/checkpoints/last.ckpt\"),\n",
    "#     dit_config_filepath=\"/home/georg/notebooks/gpu_nb/tmp/1b_25hz_30s.json\",\n",
    "    # 2min\n",
    "#     dit_model_filepath=os.path.join(CKPT_DIR, \"wz8fc0h9/checkpoints/last.ckpt\"),\n",
    "#     dit_config_filepath=\"/home/georg/notebooks/gpu_nb/tmp/1b_2m.json\",\n",
    "\n",
    "    ### 100hz ###\n",
    "    codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/100hz_vae_peaq_kl_0.005.pth\",\n",
    "    \n",
    "    # 30s (overlap newest)\n",
    "    #dit_model_filepath=os.path.join(CKPT_DIR, \"zm0i9ccq/checkpoints/last.ckpt\"),\n",
    "    dit_model_filepath=os.path.join(CKPT_DIR, \"5jvuj882/checkpoints/last.ckpt\"), \n",
    "    #dit_model_filepath=os.path.join(CKPT_DIR, \"1ergguj2/checkpoints/last.ckpt\"),\n",
    "    #dit_model_filepath=\"/home/christian/code/christian/checkpoints/1ergguj2.ckpt\", # before blow up\n",
    "    #dit_model_filepath=\"/home/christian/code/christian/checkpoints/5jvuj882-step=430k.ckpt\", # before blow up\n",
    "\n",
    "    #dit_model_filepath=\"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_context_alignments_30s_step=160000.ckpt\",\n",
    "#     dit_model_filepath=\"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_context_30s_step=360000.ckpt\",\n",
    "    dit_config_filepath=\"/home/georg/notebooks/gpu_nb/tmp/1b_30s_ctx_2.json\",\n",
    "    \n",
    "    # 30s\n",
    "#     dit_model_filepath=os.path.join(CKPT_DIR, \"v24pfvep/checkpoints/last.ckpt\"),\n",
    "#     dit_config_filepath=\"/home/georg/notebooks/gpu_nb/tmp/1b_30s_old.json\",\n",
    "    \n",
    "    # 30s (overlap)\n",
    "#     dit_model_filepath=os.path.join(CKPT_DIR, \"8opy0my4/checkpoints/last.ckpt\"),\n",
    "#     dit_config_filepath=\"/home/georg/notebooks/gpu_nb/tmp/1b_30s_ctx.json\",\n",
    "   \n",
    "    # 2mins\n",
    "#     dit_model_filepath=os.path.join(CKPT_DIR, \"3fhnp5mz/checkpoints/last.ckpt\"),\n",
    "#     dit_config_filepath=\"/home/georg/notebooks/gpu_nb/tmp/1b_2m_100hz_2.json\",\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": "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 = \"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",
    "\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",
    "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",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6e8953",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "seeds = np.random.randint(0, 2**32 - 1, 3)\n",
    "print(seeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0c34cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = 32\n",
    "text_cfg_coef = 6.0\n",
    "ctx_cfg_coef = 2.5\n",
    "\n",
    "audios = []\n",
    "for n in seeds:\n",
    "    pred_audio = generate(\n",
    "        audio, \n",
    "#         lyrics=lyrics,\n",
    "        aligned_lyrics=aligned_lyrics,\n",
    "        tags=\"rock, pop, bass, 2020s\", \n",
    "        text_cfg_coef=text_cfg_coef, \n",
    "        ctx_cfg_coef=ctx_cfg_coef,\n",
    "        steps=steps,\n",
    "        seed=n,\n",
    "    ).normalize_volume()\n",
    "    audios.append(pred_audio)\n",
    "    pred_audio.play()\n",
    "#     break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1969b8bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_cfg_coef = 4.0\n",
    "ctx_cfg_coef = 1.5\n",
    "song_name = \"golden\"\n",
    "\n",
    "os.makedirs(\"outputs/step-test2\", exist_ok=True)\n",
    "\n",
    "for steps in [2, 4, 8, 16, 32, 64]:\n",
    "    for n in [1, 2, 3]:\n",
    "        pred_audio = generate(\n",
    "            audio, \n",
    "            aligned_lyrics=aligned_lyrics,\n",
    "            tags=\"pop rock, bass\", \n",
    "            text_cfg_coef=text_cfg_coef, \n",
    "            ctx_cfg_coef=ctx_cfg_coef,\n",
    "            steps=steps,\n",
    "            seed=n,\n",
    "        ).normalize_volume()\n",
    "        pred_audio.play()\n",
    "        pred_audio.write_mp3(f\"outputs/step-test2/{song_name}-steps={steps}-seed={n}.mp3\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa02303a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "audios[-1].write_mp3(\"test-64-steps.mp3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "441ef260",
   "metadata": {},
   "source": [
    "# Testbench"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2115a8e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from suno_utils.utils.s3 import read_from_s3\n",
    "import shutil\n",
    "output_dir = \"outputs/custom_infer_stable_audio-10102024\"\n",
    "if os.path.exists(output_dir):\n",
    "    shutil.rmtree(output_dir)\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "gen_ids = [\n",
    "    (\"99bffa17-7e59-47b4-a048-5528cbda05d5\", \"country\"), # sister\n",
    "    (\"081d73c4-7805-4212-9c80-8db1137ca3c4\", \"rock\"), # friends\n",
    "    (\"562f762d-6ced-4080-9af1-910ee3d0a5dc\", \"rock\"), # something real\n",
    "    (\"23c15c62-494d-422d-8a60-8b0454044322\", \"shanty\"), # rubber duck\n",
    "    (\"4b140a9e-964b-422c-85b5-5861ad1a9d38\", \"indie rock\"), # once\n",
    "    (\"7b214347-fa38-4e9b-96f4-f7ec65adea45\", \"rock n roll\"), # rock n roll\n",
    "    (\"7f5e7819-b7a1-449d-a76c-9c38950c4274\", \"pop\"), # canada\n",
    "    (\"27135940-957f-4152-9e40-72e5c10a46bd\", \"pop\"), # drunk\n",
    "    (\"caf1d7ae-ce81-4c75-b737-fd41e47a38ba\", \"pop\"), # linger\n",
    "    (\"ac5aa29a-f18f-406f-a904-c47493f1beec\", \"pop\"), # linger vox\n",
    "    (\"b04719c3-012a-4a4a-9b02-3cb2875d1744\", \"pop\"), # courtney\n",
    "    (\"7407c7b3-8aa2-4f0a-82e8-fa19933f6e94\", \"opera\"), # opera\n",
    "    (\"5b9624ef-b416-4095-a5f3-4945f09bba86\", \"math rock\"), # butterflies\n",
    "]\n",
    "\n",
    "for gen_id, tags in gen_ids:\n",
    "    mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
    "    print(tags, mp3_filepath)\n",
    "    audio = Audio.from_s3(mp3_filepath, n_channels=2)\n",
    "    audio.normalize_volume()\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",
    "\n",
    "    steps = 32\n",
    "    text_cfg_coef = 4.0\n",
    "    ctx_cfg_coef = 2.0\n",
    "    seeds = np.random.randint(0, 2**32 - 1, 3)\n",
    "\n",
    "    for n in seeds:\n",
    "        pred_audio = generate(\n",
    "            audio, \n",
    "            aligned_lyrics=aligned_lyrics,\n",
    "            tags=tags, \n",
    "            text_cfg_coef=text_cfg_coef, \n",
    "            ctx_cfg_coef=ctx_cfg_coef,\n",
    "            steps=steps,\n",
    "            seed=n,\n",
    "        ).normalize_volume()\n",
    "\n",
    "        # save audio to disk\n",
    "        pred_audio.write_mp3(f\"{output_dir}/pred_{gen_id}_{n}.mp3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb494e8a",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ce62b93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import json\n",
    "# with open(\"walking_aligned_lyrics.json\") as f:\n",
    "#     aligned_lyrics = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75675270",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # TODO: pad with silence for integer multiple of sliding\n",
    "# if models[\"dit_model\"].ctx_len is not None:\n",
    "#     window_s = int(round((models[\"dit_model\"].block_size - models[\"dit_model\"].ctx_len) / models[\"dit_model\"].io_hz))\n",
    "#     if audio.duration_s % window_s > 0.1:\n",
    "#         audio = Audio.concatenate([\n",
    "#             audio, \n",
    "#             Audio.from_silence(window_s - audio.duration_s % window_s + 0.01, sample_rate=audio.sample_rate, n_channels=audio.n_channels)\n",
    "#         ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d70822",
   "metadata": {},
   "outputs": [],
   "source": [
    "# # get latents for preconditioning\n",
    "# import torch\n",
    "# from suno_utils.tasks.dac_vae_100hz_peaq import encode as codec_encode\n",
    "# # latent_audio = audio\n",
    "# latent_audio = Audio.from_file(\"test.mp3\", sample_rate=48_000, n_channels=2).get_segment(from_s=offset_s, to_s=offset_s+duration_s)\n",
    "# init_latents = torch.from_numpy(codec_encode(latent_audio).T[None]) * 2.5\n",
    "# latent_audio.play()"
   ]
  },
  {
   "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": "818cf6a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# PAG\n",
    "#  https://arxiv.org/pdf/2403.17377\n",
    "#  https://huggingface.co/hyoungwoncho/sd_perturbed_attention_guidance/blob/main/pipeline.py\n",
    "#  https://github.com/KU-CVLAB/Perturbed-Attention-Guidance?tab=readme-ov-file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe70d19",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56142e0c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d56f346",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b12de9ec",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acda5fe8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}