{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 3,
   "id": "88c3355d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13171812",
   "metadata": {},
   "outputs": [],
   "source": [
    "# no semantic finetine\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2024-10-30_15-14-54_s7904/last_ckpt_infer.pt\"\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",
    "codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\"\n",
    "\n",
    "\n",
    "if \"dac_vae_fixed\" in codec_filepath:\n",
    "    patch_size = 1\n",
    "    codec_scale_factor = 0.4\n",
    "else:\n",
    "    patch_size = 1\n",
    "    codec_scale_factor = 2.5\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": "code",
   "execution_count": null,
   "id": "98fb6f88",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import IPython\n",
    "\n",
    "empty_semantic_codes = torch.ones([1500], device=\"cuda\").long() * 4000\n",
    "\n",
    "yum_lyrics = \"\"\"\n",
    "[verse]\n",
    "A is for the amazing grace that we recieve\n",
    "B is for the blessings, every day we blieve\n",
    "C is for the chorus, we sing it loud and clear\n",
    "D is for the devotion that we hold dear\n",
    "\n",
    "[Chorus]\n",
    "B is for Buttocks, ripe and slightly damp\n",
    "Yum Yum, boy oh boy do I like God\n",
    "Every day is a gift, when you won a skateboard ramp.\n",
    "Yum yum, boy oh boy do I like God.\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "\n",
    "inst_lyrics = \"[Instrumental]\"\n",
    "max_sigma = 50.0\n",
    "\n",
    "seeds = [np.random.randint(0, 2**32-1) for _ in range(3)]\n",
    "#seeds = [0]\n",
    "\n",
    "param_configs = [\n",
    "    {\"text_cfg_coef\": 6.0, \"ctx_cfg_coef\": 1.0, \"sem_cfg_coef\": 1.0, \"steps\": 32, \"tags\": \"gospel\"}, # improved\n",
    "]\n",
    "\n",
    "for param_config in param_configs:\n",
    "    print(param_config)\n",
    "    for seed in seeds:\n",
    "        pred_audio = generate(\n",
    "            empty_semantic_codes, \n",
    "            lyrics=yum_lyrics,\n",
    "            tags=param_config[\"tags\"], \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_max=max_sigma,\n",
    "            scale_ctx_vector=True\n",
    "        )\n",
    "        pred_audio.play()\n",
    "\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fedf808d",
   "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
}