{ "cells": [ { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n", "import sys\n", "import torch\n", "import torch.nn.functional as F\n", "from tqdm import tqdm\n", "sys.path.insert(0, \"/home/christian/code/christian/scripts\")\n", "\n", "from train_gpt import GPTConfig, GPTModel, load_tokenizer, INFER_TOKEN, SEMANTIC_PAD_TOKEN, SEMANTIC_EOS_TOKEN" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "loading semantic model...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/christian/code/glockenspiel/suno_utils/suno_utils/tasks/mert_25.py:148: 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", " sd = _torch_load_p(checkpoint_filepath)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "loading codec model...\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "2025-01-17 20:46:09.130935: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n", "2025-01-17 20:46:09.956805: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n", "2025-01-17 20:46:10.239431: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n", "2025-01-17 20:46:10.321856: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n", "2025-01-17 20:46:10.923975: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n", "To enable the following instructions: AVX2 AVX512F AVX512_VNNI AVX512_BF16 AVX512_FP16 AVX_VNNI AMX_TILE AMX_INT8 AMX_BF16 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n", "2025-01-17 20:46:15.319060: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n", "/home/christian/code/glockenspiel/suno_utils/suno_utils/tasks/dac_vae_100hz_peaq.py:73: 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", " sd = load_f(checkpoint_filepath)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "loading diffusion model...\n", "loading weights...\n", "converting model to precision torch.bfloat16...\n", "loading tokenizer...\n", "done!\n" ] } ], "source": [ "# also load diffusion modelimport sys\n", "sys.path.insert(0, \"/home/christian/code/neon/sunoDiff/\")\n", "dit_model_filepath = \"s3://suno-data/georg/tmp/2b_prefix_ft.pt\"\n", "\n", "# mert25\n", "semantic_model_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25.pt\"\n", "semantic_clusters_filepath=\"/home/georg/notebooks/gpu_nb/tmp/mert_25_2x4k.npy\"\n", "codec_filepath=\"/home/georg/notebooks/gpu_nb/tmp/25hz_vae_peaq_kl_0.005.pth\"\n", "\n", "patch_size = 1\n", "codec_scale_factor = 2.5\n", "\n", "import torch\n", "from suno_utils.audio import Audio\n", "from generation import preload_models, generate, _retrieve_models\n", "_ = preload_models(\n", " tokenizer_filepath=\"/home/georg/notebooks/gpu_nb/tmp/tokenizer_60k.json\",\n", " semantic_model_filepath=semantic_model_filepath,\n", " semantic_clusters_filepath=semantic_clusters_filepath,\n", " codec_filepath=codec_filepath,\n", " dit_model_filepath=dit_model_filepath,\n", " weights_precision=torch.bfloat16,\n", " model_type=\"prefix\",\n", " codec_scale_factor=codec_scale_factor,\n", ")\n", "\n", "models = _retrieve_models()\n", "model_duration_s = 30\n", "if models[\"dit_model\"].ctx_len is not None:\n", " model_duration_s = 6 * 60\n", "else:\n", " model_duration_s = models[\"dit_model\"].block_size // models[\"dit_model\"].io_hz\n", "duration_s = 2*60 if model_duration_s >= 2*60 else 30\n" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/tmp/ipykernel_2111198/182291805.py:4: 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(checkpoint_filepath, map_location=\"cpu\")\n" ] } ], "source": [ "# load checkpoint\n", "#checkpoint_filepath = \"/app/suno/christian/checkpoints/gpt/2025-01-16_22-47-13_s5202/last_ckpt.pth\"\n", "checkpoint_filepath = \"/app/suno/christian/checkpoints/gpt/2025-01-17_20-43-36_s9865/last_ckpt.pth\"\n", "checkpoint = torch.load(checkpoint_filepath, map_location=\"cpu\")\n", "\n", "# load config\n", "config = checkpoint[\"config\"]\n", "\n", "# load model\n", "model = GPTModel(config)\n", "new_state_dict = {k.replace(\"module.\", \"\"): v for k, v in checkpoint[\"model\"].items()}\n", "model.load_state_dict(new_state_dict)\n", "\n", "# load tokenizer\n", "tokenizer = load_tokenizer()\n" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [], "source": [ "\n", "\n", "def generate(\n", " model, \n", " text_input_ids, \n", " infer_token_id,\n", " semantic_vocab_size,\n", " attention_mask=None, \n", " max_new_tokens=50, \n", " temperature=1.0, \n", " top_k=50, \n", " top_p=0.9\n", "):\n", " model.eval()\n", " \n", " # Make sure input_ids is 2D\n", " if text_input_ids.dim() == 1:\n", " text_input_ids = text_input_ids.unsqueeze(0)\n", " \n", " # Add infer token to signal start of semantic generation\n", " input_ids = torch.cat([text_input_ids, torch.tensor([[infer_token_id]])], dim=-1)\n", " \n", " # If no attention mask provided, create one\n", " if attention_mask is None:\n", " attention_mask = torch.ones_like(input_ids).bool()\n", " else:\n", " attention_mask = torch.cat([attention_mask, torch.ones(1, 1).bool()], dim=-1)\n", " \n", " # Move to same device as model\n", " device = next(model.parameters()).device\n", " input_ids = input_ids.to(device)\n", " attention_mask = attention_mask.to(device)\n", "\n", " with torch.no_grad():\n", " pbar = tqdm(range(max_new_tokens))\n", " for _ in pbar:\n", " # Get model predictions\n", " logits = model(input_ids, attention_mask=attention_mask)\n", " \n", " # Get next token logits and apply temperature\n", " next_token_logits = logits[:, -1, :] / temperature\n", " \n", " # Only allow sampling from semantic vocabulary for generation\n", " # Zero out logits for text vocabulary\n", " next_token_logits[:, :semantic_vocab_size] = float('-inf')\n", " \n", " # Apply top-k filtering\n", " if top_k > 0:\n", " values, _ = torch.topk(next_token_logits, top_k)\n", " min_values = values[:, -1].unsqueeze(-1).expand_as(next_token_logits)\n", " next_token_logits = torch.where(\n", " next_token_logits < min_values,\n", " torch.full_like(next_token_logits, float('-inf')),\n", " next_token_logits\n", " )\n", " \n", " # Apply top-p (nucleus) filtering\n", " if top_p < 1.0:\n", " probs = F.softmax(next_token_logits, dim=-1)\n", " sorted_probs, sorted_indices = torch.sort(probs, descending=True)\n", " cumulative_probs = torch.cumsum(sorted_probs, dim=-1)\n", " \n", " sorted_indices_to_remove = cumulative_probs > top_p\n", " sorted_indices_to_remove[..., 0] = 0\n", " \n", " indices_to_remove = sorted_indices_to_remove.scatter(\n", " dim=-1,\n", " index=sorted_indices,\n", " src=sorted_indices_to_remove\n", " )\n", " next_token_logits = torch.where(\n", " indices_to_remove,\n", " torch.full_like(next_token_logits, float('-inf')),\n", " next_token_logits\n", " )\n", " \n", " # Sample next token\n", " probs = F.softmax(next_token_logits, dim=-1)\n", " next_token = torch.multinomial(probs, num_samples=1)\n", " \n", " # Append new token\n", " input_ids = torch.cat([input_ids, next_token], dim=-1)\n", " attention_mask = torch.cat([attention_mask, torch.ones_like(next_token)], dim=-1).bool()\n", " \n", " # Optional: Stop if semantic EOS token is generated\n", " if next_token[0, 0].item() == SEMANTIC_EOS_TOKEN_ID:\n", " break\n", " \n", " return input_ids\n", "\n", "# Example usage:\n", "@torch.no_grad()\n", "def generate_semantic(\n", " model,\n", " text_tokenizer,\n", " prompt,\n", " infer_token_id,\n", " semantic_vocab_size,\n", " **gen_kwargs\n", "):\n", " # Encode text prompt\n", " text_input_ids = text_tokenizer.encode(prompt).ids\n", " text_input_ids = torch.tensor(text_input_ids)\n", " \n", " # Generate\n", " output_ids = generate(\n", " model, \n", " text_input_ids,\n", " infer_token_id=infer_token_id,\n", " semantic_vocab_size=semantic_vocab_size,\n", " **gen_kwargs\n", " )\n", " \n", " # Split output into text and semantic parts\n", " infer_token_pos = (output_ids == infer_token_id).nonzero()[0, 1]\n", " semantic_tokens = output_ids[0, infer_token_pos + 1:]\n", " \n", " return semantic_tokens" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ " 0%| | 0/50 [00:00 3\u001b[0m semantic_tokens \u001b[38;5;241m=\u001b[39m \u001b[43mgenerate_semantic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtokenizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mprompt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mINFER_TOKEN\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msemantic_vocab_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m4001\u001b[39;49m\u001b[43m)\u001b[49m\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/utils/_contextlib.py:116\u001b[0m, in \u001b[0;36mcontext_decorator..decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 113\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m 114\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m 115\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 116\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", "Cell \u001b[0;32mIn[20], line 104\u001b[0m, in \u001b[0;36mgenerate_semantic\u001b[0;34m(model, text_tokenizer, prompt, infer_token_id, semantic_vocab_size, **gen_kwargs)\u001b[0m\n\u001b[1;32m 101\u001b[0m text_input_ids \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(text_input_ids)\n\u001b[1;32m 103\u001b[0m \u001b[38;5;66;03m# Generate\u001b[39;00m\n\u001b[0;32m--> 104\u001b[0m output_ids \u001b[38;5;241m=\u001b[39m \u001b[43mgenerate\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 105\u001b[0m \u001b[43m \u001b[49m\u001b[43mmodel\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[1;32m 106\u001b[0m \u001b[43m \u001b[49m\u001b[43mtext_input_ids\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 107\u001b[0m \u001b[43m \u001b[49m\u001b[43minfer_token_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43minfer_token_id\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 108\u001b[0m \u001b[43m \u001b[49m\u001b[43msemantic_vocab_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msemantic_vocab_size\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 109\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mgen_kwargs\u001b[49m\n\u001b[1;32m 110\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 112\u001b[0m \u001b[38;5;66;03m# Split output into text and semantic parts\u001b[39;00m\n\u001b[1;32m 113\u001b[0m infer_token_pos \u001b[38;5;241m=\u001b[39m (output_ids \u001b[38;5;241m==\u001b[39m infer_token_id)\u001b[38;5;241m.\u001b[39mnonzero()[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m]\n", "Cell \u001b[0;32mIn[20], line 36\u001b[0m, in \u001b[0;36mgenerate\u001b[0;34m(model, text_input_ids, infer_token_id, semantic_vocab_size, attention_mask, max_new_tokens, temperature, top_k, top_p)\u001b[0m\n\u001b[1;32m 33\u001b[0m pbar \u001b[38;5;241m=\u001b[39m tqdm(\u001b[38;5;28mrange\u001b[39m(max_new_tokens))\n\u001b[1;32m 34\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m pbar:\n\u001b[1;32m 35\u001b[0m \u001b[38;5;66;03m# Get model predictions\u001b[39;00m\n\u001b[0;32m---> 36\u001b[0m logits \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_ids\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mattention_mask\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mattention_mask\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 38\u001b[0m \u001b[38;5;66;03m# Get next token logits and apply temperature\u001b[39;00m\n\u001b[1;32m 39\u001b[0m next_token_logits \u001b[38;5;241m=\u001b[39m logits[:, \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, :] \u001b[38;5;241m/\u001b[39m temperature\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1736\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1734\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs) \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m 1735\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1736\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", "File \u001b[0;32m~/miniconda3/envs/suno_env/lib/python3.10/site-packages/torch/nn/modules/module.py:1747\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 1742\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m 1743\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m 1744\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m 1745\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m 1746\u001b[0m \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1747\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1749\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 1750\u001b[0m called_always_called_hooks \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mset\u001b[39m()\n", "\u001b[0;31mTypeError\u001b[0m: GPTModel.forward() missing 1 required positional argument: 'semantic_input_ids'" ] } ], "source": [ "# test\n", "prompt = \"[rock]\"\n", "semantic_tokens = generate_semantic(model, tokenizer, prompt, INFER_TOKEN, semantic_vocab_size=4001)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "suno_env", "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.9" } }, "nbformat": 4, "nbformat_minor": 2 }