{ "cells": [ { "cell_type": "code", "execution_count": 1, "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": 2, "id": "3c7d0311", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Failed to import flash_attn_3.\n", "No module named 'flash_attn_interface'\n", "WILL USE FLASH ATTN: True\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/tony/anaconda3/envs/suno_env_dev/lib/python3.10/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n", " warnings.warn(\n", "/home/tony/Work/glockenspiel/suno_utils/suno_utils/gpt/generation.py:94: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " checkpoint = torch.load(str(ckpt_path), mmap=True)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Loading GPT model...\n", "model loaded: /app2/suno/data/dpo/models/sem_4rvq_6b_11141927_infer.pt\n", "model loaded: 5602.8M params\n", "GPT model config: use_vae_output=False\n" ] } ], "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-04_19-47-20/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": "code", "execution_count": 3, "id": "ae7b7934", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Failed to import flash_attn_3.\n", "Running on non-Modal\n", "Dir1: 648 files\n", "Dir2: 311 files\n", "Common files: 311\n", "All unique files: 648\n", "Match percentage: 47.99%\n", "Similarity: 47.99%\n", "Directories are not identical. Use disk model instead, if you want to use ssd models. run python -m suno_utils.worker.modal_manager_model_store sync\n", "loading tokenizer...\n", "loading diffusion model...\n", "loading weights...\n", "converting model to precision torch.bfloat16...\n", "Block size: 750\n" ] } ], "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 = \"/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(tokenizer_path)\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": "markdown", "id": "1a1d5952", "metadata": {}, "source": [ "## Inference" ] }, { "cell_type": "code", "execution_count": 4, "id": "f2e7e428", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[19333, 60000, 60000, 60000, 60, 21366, 62, 60000, 19249, 13, 19111, 19207, 60000, 19483, 19874, 19078, 19182, 60000, 19108, 8, 84, 19420, 66, 20623, 60000, 23413, 21126, 19136, 19078, 60000, 19315, 19079, 20041, 21335, 19089, 21086, 60000, 60000, 60, 20750, 62, 60000, 42, 8, 78, 19194, 21843, 19136, 28064, 60000, 28064, 19462, 19785, 19221, 60000, 74, 19182, 19108, 19315, 19262, 19252, 74, 19092, 19173, 60000, 74, 19365, 19408, 21064, 19092, 21947, 13, 60000, 19194, 19639, 19089, 19408, 21084, 60000, 60000, 60, 21366, 62, 60000, 19249, 13, 19111, 19207, 60000, 19078, 19421, 66, 19475, 19136, 24004, 60000, 19713, 21639, 74, 8, 19276, 19431, 19820, 60000, 74, 8, 19276, 19335, 21718, 19194, 20051, 60000, 19315, 19079, 20041, 21335, 19089, 21086, 60000, 60002]\n", "BlockSequence(\n", " \u001b[38;2;28;178;81mtext\u001b[0m (causal) (120)\n", " text_input: \u001b[2m[[[19333, 60000, 60000, 60000, 60, 21366, 62, 60000, 19249, 13, 19111, 19207, 60000, 19483, 19874, 19078, 19182, 60000, 19108, 8, 84, 19420, 66, 20623, 60000, 23413, 21126, 19136, 19078, 60000, 19315, 19079, 20041, 21335, 19089, 21086, 60000, 60000, 60, 20750, 62, 60000, 42, 8, 78, 19194, 21843, 19136, 28064, 60000, 28064, 19462, 19785, 19221, 60000, 74, 19182, 19108, 19315, 19262, 19252, 74, 19092, 19173, 60000, 74, 19365, 19408, 21064, 19092, 21947, 13, 60000, 19194, 19639, 19089, 19408, 21084, 60000, 60000, 60, 21366, 62, 60000, 19249, 13, 19111, 19207, 60000, 19078, 19421, 66, 19475, 19136, 24004, 60000, 19713, 21639, 74, 8, 19276, 19431, 19820, 60000, 74, 8, 19276, 19335, 21718, 19194, 20051, 60000, 19315, 19079, 20041, 21335, 19089, 21086, 60000, 60002]]]\u001b[0m\n", " \u001b[38;2;245;139;226msemantic\u001b[0m (causal) (1)\n", " semantic_input_0: \u001b[2m[[[4001]]]\u001b[0m\n", " semantic_input_1: \u001b[2m[[[4001]]]\u001b[0m\n", " semantic_input_2: \u001b[2m[[[4001]]]\u001b[0m\n", " semantic_input_3: \u001b[2m[[[4001]]]\u001b[0m\n", ")\n", "KV Cache - Allocated: 1.95GB, Blocks: 4000, Tokens supported: 1,024,000\n", "KV Cache - Allocated: 1.95GB, Blocks: 4000, Tokens supported: 1,024,000\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "100%|██████████| 3000/3000 [01:02<00:00, 47.82it/s]" ] }, { "name": "stdout", "output_type": "stream", "text": [ "0\n", "1\n", "2\n", "3\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\n" ] } ], "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", "\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", "tags = [\"pop\"]\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 * 120,\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": [ { "name": "stderr", "output_type": "stream", "text": [ "Engine step: 0it [00:00, ?it/s]" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "e676ab23c5194f9fa8a740b884edf608", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/12 [00:00\n", " \n", " Your browser does not support the audio element.\n", "\n" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "ename": "", "evalue": "", "output_type": "error", "traceback": [ "\u001b[1;31mThe Kernel crashed while executing code in the current cell or a previous cell. \n", "\u001b[1;31mPlease review the code in the cell(s) to identify a possible cause of the failure. \n", "\u001b[1;31mClick here for more info. \n", "\u001b[1;31mView Jupyter log for further details." ] } ], "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": "996e1324", "metadata": {}, "outputs": [], "source": [] }, { "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": "suno_env_dev", "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 }