{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a0d83737",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !nvidia-smi\n",
    "# !echo $HOSTNAME"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "88c3355d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, \"/home/georg/code/neon/sunoDiff/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cd1bb050",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 30s\n",
    "dit_model_filepath = \"/app/suno/checkpoints/2024-10-06_05-06-29_s6612/last_ckpt_infer.pt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a12b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "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=\"/home/georg/notebooks/gpu_nb/tmp/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy\",\n",
    "    ## 25hz new train ##\n",
    "    weights_precision=torch.bfloat16,\n",
    "    codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\",\n",
    "    dit_model_filepath=dit_model_filepath,\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": "197bd22f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "offset_s = 0\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/walking_down.mp3\"\n",
    "# tags = \"Country\"\n",
    "# if duration_s <= 30.1:\n",
    "#     lyrics = \"Walking down the street\\nFeeling so alive\\nGot my head in the clouds\\nGot a gleam in my eye\\nEvery step that i take\\nIs like a brand new start\\nNo matter where I'm going\\nI'll always find my part\\n\\nOooo\\n\\nLife is like a high\"\n",
    "# else:\n",
    "#     lyrics = \"Walking down the street\\nFeeling so alive\\nGot my head in the clouds\\nGot a gleam in my eye\\nEvery step that i take\\nIs like a brand new start\\nNo matter where I'm going\\nI'll always find my part\\n\\nOooo\\n\\nLife is like a high wire act\\nwe're dancing in the sky\\nIt's a worry no need to ask why\\nWith a little bit or courage\\nWe can chase our dreams\\nNo matter what comes away\\nWe'll always be a team\\n\\nDon't stop oooo, yeah\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/md_mirror_baby.mp3\"\n",
    "# tags = \"reggae uplifting playful\"\n",
    "# lyrics = \"[Verse]\\nLittle Leo looks so fine\\nMirror baby by his side\\nThinks he’s seeing someone new\\nEvery time he smiles too\\n\\n[Chorus]\\nMirror baby best friend\\nReflections that don't end\\nLittle Leo's world so bright\\nEvery day a new delight\\n\\n[Verse 2]\\nChubby cheeks and laughing eyes\\nMirror baby never lies\\nEvery nod and every wink\\nLeo sees the bond they link\\n\\n[Chorus]\\nMirror baby best friend\\nReflections that don't end\\nLittle Leo's world so bright\\nEvery day a new delight\\n\\n[Bridge]\\nEvery giggle every glance\\nTiny feet that love to dance\\nMirror cheers him all day long\\nWith a smile and happy song\\n\\n[Verse 3]\\nCribs and rattles filled with glee\\nMirror baby laughs with me\\nLeo finds a friend so true\\nIn the mirror just like new\"\n",
    " \n",
    "# audio_fp = \"/home/georg/notebooks/samples/norah.wav\"\n",
    "# tags = \"jazz, pop, female vocal\"\n",
    "# lyrics = \"[Verse 1]\\nI waited 'til I saw the sun\\nI don't know why I didn't come\\nI left you by the house of fun\\nI don't know why I didn't come\\nI don't know why I didn't come\\n\\n[Verse 2]\\nWhen I saw the break of day\\nI wished that I could fly away\\nInstead of kneeling in the sand\\nCatching teardrops in my hand\\n\\n[Refrain]\\nMy heart is drenched in wine\\nBut you'll be on my mind forever\\n\\n[Verse 3]\\nOut across the endless sea\\nI will die in ecstasy\\nBut I'll be a bag of bones\\nDriving down the road alone\\n\\n[Refrain]\\nMy heart is drenched in wine\\nBut you'll be on my mind\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/turkey_time.mp3\"\n",
    "# tags = \"Country\"\n",
    "# lyrics = \"[Verse]\\nGather 'round now, it's turkey time\\nCatchin' up with all your kin and dimes\\nFeastin' on that golden bird, so fine (so fine)\\nPass the gravy, pour some wine\\n\\n[Chorus]\\nTurkey time, it's a country rhyme\\nGrab your plate, fill it up, take your time\\nAnd when the sun starts settin', we'll gather 'round\\nSingin' songs, sippin' cider, dancin' all night sound (oooh)\\nTurkey time, it's a country rhyme\\n\\n[Verse]\\nWell, it's that time of year, with family near\\nGatherin' 'round the table, filled with cheer\\nTurkey's roasting, gravy's flowin'\\nMashed potatoes piled high, it's Thanksiving time\\n\\n[Chorus]\\nTurkey time, it's a country rhyme\\nGrab your plate, fill it up, take your time\\nAnd when the sun starts settin', we'll gather 'round\\nSingin' songs, sippin' cider, dancin' all night sound (oooh)\\nTurkey time, it's a country rhyme\\n\\n[Bridge]  \\nTurkey time is here again\\nTurkey time is here again\\nTurkey time is here again\\nTurkey tiiiiime!\"\n",
    "\n",
    "# audio_fp = \"/home/georg/notebooks/samples/mistletoe.mp3\"\n",
    "# tags = \"Country\"\n",
    "# lyrics = \"[intro]\\nooooh\\n(what is life, what is love?)\\n[Verse]\\nAnother year has come and gone, and I'm not where I long to be\\nSnowfall on the city streets, but it ain't Christmas to me\"\n",
    "\n",
    "audio_fp = \"/home/georg/notebooks/samples/halo.wav\"\n",
    "tags = \"R&B, Ballad, Pop\"\n",
    "if duration_s <= 30.1:\n",
    "    offset_s = 20\n",
    "    lyrics = \"[Verse 1]\\nRemember those walls I built?\\nWell, baby, they're tumblin' down\\nAnd they didn't even put up a fight\\nThey didn't even make a sound\\nI found a way to let you in\\nBut I never really had a doubt\\nStandin' in the light of your halo\\nI got my angel now\\n\\n[Pre-Chorus]\\nIt's like\"\n",
    "else:\n",
    "    lyrics = \"[Verse 1]\\nRemember those walls I built?\\nWell, baby, they're tumblin' down\\nAnd they didn't even put up a fight\\nThey didn't even make a sound\\nI found a way to let you in\\nBut I never really had a doubt\\nStandin' in the light of your halo\\nI got my angel now\\n\\n[Pre-Chorus]\\nIt's like I've been awakened\\nEvery rule, I had you breakin'\\nIt's the risk that I'm takin'\\nI ain't ever gonna shut you out\\n\\n[Chorus]\\nEverywhere I'm lookin' now\\nI'm surrounded by your embrace\\nBaby, I can see your halo\\nYou know you're my savin' grace\\nYou're everything I need and more\\nIt's written all over your face\\nBaby, I can feel your halo\\nPray it won't fade away\\n\\n[Post-Chorus]\\nI can feel your halo, halo, halo\\nI can see your halo, halo, halo\\nI can feel your halo, halo, halo\\nI can see your halo, halo, halo, ooh\\n\\n[Verse 2]\\nHit me like a ray of sun\\nBurnin' through my darkest night\\nYou're the only one that I want\\nThink I'm addicted to your light\\nI swore I'd never fall again\\nBut this don't even feel like fallin'\\nGravity can't begin\\nTo pull me back to\"\n",
    "\n",
    "audio = Audio.from_file(audio_fp, n_channels=2).get_segment(from_s=offset_s, to_s=offset_s+duration_s+0.01)\n",
    "audio.normalize_volume().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c95f9978",
   "metadata": {},
   "outputs": [],
   "source": [
    "steps = 100\n",
    "cfg_coef = 4\n",
    "\n",
    "audios = []\n",
    "for n in range(3):\n",
    "    pred_audio = generate(\n",
    "        audio, \n",
    "        lyrics=lyrics,\n",
    "        tags=tags, \n",
    "        cfg_coef=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": "a62665b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# steps = 100\n",
    "# cfg_coef = 1\n",
    "\n",
    "# audios = []\n",
    "# for n in range(3):\n",
    "#     pred_audio = generate(\n",
    "#         audio, \n",
    "#         lyrics=\"\",\n",
    "#         tags=tags, \n",
    "#         cfg_coef=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": "b86eeb90",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c014544",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "908498d4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "bb494e8a",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9843759c",
   "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": 29,
   "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": "abe70d19",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56142e0c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c838e916",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3d28c6a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acda5fe8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}