{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: gpt was finetuned with this code\n",
    "#  https://github.com/suno-ai/neon/pull/129/files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "sys.path.insert(0, \"/home/sara/neon/sunoDiff/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\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",
    "\n",
    "from generation import preload_models as preload_diff_models, generate\n",
    "from generation import encode_semantic\n",
    "\n",
    "\n",
    "gpt_model_fp = (\n",
    "    \"/app/suno/checkpoints/2025-06-04_05-27-45/last_ckpt_infer.pt\"  # bluejay infill-dpo\n",
    ")\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-05-21_22-39-00/last_ckpt_infer.pt\"  # bluejay dpo\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-05-31_07-05-35/last_ckpt_infer.pt\"  # bluejay dpo-light\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-05-19_13-46-02/last_ckpt_infer.pt\"  # bluejay sft\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-05-17_02-36-28/last_ckpt_infer.pt\"  # bluejay\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-02-04_21-04-31/last_ckpt_infer.pt\"  # base\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-04-12_03-10-00/last_ckpt_infer.pt\"  # sft\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-04-15_19-56-08/last_ckpt_infer.pt\"  # sft new\n",
    "# gpt_model_fp = \"/app/suno/checkpoints/2025-04-08_19-36-46/last_ckpt_infer.pt\"  # dpo\n",
    "\n",
    "\n",
    "diff_model_fp = \"/app/suno/tmp/diff_v2_dpo_apr12.pt\"  # v2\n",
    "# diff_model_fp = \"/app/suno/checkpoints/2025-03-28_18-46-17_s894/last_ckpt_infer.pt\"  # v2 (+infill, different code)\n",
    "\n",
    "\n",
    "# GPT\n",
    "N_BATCH = 2\n",
    "engine = Engine(\n",
    "    gpt_model_fp,\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",
    "# Diffusion\n",
    "downscale_ctx_vector = False\n",
    "extra_gen_params = {\n",
    "    \"noise_ctx_vector\": 0.75,\n",
    "    \"noise_ctx_vector_pad_size\": 0,\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/dac_vae_fixed_25hz_2.pth\",\n",
    "    codec_scale_factor=0.4,\n",
    "    dit_model_filepath=diff_model_fp,\n",
    "    model_type=\"prefix\",\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse 1]\n",
    "Another silo full, another dark dawn\n",
    "Bending the air, love is so small\n",
    "Visualize dark thoughts, it was in a dream\n",
    "Internet as a self-atomizing machine\n",
    "I designed it, reset, she's the cure\n",
    "Underground, search for more\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\n",
    "[Verse 2]\n",
    "Another silo full another dark dawn\n",
    "Bending the air, love is so full\n",
    "My desire to see the light, my desire to be pure\n",
    "Underground I visualize sighing ocean dream\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\"\"\"\n",
    "text_new = \"\"\"\n",
    "[Verse 1]\n",
    "Another silo full, another dark dawn\n",
    "Bending the air, love is so small\n",
    "Visualize dark thoughts, it was in a dream\n",
    "Internet as a self-atomizing machine\n",
    "I designed it, reset, she's the cure\n",
    "Underground, search for more\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, give that weasel some wine\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\n",
    "[Verse 2]\n",
    "Another silo full another dark dawn\n",
    "Bending the air, love is so full\n",
    "My desire to see the light, my desire to be pure\n",
    "Underground I visualize sighing ocean dream\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\"\"\"\n",
    "tags = \"slow, incredibly sad, face melting, epic, beautifully sung intimate anthem for the heat death of the universe\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# a = Audio.from_file(\"/home/georg/notebooks/samples/halo.wav\", sample_rate=48_000, n_channels=2)\n",
    "# a.play()\n",
    "# semantic_arr = encode_semantic(a)[:,:1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.utils.clip import SunoClip\n",
    "\n",
    "sc = SunoClip(\"37f34bc3-a6e9-40b9-8a51-635c268eddb4\")\n",
    "a = sc.audio()\n",
    "a.play()\n",
    "semantic_arr = sc.full_arr()[:, :1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# let's do an easy case where we re-do a verse\n",
    "# a_left_idx = 0\n",
    "# b_left_idx = 25*97\n",
    "# b_right_idx = 25*121\n",
    "# c_right_idx = len(semantic_arr)\n",
    "# history_arr = semantic_arr[a_left_idx:b_left_idx].copy()\n",
    "# future_arr = semantic_arr[b_right_idx:c_right_idx].copy()\n",
    "# infill_text = \"\"\"\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",
    "# \"\"\".strip()\n",
    "# sample_duration_s = int(round((c_right_idx - a_left_idx) / 25))\n",
    "b_left_idx = 25 * 60\n",
    "a_left_idx = b_left_idx - 350\n",
    "b_right_idx = int(25 * 100)\n",
    "c_right_idx = b_right_idx + 350\n",
    "history_arr = semantic_arr[a_left_idx:b_left_idx].copy()\n",
    "future_arr = semantic_arr[b_right_idx:c_right_idx].copy()\n",
    "infill_text = \"\"\"\n",
    "I designed it, reset, she's the cure\n",
    "Underground, search for more\n",
    "\n",
    "[Chorus]\n",
    "Returnal\n",
    "You've never left, give that weasel some wine\n",
    "Returnal\n",
    "You've never left, you've been here the whole time\n",
    "\n",
    "[Verse 2]\n",
    "Another silo full another dark dawn\n",
    "Bending the air, love is so small\n",
    "\"\"\".strip()\n",
    "# sample_duration_s = int(round((c_right_idx - a_left_idx) / 25))\n",
    "sample_duration_s = int(round((b_right_idx - b_left_idx) / 25))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.generation import CfgGenerationConfig\n",
    "from suno_utils.gpt.generation_prompt import ALL_AUDIO_PROMPTS\n",
    "\n",
    "use_alternte_tags = False\n",
    "\n",
    "extra_fields = {}\n",
    "if use_alternte_tags:\n",
    "    extra_fields[\"cfg_streams\"] = [\n",
    "        #         CfgGenerationConfig(\n",
    "        #             stream_type=\"tag\",\n",
    "        #             prompts=[\"tag\", \"lyrics\"] + [\"history\", \"future\"],\n",
    "        #             null_prompts=[\"lyrics\"] + [\"history\", \"future\"],\n",
    "        #             weight=2.5,\n",
    "        #             max_steps=None,\n",
    "        #         ),\n",
    "        CfgGenerationConfig(\n",
    "            stream_type=\"custom\",\n",
    "            prompts=[\"tag\", \"lyrics\"] + [\"future\", \"history\"],\n",
    "            null_prompts=[\"tag\", \"lyrics\"] + [\"history\"],\n",
    "            weight=0.5,\n",
    "            max_steps=None,\n",
    "        ),\n",
    "    ]\n",
    "#     extra_fields[\"custom_null_fields\"] = ALL_AUDIO_PROMPTS\n",
    "\n",
    "n_skip_semantic = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_gpt = []\n",
    "for nnn in range(3):\n",
    "    gconf = GenerationConfig(\n",
    "        #     text=text,\n",
    "        #         text=infill_text,\n",
    "        text=text_new,\n",
    "        text_tags=tags,\n",
    "        history_arr=history_arr,\n",
    "        future_arr=future_arr,\n",
    "        cfg_coef=1.2,\n",
    "        # todo: sample duration for skip is wrong cause bug in train\n",
    "        text_start_control_tags=\"{duration:\"\n",
    "        + str(int(round(sample_duration_s / n_skip_semantic)))\n",
    "        + \"}\",\n",
    "        #     min_eos_p=0.1,\n",
    "        cfg_coef_tags=0.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",
    "        n_skip_semantic=n_skip_semantic,\n",
    "        cfg_coef_neg_tags=0.0,\n",
    "        text_neg_tags=None,\n",
    "        temp_semantic=0.9,\n",
    "        top_k_semantic=None,\n",
    "        top_p_semantic=None,\n",
    "        min_p_semantic=0.005,\n",
    "        n_batch=1,\n",
    "        min_text_offset=0,\n",
    "        eos_pad_duration_s=0,\n",
    "        max_gen_duration_s=int(1 * 60),  # just incase\n",
    "        random_seed=nnn,\n",
    "        **extra_fields,\n",
    "    )\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",
    "    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",
    "        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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "# upsample with 10s pad on each side just incase\n",
    "# note: for simplicity let's upsample the entire clip rather than just the infill with context\n",
    "out_audio = []\n",
    "for in_sem_arr in out_gpt:\n",
    "    a = generate(\n",
    "        torch.cat(\n",
    "            [\n",
    "                torch.from_numpy(history_arr[-25 * 10 :, 0].astype(np.int32)),\n",
    "                in_sem_arr,\n",
    "                torch.from_numpy(future_arr[: 25 * 10, 0].astype(np.int32)),\n",
    "            ],\n",
    "            dim=0,\n",
    "        ),\n",
    "        #         lyrics=text,\n",
    "        lyrics=text_new,\n",
    "        tags=tags,\n",
    "        text_cfg_coef=2.0,\n",
    "        downscale_ctx_vector=downscale_ctx_vector,\n",
    "        steps=16,\n",
    "        seed=0,\n",
    "        **extra_gen_params,\n",
    "    )\n",
    "    out_audio.append(a)\n",
    "    a.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# I designed it, reset, she's the cure\n",
    "# Underground, search for more\n",
    "\n",
    "# [Chorus]\n",
    "# Returnal\n",
    "# You've never left, give that weasel some wine\n",
    "# Returnal\n",
    "# You've never left, you've been here the whole time\n",
    "\n",
    "# [Verse 2]\n",
    "# Another silo full another dark dawn\n",
    "# Bending the air, love is so full"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: give actual vae context in upsample\n",
    "#   see more info here https://github.com/suno-ai/glockenspiel/blob/main/suno_utils/notebooks/diffusion_engine/diffusion_engine.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize prompt\n",
    "import numpy as np\n",
    "from suno_utils.gpt.prompt import Prompt\n",
    "\n",
    "in_arr = requests[0].streams[0].prompt\n",
    "prompt = Prompt(\"\", engine.model.config)\n",
    "prompt.visualize(in_arr, compress=True)\n",
    "prompt.visualize(in_arr, compress=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Playground"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "from suno_utils.utils.text import read_jsonl\n",
    "\n",
    "metas_tr = read_jsonl(\"/app/suno/data/chirp_v5/v2/metas_tr.jsonl\")\n",
    "lll = [m for m in metas_tr if \"text_lines\" in m]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_tr = np.memmap(\"/app/suno/data/chirp_v5/v2/data_tr.bin\", dtype=np.uint16, mode=\"r\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = random.choice(lll)\n",
    "# m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "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
}