{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f130bb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: use branch gk_prefix from neon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a0d83737",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "88c3355d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "cd1bb050",
   "metadata": {},
   "outputs": [],
   "source": [
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-23_14-12-03_s7273/step_500000_infer.pt\"\n",
    "#dit_model_filepath = \"/app/suno/checkpoints/2024-10-30_20-11-31_s2639/last_ckpt_infer.pt\" # first chunk only\n",
    "dit_model_filepath = \"s3://suno-data/georg/tmp/2b_prefix_ft.pt\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cebbeee8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.audio import Audio\n",
    "from generation import preload_models, generate, _retrieve_models, codec_encode, codec_decode\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",
    "    codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\",  # 25hz\n",
    "    dit_model_filepath=dit_model_filepath,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\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": "115cb6b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "#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",
    "\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": "11e02695",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "\n",
    "#audio_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "audio_filepath = \"/home/christian/audio/50_genre_songs/BTS (방탄소년단) 'Dynamite' Official MV.mp3\"\n",
    "audio = Audio.from_file(audio_filepath, n_channels=2).get_slice(from_s=60, to_s=160.02)\n",
    "print(\"source\")\n",
    "print(audio.sample_rate)\n",
    "audio.normalize_volume().play()\n",
    "\n",
    "lyrics = \"\"\"\n",
    "'Cause I-I-I'm in the stars tonight\n",
    "So watch me bring the fire and set the night alight\n",
    "Shoes on, get up in the morn'\n",
    "Cup of milk, let's rock and roll\n",
    "King Kong, kick the drum, rolling on like a Rolling Stone\n",
    "Sing song when I'm walking home\n",
    "Jump up to the top, LeBron\n",
    "Ding dong, call me on my phone\n",
    "Ice tea and a game of ping pong, huh\n",
    "This is getting heavy\n",
    "Can you hear the bass boom? I'm ready (woo hoo)\n",
    "Life is sweet as honey\n",
    "Yeah, this beat cha-ching like money, huh\n",
    "Disco overload, I'm into that, I'm good to go\n",
    "I'm diamond, you know I glow up\n",
    "Hey, so let's go\n",
    "'Cause I-I-I'm in the stars tonight\n",
    "So watch me bring the fire and set the night alight (hey)\n",
    "Shining through the city with a little funk and soul\n",
    "So I'ma light it up like dynamite, whoa oh oh\n",
    "Bring a friend, join the crowd\n",
    "Whoever wanna come along\n",
    "Word up, talk the talk\n",
    "Just move like we off the wall\n",
    "Day or night, the sky's alight\n",
    "So we dance to the break of dawn (hey)\n",
    "Ladies and gentlemen, I got the medicine\n",
    "So you should keep ya eyes on the ball, huh\n",
    "This is getting heavy\n",
    "Can you hear the bass boom? I'm ready (woo hoo)\n",
    "Life is sweet as honey\n",
    "Yeah, this beat cha-ching like money\n",
    "Disco overload, I'm into that, I'm good to go\n",
    "I'm diamond, you know I glow up\n",
    "Let's go\n",
    "'Cause I-I-I'm in the stars tonight\n",
    "So watch me bring the fire and set the night alight (hey)\n",
    "Shining through the city with a little funk and soul\n",
    "So I'ma light it up like dynamite, whoa oh oh\n",
    "Dy-na-na-na, na-na, na-na-na, na-na-na, life is dynamite\n",
    "Dy-na-na-na, na-na, na-na-na, na-na-na, life is dynamite\n",
    "Shining through the city with a little funk and soul\n",
    "So I'ma light it up like dynamite, whoa oh oh\n",
    "Dy-na-na-na, na-na, na-na, ayy\n",
    "Dy-na-na-na, na-na, na-na, ayy\n",
    "Dy-na-na-na, na-na, na-na, ayy\n",
    "Light it up like dynamite\n",
    "Dy-na-na-na, na-na, na-na, ayy\n",
    "Dy-na-na-na, na-na, na-na, ayy\n",
    "Dy-na-na-na, na-na, na-na, ayy\n",
    "Light it up like dynamite\n",
    "'Cause I-I-I'm in the stars tonight\n",
    "So watch me bring the fire and set the night alight\n",
    "Shining through the city with a little funk and soul\n",
    "So I'ma light it up like dynamite (this is ah)\n",
    "'Cause I-I-I'm in the stars tonight\n",
    "So watch me bring the fire and set the night alight (alight, oh)\n",
    "Shining through the city with a little funk and soul\n",
    "So I'ma light it up like dynamite, whoa (light it up like dynamite)\n",
    "Dy-na-na-na, na-na, na-na-na, na-na-na, life is dynamite\n",
    "Dy-na-na-na, na-na, na-na-na, na-na-na, life is dynamite\n",
    "Shining through the city with a little funk and soul\n",
    "So I'ma light it up like dynamite, whoa oh oh\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "# extract vae latents to use as context oracle\n",
    "codec_codes = codec_encode(audio)\n",
    "print(codec_codes.shape)\n",
    "\n",
    "print(\"decoded\")\n",
    "audio_decoded = codec_decode(codec_codes)\n",
    "audio_decoded.normalize_volume().play()\n",
    "\n",
    "tags = \"rock\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b38cf75c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "\n",
    "#audio_filepath = \"/home/christian/audio/reference-audio-wav/03 Guess I_m Doing Fine.wav\"\n",
    "audio_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n",
    "\n",
    "audio = Audio.from_file(audio_filepath, n_channels=2).get_slice(from_s=0, to_s=60.02).convert(sample_rate=48000, n_channels=2, byte_width=2)\n",
    "print(audio.sample_rate)\n",
    "audio.normalize_volume().play()\n",
    "\n",
    "audio_decoded = codec_decode(codec_codes)\n",
    "audio_decoded.normalize_volume().play()\n",
    "\n",
    "lyrics = \"\"\"\n",
    "[Verse 1]\n",
    "There's a bluebird at my window\n",
    "I can't hear the songs he sings\n",
    "All the jewels in Heaven\n",
    "They don't look the same to me\n",
    "\n",
    "[Pre-Chorus]\n",
    "I just wade the tides that turned\n",
    "'Til I learn to leave the past behind\n",
    "\n",
    "[Chorus]\n",
    "It's only lies that I'm living\n",
    "It's only tears that I'm crying\n",
    "It's only you that I'm losing\n",
    "Guess I'm doing fine\n",
    "\n",
    "[Verse 2]\n",
    "All the battlements are empty\n",
    "And the moon is laying low\n",
    "Yellow roses in the graveyard\n",
    "Got no time to watch them grow\n",
    "\n",
    "[Pre-Chorus]\n",
    "Now I bade a friend farewell\n",
    "I can do whatever pleases me\n",
    "\n",
    "[Chorus]\n",
    "It's only lies that I'm living\n",
    "It's only tears that I'm crying\n",
    "It's only you that I'm losing\n",
    "Guess I'm doing fine\n",
    "\n",
    "[Bridge]\n",
    "Press my face up to the window\n",
    "To see how warm it is inside\n",
    "See the things that I've been missing\n",
    "Missing all this time\n",
    "\n",
    "[Chorus]\n",
    "It's only lies that I'm living\n",
    "It's only tears that I'm crying\n",
    "It's only you I'm losing\n",
    "Guess I'm doing fine\n",
    "\n",
    "[Outro]\n",
    "Guess I'm doing fine\n",
    "\"\"\"\n",
    "\n",
    "tags = \"rock\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2638a647",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = 32\n",
    "text_cfg_coef = 2.0\n",
    "\n",
    "import IPython\n",
    "\n",
    "audios = []\n",
    "for oracle_first_chunk in [True, False]:    \n",
    "    for n in range(1):\n",
    "        pred_audio = generate(\n",
    "        audio, \n",
    "        lyrics=lyrics,\n",
    "        tags=tags, \n",
    "            text_cfg_coef=text_cfg_coef,\n",
    "            ctx_cfg_coef=1.0,\n",
    "            oracle_first_chunk=oracle_first_chunk,\n",
    "            steps=steps,\n",
    "            seed=n,\n",
    "    ).normalize_volume().play()\n",
    "    #audios.append(pred_audio)\n",
    "        #pred_audio.play()\n",
    "        #IPython.display.display(IPython.display.Audio(pred_audio, rate=48000))\n",
    "    #     break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74289d4b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ec79526",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943a5c86",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91bc6c4b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "908498d4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acda5fe8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_env2",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}