{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"7\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c62b25",
   "metadata": {},
   "source": [
    "# BCT Inference\n",
    "\n",
    "This notebook is a hackable place to test BCT model inference."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23f330f5",
   "metadata": {},
   "source": [
    "## Load model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3dd2296",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch._C._cuda_getDeviceCount()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7d0311",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.generation import load_model, GPT\n",
    "from suno_utils.utils.s3 import download_s3_file_if_needed\n",
    "import torch\n",
    "\n",
    "N_BATCH = 2\n",
    "\n",
    "# normal\n",
    "gpt_model_path_diffusion = \"/app2/suno/checkpoints/2025-08-18_20-49-50/last_ckpt_infer.pt\"\n",
    "# routed\n",
    "gpt_model_path_diffusion = \"/app2/suno/checkpoints/2025-08-19_03-28-37/step_300000_infer.pt\"\n",
    "# 20s\n",
    "gpt_model_path_diffusion = \"/app2/suno/checkpoints/2025-08-21_02-09-32/last_ckpt_infer.pt\"\n",
    "# 60s fix crop\n",
    "gpt_model_path_diffusion = \"/app2/suno/checkpoints/2025-08-21_02-08-27/last_ckpt_infer.pt\"\n",
    "# 8 node\n",
    "gpt_model_path_diffusion = \"/app2/suno/checkpoints/2025-09-01_12-56-36/last_ckpt_infer.pt\"\n",
    "tokenizer_path = \"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\"\n",
    "\n",
    "\n",
    "model_container = load_model(\n",
    "    ckpt_path=download_s3_file_if_needed(gpt_model_path_diffusion),\n",
    "    tokenizer_path=download_s3_file_if_needed(tokenizer_path),\n",
    ")\n",
    "model: GPT = model_container[\"model\"]\n",
    "assert isinstance(model, GPT)\n",
    "\n",
    "\n",
    "cfg = model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3118c1b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.diffusion import generation as diffusion_gen\n",
    "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request\n",
    "from suno_utils.audio import Audio\n",
    "\n",
    "# dit_model_filepath = \"/app/suno/checkpoints/2025-02-17_16-54-01_s7787/last_ckpt.pt\"  # base\n",
    "# diffusion_gen.preload_models(\n",
    "#     dit_model_filepath=dit_model_filepath,\n",
    "# )\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import preload_models as preload_codec_models, decode\n",
    "\n",
    "preload_codec_models(\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\")\n",
    "\n",
    "# diffusion_engine = UpsampleEngine(compile=False)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a1d5952",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1e21a6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Block type name to ID mapping (must match GPTConfig.block_type_vocab_size)\n",
    "BLOCK_TYPE_NAME_TO_ID = {\n",
    "    \"text\": 0,\n",
    "    \"mt5_text\": 1,\n",
    "    \"hoot_text\": 2,\n",
    "    \"ditto\": 3,\n",
    "    \"semantic\": 4,\n",
    "    \"artist\": 5,\n",
    "    \"playlist\": 6,\n",
    "    \"underpaint\": 7,\n",
    "    \"overpaint\": 8,\n",
    "    \"vox\": 9,\n",
    "    \"stem\": 10,\n",
    "    \"sample\": 11,\n",
    "    \"cover\": 12,\n",
    "    \"prefix\": 13,\n",
    "    \"suffix\": 14,\n",
    "    \"non_causal_semantic\": 15,\n",
    "    \"diffusion\": 16,\n",
    "    \"vae_artist\": 17,\n",
    "    \"vae_playlist\": 18,\n",
    "    \"vae_underpaint\": 19,\n",
    "    \"vae_overpaint\": 20,\n",
    "    \"vae_vox\": 21,\n",
    "    \"vae_stem\": 22,\n",
    "    \"vae_sample\": 23,\n",
    "    \"vae_cover\": 24,\n",
    "    \"vae_prefix\": 25,\n",
    "    \"vae_suffix\": 26,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e7e428",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.gpt.bct.bct_generation_simple import BCTGenerationConfig, BlockSequence, generate_block\n",
    "from suno_utils.gpt.bct.bct import Block, BlockType\n",
    "from suno_utils.utils.clip import SunoClip\n",
    "\n",
    "TextBlockType = BlockType(\n",
    "    name=\"text\",\n",
    "    is_causal=True,\n",
    ")\n",
    "CausalSemanticBlockType = BlockType(\n",
    "    name=\"semantic\",\n",
    "    is_causal=True,\n",
    ")\n",
    "DiffusionBlockType = BlockType(\n",
    "    name=\"diffusion\",\n",
    "    is_causal=False,\n",
    ")\n",
    "\n",
    "lyrics = \"\"\"{activity:100%;activity>10%}\n",
    "[verse]\n",
    "oh, my love\n",
    "My friend you know\n",
    "it's been a while\n",
    "Without thinking of you\n",
    "but the thought makes me smile\n",
    "\n",
    "[chorus]\n",
    "I'm so tired of wanting\n",
    "wanting more than this\n",
    "i know it but what am i to do\n",
    "i need some space to breathe,\n",
    "so give me some room\n",
    "\n",
    "[verse]\n",
    "oh, my love\n",
    "you have a heart of stone\n",
    "cause since i've come home\n",
    "i've never felt so alone\n",
    "but the thought makes me smile\"\"\"\n",
    "lyrics_small = \"\"\"{activity:100%;activity>10%;min_duration:0}\n",
    "[verse]\n",
    "oh, my love\n",
    "My friend you know\n",
    "it's been a while\n",
    "but the thought makes me smile\"\"\"\n",
    "text = lyrics_small\n",
    "text_tokens = model_container[\"tokenizer\"].encode(\n",
    "    text,\n",
    "    add_special_tokens=False,\n",
    ")\n",
    "# print(text_tokens)\n",
    "\n",
    "text_block_type_ids = torch.full((1, 1, len(text_tokens)), BLOCK_TYPE_NAME_TO_ID[\"text\"])\n",
    "text_block = Block(\n",
    "    TextBlockType,\n",
    "    inputs={\n",
    "        \"text_input\": torch.tensor(text_tokens).reshape(1, 1, -1),\n",
    "        \"block_type_input\": text_block_type_ids,\n",
    "    },\n",
    ")\n",
    "gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"\n",
    "# gen_id = \"961c62b6-5469-4e4b-b654-1e3eaac9bd94\"\n",
    "clip = SunoClip(gen_id)\n",
    "# clip.audio().play()\n",
    "npz = clip._get_npz()\n",
    "print(list(npz.keys()))\n",
    "\n",
    "T = 20 * 25\n",
    "\n",
    "sem_codes = npz[\"v3.5_raw\"]\n",
    "# print(sem_codes.shape)\n",
    "sem_codes = sem_codes[:, 0]\n",
    "sem_codes = torch.from_numpy(sem_codes)[..., :T]\n",
    "# add pad token\n",
    "sem_codes = torch.cat([sem_codes, torch.full((1,), cfg.semantic_pad_token)])\n",
    "sem_codes = sem_codes.reshape(1, 1, -1)\n",
    "# print(sem_codes.shape)\n",
    "# print(sem_codes)\n",
    "\n",
    "sem_block_type_ids = torch.full((1, 1, T), BLOCK_TYPE_NAME_TO_ID[\"semantic\"])\n",
    "sem_block = Block(\n",
    "    CausalSemanticBlockType,\n",
    "    inputs={\n",
    "        \"semantic_input\": sem_codes,\n",
    "        \"block_type_input\": sem_block_type_ids,\n",
    "    },\n",
    ")\n",
    "\n",
    "diffusion_block_type_ids = torch.full((1, 1, T), BLOCK_TYPE_NAME_TO_ID[\"diffusion\"])\n",
    "diffusion_block = Block(\n",
    "    DiffusionBlockType,\n",
    "    inputs={\n",
    "        \"vae_input\": torch.randn((1, 128, T), device=model.device, dtype=torch.bfloat16),\n",
    "        \"block_type_input\": diffusion_block_type_ids,\n",
    "    },\n",
    ")\n",
    "\n",
    "\n",
    "no_text_blocks = BlockSequence([sem_block, diffusion_block])\n",
    "blocks = BlockSequence([text_block, sem_block, diffusion_block])\n",
    "print(blocks)\n",
    "\n",
    "text_cfg_boost = 1\n",
    "gconf = BCTGenerationConfig(\n",
    "    # [(1, no_text_blocks)],\n",
    "    [(1 + text_cfg_boost, blocks), (-text_cfg_boost, no_text_blocks)],\n",
    "    max_autoregressive_steps=16,\n",
    "    compile=False,\n",
    ")\n",
    "print(len(diffusion_block))\n",
    "\n",
    "block = generate_block(model, gconf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc4ad5ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = decode(block.inputs[\"vae_input\"][0].T / 0.4).normalize_volume()\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0e4747f",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio.write_mp3(\"stone_but_victor_trained_the_model.mp3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f3e460b",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_container[\"tokenizer\"].encode(\n",
    "    \"{start_} fesfes\",\n",
    "    add_special_tokens=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe209a25",
   "metadata": {},
   "source": [
    "## chunk it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "656f7c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def diffuse(sem_codes, lyrics=\"\", text_cfg_boost=1.0):\n",
    "    text_tokens = model_container[\"tokenizer\"].encode(lyrics)\n",
    "\n",
    "    # split into chunks of 30s\n",
    "    print(sem_codes.shape)\n",
    "    T = sem_codes.shape[0]\n",
    "    sem_codes = torch.from_numpy(sem_codes)\n",
    "    outputs = []\n",
    "\n",
    "    for i in range(0, T, 30 * 25):\n",
    "        sem_codes_chunk = sem_codes[i : i + 30 * 25]\n",
    "        sem_codes_chunk = torch.cat([sem_codes_chunk, torch.full((1,), cfg.semantic_pad_token)])\n",
    "        sem_codes_chunk = sem_codes_chunk.reshape(1, 1, -1)\n",
    "\n",
    "        text_block_type_ids = torch.full((1, 1, len(text_tokens)), BLOCK_TYPE_NAME_TO_ID[\"text\"])\n",
    "        text_block = Block(\n",
    "            TextBlockType,\n",
    "            inputs={\n",
    "                \"text_input\": torch.tensor(text_tokens).reshape(1, 1, -1),\n",
    "                \"block_type_input\": text_block_type_ids,\n",
    "            },\n",
    "        )\n",
    "        sem_block_type_ids = torch.full(\n",
    "            (1, 1, sem_codes_chunk.shape[-1]), BLOCK_TYPE_NAME_TO_ID[\"semantic\"]\n",
    "        )\n",
    "        sem_block = Block(\n",
    "            CausalSemanticBlockType,\n",
    "            inputs={\n",
    "                \"semantic_input\": sem_codes_chunk,\n",
    "                \"block_type_input\": sem_block_type_ids,\n",
    "            },\n",
    "        )\n",
    "        diffusion_block_type_ids = torch.full(\n",
    "            (1, 1, sem_codes_chunk.shape[-1]), BLOCK_TYPE_NAME_TO_ID[\"diffusion\"]\n",
    "        )\n",
    "        diffusion_block = Block(\n",
    "            DiffusionBlockType,\n",
    "            inputs={\n",
    "                \"vae_input\": torch.randn((1, 128, 30 * 25), device=model.device, dtype=torch.bfloat16),\n",
    "                \"block_type_input\": diffusion_block_type_ids,\n",
    "            },\n",
    "        )\n",
    "        blocks = BlockSequence([text_block, sem_block, diffusion_block])\n",
    "\n",
    "        if i > 0:\n",
    "            history_block_type_ids = torch.full(\n",
    "                (1, 1, sem_codes_chunk.shape[-1]), BLOCK_TYPE_NAME_TO_ID[\"prefix\"]\n",
    "            )\n",
    "            vae_input = outputs[-1].reshape(1, 128, -1)\n",
    "            vae_input = vae_input + torch.rand_like(vae_input) * 1\n",
    "            history_block = Block(\n",
    "                DiffusionBlockType,\n",
    "                inputs={\n",
    "                    \"vae_input\": vae_input,\n",
    "                    \"block_type_input\": history_block_type_ids,\n",
    "                },\n",
    "            )\n",
    "            blocks = BlockSequence([text_block, history_block, sem_block, diffusion_block])\n",
    "        no_text_blocks = BlockSequence(blocks[1:])\n",
    "\n",
    "        gconf = BCTGenerationConfig(\n",
    "            [(1 + text_cfg_boost, blocks), (-text_cfg_boost, no_text_blocks)],\n",
    "            max_autoregressive_steps=16,\n",
    "            compile=False,\n",
    "        )\n",
    "        block = generate_block(model, gconf)\n",
    "        outputs.append(block.inputs[\"vae_input\"][0])\n",
    "        print(outputs[-1].shape)\n",
    "    audio = decode(torch.cat(outputs, dim=1).T / 0.4).get_segment(to_s=T / 25)\n",
    "    return audio\n",
    "\n",
    "\n",
    "sem_codes = npz[\"v3.5_raw\"][:, 0]\n",
    "diffuse(sem_codes, lyrics).play()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be22284f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}