{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from suno_utils.utils.clip import SunoClip\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",
    "import numpy as np\n",
    "from suno_utils.gpt.generation_prompt import ALL_AUDIO_PROMPTS\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\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_model_fp = \"/app2/suno/modal/models/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt\"\n",
    "\n",
    "\n",
    "_ = preload_diff_models(\n",
    "    tokenizer_filepath=\"/app/suno/data/dpo/models/tokenizer_60k.json\",\n",
    "    semantic_model_filepath=\"/app/suno/data/dpo/models/mert_25.pt\",\n",
    "    semantic_clusters_filepath=\"/app/suno/data/dpo/models/mert_25_2x4k.npy\",\n",
    "    codec_filepath=\"/app/suno/data/dpo/models/dac_vae_tuned_25hz.pth\",\n",
    "    dit_model_filepath=diff_model_fp,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=True,\n",
    "    codec_scale_factor=0.4,\n",
    ")\n",
    "\n",
    "N_BATCH = 2\n",
    "n_skip_semantic = 1\n",
    "\n",
    "engine = Engine(\n",
    "    # \"/app2/suno/modal/models/tony/sem/model_45_6b_apr15_sft_dpo_v10_11.pt\",\n",
    "    \"/app2/suno/checkpoints/2025-07-14_05-30-11/last_ckpt_infer.pt\",\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=5 * N_BATCH,\n",
    "    compile=False,\n",
    ")\n",
    "cfg = engine.model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "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",
    "            steps=16,\n",
    "            semantic_skip_factor=1,\n",
    "            downscale_ctx_vector=False,\n",
    "            noise_ctx_vector=0.5,\n",
    "            noise_ctx_vector_pad_size=25,\n",
    "        )\n",
    "        out_diff.append(output)\n",
    "\n",
    "    return out_diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc = SunoClip(\"cbea9ff8-01e3-42e1-90b4-12eb45f73e98\")\n",
    "\n",
    "sc.audio().play()\n",
    "\n",
    "semantic_arr = sc.full_arr()[:, :1]\n",
    "\n",
    "print(semantic_arr.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.audio().get_segment(from_s=35, to_s=50).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.audio().get_segment(from_s=21, to_s=64).play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "CONTEXT_TOKENS = 350\n",
    "b_left_idx = int(25 * 35)\n",
    "a_left_idx = b_left_idx - CONTEXT_TOKENS\n",
    "b_right_idx = int(25 * 50)\n",
    "c_right_idx = b_right_idx + CONTEXT_TOKENS\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",
    "tags = \"Folktronica, acoustic-pop, experimental, Heavy beats, synth textures, slick rhythm & seductive vocals, hip hop, pop, catchy, phonk, edm\"\n",
    "infill_text = \"\"\"\n",
    "let it combust in it\n",
    "It’s getting  strong don’t get  lost in it\n",
    "\n",
    "[Bridge]\n",
    "Got me caught up in the games you play\n",
    "Every word a trap gon make me stray \n",
    "Your touch is fire\n",
    "\"\"\"\n",
    "sample_duration_s = int(round((c_right_idx - a_left_idx) / 25))\n",
    "infill_duration_s = int(round((b_right_idx - b_left_idx) / 25))\n",
    "print(sample_duration_s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "context_text = \"\"\"\n",
    "I'm whole alone with or without you\n",
    "But this right here fills me with doubt too\n",
    "\n",
    "[Chorus]\n",
    "Trust trust put your trust in it\n",
    "Watch you move with the lust in it\n",
    "Hold it back don’t let it combust in it\n",
    "It’s getting  strong don’t get  lost in it\n",
    "\n",
    "[Bridge]\n",
    "Got me caught up in the games you play\n",
    "Every word a trap gon make me stray \n",
    "Your touch is fire but your heart is cold\n",
    "This love is new but feels so old\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_text = \"\"\"\n",
    "[Verse]\\nDo we let love live or let it slide\\nToo much pride and love don't hide\\n\n",
    "You ghost on me then you never call\\nRun back to them that's the worst of all\\n\\n[Verse 2]\n",
    "\\nI'm worth more how you make me stall\\nA love like ours ain't nothing small\\nI'm whole alone \n",
    "with or without you\\nBut this right here fills me with doubt too\\n\\n[Chorus]\\nTrust trust put your \n",
    "trust in it\\nWatch you move with the lust in it\\nHold it back don’t let it combust in it\\nIt’s getting  \n",
    "strong don’t get  lost in it\\n\\n[Bridge]\\nGot me caught up in the games you play\\nEvery word a trap gon \n",
    "make me stray \\nYour touch is fire but your heart is cold\\nThis love is new but feels so old\\n\\n[Verse 3]\\\n",
    "    nHow you leave me hanging in the dark\\nWhen you know I’m the one who sparks your heart\\nYou play it cool \n",
    "    but I see through\\nThe walls you build to hide the truth\\n\\n[Chorus]\\nTrust trust put your trust in it\\nWatching \n",
    "    you move your  lust in it\\nHold it back don’t let it combust in it\\nIt’s getting  strong don’t get  lost in it\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=full_text,\n",
    "    text_tags=tags,\n",
    "    text_neg_tags=\"repetitive, loop\",\n",
    "    cover_arr=semantic_arr,\n",
    "    cfg_coef=1.2,\n",
    "    cfg_coef_max_steps=200,\n",
    "    cfg_coef_tags=1.0,\n",
    "    cfg_coef_neg_tags=-1.0,\n",
    "    min_eos_p=0.1,\n",
    "    use_whisper=False,\n",
    "    min_text_offset=0,\n",
    "    max_tag_len=512,\n",
    "    n_repeat_tags=3,\n",
    "    temp_semantic=0.80,\n",
    "    top_k_semantic=1500,\n",
    "    top_p_semantic=None,\n",
    "    min_p_semantic=0.005,\n",
    "    eos_pad_duration_s=0,\n",
    "    cfg_coef_tags_max_steps=300,\n",
    "    max_gen_duration_s=int(8 * 60 / n_skip_semantic),\n",
    "    # custom_null_fields=ALL_AUDIO_PROMPTS,\n",
    "    n_batch=1,\n",
    "    # cfg_streams=[\n",
    "    #    {\n",
    "    #        \"stream_type\": \"custom\",\n",
    "    #        \"weight\": 2,\n",
    "    #        \"null_prompts\": [ \"tag\", \"lyrics\", \"future\", \"history\" ],\n",
    "    #        \"max_steps\": 200,\n",
    "    #        \"prompts\": [ \"tag\", \"lyrics\", \"vox\", \"playlist\", \"multi_artist\", \"artist\", \"cover\", \"overpaint\", \"underpaint\", \"future\", \"history\"]\n",
    "    #    }\n",
    "    # ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_gpt = run_gpt(gconf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "full_out = []\n",
    "for out in out_gpt:\n",
    "    full = torch.cat(\n",
    "        [\n",
    "            torch.from_numpy(history_arr[:, 0].astype(np.int32)),\n",
    "            out,\n",
    "            torch.from_numpy(future_arr[:, 0].astype(np.int32)),\n",
    "        ]\n",
    "    )\n",
    "    full_out.append(full)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_diff = run_diffusion(full_out, text=context_text, tags=tags)\n",
    "\n",
    "for output in out_diff:\n",
    "    output.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "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
}
