{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"5\""
   ]
  },
  {
   "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 = \"/app2/suno/modal/models/tony/sem/model_45_6b_crow_sft0828_t14_r15_d87.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": "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_infer.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": "f171a38c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.ditto_v2 import preload_models as preload_ditto_models\n",
    "\n",
    "preload_ditto_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0280920",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.tasks.ditto_v2 import encode as encode_ditto, encode_text\n",
    "\n",
    "ditto_emb = torch.tensor(encode_text(\"bluegrass, female vocalist\"), dtype=torch.bfloat16)\n",
    "ditto_emb.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "z8n0mas9ev",
   "metadata": {},
   "source": [
    "## Ditto Block with Text Diff + SLERP\n",
    "\n",
    "Create a Ditto block by:\n",
    "1. Computing text difference: `text_diff = new_text_emb - generic_music_emb`\n",
    "2. Normalizing the text diff to unit length\n",
    "3. Using SLERP to interpolate between audio embedding and normalized text diff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rz8fsha36j",
   "metadata": {},
   "outputs": [],
   "source": [
    "# For reference: play the original stone_clip audio\n",
    "from suno_utils.gpt.bct.bct_test_utils import stone_clip\n",
    "\n",
    "print(\"Original stone_clip:\")\n",
    "print(f\"Lyrics: {stone_clip.lyrics()[:100]}...\")\n",
    "stone_clip.audio().play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "z69efb7n4wf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import utilities\n",
    "from suno_utils.tasks.ditto_interpolation import interpolate_latents\n",
    "from suno_utils.gpt.bct.bct_test_utils import stone_clip\n",
    "from suno_utils.gpt.bct.bct import TensorDict\n",
    "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",
    "\n",
    "# Get text encoding from \"bluegrass\"\n",
    "text_ditto_emb = encode_text(\n",
    "    \"Heavy, syncopated drums and jagged metal guitar riffs lead, layered with dark, pulsing synths and deep sub-bass, Trap hi-hats accent rapid rhythms, while male and female vocals alternate: visceral rap in verses, aggressive sung choruses, Beat drops unleash witch house atmospheres and bounce house grooves, intensifying the dance-driven texture\"\n",
    ").cpu()\n",
    "print(f\"Text ditto embedding shape: {text_ditto_emb.shape}\")\n",
    "\n",
    "# Get audio encoding from stone_clip\n",
    "audio_ditto_emb = torch.tensor(encode_ditto(stone_clip.audio()), dtype=torch.float32).cpu()\n",
    "print(f\"Audio ditto embedding shape: {audio_ditto_emb.shape}\")\n",
    "\n",
    "# Get generic \"music\" text embedding as baseline\n",
    "generic_music_emb = encode_text(\"music\").cpu()\n",
    "print(f\"Generic music embedding shape: {generic_music_emb.shape}\")\n",
    "\n",
    "# Calculate text difference: D_newtext - D_generic\n",
    "text_diff = text_ditto_emb  # - generic_music_emb\n",
    "print(f\"Text diff shape: {text_diff.shape}\")\n",
    "print(f\"Text diff norm before normalization: {torch.norm(text_diff).item():.4f}\")\n",
    "\n",
    "# Normalize the text diff to unit length\n",
    "text_diff_normalized = text_diff / (torch.norm(text_diff) + 1e-8)\n",
    "print(f\"Text diff norm after normalization: {torch.norm(text_diff_normalized).item():.4f}\")\n",
    "\n",
    "# SLERP between audio embedding and normalized text diff (70% audio, 30% text diff)\n",
    "final_emb = interpolate_latents([audio_ditto_emb, text_diff_normalized], weights=[0.7, 0.3])\n",
    "print(f\"Final embedding shape: {final_emb.shape}\")\n",
    "\n",
    "# Convert to bfloat16 for model compatibility\n",
    "final_emb = final_emb.to(dtype=torch.bfloat16)\n",
    "\n",
    "# Create Ditto block with the modified embedding\n",
    "DittoBlockType = BlockType(\n",
    "    name=\"ditto\",\n",
    "    is_causal=True,\n",
    ")\n",
    "\n",
    "ditto_block = Block(\n",
    "    spec=DittoBlockType,\n",
    "    inputs=TensorDict({\"ditto_input\": final_emb.reshape(1, -1, 1)}),\n",
    ")\n",
    "\n",
    "print(f\"Ditto block created with input shape: {ditto_block.inputs['ditto_input'].shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "iayhgsxfsjq",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate using the interpolated Ditto block\n",
    "\n",
    "\n",
    "TextBlockType = BlockType(\n",
    "    name=\"text\",\n",
    "    is_causal=True,\n",
    ")\n",
    "CausalSemanticBlockType = BlockType(\n",
    "    name=\"semantic\",\n",
    "    is_causal=True,\n",
    ")\n",
    "# Create mumble mode text block\n",
    "mumble_text = \"[mumble mode]\"\n",
    "mumble_tokens = model_container[\"tokenizer\"].encode(mumble_text)\n",
    "mumble_block = Block(\n",
    "    TextBlockType,\n",
    "    inputs=TensorDict(\n",
    "        {\n",
    "            \"text_input\": torch.tensor(mumble_tokens).reshape(1, 1, -1),\n",
    "        }\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Create semantic inference block\n",
    "sem_infer_block = Block(\n",
    "    CausalSemanticBlockType,\n",
    "    inputs=TensorDict(\n",
    "        {\n",
    "            \"semantic_input\": torch.full((1, 1, 1), cfg.semantic_infer_token),\n",
    "        }\n",
    "    ),\n",
    ")\n",
    "\n",
    "# Create block sequence with interpolated ditto block\n",
    "ditto_blocks = BlockSequence([ditto_block, mumble_block, sem_infer_block])\n",
    "\n",
    "# Generate with the interpolated ditto block\n",
    "ditto_gconf = BCTGenerationConfig(\n",
    "    [(1, ditto_blocks)],\n",
    "    max_autoregressive_steps=25 * 60,\n",
    "    eos_token=cfg.semantic_pad_token,\n",
    "    temperature=0.9,\n",
    ")\n",
    "\n",
    "print(\"Generating with interpolated Ditto block...\")\n",
    "ditto_gen_block = generate_block(model, ditto_gconf)\n",
    "ditto_sem_codes = ditto_gen_block.inputs[\"semantic_input\"][0, 0, 1 : len(ditto_gen_block)]\n",
    "print(f\"Generated semantic codes shape: {ditto_sem_codes.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "rboadisgeoh",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Render audio from the generated semantic codes\n",
    "ditto_gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n",
    "    lyrics=mumble_text,\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",
    "\n",
    "ditto_request = Request(\n",
    "    id=\"ditto_interpolated\",\n",
    "    generation_config=ditto_gen_cfg,\n",
    "    tokens=ditto_sem_codes.cpu(),\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "\n",
    "print(\"Rendering audio...\")\n",
    "ditto_result = diffusion_engine.run_request(ditto_request)\n",
    "ditto_vae_latents = torch.concat(ditto_result.vae_latents)\n",
    "ditto_audio = decode(ditto_vae_latents)\n",
    "ditto_audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e7e428",
   "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",
    "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",
    "sem_block = Block(\n",
    "    CausalSemanticBlockType,\n",
    "    inputs={\n",
    "        \"semantic_input\": torch.full((1, 1, 1), cfg.semantic_infer_token),\n",
    "    },\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 * 60,\n",
    "    eos_token=cfg.semantic_pad_token,\n",
    "    temperature=0.9,\n",
    ")\n",
    "\n",
    "block = generate_block(model, gconf)\n",
    "sem_codes = block.inputs[\"semantic_input\"][0, 0, 1 : len(block)]\n",
    "sem_codes.shape\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",
    "request = Request(\n",
    "    id=\"dummy\",\n",
    "    generation_config=gen_cfg,\n",
    "    tokens=sem_codes.cpu(),\n",
    "    input_tokens_finished=True,\n",
    ")\n",
    "result = diffusion_engine.run_request(request)\n",
    "vae_latents = torch.concat(result.vae_latents)\n",
    "audio = decode(vae_latents)\n",
    "audio.play()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}