{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.gpt.generation import GenerationConfig, CfgGenerationConfig\n",
    "from suno_utils.gpt.engine import Engine\n",
    "from suno_utils.gpt.generation_engine import make_request\n",
    "from suno_utils.gpt.generation_prompt import ALL_AUDIO_PROMPTS\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.tasks.demucs import split_vocals\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"/home/sara/neon/sunoDiff/\")\n",
    "from generation import preload_models as preload_diff_models, generate, encode_semantic\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_BATCH = 2\n",
    "n_skip_semantic = 1\n",
    "\n",
    "engine = Engine(\n",
    "    \"/app/suno/checkpoints/2025-01-28_12-59-39/last_ckpt_infer.pt\",  # \"/app/suno/checkpoints/2025-02-04_21-04-31/last_ckpt_infer.pt\", # ckpt before finetune\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=4 * N_BATCH,\n",
    "    compile=False,\n",
    ")\n",
    "cfg = engine.model.config\n",
    "\n",
    "_ = preload_diff_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\",\n",
    "    dit_model_filepath=\"/app/suno/data/dpo/models/diff_vae_25_peaq_v4_jan28.pt\",\n",
    "    model_type=\"prefix\",\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "codec_filepath = \"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\"  # prod\n",
    "codec_scale_factor = 2.5\n",
    "downscale_ctx_vector = True\n",
    "\n",
    "# use new vae\n",
    "if True:\n",
    "    diff_model_fp = \"/app/suno/checkpoints/2025-02-17_16-54-01_s7787/last_ckpt_infer.pt\"  # new\n",
    "    codec_filepath = \"/home/georg/notebooks/gpu_nb/tmp/dac_vae_fixed_25hz_2.pth\"\n",
    "    codec_scale_factor = 0.4\n",
    "    downscale_ctx_vector = False  \n",
    "\n",
    "from generation import preload_models as preload_diff_models, generate, encode_semantic\n",
    "_ = preload_diff_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=codec_filepath,\n",
    "    codec_scale_factor=codec_scale_factor,\n",
    "    dit_model_filepath=diff_model_fp,\n",
    "    model_type=\"prefix\",\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    ")\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_gpt(gconf):\n",
    "    requests = [\n",
    "        make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer)\n",
    "        for i in range(N_BATCH)\n",
    "    ]\n",
    "    jobs = engine.run_request(requests, tqdm_enabled=True)\n",
    "    out_gpt = []\n",
    "    for n, job in enumerate(jobs):\n",
    "        stream = engine.token_generator(job)\n",
    "        arr = torch.stack(list(stream))[:, 1]\n",
    "        if arr[-1] == 4000:\n",
    "            arr = arr[:-1]\n",
    "        print(f\"{round(arr.shape[-1]/25*n_skip_semantic)}s for track {n}\")\n",
    "        # do stuff incase skip\n",
    "        arr2 = (\n",
    "            torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype)\n",
    "            + cfg.semantic_pad_token\n",
    "        )\n",
    "        arr2[::n_skip_semantic] = arr\n",
    "        # add\n",
    "        out_gpt.append(arr2)\n",
    "\n",
    "    return out_gpt\n",
    "\n",
    "\n",
    "def run_diffusion(out_gpt, text, tags):\n",
    "    out_diff = []\n",
    "    for in_sem_arr in out_gpt:\n",
    "        output = generate(\n",
    "            in_sem_arr,\n",
    "            lyrics=text,\n",
    "            tags=tags,\n",
    "            text_cfg_coef=1.2,\n",
    "            steps=16,\n",
    "            seed=0,\n",
    "        )\n",
    "        out_diff.append(output)\n",
    "\n",
    "    return out_diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "mm = np.memmap(\"/app/suno/data/chirp_v5/v1/data_val.bin\", dtype=np.uint16, mode=\"r\")\n",
    "metas = read_jsonl(\"/app/suno/data/chirp_v5/v1/metas_val.jsonl\")\n",
    "overpaint_metas = [\n",
    "    m\n",
    "    for m in metas\n",
    "    if m[\"dataset\"] == \"musdb_stems_overpaint\" and m[\"type\"] == \"instrumental\"\n",
    "]\n",
    "underpaint_metas = [\n",
    "    m\n",
    "    for m in metas\n",
    "    if m[\"dataset\"] == \"musdb_stems_underpaint\" and m[\"type\"] == \"vocals\"\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "text1 = \"\"\"\n",
    "[Verse 1]\n",
    "Remember those walls I built?\n",
    "Well, baby, they're tumblin' down\n",
    "And they didn't even put up a fight\n",
    "They didn't even make a sound\n",
    "I found a way to let you in\n",
    "But I never really had a doubt\n",
    "Standin' in the light of your halo\n",
    "I got my angel now\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "It's the risk that I'm takin'\n",
    "I ain't ever gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\n",
    "[Verse 2]\n",
    "Hit me like a ray of sun\n",
    "Burnin' through my darkest night\n",
    "You're the only one that I want\n",
    "Think I'm addicted to your light\n",
    "I swore I'd never fall again\n",
    "But this don't even feel like fallin'\n",
    "Gravity can't begin\n",
    "To pull me back to the ground again\n",
    "\n",
    "[Pre-Chorus]\n",
    "It's like I've been awakened\n",
    "Every rule, I had you breakin'\n",
    "The risk that I'm takin'\n",
    "I'm never gonna shut you out\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\n",
    "[Bridge]\n",
    "Halo, ooh\n",
    "Halo, ooh\n",
    "Ooh\n",
    "\n",
    "[Chorus]\n",
    "Everywhere I'm lookin' now\n",
    "I'm surrounded by your embrace\n",
    "Baby, I can see your halo\n",
    "You know you're my savin' grace\n",
    "You're everything I need and more\n",
    "It's written all over your face\n",
    "Baby, I can feel your halo\n",
    "Pray it won't fade away\n",
    "\n",
    "[Post-Chorus]\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo\n",
    "I can feel your halo, halo, halo\n",
    "I can see your halo, halo, halo, ooh\n",
    "\"\"\"\n",
    "\n",
    "text2 = \"\"\"\n",
    "[Verse 1]\n",
    "Almost Heaven, West Virginia\n",
    "Blue Ridge Mountains, Shenandoah River\n",
    "Life is old there, older than the trees\n",
    "Younger than the mountains, growing like a breeze\n",
    "\n",
    "[Chorus]\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\n",
    "[Verse 2]\n",
    "All my memories gather 'round her\n",
    "Miner's lady, stranger to blue water\n",
    "Dark and dusty, painted on the sky\n",
    "Misty taste of moonshine, teardrop in my eye\n",
    "\n",
    "[Chorus]\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\n",
    "[Bridge]\n",
    "I hear her voice in the morning hour, she calls me\n",
    "The radio reminds me of my home far away\n",
    "Driving down the road, I get a feeling\n",
    "That I should have been home yesterday, yesterday\n",
    "\n",
    "[Chorus]\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "Country roads, take me home\n",
    "To the place I belong\n",
    "West Virginia, mountain mama\n",
    "Take me home, country roads\n",
    "\n",
    "[Outro]\n",
    "Take me home, (Down) country roads\n",
    "Take me home, (Down) country roads\n",
    "\"\"\"\n",
    "\n",
    "text3 = \"\"\"\n",
    "Well, you only need the light when it's burning low\n",
    "Only miss the sun when it starts to snow\n",
    "Only know you love her when you let her go\n",
    "Only know you've been high when you're feeling low\n",
    "Only hate the road when you're missing home\n",
    "Only know you love her when you let her go\n",
    "\n",
    "And you let her go\n",
    "\"\"\"\n",
    "\n",
    "text = text1\n",
    "tags = \"piano ballad, pop\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = Audio.from_file(\"/home/sara/samples/halo.wav\", sample_rate=44_100, n_channels=2)\n",
    "# # a = Audio.from_file(\"../samples/calm_down.mp3\", sample_rate=44_100, n_channels=2)\n",
    "# a = Audio.from_file(\"../samples/martin_vocals.m4a\", sample_rate=44_100, n_channels=2)\n",
    "\n",
    "a_vocals, a_other = split_vocals(a.convert(44_100, 2, 2))\n",
    "vocals_arr = encode_semantic(a_vocals.convert(44_100, 2, 2).normalize_volume())[:, :1]\n",
    "instrumental_arr = encode_semantic(a_other.convert(44_100, 2, 2).normalize_volume())[\n",
    "    :, :1\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    # underpaint_arr=vocals_arr,\n",
    "    overpaint_arr=instrumental_arr,\n",
    "    cfg_coef=1.05,\n",
    "    cfg_coef_tags=2.0,\n",
    "    cfg_coef_max_steps=None,\n",
    "    cfg_coef_tags_max_steps=None,\n",
    "    n_repeat_tags=1,\n",
    "    n_repeat_neg_tags=1,\n",
    "    text_start_control_tags=\"{min_duration:60}\",  # ;max_duration:360}\",\n",
    "    n_batch=1,\n",
    "    cfg_coef_neg_tags=0.0,  # -1.0\n",
    "    text_neg_tags=None,  # \"repetitive, loop\",\n",
    "    temp_semantic=0.92,\n",
    "    top_k_semantic=None,\n",
    "    top_p_semantic=None,\n",
    "    min_p_semantic=0.001,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=2 * 60,\n",
    "    random_seed=0,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_gpt = run_gpt(gconf)\n",
    "out_diff = run_diffusion(out_gpt, text=text, tags=tags)\n",
    "for output in out_diff:\n",
    "    output.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "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
}