{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"import torch\n",
"\n",
"# test the conductor\n",
"import json\n",
"import modal\n",
"\n",
"import numpy as np\n",
"from suno_utils.diffusion.generation import (\n",
" generate,\n",
" preload_models,\n",
" DiffusionGenerationConfig,\n",
" generate_stream\n",
")\n",
"from suno_utils.tasks.dac_vae_100hz_peaq import (\n",
" preload_models as preload_vae_models,\n",
" decode_stream_to_full_audio as decode_vae_to_audio,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"# input = json.dumps(\n",
"# {\n",
"# \"id\": \"f3fb4b05-49ab-45f2-9591-1bcdcef6ece6\",\n",
"# \"prompt_audio\": \"99c4280e-c26b-4c8b-87a8-494512a24484.mp3\",\n",
"# \"prompt_npz\": None,\n",
"# \"prompt_text\": \"[verse]\\n\\n\",\n",
"# \"metadata\": {\n",
"# \"tags\": \"shoegaze, ethereal, art, ego\",\n",
"# \"gpt_prompt\": None,\n",
"# \"gpt_description_prompt\": None,\n",
"# \"stream\": True,\n",
"# \"make_instrumental\": False,\n",
"# \"priority\": 10,\n",
"# \"user_id\": 172918,\n",
"# \"is_bot\": False,\n",
"# \"task\": \"extend\",\n",
"# \"continued_from_prompt\": \"[melodic hook]\\n\\n\",\n",
"# \"continued_aligned_prompt\": None,\n",
"# \"history\": [\n",
"# {\n",
"# \"id\": \"99c4280e-c26b-4c8b-87a8-494512a24484\",\n",
"# \"continue_at\": 160.0,\n",
"# \"type\": \"gen\",\n",
"# \"source\": \"web\",\n",
"# \"infill\": False,\n",
"# }\n",
"# ],\n",
"# \"continue_at\": 160.0,\n",
"# \"options\": None,\n",
"# \"edit_session_id\": None,\n",
"# \"infill\": False,\n",
"# \"type\": \"gen\",\n",
"# \"lang\": \"English\",\n",
"# },\n",
"# \"gen_duration\": 12,\n",
"# \"callback_url\": \"https://studio-api.prod.suno.com/api/generate/finish-clip/\",\n",
"# \"model_name\": \"chirp-v4-xxx\",\n",
"# \"title\": \"\",\n",
"# \"ids\": [\n",
"# \"e72a78db-a9a2-43a3-8b97-597db7f76c12\",\n",
"# \"3c0160e0-0924-45a3-9c95-fb2e8ffaaddc\",\n",
"# ],\n",
"# }\n",
"# )"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"# # f = modal.Function.lookup(\"history_encoder-dev\", \"HistoryEncoderStub.encode_history\")\n",
"## fn_call = f.spawn(input)\n",
"# result = fn_call.get()\n",
"# result.history_latents.shape"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"CODEC_FILEPATH = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\""
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/tony/Work/glockenspiel/suno_utils/suno_utils/utils/s3.py:325: 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",
" data = read_f(tmp_filepath)\n"
]
}
],
"source": [
"preload_vae_models(CODEC_FILEPATH)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [],
"source": [
"\n",
"# print(result.history_latents.dtype)\n",
"# decoded_audio = decode_vae_to_audio(\n",
"# torch.from_numpy(result.history_latents).to(torch.bfloat16)\n",
"# )\n",
"# decoded_audio.play()\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"loading semantic model...\n",
"loading codec model...\n",
"loading diffusion model...\n",
"loading weights...\n",
"converting model to precision torch.bfloat16...\n",
"dit_model loaded onto cuda:0\n",
"loading tokenizer...\n",
"done!\n"
]
}
],
"source": [
"preload_models(\n",
" dit_model_filepath=\"/app/suno/data/dpo/models/diff_vae_25_peaq_v1_nov7.pt\"\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"parent_gen_id = \"66636a2e-77da-4162-85f5-521916225898\" # artist condition\n",
"gen_id = \"fb177587-f665-412a-b768-3e27d2f04285\" # infill \n",
"# gen_id = \"6b2bcada-d91a-46c0-94bf-e0682578cee7\" # full\n",
"# gen_id = \"68817224-d5fc-415e-80d7-236ed6576228\""
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"source": [
"from suno_utils.utils.s3 import read_from_s3"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"s3://suno-data-uploads/studio/uploads/fb177587-f665-412a-b768-3e27d2f04285.npz\n",
"sem_codes.shape torch.Size([632]) 25.28\n",
"vae_codes.shape (627, 128) 25.08\n"
]
}
],
"source": [
"s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n",
"mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n",
"print(s3_filepath)\n",
"data = read_from_s3(s3_filepath, read_f=np.load)\n",
"\n",
"if \"v3.0_raw\" in data:\n",
" codes = data[\"v3.0_raw\"]\n",
"elif \"v3.5_raw\" in data:\n",
" codes = data[\"v3.5_raw\"]\n",
"elif \"v4.0_raw\" in data:\n",
" codes = data[\"v4.0_raw\"]\n",
"else:\n",
" raise ValueError(\"No codes found\")\n",
"sem_codes = torch.tensor(codes[:, 0]).to(torch.long)\n",
"print(\"sem_codes.shape\", sem_codes.shape, sem_codes.shape[0] / 25)\n",
"\n",
"vae_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}_vae.npz\"\n",
"vae_data = read_from_s3(vae_filepath, read_f=np.load)\n",
"if \"vae_latents\" in vae_data:\n",
" infill_with_pad_latents = vae_data[\"vae_latents\"]\n",
"else:\n",
" raise ValueError(\"No vae codes found\")\n",
"print(\"vae_codes.shape\", infill_with_pad_latents.shape, infill_with_pad_latents.shape[0] / 25)"
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"s3://suno-data-uploads/studio/uploads/66636a2e-77da-4162-85f5-521916225898.npz\n",
"parent_sem_codes.shape torch.Size([5295]) 211.8\n",
"parent_vae_codes.shape (5395, 128) 215.8\n"
]
}
],
"source": [
"parent_s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{parent_gen_id}.npz\"\n",
"print(parent_s3_filepath)\n",
"parent_data = read_from_s3(parent_s3_filepath, read_f=np.load)\n",
"\n",
"if \"v3.0_raw\" in parent_data:\n",
" parent_codes = parent_data[\"v3.0_raw\"]\n",
"elif \"v3.5_raw\" in parent_data:\n",
" parent_codes = parent_data[\"v3.5_raw\"]\n",
"elif \"v4.0_raw\" in parent_data:\n",
" parent_codes = parent_data[\"v4.0_raw\"]\n",
"else:\n",
" raise ValueError(\"No codes found\")\n",
"parent_sem_codes = torch.tensor(parent_codes[:, 0]).to(torch.long)\n",
"print(\"parent_sem_codes.shape\", parent_sem_codes.shape, parent_sem_codes.shape[0] / 25)\n",
"\n",
"parent_vae_filepath = f\"s3://suno-data-uploads/studio/uploads/{parent_gen_id}_vae.npz\"\n",
"parent_vae_data = read_from_s3(parent_vae_filepath, read_f=np.load)\n",
"if \"vae_latents\" in parent_vae_data:\n",
" parent_latents = parent_vae_data[\"vae_latents\"]\n",
"else:\n",
" raise ValueError(\"No vae codes found\")\n",
"print(\"parent_vae_codes.shape\", parent_latents.shape, parent_latents.shape[0] / 25)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"parent_from_start_index 0\n",
"parent_from_end_index 0\n"
]
}
],
"source": [
"parent_from_start_index = vae_data.get(\"parent_from_start_index\", 0)\n",
"parent_from_end_index = vae_data.get(\"parent_from_end_index\", 0)\n",
"print(\"parent_from_start_index\", parent_from_start_index)\n",
"print(\"parent_from_end_index\", parent_from_end_index)"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"retrieved_latents_shape 627\n",
"total_latents.shape (6022, 128)\n"
]
}
],
"source": [
"retrieved_latents_shape = (\n",
" parent_from_start_index\n",
" + infill_with_pad_latents.shape[0]\n",
" + parent_from_end_index\n",
")\n",
"print(\"retrieved_latents_shape\", retrieved_latents_shape)\n",
"history_latents = parent_latents[:parent_from_start_index]\n",
"future_latents = parent_latents[-parent_from_end_index:]\n",
"total_latents = np.concatenate(\n",
" [history_latents, infill_with_pad_latents, future_latents], axis=0\n",
")\n",
"print(\"total_latents.shape\", total_latents.shape)"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"decoded_audio = decode_vae_to_audio(\n",
" torch.from_numpy(infill_with_pad_latents).to(torch.bfloat16)\n",
")\n",
"decoded_audio.play()"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"decoded_audio = decode_vae_to_audio(\n",
" torch.from_numpy(parent_latents).to(torch.bfloat16)\n",
")\n",
"decoded_audio.play()"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [],
"source": [
"try:\n",
" history_text = str(data[\"history_text\"])\n",
"except:\n",
" history_text = None"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [],
"source": [
"text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n",
"aligned_lyrics = json.loads(text_data)\n",
"# print(aligned_lyrics)\n",
"lyrics = \"\".join(w[\"word\"] for w in aligned_lyrics if \"word\" in w)"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [],
"source": [
"test_generation_config = DiffusionGenerationConfig(\n",
" audio=sem_codes[200:],\n",
" lyrics=lyrics,\n",
" history_latents=None,\n",
" history_lyrics=None,\n",
" steps=12,\n",
" text_cfg_coef=3.0,\n",
" seed=422,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"0it [00:00, ?it/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"cross_attn_cond: torch.Size([1, 3036, 2048]) tensor([[[-0.2910, -0.4062, 0.1562, ..., 0.0879, -0.1650, -0.8711],\n",
" [-0.0186, 0.0811, -0.5625, ..., 0.0272, 0.5820, 0.0474],\n",
" [ 0.1230, -0.4961, -0.1641, ..., 0.7109, -0.1719, -0.3457],\n",
" ...,\n",
" [ 0.0225, 0.0123, 0.0084, ..., -0.0090, 0.0310, -0.0298],\n",
" [ 0.0225, 0.0123, 0.0084, ..., -0.0090, 0.0310, -0.0298],\n",
" [ 0.0225, 0.0123, 0.0084, ..., -0.0090, 0.0310, -0.0298]]],\n",
" device='cuda:0', dtype=torch.bfloat16)\n",
"no_text_cross_attn_cond: torch.Size([1, 3036, 2048]) tensor([[[-0.1543, 0.3242, -0.1973, ..., -0.3652, -0.5078, 0.2412],\n",
" [-0.1543, 0.3242, -0.1973, ..., -0.3652, -0.5078, 0.2412],\n",
" [-0.1543, 0.3242, -0.1973, ..., -0.3652, -0.5078, 0.2412],\n",
" ...,\n",
" [ 0.0225, 0.0123, 0.0084, ..., -0.0090, 0.0310, -0.0298],\n",
" [ 0.0225, 0.0123, 0.0084, ..., -0.0090, 0.0310, -0.0298],\n",
" [ 0.0225, 0.0123, 0.0084, ..., -0.0090, 0.0310, -0.0298]]],\n",
" device='cuda:0', dtype=torch.bfloat16)\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"1it [00:01, 1.65s/it]\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"gen_audio = generate(test_generation_config)\n",
"gen_audio.play()"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [],
"source": [
"test_generation_config_hist = DiffusionGenerationConfig(\n",
" audio=sem_codes[200:],\n",
" lyrics=lyrics,\n",
" history_latents=torch.ones(750, 128), # torch.from_numpy(parent_latents)[:3000],\n",
" history_lyrics=\"\",\n",
" steps=12,\n",
" text_cfg_coef=3.0,\n",
" seed=422,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"0it [00:00, ?it/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loaded history latents of shape torch.Size([1, 750, 128])\n",
"History latents: tensor([[[1., 1., 1., ..., 1., 1., 1.],\n",
" [1., 1., 1., ..., 1., 1., 1.],\n",
" [1., 1., 1., ..., 1., 1., 1.],\n",
" ...,\n",
" [1., 1., 1., ..., 1., 1., 1.],\n",
" [1., 1., 1., ..., 1., 1., 1.],\n",
" [1., 1., 1., ..., 1., 1., 1.]]])\n",
"cross_attn_cond: torch.Size([1, 3036, 2048]) tensor([[[-0.2910, -0.4062, 0.1562, ..., 0.0879, -0.1650, -0.8711],\n",
" [-0.0186, 0.0811, -0.5625, ..., 0.0272, 0.5820, 0.0474],\n",
" [ 0.1230, -0.4961, -0.1641, ..., 0.7109, -0.1719, -0.3457],\n",
" ...,\n",
" [-0.1709, -0.0071, -0.8359, ..., -0.6250, 0.1177, -0.7305],\n",
" [-0.1709, -0.0071, -0.8359, ..., -0.6250, 0.1177, -0.7305],\n",
" [-0.1709, -0.0071, -0.8359, ..., -0.6250, 0.1177, -0.7305]]],\n",
" device='cuda:0', dtype=torch.bfloat16)\n",
"no_text_cross_attn_cond: torch.Size([1, 3036, 2048]) tensor([[[-0.1543, 0.3242, -0.1973, ..., -0.3652, -0.5078, 0.2412],\n",
" [-0.1543, 0.3242, -0.1973, ..., -0.3652, -0.5078, 0.2412],\n",
" [-0.1543, 0.3242, -0.1973, ..., -0.3652, -0.5078, 0.2412],\n",
" ...,\n",
" [-0.1709, -0.0071, -0.8359, ..., -0.6250, 0.1177, -0.7305],\n",
" [-0.1709, -0.0071, -0.8359, ..., -0.6250, 0.1177, -0.7305],\n",
" [-0.1709, -0.0071, -0.8359, ..., -0.6250, 0.1177, -0.7305]]],\n",
" device='cuda:0', dtype=torch.bfloat16)\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"1it [00:01, 1.51s/it]\n"
]
},
{
"data": {
"text/html": [
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"gen_audio_hist = generate(test_generation_config_hist)\n",
"gen_audio_hist.play()"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [
{
"ename": "NameError",
"evalue": "name 'BREAKPOINT' is not defined",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[22], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mBREAKPOINT\u001b[49m\n",
"\u001b[0;31mNameError\u001b[0m: name 'BREAKPOINT' is not defined"
]
}
],
"source": [
"BREAKPOINT"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"test_generation_config = DiffusionGenerationConfig(\n",
" audio=sem_codes,\n",
" lyrics=lyrics,\n",
" history_latents=torch.from_numpy(parent_latents),\n",
" history_lyrics=history_text,\n",
" steps=16,\n",
" text_cfg_coef=3.0,\n",
" seed=42,\n",
")\n",
"gen_audio = generate(test_generation_config)\n",
"gen_audio.play()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"test_generation_config = DiffusionGenerationConfig(\n",
" audio=sem_codes,\n",
" lyrics=lyrics,\n",
" history_latents=torch.from_numpy(result.history_latents),\n",
" history_lyrics=history_text,\n",
" steps=16,\n",
" text_cfg_coef=3.0,\n",
" seed=412232,\n",
")\n",
"gen_audio = generate(test_generation_config)\n",
"gen_audio.play()"
]
},
{
"cell_type": "code",
"execution_count": 54,
"metadata": {},
"outputs": [],
"source": [
"if len(sem_codes.shape) == 1:\n",
" test_sem_codes = sem_codes[None]\n",
"assert test_sem_codes.shape[0] == 1 # single channel\n"
]
},
{
"cell_type": "code",
"execution_count": 58,
"metadata": {},
"outputs": [],
"source": [
"code_stream = generate_stream(test_sem_codes.T, test_generation_config)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"latents = []\n",
"for chunk in code_stream:\n",
" latents.append(chunk)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"latents[0].shape\", latents[0].shape)\n",
"vae_codes = np.concatenate(latents, axis=0)\n",
"print(\"vae_codes.shape\", vae_codes.shape)"
]
},
{
"cell_type": "code",
"execution_count": 65,
"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.15"
}
},
"nbformat": 4,
"nbformat_minor": 2
}