{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\""
   ]
  },
  {
   "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 models\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48588d23-a088-4354-9bf2-319d394d2c07",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "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 = \"/app2/suno/checkpoints/2025-11-02_15-41-25_s2081/last_ckpt.pt\"\n",
    "dit_model_filepath = (\n",
    "    \"/app2/suno/checkpoints/2025-11-11_20-38-05_s6962/last_ckpt.pt\"  #  newer run\n",
    ")\n",
    "diffusion_gen.preload_tokenizer(\"/app/suno/models/chirp_v2/tokenizer_60k.json\")\n",
    "diffusion_gen.preload_dit_model(\n",
    "    dit_model_filepath=dit_model_filepath,\n",
    "    use_ema_if_exists=True,\n",
    "    compile=False,\n",
    "    weights_precision=torch.bfloat16,\n",
    "    use_vox=True,\n",
    ")\n",
    "\n",
    "diffusion_engine = UpsampleEngine(compile=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7d0311",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from einops import rearrange\n",
    "from suno_utils.gpt.generation import (\n",
    "    _load_model,\n",
    "    GPT,\n",
    ")  # Use _load_model to bypass global caching\n",
    "from suno_utils.utils.s3 import download_s3_file_if_needed\n",
    "from suno_utils.tasks.dac_vae_fixed_25hz import (\n",
    "    preload_models as preload_codec_models,\n",
    "    decode,\n",
    ")\n",
    "from suno_utils.tasks.musicfm_v3 import ClusterModel\n",
    "from suno_utils.gpt.bct.bct_generation_simple import (\n",
    "    BCTGenerationConfig,\n",
    "    BlockSequence,\n",
    "    generate_block,\n",
    ")\n",
    "from suno_utils.gpt.bct.bct import Block, BlockType\n",
    "from transformers import PreTrainedTokenizerFast\n",
    "from tokenizers import AddedToken\n",
    "\n",
    "cluster_model = ClusterModel(np.load(\"/home/minz/temp/mert_768d_centroids_4000_50.npy\"))\n",
    "preload_codec_models(\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\")\n",
    "\n",
    "N_BATCH = 2\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "# Load tokenizer separately to reuse\n",
    "tokenizer_path = \"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\"\n",
    "local_tokenizer_path = download_s3_file_if_needed(tokenizer_path)\n",
    "\n",
    "tokenizer = PreTrainedTokenizerFast(\n",
    "    tokenizer_file=local_tokenizer_path,\n",
    "    unk_token=\"[UNK]\",\n",
    "    pad_token=\"[PAD]\",\n",
    ")\n",
    "tokenizer.add_special_tokens({\"additional_special_tokens\": [AddedToken(\"\\n\")]})\n",
    "\n",
    "# Load GPT model (first model)\n",
    "n_rvq = 4\n",
    "# gpt_model_path = \"/app2/suno/checkpoints/2025-10-30_18-43-30/last_ckpt_infer.pt\"  # 8 nodes, shift 1\n",
    "gpt_model_path = (\n",
    "    \"/app2/suno/checkpoints/2025-11-15_10-02-37/last_ckpt_infer.pt\"  # 32 nodes, shift 5\n",
    "    # \"/app2/suno/data/dpo/models/sem_4rvq_6b_11141927_infer.pt\"\n",
    ")\n",
    "print(\"Loading GPT model...\")\n",
    "model, _ = _load_model(\n",
    "    download_s3_file_if_needed(gpt_model_path), local_tokenizer_path, device\n",
    ")\n",
    "cfg = model.config\n",
    "print(f\"GPT model config: use_vae_output={cfg.use_vae_output}\")\n",
    "\n",
    "model_container = {\"model\": model, \"tokenizer\": tokenizer}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a1d5952",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd2f70ba-c9d1-4662-95ce-b7b5037ff421",
   "metadata": {},
   "outputs": [],
   "source": [
    "lyrics = \"\"\"\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",
    "\"\"\"\n",
    "tags = [\"pop\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46e7a399-7f55-4a89-a56f-3030c23ba3ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "lyrics = \"\"\"\n",
    "[intro]\n",
    "\n",
    "[verse]\n",
    "Walking down the street, feeling so alive\n",
    "Got my head in the clouds, got a gleam in my eye\n",
    "Every step I take, it's like a brand new start\n",
    "No matter where I'm going, I'll always find my part\n",
    "(oh-oh-oh)\n",
    "\n",
    "[chorus]\n",
    "Life is like a high-wire act, we're dancing in the sky\n",
    "No need to worry, no need to ask why\n",
    "With a little bit of courage, we can chase our dreams\n",
    "No matter what comes our way, we'll always be a team\n",
    "(we're unstoppable, yeah)\n",
    "\n",
    "[outro]\n",
    "\"\"\"\n",
    "tags = [\"epic film orchestral\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826dbf08-4612-4553-b5b5-ba69af2eaf38",
   "metadata": {},
   "outputs": [],
   "source": [
    "# lyrics = \"\"\"\n",
    "# \"\"\"\n",
    "# tags = [\"Mozart, Mozart, symphony, romantic, orchestra\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e7e428",
   "metadata": {},
   "outputs": [],
   "source": [
    "TextBlockType = BlockType(\n",
    "    name=\"text\",\n",
    "    is_causal=True,\n",
    ")\n",
    "CausalSemanticBlockType = BlockType(\n",
    "    name=\"semantic\",\n",
    "    is_causal=True,\n",
    ")\n",
    "CausalSemanticContinuousBlockType = BlockType(\n",
    "    name=\"semantic_continuous\",\n",
    "    is_causal=True,\n",
    ")\n",
    "DiffusionBlockType = BlockType(\n",
    "    name=\"diffusion\",\n",
    "    is_causal=False,\n",
    ")\n",
    "text = \", \".join(tags) + \"\\n\\n\" + lyrics\n",
    "# text = lyrics\n",
    "text_tokens = model_container[\"tokenizer\"].encode(text) + [cfg.text_infer_token]\n",
    "print(text_tokens)\n",
    "\n",
    "text_block = Block(\n",
    "    TextBlockType,\n",
    "    inputs={\n",
    "        \"text_input\": torch.tensor(text_tokens).reshape(1, 1, -1),\n",
    "    },\n",
    ")\n",
    "\n",
    "inputs = {}\n",
    "for i in range(n_rvq):\n",
    "    inputs[\"semantic_input_%d\" % i] = torch.full((1, 1, 1), cfg.semantic_infer_token)\n",
    "sem_block = Block(\n",
    "    CausalSemanticBlockType,\n",
    "    inputs=inputs,\n",
    ")\n",
    "\n",
    "no_text_blocks = BlockSequence([sem_block])\n",
    "blocks = BlockSequence([text_block, sem_block])\n",
    "\n",
    "text_cfg_boost = 0.5\n",
    "gconf = BCTGenerationConfig(\n",
    "    [(1 + text_cfg_boost, blocks), (-text_cfg_boost, no_text_blocks)],\n",
    "    # [(1, blocks)],\n",
    "    max_autoregressive_steps=25 * 4 * 60,\n",
    "    eos_token=cfg.semantic_pad_token,\n",
    "    temperature=0.9,\n",
    "    compile=False,\n",
    ")\n",
    "print(blocks)\n",
    "\n",
    "block = generate_block(model, gconf)\n",
    "\n",
    "ll = []\n",
    "for i in range(n_rvq):\n",
    "    ll.append(block.inputs[\"semantic_input_%d\" % i])\n",
    "sem_codes = torch.stack(ll)\n",
    "sem_codes = sem_codes[:, 0, 0, 1 : len(block)]\n",
    "\n",
    "# unshift delay pattern (shift each channel by 5*i)\n",
    "shift_amount = 5\n",
    "max_shift = (n_rvq - 1) * shift_amount\n",
    "unshifted = torch.full((n_rvq, sem_codes.shape[-1] - max_shift), -1)\n",
    "for i in range(n_rvq):\n",
    "    print(i)\n",
    "    shift = i * shift_amount\n",
    "    if i == n_rvq - 1:\n",
    "        unshifted[i] = sem_codes[i, shift:]\n",
    "    else:\n",
    "        unshifted[i] = sem_codes[i, shift : shift - max_shift]\n",
    "unshifted[unshifted > 4000] = 4000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ee67009",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "    lyrics=lyrics,\n",
    "    text_cfg_coef=2.0,\n",
    "    ctx_cfg_coef=1.0,\n",
    "    steps=12,\n",
    "    codec_scale_factor=0.4,\n",
    "    scale_ctx_vector=True,\n",
    "    objective=\"rectified_flow\",\n",
    ")\n",
    "request = Request(\n",
    "    id=\"dummy\",\n",
    "    generation_config=gen_cfg,\n",
    "    tokens=unshifted.T,\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "result = diffusion_engine.run_request(request)\n",
    "\n",
    "vae_latents = torch.concat(result.vae_latents)\n",
    "audio = decode(vae_latents)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce349010",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f391977",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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
}