{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\""
   ]
  },
  {
   "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": "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",
    "gpt_model_path = \"/app/suno/checkpoints/2025-04-23_16-23-43/last_ckpt_infer.pt\"\n",
    "# gpt_model_path_asr = \"/app/suno/checkpoints/2025-04-15_23-31-50/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),\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": "8bde46df",
   "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",
    "dit_model_filepath = \"/app/suno/tmp/diff_v2_dpo_apr12.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": "33c1eb9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# from suno_utils.gpt.bct.bct_generation_simple import BCTGenerationConfig, BlockSequence, generate_block\n",
    "from suno_utils.gpt.bct.bct_engine import BCTEngine, BlockSequence, Config\n",
    "from suno_utils.gpt.bct.bct import Block, BlockType, TensorDict, SamplingParams\n",
    "\n",
    "engine_config = Config()\n",
    "engine = BCTEngine(engine_config, model, compile=False)\n"
   ]
  },
  {
   "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",
    "\n",
    "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",
    "text = lyrics\n",
    "text_tokens = model_container[\"tokenizer\"].encode(text)\n",
    "print(text_tokens)\n",
    "\n",
    "text_block = Block(\n",
    "    TextBlockType,\n",
    "    inputs=TensorDict(\n",
    "        text_input=torch.tensor(text_tokens).reshape(1, 1, -1),\n",
    "    ),\n",
    ")\n",
    "\n",
    "sem_block = Block(\n",
    "    CausalSemanticBlockType,\n",
    "    inputs=TensorDict(\n",
    "        semantic_input=torch.full((1, 1, 1), cfg.semantic_infer_token),\n",
    "    ),\n",
    ")\n",
    "\n",
    "sampling_params = SamplingParams(\n",
    "    stop_token_ids=[cfg.semantic_pad_token],\n",
    "    temperature=0.9,\n",
    "    max_tokens=1000,\n",
    ")\n",
    "no_text_blocks = BlockSequence([sem_block], sampling_params=sampling_params)\n",
    "blocks = BlockSequence([text_block, sem_block], sampling_params=sampling_params)\n",
    "\n",
    "batch = [BlockSequence([text_block, sem_block], sampling_params=sampling_params) for _ in range(1)]\n",
    "engine.generate(batch, use_tqdm=True)\n",
    "\n",
    "sem_codes = [TensorDict.concatenate(seq.generated_tokens)[\"semantic_input\"] for seq in batch]\n",
    "# sem_codes = TensorDict.concatenate(blocks.generated_tokens)[\"semantic_input\"]\n",
    "# sem_codes.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0dd8bc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "sem_codes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ff662f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.generation import GenerationConfig\n",
    "from suno_utils.gpt.generation_engine import make_request\n",
    "\n",
    "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",
    "gconf = GenerationConfig(\n",
    "    text=lyrics,\n",
    "    text_tags=\"\",\n",
    "    cfg_coef=1.0,\n",
    "    # cfg_coef_tags=1.5,\n",
    "    n_batch=1,\n",
    "    min_text_offset=0,\n",
    "    eos_pad_duration_s=0,\n",
    "    max_gen_duration_s=120,\n",
    ")\n",
    "request = make_request(\"yo\", gconf, engine.model.config, model_container[\"tokenizer\"])\n",
    "request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7517835",
   "metadata": {},
   "outputs": [],
   "source": [
    "request.streams[0].prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1856da70",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.job import Request as GPTRequest\n",
    "\n",
    "\n",
    "def request_to_block_sequence(request: GPTRequest):\n",
    "    streams = []\n",
    "    for i, stream in enumerate(request.streams):\n",
    "        req_block = Block(\n",
    "            CausalSemanticBlockType,\n",
    "            inputs=TensorDict(\n",
    "                text_input=stream.prompt[0].reshape(1, 1, -1),\n",
    "                semantic_input=stream.prompt[1].reshape(1, 1, -1),\n",
    "            ),\n",
    "        )\n",
    "        if stream.max_n_steps is not None:\n",
    "            max_tokens = stream.max_n_steps[-1]\n",
    "        else:\n",
    "            max_tokens = request.max_gen_duration_s * 25\n",
    "\n",
    "        sampling_params = SamplingParams(\n",
    "            stop_token_ids=[cfg.semantic_pad_token] if request.allow_eos else [],\n",
    "            temperature=request.temp_semantic,\n",
    "            max_tokens=max_tokens,\n",
    "        )\n",
    "        if request.top_k_semantic is not None:\n",
    "            sampling_params.top_k = request.top_k_semantic\n",
    "        if request.top_p_semantic is not None:\n",
    "            sampling_params.top_p = request.top_p_semantic\n",
    "        if request.min_p_semantic is not None:\n",
    "            sampling_params.min_p = request.min_p_semantic\n",
    "\n",
    "        seq = BlockSequence([req_block], sampling_params=sampling_params)\n",
    "        seq.id = str(i)\n",
    "        seq.weight = stream.weight[0]\n",
    "        seq.main_cfg_stream_id = \"0\"\n",
    "\n",
    "        streams.append(seq)\n",
    "\n",
    "    return streams\n",
    "\n",
    "\n",
    "batch = request_to_block_sequence(request)\n",
    "engine.generate(batch, use_tqdm=True)\n",
    "sem_codes = [TensorDict.concatenate(seq.generated_tokens)[\"semantic_input\"] for seq in batch]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68a2b9b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08f47497",
   "metadata": {},
   "source": [
    "## Diffusion Engine\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc4ad5ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "    # audio=vae_latents,\n",
    "    # tags=\"extract Lead Vocal\",\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",
    ")\n",
    "for sem_code in sem_codes:\n",
    "    diff_request = Request(\n",
    "        id=\"dummy\",\n",
    "        generation_config=gen_cfg,\n",
    "        tokens=sem_code.cpu()[0, 0, :-1],\n",
    "        input_tokens_finished=True,\n",
    "    )\n",
    "    result = diffusion_engine.run_request(diff_request)\n",
    "    vae_latents = torch.concat(result.vae_latents)\n",
    "    audio = decode(vae_latents)\n",
    "    audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c60080cf",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}