{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "fed7c348", "metadata": {}, "outputs": [], "source": [ "import os\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"" ] }, { "cell_type": "code", "execution_count": null, "id": "c046f7b4", "metadata": {}, "outputs": [], "source": [ "from suno_utils.diffusion import generation as diffusion_gen\n", "from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job\n", "from suno_utils.tasks.dac_vae_fixed_25hz import decode_stream_to_full_audio\n", "import torch\n", "from suno_utils.audio import Audio" ] }, { "cell_type": "code", "execution_count": null, "id": "a718fb2b", "metadata": {}, "outputs": [], "source": [ "diffusion_gen.preload_models(\n", " # dit_model_filepath=\"/app/suno/victor/2b_prefix_ft.pt\",\n", " # codec_filepath=\"/app/suno/victor/25hz_vae_peaq_kl_0.005.pth\",\n", " # dit_model_filepath=\"/app/suno/checkpoints/2025-03-24_22-31-58_s4919/last_ckpt_infer.pt\", # finetune infill fixed\n", " # dit_model_filepath=\"/app/suno/checkpoints/2025-03-28_18-46-17_s894/last_ckpt_infer.pt\", # base model\n", " dit_model_filepath=\"/app/suno/data/dpo/models/diff_v2_2b_2mil_ft_infill_20250421_v2.pt\", # base model\n", " codec_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "cf696991", "metadata": {}, "outputs": [], "source": [ "import json\n", "import numpy as np\n", "from suno_utils.utils.s3 import read_from_s3\n", "\n", "# gen_id = \"99bffa17-7e59-47b4-a048-5528cbda05d5\" # sister\n", "# gen_id = \"081d73c4-7805-4212-9c80-8db1137ca3c4\" # friends\n", "# gen_id = \"562f762d-6ced-4080-9af1-910ee3d0a5dc\" # something real\n", "# gen_id = \"23c15c62-494d-422d-8a60-8b0454044322\" # rubber duck\n", "# gen_id = \"4b140a9e-964b-422c-85b5-5861ad1a9d38\" # once\n", "# gen_id = \"7b214347-fa38-4e9b-96f4-f7ec65adea45\" # rock n roll\n", "# gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\" # stone\n", "# gen_id = \"05ffc918-0bde-4418-8828-5105cd5717e1\" # sara bad click\n", "# gen_id = \"7f5e7819-b7a1-449d-a76c-9c38950c4274\" # canada\n", "# gen_id = \"27135940-957f-4152-9e40-72e5c10a46bd\" # drunk\n", "# gen_id = \"caf1d7ae-ce81-4c75-b737-fd41e47a38ba\" # linger\n", "# gen_id = \"ac5aa29a-f18f-406f-a904-c47493f1beec\" # linger vox\n", "# gen_id = \"b04719c3-012a-4a4a-9b02-3cb2875d1744\" # courtney\n", "# gen_id = \"7407c7b3-8aa2-4f0a-82e8-fa19933f6e94\" # opera\n", "# gen_id = \"5b9624ef-b416-4095-a5f3-4945f09bba86\" # butterflies\n", "# gen_id = \"e757aa19-5c33-4945-aaa1-1519b7890f0b\" # bubblegum\n", "# gen_id = \"7f774078-1672-4858-a37f-acad373c5a84\" # golden\n", "# gen_id = \"1cee79db-a0ee-44e4-9646-d312cd620997\" # chinese\n", "# gen_id = \"e404a4f0-64e0-423a-93d9-486b64c606ce\" # mikey click\n", "# gen_id = \"ebb4ce1d-8f79-4a12-8aa8-858bcc65b44a\" # optimistic\n", "\n", "# v2 codec clips\n", "gen_id = \"5ed9d01d-6e8b-41a0-968d-70ebbe51ba24\" #\n", "# gen_id = \"f497991a-ad42-45be-89f9-5fc324252f24\"\n", "\n", "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", "vae_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}_vae.npz\"\n", "print(s3_filepath)\n", "data = read_from_s3(s3_filepath, read_f=np.load)\n", "vae_data = read_from_s3(vae_filepath, read_f=np.load)\n", "audio = Audio.from_s3(mp3_filepath, n_channels=2)\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", "\n", "os.system(f\"aws s3 cp s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json text_data.json\")\n", "text_data = open(\"text_data.json\", \"r\", encoding=\"utf-8\").read()\n", "aligned_lyrics = json.loads(text_data)\n", "\n", "tags = \"pop\"\n", "\n", "lyrics = \"\"\n", "for elem in aligned_lyrics:\n", " if \"word\" in elem:\n", " lyrics += elem[\"word\"]" ] }, { "cell_type": "code", "execution_count": null, "id": "47656293", "metadata": {}, "outputs": [], "source": [ "# codes = torch.from_numpy(np.load(\"degrade.npy\")).unsqueeze(-1)\n", "# codes.shape" ] }, { "cell_type": "code", "execution_count": null, "id": "c0c34cd3", "metadata": {}, "outputs": [], "source": [ "engine = UpsampleEngine(min_chunk_size=25 * 30)" ] }, { "cell_type": "code", "execution_count": null, "id": "37b3fcef", "metadata": {}, "outputs": [], "source": [ "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n", " audio=audio,\n", " steps=12,\n", " lyrics=lyrics,\n", " tags=tags,\n", " text_cfg_coef=2.0,\n", " codec_scale_factor=0.4,\n", " scale_ctx_vector=True,\n", " noise_ctx_level=0.5,\n", " noise_ctx_pad_len=25,\n", " drop_semantic_tokens=False,\n", ")\n", "request = Request(\n", " id=\"dummy\",\n", " generation_config=gen_cfg,\n", " tokens=codes[:, 0],\n", " input_tokens_finished=True,\n", ")\n", "result = engine.run_request(request)" ] }, { "cell_type": "code", "execution_count": null, "id": "d0aeca0a", "metadata": {}, "outputs": [], "source": [ "vae_latents = torch.concat(result.vae_latents)\n", "out = decode_stream_to_full_audio(vae_latents)\n", "out.play()" ] }, { "cell_type": "code", "execution_count": null, "id": "b5bad6bc", "metadata": {}, "outputs": [], "source": [ "print(vae_latents.shape)" ] }, { "cell_type": "markdown", "id": "55e4f970", "metadata": {}, "source": [ "# Infill" ] }, { "cell_type": "code", "execution_count": null, "id": "1194f256", "metadata": {}, "outputs": [], "source": [ "list(vae_data.keys())\n", "ctx_vae_latents = vae_data[\"vae_latents\"]\n", "ctx_vae_latents.shape\n", "print(ctx_vae_latents.shape)" ] }, { "cell_type": "code", "execution_count": null, "id": "8628f293", "metadata": {}, "outputs": [], "source": [ "start_idx = 25 * 30\n", "dur = 25 * 5 # can change this to anything\n", "gen_cfg = diffusion_gen.DiffusionGenerationConfig(\n", " audio=audio,\n", " steps=32,\n", " lyrics=lyrics.replace(\"Things\", \"Sunos\"),\n", " tags=tags,\n", " text_cfg_coef=4.0,\n", " codec_scale_factor=0.4,\n", " scale_ctx_vector=True,\n", " # infill_prefix_latents=torch.from_numpy(ctx_vae_latents[start_idx : start_idx + 25 * 5, :]),\n", " generation_history_latents=torch.from_numpy(\n", " ctx_vae_latents[max(0, start_idx - 5 * 25) : start_idx, :]\n", " ),\n", " infill_suffix_latents=torch.from_numpy(\n", " ctx_vae_latents[start_idx + dur : start_idx + dur + 25 * 5, :]\n", " ),\n", " noise_ctx_level=0.0,\n", " drop_semantic_tokens=False,\n", ")\n", "request = Request(\n", " id=\"dummy\",\n", " generation_config=gen_cfg,\n", " tokens=codes[start_idx : start_idx + dur, 0],\n", " input_tokens_finished=True,\n", ")\n", "result = engine.run_request(request)" ] }, { "cell_type": "code", "execution_count": null, "id": "7dc575bb", "metadata": {}, "outputs": [], "source": [ "vae_latents = torch.concat(result.vae_latents)\n", "decode_stream_to_full_audio(vae_latents).play()" ] }, { "cell_type": "code", "execution_count": null, "id": "9f6bc7c4", "metadata": {}, "outputs": [], "source": [ "full_vae_latents = torch.from_numpy(ctx_vae_latents.copy())\n", "full_vae_latents[start_idx : start_idx + dur] = vae_latents\n", "decode_stream_to_full_audio(full_vae_latents[: start_idx + dur + 25 * 5]).play()" ] }, { "cell_type": "code", "execution_count": null, "id": "34f230df", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }