{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\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",
    "from suno_utils.gpt.generation_prompt import ALL_AUDIO_PROMPTS\n",
    "from suno_utils.tasks.mert_25 import encode as encode_semantic\n",
    "from suno_utils.diffusion.generation import generate, DiffusionGenerationConfig\n",
    "from suno_utils.diffusion.generation import preload_models as preload_diff_models\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74d60c8a",
   "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\",\n",
    "    \"/app/suno/models/chirp_v2/tokenizer_60k.json\",\n",
    "    max_sequences=5 * 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",
    "    weights_precision=torch.bfloat16,\n",
    "    compile=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9fdb7a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_gpt(gconf):\n",
    "    requests = [\n",
    "        make_request(f\"{i}\", gconf, engine.model.config, engine.tokenizer) 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 = torch.zeros(arr.shape[0] * n_skip_semantic, dtype=arr.dtype) + cfg.semantic_pad_token\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",
    "            DiffusionGenerationConfig(\n",
    "                audio=in_sem_arr,\n",
    "                lyrics=text,\n",
    "                tags=tags,\n",
    "                text_cfg_coef=1.0,\n",
    "                steps=16,\n",
    "                seed=0,\n",
    "            )\n",
    "        )\n",
    "        out_diff.append(output)\n",
    "\n",
    "    return out_diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40e99b14",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Verse]\n",
    "By the river where currents flow\n",
    "In the moonlight my heartache shows\n",
    "Herbs I gather with trembling hands\n",
    "My tears mingle with the sand\n",
    "\n",
    "[Chorus]\n",
    "Oh why do you deceive me\n",
    "In this world so wild and free\n",
    "Your heart belongs to another\n",
    "While I wait by the river\n",
    "\n",
    "[Verse 3]\n",
    "Herbs of healing in my grasp\n",
    "Cannot mend this broken past\n",
    "Nature’s balm cannot restore\n",
    "Trust that walked out the door\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "tags = \"renaissance, ethereal, dorian folk, female voice, mandolin, dulcimer, danceable\"\n",
    "neg_tags = \"pop, rock, male\"\n",
    "\n",
    "n_skip_semantic = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b872ee",
   "metadata": {},
   "source": [
    "## Main Stream with Custom Null\n",
    "\n",
    "This example has one main stream and one null stream. Tag and Neg Tag CFG is turned off. By default, the main stream uses a null of just `[\"history\"]`, but this can be customized with the `custom_null_fields` parameter. The main stream is always active, and always has all available audio and text prompts in it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77176820",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    cfg_coef=1.2,  # by default the main stream has all available text and audio\n",
    "    cfg_coef_tags=0.0,  # off\n",
    "    cfg_coef_neg_tags=0.0,  # off\n",
    "    n_repeat_tags=3,\n",
    "    n_skip_semantic=n_skip_semantic,\n",
    "    text_start_control_tags=\"{min_duration:60}\",\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=int(2 * 60 / n_skip_semantic),\n",
    "    random_seed=0,\n",
    "    custom_null_fields=ALL_AUDIO_PROMPTS,  # define specific prompts to use in the main null stream\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c4772e",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_gpt = run_gpt(gconf)\n",
    "out_diff = run_diffusion(out_gpt, text=text, tags=tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0db8d7df",
   "metadata": {},
   "outputs": [],
   "source": [
    "for output in out_diff:\n",
    "    output.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43e5d68b",
   "metadata": {},
   "source": [
    "## Custom Tag Stream\n",
    "\n",
    "This example adds on a custom tag CFG stream. When a custom stream of type \"tag\" is defined, it takes precedence over cfg_coef_tags and cfg_coef_tags_max_steps. Within the custom tag stream we define the prompts we want in the positive stream and in the null stream. This null stream differs from the main null stream, so TWO null streams will be run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fff7b49",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    cfg_coef=1.2,  # by default the main stream has all available text and audio\n",
    "    n_repeat_tags=3,\n",
    "    n_skip_semantic=n_skip_semantic,\n",
    "    text_start_control_tags=\"{min_duration:60}\",\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=int(2 * 60 / n_skip_semantic),\n",
    "    random_seed=0,\n",
    "    cfg_streams=[\n",
    "        CfgGenerationConfig(\n",
    "            stream_type=\"tag\",\n",
    "            prompts=[\"tag\", \"lyrics\"] + ALL_AUDIO_PROMPTS,\n",
    "            null_prompts=[\"lyrics_no_tags\"] + ALL_AUDIO_PROMPTS,\n",
    "            weight=2.5,\n",
    "            max_steps=250,\n",
    "        )\n",
    "    ],\n",
    "    custom_null_fields=ALL_AUDIO_PROMPTS,  # define specific prompts to use in the main null stream\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0335bb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_gpt = run_gpt(gconf)\n",
    "out_diff = run_diffusion(out_gpt, text=text, tags=tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "158b333e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for output in out_diff:\n",
    "    output.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85080785",
   "metadata": {},
   "source": [
    "## Custom Tag and Neg Tag Streams\n",
    "\n",
    "This example adds a third negative tag stream. Like before, defining a custom stream of type \"neg_tag\" will take precedence over cfg_coef_neg_tag. Since the tag and neg_tag streams have matching null prompts, we can combine their null streams. This config will result in 5 total streams: main positive, tag positive, neg_tag positive, main null, tag&neg_tag null. Null streams are automatically stopped when their last associated positive stream ends (step 250 in this case)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aa30c2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    text_neg_tags=neg_tags,\n",
    "    cfg_coef=1.2,  # by default the main stream has all available text and audio\n",
    "    n_repeat_tags=3,\n",
    "    n_skip_semantic=n_skip_semantic,\n",
    "    text_start_control_tags=\"{min_duration:60}\",\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=int(2 * 60 / n_skip_semantic),\n",
    "    random_seed=0,\n",
    "    cfg_streams=[\n",
    "        CfgGenerationConfig(\n",
    "            stream_type=\"tag\",\n",
    "            prompts=[\"tag\", \"lyrics\"] + ALL_AUDIO_PROMPTS,\n",
    "            null_prompts=[\"lyrics_no_tags\"] + ALL_AUDIO_PROMPTS,\n",
    "            weight=2.5,\n",
    "            max_steps=250,\n",
    "        ),\n",
    "        CfgGenerationConfig(\n",
    "            stream_type=\"neg_tag\",\n",
    "            prompts=[\"neg_tag\", \"lyrics\"] + ALL_AUDIO_PROMPTS,\n",
    "            null_prompts=[\"lyrics_no_tags\"] + ALL_AUDIO_PROMPTS,\n",
    "            weight=-1.0,\n",
    "            max_steps=250,\n",
    "        ),\n",
    "    ],\n",
    "    custom_null_fields=ALL_AUDIO_PROMPTS,  # define specific prompts to use in the main null stream\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "041113fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_gpt = run_gpt(gconf)\n",
    "out_diff = run_diffusion(out_gpt, text=text, tags=tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "735e028e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for output in out_diff:\n",
    "    output.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bc03413",
   "metadata": {},
   "source": [
    "## Cover"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7be1ae8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"\"\"\n",
    "[Intro]\n",
    "Welcome to nowhere,\n",
    "\n",
    "[Verse 1]\n",
    "We met in the corner of the Sycamore\n",
    "dancing to the ones no one plays anymore.You sang the wrong word to \"London Calling,\"I joined in, and we both started laughing.\n",
    "\n",
    "You pulled out a mixtape, held it like gold,“Skip to Side B, that’s the best part, I’m told.\"I said, “It’s like life—rough in the middle.”You smiled, and my heart got caught in your riddle.\n",
    "\n",
    "[Chorus]\n",
    "welcome to nowhere, headlights on,Living like the punchline of a half-written song.We’re stuck on repeat, stuck on repeat, stuckBut somehow all the stars line up(We’re stuck on repeat, stuck on repeat, stuck)(But somehow all the stars line up)\n",
    "\n",
    "[Verse 2]\n",
    "We snuck onto rooftops, just to kill time,Counting satellite trails and graffiti signs.You said, “Do you think we’ll remember this night?”I said, “If we don’t, it’s probably right.”\n",
    "We played dumb games on a gameboy you found,The beeps and your silence both equally loud.But then you paused, looked up, and said,\"Isn’t it wild how this feels like the end?\"\n",
    "\n",
    "[Chorus]\n",
    "Welcome to nowhere, headlights on,Living like the punchline of a half-written song.We’re stuck on repeat, stuck on repeat, stuckBut somehow all the stars line up(We’re stuck on repeat, stuck on repeat, stuck)(But somehow all the stars line up)\n",
    "\n",
    "[Bridge]\n",
    "We burned the time we thought we’d keep,Let the ash settle like snow on the street.Who cares if it’s fleeting, who cares if it’s wrong?Not every chorus needs a verse to belong.\n",
    "\n",
    "[Chorus]\n",
    "welcome to nowhere, headlights on,Living like the punchline of a half-written song.We’re stuck on repeat, stuck on repeat, stuckBut somehow all the stars line up(We’re stuck on repeat, stuck on repeat, stuck)(But somehow all the stars line up)\n",
    "\n",
    "Stuck. We’re stuck on repeat\n",
    "Stuck. We’re stuck on repeat\n",
    "welcome to nowhere\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "tags = \"techno, dubstep, glitchy electronic, heavy bass, 808 drum machine\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d86e6c66",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_filepath = \"/home/sara/samples/welcome_to_nowhere.wav\"\n",
    "\n",
    "a = Audio.from_file(audio_filepath, sample_rate=44_100, n_channels=2)\n",
    "cover_arr = encode_semantic(a.normalize_volume())[:, :1]\n",
    "a.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5837b06b",
   "metadata": {},
   "outputs": [],
   "source": [
    "gconf = GenerationConfig(\n",
    "    text=text,\n",
    "    text_tags=tags,\n",
    "    cover_arr=cover_arr,\n",
    "    cfg_coef=1.2,  # by default the main stream has all available text and audio\n",
    "    n_repeat_tags=3,\n",
    "    n_skip_semantic=n_skip_semantic,\n",
    "    text_start_control_tags=\"{min_duration:60}\",\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=int(2 * 60 / n_skip_semantic),\n",
    "    random_seed=0,\n",
    "    cfg_streams=[\n",
    "        CfgGenerationConfig(\n",
    "            stream_type=\"tag\",\n",
    "            prompts=[\"tag\", \"lyrics\"] + [\"history\", \"future\"],\n",
    "            null_prompts=[\"lyrics_no_tags\"] + [\"history\", \"future\"],\n",
    "            weight=2.5,\n",
    "            max_steps=750,\n",
    "        )\n",
    "    ],\n",
    "    custom_null_fields=ALL_AUDIO_PROMPTS,  # define specific prompts to use in the main null stream\n",
    ")\n",
    "\n",
    "out_gpt = run_gpt(gconf)\n",
    "out_diff = run_diffusion(out_gpt, text=text, tags=tags)\n",
    "\n",
    "for output in out_diff:\n",
    "    output.play()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de695c8a",
   "metadata": {},
   "source": [
    "See the `cfg_over_time_infill.ipynb` example notebook for details on ramping CFG weights over time"
   ]
  }
 ],
 "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
}