{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# autoreload\n", "%load_ext autoreload\n", "%autoreload 2\n", "\n", "import os\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n", "import json\n", "import glob\n", "import torch\n", "import funcy\n", "import IPython\n", "import torchaudio\n", "import numpy as np\n", "from stable_audio_tools.inference.generation import (\n", " upsample_diffusion_from_semantic_and_text,\n", " upsample_diffusion_from_semantic_and_text_with_phonemes\n", ")\n", "from dac.model.dac4 import DAC\n", "\n", "from stable_audio_tools.interface.gradio import load_model\n", "from stable_audio_tools.models.utils import apply_normalization\n", "\n", "from suno_utils.utils.s3 import read_from_s3\n", "#from suno_utils.tasks.dac_2c_12cb import load_model as load_vae_model\n", "\n", "# VAE\n", "from suno_utils.tasks.dac_vae_peaq import (\n", " preload_models as preload_vae_models,\n", " load_model as load_vae_model,\n", " encode as vae_encode,\n", " decode as vae_decode,\n", ")\n", "\n", "# MERT\n", "from suno_utils.tasks.mert_25 import (\n", " preload_models as preload_semantic_models,\n", " encode as semantic_encode,\n", " encode_files as semantic_encode_files,\n", ")\n", "\n", "_ = preload_semantic_models(\n", " checkpoint_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n", " centroids_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\",\n", " device=\"cuda\",\n", ")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# 2min ckpt\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_epoch=14-step=580000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_epoch=15-step=620000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_no_val_step=740000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_vlg_scale=2.5_qk_norm_adaln_2min_step=130000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_1node_epoch=22-step=1200000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_vlg_scale=2.5_qk_norm_adaln_2min_step=190000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_2min_step=90000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_legacy_2min_step=220000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_2min_step=250000.ckpt\"\n", "\n", "# 25hz ckpt\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_no_val_step=1400000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_no_val_step=1440000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_no_val_step=1500000.ckpt\"\n", "ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_no_val_step=2400000.ckpt\"\n", "\n", "# 100hz ckpt\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_2min_step=350000.ckpt\"\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_2min_step=600000.ckpt\"\n", "\n", "# 2min ckpt local attent\n", "#ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_local_epoch=9-step=740000.ckpt\"\n", "\n", "config_dir = \"/home/christian/code/neon/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio\"\n", "\n", "ckpt_name = os.path.basename(ckpt_path).replace(\".ckpt\", \"\")\n", "\n", "# check if ckpt exists\n", "if not os.path.exists(ckpt_path):\n", " raise ValueError(f\"Checkpoint not found: {ckpt_path}\")\n", "\n", "vae_frame_rate = 25\n", "device = \"cuda:0\"\n", "\n", "if vae_frame_rate == 25:\n", " # load VAE model\n", " checkpoint_filepath = \"s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth\"\n", "elif vae_frame_rate == 100:\n", " checkpoint_filepath = \"s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth\"\n", "else:\n", " raise ValueError(f\"VAE frame rate not supported: {vae_frame_rate}\")\n", "\n", "#checkpoint_filepath = \"/home/christian/code/christian/checkpoints/25hz_vae_peaq_kl_0.005.pth\"\n", "load_f = funcy.partial(torch.load, map_location=\"cpu\")\n", "\n", "if checkpoint_filepath.startswith(\"s3://\"):\n", " sd = read_from_s3(checkpoint_filepath, read_f=load_f)\n", "else:\n", " sd = load_f(checkpoint_filepath)\n", "\n", "sd[\"metadata\"][\"kwargs\"] = {\n", " k: v\n", " for k, v in sd[\"metadata\"][\"kwargs\"].items()\n", " if k in DAC.__init__.__code__.co_varnames\n", "}\n", "vae_model = DAC(**sd[\"metadata\"][\"kwargs\"])\n", "vae_model.load_state_dict(sd[\"state_dict\"])\n", "vae_model.eval()\n", "vae_model.to(device)\n", "\n", "# setup config\n", "if \"diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_legacy_2min\" in ckpt_path:\n", " model_type = \"diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_legacy_2min\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm_legacy.json\"\n", " )\n", "elif \"diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_2min\" in ckpt_path:\n", " model_type = \"diffusion_semantic+text_100hz_lg_scale=2.5_qk_norm_2min\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm.json\"\n", " )\n", "elif \"diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_adaln_phonemes_2min\" in ckpt_path:\n", " model_type = \"diffusion_semantic+text_25hz_lg_scale=2.5_qk_norm_adaln_phonemes_2min\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes.json\"\n", " )\n", "elif \"semantic+text_48khz_vlg_scale=2.5_qk_norm_adaln\" in ckpt_path:\n", " model_type = \"semantic+text_48khz_vlg_scale=2.5_qk_norm_adaln\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_vlg_scale=2.5_qk_norm_adaln.json\"\n", " )\n", "elif \"semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_local\" in ckpt_path:\n", " model_type = \"semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_2min_local\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes_local.json\"\n", " )\n", "elif \"semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes\" in ckpt_path:\n", " model_type = \"semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5_qk_norm_adaln_phonemes.json\"\n", " )\n", "else:\n", " raise ValueError(f\"Model type not found: {ckpt_path}\")\n", "\n", "#model_type = \"semantic+text_100hz_1b\"\n", "#model_config_path = os.path.join(\n", "# config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5.json\"\n", "#)\n", "\n", "print(model_config_path)\n", "\n", "# load model from checkpoint\n", "if model_config_path is not None:\n", " # Load config from json file\n", " with open(model_config_path) as f:\n", " model_config = json.load(f)\n", "else:\n", " model_config = None\n", "\n", "for key, val in model_config.items():\n", " print(f\"{key}: {val}\")\n", "\n", "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n", "model, model_config = load_model(\n", " model_config,\n", " ckpt_path,\n", " # pretrained_name=pretrained_name,\n", " # pretransform_ckpt_path=pretransform_ckpt_path,\n", " # model_half=model_half,\n", " device=\"cuda\",\n", ")\n", "\n", "scale_factor = model_config[\"training\"][\"scale_factor\"]\n", "print(f\"scale_factor: {scale_factor}\")\n", "\n", "max_dur_s = 120\n", "#max_dur_s = 30.0\n", "semantic_n_tokens = int(25 * max_dur_s)\n", "n_vae_tokens = int(vae_frame_rate * max_dur_s)\n", "\n", "print(semantic_n_tokens, n_vae_tokens)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## From val set\n", "Load random example of semantic + text from val set for inference." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.text import read_jsonl\n", "\n", "vae_memmap_path = \"/app/suno/data/chirp_v4/vae_v5/data_vae_val.bin\"\n", "semantic_memmap_path = \"/app/suno/data/chirp_v4/vae_v5/data_semantic_val.bin\"\n", "\n", "vae_n_tokens = 36000\n", "semantic_n_tokens = 9000\n", "vae_dim = 128\n", "\n", "# load memmap and get semantic\n", "vae_data = np.memmap(vae_memmap_path, dtype=np.float16, mode=\"r\")\n", "vae_data = vae_data.reshape(-1, vae_n_tokens, vae_dim)\n", "vae_data = vae_data[:,:12000, :]\n", "\n", "# open semantic memmap\n", "semantic_data = np.memmap(semantic_memmap_path, dtype=np.uint16, mode=\"r\")\n", "semantic_data = semantic_data.reshape(-1, semantic_n_tokens, 1)\n", "semantic_data = semantic_data[:, :, 0]\n", "\n", "# open metas\n", "metas = read_jsonl(\"/app/suno/data/chirp_v4/vae_v5/metas_val.jsonl\")\n", "\n", "print(len(metas), vae_data.shape, semantic_data.shape)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# country - 95, 98\n", "# pop - 292\n", "# folk acoustic - 527\n", "# jazz - 322\n", "\n", "rand_idx = np.random.randint(len(metas))\n", "#rand_idx = 700\n", "meta = metas[rand_idx]\n", "tags = meta.get(\"tags\", [])\n", "lyrics = meta.get(\"text\", \"\")\n", "semantic_codes = torch.from_numpy(semantic_data[rand_idx].copy()).long().cuda()\n", "vae_latents = torch.from_numpy(vae_data[rand_idx].copy()).float().unsqueeze(0).cuda()\n", "n_vae_tokens = meta.get(\"n_vae_tokens\", None)\n", "padding_mask = torch.ones(vae_n_tokens)\n", "padding_mask[n_vae_tokens:] = 0\n", "padding_mask = padding_mask.unsqueeze(0).cuda()\n", "padding_mask = padding_mask.bool()\n", "print(rand_idx, meta.get(\"dataset\"))\n", "print(n_vae_tokens, vae_n_tokens)\n", "\n", "print(tags)\n", "print(lyrics)\n", "print(semantic_codes.shape, vae_latents.shape, padding_mask.shape)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "with torch.no_grad():\n", " audio = vae_model.decode(vae_latents.permute(0, 2, 1))[0].detach().cpu() \n", "audio /= audio.abs().max().clamp(1e-8)\n", "print(audio.mean())\n", "\n", "IPython.display.display(IPython.display.Audio(audio.numpy(), rate=48000))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Raw audio\n", "Read audio file from disk, semantic encode, have to manually provide lyrics and tags" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.tasks.mert_25 import preload_models\n", "\n", "mert_filepath = \"s3://suno-data/georg/models/semantic/mert_25.pt\"\n", "centroids_filepath = \"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n", "\n", "_ = preload_models(\n", " checkpoint_filepath=mert_filepath,\n", " centroids_filepath=centroids_filepath,\n", ")\n", "\n", "from suno_utils.tasks.mert_25 import encode, SAMPLE_RATE, EMBEDDING_RATE" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# reconstruction from generation (real audio)\n", "from suno_utils.tasks.data_loader import load_audio_mp\n", "\n", "audio_path = \"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/04 Fuckwithmeyouknowigotit.wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/Speak For Me [omeNvD8IddM].wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/09 Sounds Like Hallelujah.wav\"\n", "#audio_path = \"/home/christian/code/christian/outputs/halo_24.wav\"\n", "#audio_path = \"/home/christian/code/christian/outputs/Beyoncé - Halo (Lyrics) [wekDNXDWGjM].wav\"\n", "#audio_path = \"/home/christian/audio/bad-audio/bill-evans-intro.wav\"'\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/Crazy [CKTOvHw8qFM].wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/03 Your New Aesthetic.wav\"\n", "#audio_path = \"/home/christian/code/christian/outputs/linger-30s.wav\"\n", "#audio_path = \"/home/christian/code/christian/outputs/turkey.mp3\"\n", "#audio_path = \"/home/christian/code/neon/stable-audio-tools/outputs/Cash Cobain & J. Cole - Grippy (AUDIO) [9wn_ARfYMw4].mp3\"\n", "#audio_path = \"/home/christian/code/christian/outputs/Lost in Tokyo.m4a\"\n", "#audio_path = \"/home/christian/code/christian/outputs/i don't trust you (ben camp x suno).mp3\"\n", "\n", "audio_arrays_48khz = load_audio_mp(\n", " [audio_path],\n", " target_sample_rate=48000,\n", " normalize_volume=True,\n", " num_workers=12,\n", " n_channels=2,\n", ")\n", "\n", "if \"Fuck\" in audio_path:\n", " start_s = 45.0\n", " end_s = start_s + max_dur_s\n", "if \"Speak\" in audio_path: \n", " start_s = 25.0\n", " end_s = start_s + max_dur_s\n", "if \"i don't trust you\" in audio_path:\n", " start_s = 26.0\n", " end_s = start_s + max_dur_s\n", "else:\n", " start_s = 0.0\n", " end_s = start_s + max_dur_s\n", "\n", "print(start_s, end_s)\n", "audio_arrays_48khz = [a[:, int(start_s * 48000) : int(end_s * 48000)] for a in audio_arrays_48khz]\n", "\n", "print(audio_arrays_48khz[0].shape)\n", "IPython.display.display(IPython.display.Audio(audio_arrays_48khz[0].numpy(), rate=48000))\n", "\n", "# now load audio at 24khz\n", "audio_arrays_24khz = load_audio_mp(\n", " [audio_path],\n", " target_sample_rate=24000,\n", " normalize_volume=True,\n", " min_duration_s=30.0,\n", " num_workers=12,\n", " n_channels=2,\n", ")\n", "\n", "audio_arrays_24khz = [a[:, int(start_s * 24000) : int(end_s * 24000)] for a in audio_arrays_24khz]\n", "\n", "# semantic encode \n", "semantic_codes = encode([audio_arrays_24khz[0].mean(axis=0, keepdim=True)], SAMPLE_RATE, EMBEDDING_RATE)\n", "semantic_codes = np.array(semantic_codes)[0, :, 0]\n", "\n", "n_valid_semanitc_codes = semantic_codes.shape[0]\n", "#print(semantic_codes.shape)\n", "sementic_n_tokens = 3000\n", "if semantic_codes.shape[0] < sementic_n_tokens: # pad with 4000\n", " semantic_codes = np.pad(semantic_codes, (0, sementic_n_tokens - semantic_codes.shape[0]), mode=\"constant\", constant_values=4000)\n", "# # repeat the last code\n", "# #semantic_codes = np.pad(semantic_codes, (0, sementic_n_tokens - semantic_codes.shape[0]), mode=\"constant\", constant_values=semantic_codes[-1])\n", "semantic_codes = torch.from_numpy(semantic_codes).long().cuda() \n", "print(semantic_codes.shape)\n", "n_valid_semantic_tokens = semantic_codes.shape[0]\n", "\n", "\n", "if \"Don't Know Why\" in audio_path:\n", " lyrics = \"\"\"\n", "I waited 'til I saw the sun\n", "I don't know why I didn't come\n", "I left you by the house of fun\n", "I don't know why I didn't come\n", "I don't know why I didn't come\n", "When I saw the break of day\n", "I wished that I could fly away\n", "Instead of kneeling in the sand\n", "Catching tear-drops in my hand\n", "My heart is drenched in wine\n", "But you'll be on my mind forever\n", "Out across the endless sea\n", "I will die in ecstasy\n", "But I'll be a bag of bones\n", "Driving down the road alone\n", "My heart is drenched in wine\n", "But you'll be on my mind forever\n", "Something has to make you run\n", "I don't know why I didn't come\n", "I feel as empty as a drum\n", "I don't know why I didn't come\n", "I don't know why I didn't come\n", "I don't know why I didn't come\n", " \"\"\".strip()\n", " tags = [\"jazz\", \"pop\", \"female vocal\"]\n", "elif \"Dreams\" in audio_path:\n", " lyrics = \"[Verse 1] Now here you go again, you say you want your freedom, well who am I to keep you down?\"\n", " tags = [\"rock\", \"pop\"]\n", "elif \"Sounds Like\" in audio_path:\n", " lyrics = \"One, two, one two three four.\\n I'm just waiting on the sun\\n to close his eyes and call the night\\n so we can put all our differences aside\\n I'm just waiting on the moon, with his stars and all its gloom. \\nWe can watch it fall right back into place\\n So I won't keep my\"\n", " tags = [\"Rock\"]\n", "elif \"halo\" in audio_path:\n", " lyrics = \"\"\"\n", " [Verse 1]\n", "Remember those walls I built?\n", "Well, baby, they're tumblin' down\n", "And they didn't even put up a fight\n", "They didn't even make a sound\n", "I found a way to let you in\n", "But I never really had a doubt\n", "Standin' in the light of your halo\n", "I got my angel now\n", "\n", "[Pre-Chorus]\n", "It's like\n", " \"\"\".strip()\n", " tags = [\"Pop\", \"Ballad\", \"R&B\"]\n", "elif \"Halo\" in audio_path:\n", " tags = [\"Pop\", \"Ballad\", \"R&B\"]\n", " lyrics = \"\"\"\n", "[Pre-Chorus]\n", "It's like I've been awakened\n", "Every rule, I had you breakin'\n", "It's the risk that I'm takin'\n", "I ain't ever gonna shut you out\n", "\n", "[Chorus]\n", "Everywhere I'm lookin' now\n", "I'm surrounded by your embrace\n", "Baby, I can see your halo\n", "\"\"\".strip() \n", "elif \"Speak For Me\" in audio_path:\n", " tags = [\"Folk\", \"Rock\"]\n", " lyrics = \"\"\"\n", " Now the cover of a Rolling Stone\n", "Ain't the cover of a Rolling Stone\n", "And the music on my radio\n", "Ain't supposed to make me feel alone\n", "Would you liked to know\n", "Have to learn to let it go\n", "Show me something I can be\n", "Play a song that I can sing\n", "Make me feel as I am free\n", " \"\"\".strip()\n", "elif \"Crazy\" in audio_path:\n", " tags = [\"Pop\", \"Rock\"]\n", " lyrics = \"\"\"\n", " Crazy, I'm crazy for feeling so lonely\n", "I'm crazy, crazy for feeling so blue\n", "\n", "[Verse 2]\n", "I knew, you'd love me as long as you wanted\n", "And then someday, you'd leave me for somebody new\n", "\n", "[Verse 3]\n", " Worry, why do I let myself worry?\n", " Wondering, what in the world did I do?\"\"\".strip()\n", "elif \"Fuckwithmeyouknowigotit\" in audio_path:\n", " tags = [\"Hip-Hop\", \"Rap\"]\n", " lyrics = \"\"\"\n", " [Chorus: Rick Ross]\n", " Fuck with me, you know I got it\n", " Fuck with me, you know I got it\n", " Sexy bitch, I hope she 'bout it\n", " Come fuck with me, you know I got it\n", " Fuck with me, you know I got it\n", " Fuck with me, you know I got it\n", " Sexy bitch, I hope she 'bout it\n", " Come fuck with me, you know I got it\n", " \"\"\".strip()\n", "elif \"linger\" in audio_path:\n", " tags = [\"Pop\", \"Electronic\", \"EDM\", \"Dance\"]\n", " lyrics = \"\"\"\n", "If you, if you could return\n", "Don't let it burn\n", "Don't let it fade\n", "I'm sure I'm not being rude\n", "But it's just your attitude\n", "It's tearing me apart\n", "It's ruining every day\n", "\"\"\".strip()\n", "elif \"turkey\" in audio_path:\n", " tags = [\"pop\", \"rock\", \"indie\"]\n", " lyrics = \"\"\"\n", "[Verse]\n", "Gather 'round now, it's turkey time\n", "Catchin' up with all your kin and dimes\n", "Feastin' on that golden bird, so fine (so fine)\n", "Pass the gravy, pour some wine\n", "\n", "[Chorus]\n", "Turkey time, it's a country rhyme\n", "Grab your plate, fill it up, take your time\n", "And when the sun starts settin', we'll gather 'round\n", "Singin' songs, sippin' cider, dancin' all night sound (oooh)\n", "Turkey time, it's a country rhyme\n", "\n", "[Verse]\n", "Well, it's that time of year, with family near\n", "Gatherin' 'round the table, filled with cheer\n", "Turkey's roasting, gravy's flowin'\n", "Mashed potatoes piled high, it's Thanksiving time\n", "\n", "[Chorus]\n", "Turkey time, it's a country rhyme\n", "Grab your plate, fill it up, take your time\n", "And when the sun starts settin', we'll gather 'round\n", "Singin' songs, sippin' cider, dancin' all night sound (oooh)\n", "Turkey time, it's a country rhyme\n", "\n", "[Bridge] \n", "Turkey time is here again\n", "Turkey time is here again\n", "Turkey time is here again\n", "Turkey tiiiiime!\n", "\n", "[Chorus]\n", "Turkey time, it's a country rhyme\n", "Grab your plate, fill it up, take your time\n", "And when the sun starts settin', we'll gather 'round\n", "Singin' songs, sippin' cider, dancin' all night sound (oooh)\n", "Turkey time, it's a country rhyme\n", "\n", "Turkey time, it's a country rhyme\n", "\n", "[guitar, outro]\n", "\n", "[outro]\n", " \"\"\".strip()\n", "elif \"Grippy\" in audio_path:\n", " tags = [\"Rap\", \"Bass\", \"Pop\"]\n", " lyrics = \"\"\"\n", "She like my kick game\n", "And when you me, you don't kick game\n", "I put her front row at the Knick game\n", "Now she in my phone with a nickname\n", "It's, it's, it's, hmm (mmm)\n", "Grippy\n", "Yeah (Gvrlnd!, I'm rockin' with this one)\n", "Grippy, huh\n", "Grippy, I call her that 'cause it's grippy\n", "She thick in the hips, she a hippie\n", "And she thick in the lips, she gon' lick me\n", "And she sendin' a flick when she hit me\n", "With a kissy emoji, she miss me\n", "When she see me, she say she gon' strip me\n", "She gon' chew on this stick like it's Wrigley's\n", "Mm, yeah, believe it or not like Ripley's\n", "She said she was gay until I slayed, \n", "\"\"\"\n", "elif \"Tokyo\" in audio_path:\n", " tags = [\"Pop\", \"Electronic\", \"Dance\"]\n", " lyrics = \"\"\n", "elif \"i don't trust you\" in audio_path:\n", " tags = [\"rap\", \"pop\", \"hip-hop\"]\n", " lyrics = \"\"\"\n", "[verse]\n", "you didn't want me\n", "you only let me\n", "cause you didnt know how\n", "to set a boundary\n", "and now you pull back\n", "like a slingshot\n", "and my ego's about to catch rock\n", "\n", "[chorus]\n", "i dont trust you, oooh\n", "i dont trust you\n", "\n", "so sexy, you're so sus\n", "does yes mean yes, or should your legs be\n", " \"\"\".strip()\n", "else:\n", " tags = []\n", " lyrics = \"\"\n", "\n", "print(tags)\n", "print(lyrics)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## From npz\n", "Get npz with semantic codes from generations. Note: still need to grab tags and lyrics from db." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# given the s3 id of a generation, grab the npz, load the codes, and extract semantic codes\n", "from suno_utils.utils.s3 import read_from_s3\n", "\n", "#song_name = \"i-am-here-for-you\"#\n", "#song_name = \"on-the-ground\"\n", "#song_name = \"see-the-sky-again\"\n", "#song_name = \"stone\"\n", "#song_name = \"cat\"\n", "#song_name = \"ode-to-suno\"\n", "song_name = \"sister\"\n", "#song_name = \"turkey\"\n", "#song_name = \"hydra\"\n", "#song_name = \"friends\"\n", "#song_name = \"faded\"\n", "#song_name = \"here-for-you\"\n", "\n", "start_s = 0.0\n", "end_s = start_s + 120.0\n", "\n", "if song_name == \"cat\":\n", " gen_id = \"ee467d00-5813-4a74-9792-c9ae4a09d344\"\n", " tags = [\"electronic\"]\n", " lyrics = \"\"\"\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", "\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", " Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat Cat\n", " \"\"\".strip()\n", "elif song_name == \"stone\":\n", " gen_id = \"a5e2198a-f352-4abb-9a24-7f81b143ded3\"\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", " \"\"\".strip()\n", " tags = [\"rock\", \"indie pop\"]\n", " #tags = [\"metal\", \"metal\", \"metal\"]\n", " #tags = [\"country\", \"country\", \"country\"]\n", " #tags = [\"1950s\", \"barbershop\", \"pop\", \"vintage\"]\n", " start_s = 0.0\n", " end_s = start_s + 30.0\n", "elif song_name == \"jack\":\n", " gen_id = \"4fc20f50-866a-4066-a7ce-1bd406b8ccdb\"\n", " lyrics = \"\"\"\n", " (Verse 1)\n", " In the kitchen spotlight, when the dinner bell rings,\n", " There’s a grey tabby cat who’s the king of all things,\n", " His name is Jack, but don’t be deceived,\n", " This little bottomfeeder’s got tricks up his sleeve.\n", "\n", " (Chorus)\n", " Oh, Jack, the bottomfeeder’s got the groove,\n", " Every meal’s a chance for him to prove,\n", " He’ll swipe a bite, and lick the plate,\n", " He’s the master of the food debate.\n", " \"\"\".strip()\n", " tags = [\"broadway\", \"musical\", \"pop\", \"orchestral\"]\n", "elif song_name == \"i-am-here-for-you\":\n", " gen_id = \"a7199702-104f-42bd-b072-8b82c015551f\"\n", " lyrics = \"\"\"\n", " [Intro]\n", " Hey how are you feeling?\n", " Are you doing alright?\n", "\n", " [Instrumental]\n", "\n", " [Verse]\n", " When the night is long\n", " And you're feeling down\n", " I'm right by your side\n", " I am here for you\n", " (Just hold on tight)\n", "\n", " [Guitar interlude]\n", " \"\"\".strip()\n", " tags = [\"Blues\", \"romantic\", \"soft guitar\",\" gentle female vocals\"]\n", " tags = [\"pop\", \"electronic\", \"rock\"]\n", " start_s = 0\n", " end_s = start_s + 30.0\n", "elif song_name == \"see-the-sky-again\":\n", " gen_id = \"bae5eec7-652d-4e01-8145-af4e693b31d5\"\n", " lyrics = \"\"\"\n", "Ohhhh ohooo\n", "When I see the sky\n", "That moment when there were no clouds\n", "If you are the reason for my curses\n", "That moment\n", "Yes, at that moment\n", "Ohhhh ohooo\n", "Will I ever see the sky again?\n", "While looking into your eyes\n", "I'm drowning\n", "I'm drowning\n", "I'm drowning again\n", "Ohhhh ohoo\n", "Do apologies have any value?\n", "or pleas\n", "(It doesn't matter anymore)\n", "Every tear that flows from your eyes\n", "It pulls me deeper\n", "I'm looking for hope, but I can't find it\n", "I'm lost in the dark\n", "And every night it's the same nightmare\n", "I think of you again\n", "Ohhhh ohoo\n", "Maybe one day\n", "I'll be free from these chains\n", "And I can see the sky again\n", "My eyes are filled with tears\n", "But is there still hope?\n", "\"\"\".strip()\n", " tags = [\"soft rock\", \"pop\"]\n", "elif song_name == \"on-the-ground\":\n", " gen_id = \"3c0f052c-1500-4e54-970f-481db96739de\"\n", " tags = [\"Dance punk\", \"pop punk\", \"female singer\", \"2000s\", \"catchy\", \"rock\"]\n", " lyrics = \"\"\"\n", " [Alternate punk version of 'On the ground'.]\n", "\n", " (One, Two, Three, Four)\n", "\n", " [Intro]\n", " You’re all just starting,\n", " I’m already there,\n", " I’m stuck in this state, (going nowhere)\n", " To stop, to have it all gone\n", " Being just me, would feel so wrong\n", "\n", " (Can’t keep this up)\n", " \"\"\".strip()\n", "elif song_name == \"ode-to-suno\":\n", " gen_id = \"f6f46792-864a-41bc-b7c4-3aaf1711d498\"\n", " lyrics = \"\"\"\n", " Verse 1)\n", " I asked Suno AI to write a song for me,\n", " But what came out wasn't what I hoped it would be.\n", " Tried to capture my heart, my soul, my vibe,\n", " But it gave what I didn't describe.\n", "\n", " (Chorus)\n", " Oh Suno AI, what happened to your flow?\n", " Your words are off-key, your melody's low.\n", " You tried your best, but it's plain to see,\n", " Your songs just ain't cutting it for me.\"\"\"\n", " tags = [\"blues\", \"soul\", \"rock\"]\n", "elif song_name == \"sister\":\n", " gen_id = \"99bffa17-7e59-47b4-a048-5528cbda05d5\"\n", " lyrics = \"\"\"\n", "Verse]\n", "She grew up in the shadow of the steeple,\n", "Small town roads where the night stands still.\n", "She heard the Lord in the whispers of the people,\n", "Calling her heart to follow His will.\n", "\n", "[Verse 2]\n", "Her mama cried but knew deep down,\n", "That little girl was Heaven-sent.\n", "She traded her dresses for a plain white gown,\n", "And left behind the life she’d spent.\n", "\n", "[Chorus]\n", "Sister of the Light, with rosary beads,\n", "Praying through the night, fulfilling good deeds.\n", "She's found her calling, in the chapel's glow,\n", "A beacon of grace, wherever she goes.\n", "\n", "[Verse 3]\n", "She helps the lost and feeds the hungry,\n", "In a world that’s often cold and unkind.\n", "Her hands are worn but her spirit’s humbly,\n", "Carrying love of the most divine kind.\n", "\n", "[Verse 4]\n", "Sundays in the choir, hymns touch the sky,\n", "A voice pure as the morning dew.\n", "With faith so deep, she makes angels sigh,\n", "Her prayers reach heaven like they’re brand new.\n", "\n", "[Chorus]\n", "Sister of the Light, with rosary beads,\n", "Praying through the night, fulfilling good deeds.\n", "She's found her calling, in the chapel's glow,\n", "A beacon of grace, wherever she goes.\n", " \"\"\".strip()\n", " tags = [\"rock\", \"pop\", \"indie\"]\n", "elif song_name == \"each-page\":\n", " gen_id = \"e815230b-ab9b-4894-bb15-031d6d7ff37f\"\n", " tags = [\"emo\", \"pop\", \"acoustic\", \"melodic\", \"syncopated\"]\n", " lyrics = \"\"\"\n", "[Verse]\n", "All my tales fade\n", "Lost in the night\n", "Pages turn gray\n", "Out of sight\n", "\n", "[Verse 2]\n", "Whispers so soft\n", "Ghosts of my mind\n", "Stories I’ve lost\n", "No one will find\n", "\n", "[Chorus]\n", "Unread I stay\n", " \"\"\".strip()\n", "elif song_name == \"turkey\":\n", " gen_id = \"f5fc4429-b62f-4121-b73b-6028c81879c6\"\n", " tags = [\"pop\", \"rock\", \"indie\"]\n", " lyrics = \"\"\"\n", "[Verse]\n", "Gather 'round now, it's turkey time\n", "Catchin' up with all your kin and dimes\n", "Feastin' on that golden bird, so fine (so fine)\n", "Pass the gravy, pour some wine\n", "\n", "[Chorus]\n", "Turkey time, it's a country rhyme\n", "Grab your plate, fill it up, take your time\n", "And when the sun starts settin', we'll gather 'round\n", "Singin' songs, sippin' cider, dancin' all night sound\n", " \"\"\".strip()\n", "elif song_name == \"hydra\":\n", " gen_id = \"271a6cad-8b4f-446c-8545-8f4daa8120af\"\n", " tags = [\"metal\"]\n", " lyrics = \"\"\"\n", "Silent descent through azure mist\n", "Our ship touches down on Neptune's crust\n", "Unaware of the myth that exists\n", "In this world of frozen dust\n", "\n", "Nine heads emerge from icy depths\n", "Silken twine of tentacles reach\n", "Blinding everyone with frost-breath\n", "Our mission now beyond our reach\n", "\"\"\".strip()\n", " start_s = 20.0\n", " end_s = start_s + 30.0\n", "elif song_name == \"friends\":\n", " gen_id = \"081d73c4-7805-4212-9c80-8db1137ca3c4\"\n", " tags = [\"Rock\"]\n", " lyrics = \"\"\"\n", "What would you do if I lost my tune\n", "Would you laugh and leave me here?\n", "Lend your voice and let's sing\n", "I'll try not to get out of the melody\n", "Oh, baby, I'm going to make it\n", "(With a little help from my friends)\n", "All I need are my friends\n", "(Try it with a little help from my friends)\n", "I said I want to feel alright, I'm going to\n", "(Fly with a little help from my friends)\n", "What do I do when I'm feeling lonely\n", "(Do you worry about being alone?)\n", "How I feel at the end of the day\n", "(Are you sad to be alone?)\n", "I hope you don't feel it too\n", "Oh, baby, I'm going to make it\n", "(With a little help from my friends)\n", "All I need are my friends\n", "(Try it with a little help from my friends)\n", "I said I want to feel alright, I'm going to\n", "(Fly with a little help from my friends)\n", "(Do you need somebody?)\n", "I need someone to love\n", "(Could it be anyone?)\n", "All I need is someone who understands my way, yes\n", "(Do you need somebody?)\n", "I need someone to love\n", "(Could it be anyone?)\n", "All I need is someone who understands my way, yes\n", "Refrain:\n", "Oh, baby, I'm going to make it\n", "(With a little help from my friends)\n", "All I need are my friends\n", "(Try it with a little help from my friends)\n", "I said I want to feel alright, I'm going to\n", "(Fly with a little help from my friends)\n", "I'm going to keep trying\n", "(Fly with a little help from my friends)\n", "I'm going to make it with my friends\"\"\".strip()\n", " start_s = 0.0\n", " end_s = start_s + max_dur_s\n", "elif song_name == \"faded\":\n", " gen_id = \"55376741-0412-41d0-bffa-e7e2b290f381\"\n", " tags = [\"Pop\", \"Reverb\", \"Indie\"]\n", " lyrics = \"\"\"\n", "[Verse]\n", "We danced in the sunlight\n", "Whispers in the breeze\n", "Now those days are shadows\n", "Just memories\n", "\n", "[Verse 2]\n", "Coffee in the morning\n", "\"\"\"\n", " start_s = 0.0\n", " end_s = start_s + 30.0\n", "elif song_name == \"here-for-you\":\n", " gen_id = \"84deb8ef-8222-4486-af7b-e13d216e5454\"\n", " tags = [\"Pop\"]\n", " lyrics = \"\"\"\n", "ain't no worries when your with me,\n", "there ain't no trouble here to see,\n", "\n", "[guitar riff]\n", "\n", "[chorus]\n", "dont you worry my love,\n", "dont you worry no more,\n", "i'll be holding your hand,\n", "as the sea comes to shore,\n", " \"\"\".strip()\n", " start_s = 26.0\n", " end_s = start_s + 30.0\n", "else:\n", " start_s = 0.0\n", " end_s = start_s + 30.0\n", "\n", "start_s = 0.0\n", "end_s = 120.0\n", "\n", "start_frame = int(start_s * 25)\n", "end_frame = int(end_s * 25)\n", "\n", "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n", "print(s3_filepath)\n", "data = read_from_s3(s3_filepath, read_f=np.load)\n", "\n", "# tony\n", "if False:\n", " start_s = 90.0\n", " end_s = start_s + 30.0\n", " s3_filepath = \"/home/tony/Work/gpt_samples/30b_t3_v1_repro7/ts_pop_0.npz\"\n", " data = np.load(s3_filepath)\n", " codes = data[\"codes\"]\n", " tags = [\"pop country\", \"female vocal\", \"fast tempo\"]\n", " lyrics = \"\"\"\n", " Hey, whatcha doing with a girl like that?\n", "\n", " [Pre-Chorus]\n", " She wears high heels, I wear sneakers\n", " She's Cheer Captain and I'm on the bleachers\n", " Dreaming 'bout the day when you wake up and find\n", " That what you're looking for has been here the whole time\n", "\n", " [Chorus]\n", " If you could see that I'm the one who understands you\n", " Been here all along, so why can't you see?\n", " You belong with me\"\"\".strip()\n", " mp3_filepath = \"/home/tony/Work/gpt_samples/30b_t3_v1_repro7/ts_pop_0_seed42.mp3\"\n", " mp3_data, mp3_sr = torchaudio.load(mp3_filepath)\n", "\n", "\n", "# also read the mp3\n", "mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.mp3\"\n", "mp3_data, mp3_sr = read_from_s3(mp3_filepath, read_f=torchaudio.load)\n", "\n", "start_sample = int(start_s * 48000)\n", "end_sample = int(end_s * 48000)\n", "\n", "audio_arrays_48khz = [mp3_data[...,start_sample:end_sample].cuda()]\n", "\n", "#text_data = read_from_s3(f\"s3://suno-data-uploads/studio/uploads/{gen_id}_hoot.json\")\n", "#text_data = json.loads(text_data)\n", "\n", "#print(text_data)\n", "#print(data.keys())\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", "print(codes.shape)\n", "semantic_codes = codes[start_frame:end_frame, 0].astype(np.uint16)\n", "semantic_codes = torch.from_numpy(semantic_codes).long().cuda()\n", "print(song_name, semantic_codes.shape)\n", "n_valid_semantic_tokens = semantic_codes.shape[0]\n", "\n", "# encode audio with vae\n", "#with torch.no_grad():\n", "# latents = vae_model.encode(audio_arrays_48khz[0].unsqueeze(0).cuda())[\"z\"].detach() \n", "\n", "#print(latents.shape)\n", "\n", "#latents *= scale_factor\n", "\n", "IPython.display.display(IPython.display.Audio(audio_arrays_48khz[0].cpu().numpy(), rate=mp3_sr))\n", "print(tags)\n", "print(lyrics)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Inference\n", "\n", "Run diffusion inference with operational CFG." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_steps = 256\n", "if vae_frame_rate == 25:\n", " n_tokens_memmap_used = n_valid_semantic_tokens\n", " n_tokens_memmap = semantic_n_tokens\n", "else:\n", " n_tokens_memmap_used = n_valid_semantic_tokens * 4\n", " n_tokens_memmap = semantic_n_tokens * 4\n", "#mask = torch.ones(n_tokens_memmap).bool().cuda()\n", "#mask[n_tokens_memmap_used:] = False\n", "\n", "print(semantic_codes.shape)\n", "if semantic_codes.shape[0] < semantic_n_tokens: # pad with 4000\n", " semantic_codes = torch.cat([semantic_codes, semantic_codes.new_full((semantic_n_tokens - semantic_codes.shape[0],), 4000)])\n", "print(semantic_codes.shape)\n", "\n", "\n", "out_dir = f\"outputs/val-{ckpt_name}\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "semantic_codes_pad = (torch.ones(750) * 4000).long().cuda()\n", "\n", "cfg_scales = [2.0, 4.0, 6.0, 8.0]\n", "seeds = [np.random.randint(0, 2**32 - 1) for n in range(1)]\n", "\n", "\n", "for seed in seeds:\n", " for n in range(len(cfg_scales)):\n", " print(cfg_scales[n], seed)\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic_and_text(\n", " model,\n", " semantic_codes[:3000],\n", " [\"pop\"],\n", " lyrics,\n", " steps=num_steps,\n", " cfg_scale=cfg_scales[n],\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " seed=seed,\n", " sampler_type=\"dpmpp-2m-sde\",\n", " mask=None,\n", " compile=False,\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n", " pred_zq /= scale_factor\n", " # Wpred_zq = torch.randn_like(pred_zq)\n", " print(\"pred_zq\", pred_zq.shape)\n", "\n", " pred_zq = pred_zq[...,:n_tokens_memmap_used]\n", "\n", " pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() \n", " print(pred_audio.abs().max())\n", " pred_audio /= pred_audio.abs().max().clamp(1e-8)\n", " print(pred_audio.mean())\n", "\n", " # save audio\n", " pred_audio_filepath = os.path.join(out_dir, f\"{seed}-pred.wav\")\n", " torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n", " IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n", " \n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from stable_audio_tools.inference.sampling import get_alphas_sigmas, sample\n", "import time\n", "import k_diffusion as K\n", "\n", "from tqdm import tqdm" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def ddim_sample_step(v_t, sigma_t, sigma_t1, alpha_t, alpha_t1, z_t):\n", " return alpha_t1 * v_t + ((sigma_t1 / sigma_t) * (z_t - alpha_t * v_t))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def ddim_step(v_t, alpha_t, alpha_t1, sigma_t, sigma_t1, z_t):\n", " return alpha_t1 * v_t + ((sigma_t1 / sigma_t) * (z_t - alpha_t * v_t))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# sampling process\n", "N = 128\n", "batch_size = 1\n", "sample_size = 3000\n", "cfg_scale = 4.0\n", "device = \"cuda\"\n", "\n", "# create conditioning\n", "tags_and_lyrics_tensor = model.conditioner.conditioners[\"tags_and_lyrics\"](\n", " [(tags, lyrics)], device\n", ")\n", "semantic_codes_tensor = model.conditioner.conditioners[\"semantic_codes\"](\n", " [semantic_codes], device\n", ")\n", "\n", "conditioning_tensors = {\n", " \"tags_and_lyrics\": tags_and_lyrics_tensor,\n", " \"semantic_codes\": semantic_codes_tensor,\n", "}\n", "\n", "# create unconditional conditioning\n", "empty_tags_and_lyrics_tensor = model.conditioner.conditioners[\n", " \"tags_and_lyrics\"\n", "]([(\"\", \"\")], device)\n", "empty_semantic_codes_tensor = model.conditioner.conditioners[\"semantic_codes\"](\n", " [torch.ones_like(semantic_codes) * 4000], device\n", ")\n", "# use the empty tags but keep semantic conditioning the same\n", "conditioning_tensors[\"empty_tags_and_lyrics\"] = empty_tags_and_lyrics_tensor\n", "conditioning_tensors[\"empty_semantic_codes\"] = semantic_codes_tensor\n", "print(\"no cfg on semantic\")\n", "\n", "conditioning = model.get_conditioning_inputs(\n", " conditioning_tensors, negative=False, empty=True\n", ")\n", "print(conditioning)\n", "\n", "# create init noise\n", "eps = torch.randn([batch_size, model.io_channels, sample_size], device=device)\n", "x = torch.randn([batch_size, model.io_channels, sample_size], device=device)\n", "\n", "\n", "ts = x.new_ones([x.shape[0]])\n", "\n", "# Create the noise schedule\n", "t = torch.linspace(1, 0, N + 1)[:-1]\n", "\n", "alphas, sigmas = get_alphas_sigmas(t)\n", "print(sigmas[0])\n", "\n", "x = x * sigmas[0]\n", "\n", "eta = 0.0\n", "\n", "# The sampling loop\n", "for i in tqdm(range(N)):\n", "\n", " # Get the model output (v, the predicted velocity)\n", " with torch.no_grad():\n", " with torch.cuda.amp.autocast():\n", " v = model.model(x, ts * t[i], **conditioning, cfg_scale=cfg_scale).float()\n", "\n", " # Predict the noise and the denoised image\n", " pred = x * alphas[i] - v * sigmas[i]\n", " eps = x * sigmas[i] + v * alphas[i]\n", "\n", " # If we are not on the last timestep, compute the noisy image for the\n", " # next timestep.\n", " if i < N - 1:\n", " # If eta > 0, adjust the scaling factor for the predicted noise\n", " # downward according to the amount of additional noise to add\n", " ddim_sigma = (\n", " eta\n", " * (sigmas[i + 1] ** 2 / sigmas[i] ** 2).sqrt()\n", " * (1 - alphas[i] ** 2 / alphas[i + 1] ** 2).sqrt()\n", " )\n", " adjusted_sigma = (sigmas[i + 1] ** 2 - ddim_sigma**2).sqrt()\n", "\n", " # Recombine the predicted noise and predicted denoised image in the\n", " # correct proportions for the next step\n", " x = pred * alphas[i + 1] + eps * adjusted_sigma\n", "\n", " # Add the correct amount of fresh noise\n", " if eta:\n", " x += torch.randn_like(x) * ddim_sigma\n", "\n", "# If we are on the last timestep, output the denoised image\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "pred_zq = pred#.squeeze()#.permute(1, 0)\n", "pred_zq /= scale_factor\n", "#Wpred_zq = torch.randn_like(pred_zq)\n", "print(\"pred_zq\", pred_zq.shape)\n", "\n", "pred_zq = pred_zq[...,:3000]\n", "with torch.no_grad():\n", " pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() \n", " pred_audio /= pred_audio.abs().max().clamp(1e-8)\n", " print(pred_audio.mean())\n", "\n", "# save audio\n", "IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n", " " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "\n", "# sample from categorical distribution (1, 2, .., N)\n", "i = torch.distributions.Categorical(torch.ones(N)).sample()\n", "print(i)\n", "\n", "# convert step to continous time\n", "t = i / N\n", "print(t)\n", "\n", "# get original data example and noise\n", "x = torch.randn(1, 128, 3000)\n", "eps = torch.randn(1, 128, 3000)\n", "\n", "# get alphas and sigmas\n", "alpha_t, sigma_t = get_alphas_sigmas(t)\n", "print(alpha_t, sigma_t)\n", "\n", "# combine data and noise\n", "z_t = alpha_t * x + sigma_t * eps\n", "\n", "# two steps with ddim\n", "t1 = t - (0.5 / N)\n", "t2 = t - (1 / N)\n", "\n", "alpha_t1, sigma_t1 = get_alphas_sigmas(t1)\n", "alpha_t2, sigma_t2 = get_alphas_sigmas(t2)\n", "\n", "print(t1, t2)\n", "\n", "z_t1 = ddim_sample_step(model, sigma_t, sigma_t1, alpha_t, alpha_t1, z_t)\n", "z_t2 = ddim_sample_step(model, sigma_t1, sigma_t2, alpha_t1, alpha_t2, z_t1)\n", "\n", "x_target = (z_t2 - (sigma_t2 / sigma_t) * z_t) / (alpha_t2 - (sigma_t2 / sigma_t) * alpha_t)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def sample_ddim(model, x, sigmas):\n", " alphas, sigmas = get_alphas_sigmas(t)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# test throughput\n", "model.eval()\n", "\n", "# evaluate the teacher model at fixed time steps `\n", "N = 50 \n", "t_set = torch.linspace(0, 1, N).cuda()\n", "bs = 2\n", "extra_args = {}\n", "num_batches = 100\n", "times = []\n", "\n", "for bidx in tqdm(range(num_batches)):\n", "\n", " start_time = time.time()\n", "\n", " # dummy data\n", " reals = torch.randn(bs, 128, 3000).cuda()\n", "\n", " metadata = [{\"semantic_codes\": torch.randint(0, 4000, (3000,)).cuda(), \"tags_and_lyrics\": ([\"tags\"], \"lyrics\"), \"phonemes\" : \"\"} for n in range(bs)]\n", "\n", " # sample t \n", " t = t_set[torch.randint(0, N, (bs,))]\n", "\n", " # Calculate the noise schedule parameters for those timesteps\n", " alphas, sigmas = get_alphas_sigmas(t)\n", "\n", " diffusion_input = reals.cuda()\n", " diffusion_input = diffusion_input * scale_factor\n", "\n", " # Combine the ground truth data and the noise\n", " alphas = alphas[:, None, None]\n", " sigmas = sigmas[:, None, None]\n", " noise = torch.randn_like(diffusion_input)\n", " noised_inputs = diffusion_input * alphas + noise * sigmas\n", " targets = noise * alphas - diffusion_input * sigmas\n", "\n", " with torch.cuda.amp.autocast():\n", " conditioning = model.conditioner(metadata, \"cuda\", \"val\")\n", "\n", " with torch.no_grad():\n", " v = model(\n", " noised_inputs,\n", " t,\n", " cond=conditioning,\n", " cfg_dropout_prob=0.0,\n", " **extra_args,\n", " )\n", "\n", " elapsed_time = time.time() - start_time\n", " times.append(elapsed_time)\n", "\n", "print(np.mean(times), np.std(times))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Upsample\n", "Add some noise to the ground truth latents and run inference. This only works for audio, not npz." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_steps = 200\n", "n_tokens_memmap = 3000\n", "\n", "out_dir = f\"outputs/val-{ckpt_name}\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "semantic_codes_pad = (torch.ones(750) * 4000).long().cuda()\n", "\n", "cfg_scales = [12.0]\n", "seeds = [np.random.randint(0, 2**32 - 1) for n in range(3)]\n", "\n", "#tags = [\"Jazz\", \"Piano\", \"High-fidelity\"]\n", "#tags = [\"drums\", \"bass\", \"synth\", \"edm\", \"hard\"]\n", "\n", "for seed in seeds:\n", " for n in range(len(cfg_scales)):\n", " print(cfg_scales[n], seed)\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic_and_text(\n", " model,\n", " semantic_codes,\n", " tags,\n", " lyrics,\n", " init_audio=latents,\n", " steps=num_steps,\n", " cfg_scale=cfg_scales[n],\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " seed=seed,\n", " sampler_type=\"dpmpp-2m-sde\",\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n", " pred_zq /= scale_factor\n", " #Wpred_zq = torch.randn_like(pred_zq)\n", " print(\"pred_zq\", pred_zq.shape)\n", "\n", " pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() \n", " pred_audio /= pred_audio.abs().max().clamp(1e-8)\n", " print(pred_audio.mean())\n", "\n", " # save audio\n", " pred_audio_filepath = os.path.join(out_dir, f\"{seed}-pred.wav\")\n", " torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n", " IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n", " \n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# No semantic" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_steps = 256\n", "n_tokens_memmap = 3000\n", "\n", "out_dir = f\"outputs/val-{ckpt_name}\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "semantic_codes_pad = (torch.ones(3000) * 4000).long().cuda()\n", "lyrics_empty = \"\"\"\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", " \"\"\".strip()\n", "#lyrics_empty = \"Country roads, take me home, to the place I belong, West Virginia, mountain mama, take me home, country roads.\"\n", "tags_empty = [\"pop\", \"r&b\", \"edm\"]\n", "#tags_empty = [\"jazz\", \"pop\", \"female vocal\", \"ballad\"]\n", "#tags_empty = [\"rap\", \"hip-hop\", \"bass\"]\n", "#tags_empty = [\"pop\"]\n", "\n", "#yrics_empty = lyrics\n", "#tags_empty = tags\n", "print(tags_empty)\n", "print(lyrics_empty)\n", "\n", "#lyrics_empty = \"\"\n", "#tags_empty = []\n", "\n", "#lyrics_empty = \"\"\"\"\"\"\n", "#tags_empty = [\"southern rock\", \"country\", \"folk\"]\n", "\n", "num_seeds = 1\n", "cfg_scales = [6.0]\n", "seeds = [np.random.randint(0, 2**32 - 1) for _ in range(num_seeds)]\n", "\n", "for n in range(len(cfg_scales)):\n", " for seed in seeds:\n", " print(cfg_scales[n], seed)\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic_and_text(\n", " model,\n", " semantic_codes_pad,\n", " tags_empty,\n", " lyrics_empty,\n", " steps=num_steps,\n", " cfg_scale=cfg_scales[n],\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " seed=seed,\n", " sampler_type=\"dpmpp-3m-sde\",\n", " #init_audio=latents,\n", " compile=False,\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n", " pred_zq /= scale_factor\n", " #Wpred_zq = torch.randn_like(pred_zq)\n", " print(\"pred_zq\", pred_zq.shape)\n", "\n", " pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() \n", " pred_audio /= pred_audio.abs().max().clamp(1e-8)\n", " print(pred_audio.mean())\n", "\n", " # save audio\n", " pred_audio_filepath = os.path.join(out_dir, f\"{seed}-pred.wav\")\n", " torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n", " IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Long inference" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Load audio first" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio_path = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/Speak For Me [omeNvD8IddM].wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/09 Sounds Like Hallelujah.wav\"\n", "#audio_path = \"/home/christian/code/christian/outputs/halo_24.wav\"\n", "#audio_path = \"/home/christian/audio/bad-audio/bill-evans-intro.wav\"'\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/Crazy [CKTOvHw8qFM].wav\"\n", "#audio_path = \"/home/christian/audio/reference-audio-wav/03 Your New Aesthetic.wav\"\n", "\n", "audio_arrays_48khz = load_audio_mp(\n", " [audio_path],\n", " target_sample_rate=48000,\n", " normalize_volume=True,\n", " num_workers=12,\n", " n_channels=2,\n", ")\n", "\n", "start_s = 0.0\n", "end_s = start_s + 120.0\n", "audio_arrays_48khz = [a[:, int(start_s * 48000) : int(end_s * 48000)] for a in audio_arrays_48khz]\n", "\n", "print(audio_arrays_48khz[0].shape)\n", "IPython.display.display(IPython.display.Audio(audio_arrays_48khz[0].numpy(), rate=48000))\n", "\n", "# now load audio at 24khz\n", "audio_arrays_24khz = load_audio_mp(\n", " [audio_path],\n", " target_sample_rate=24000,\n", " normalize_volume=True,\n", " min_duration_s=30.0,\n", " num_workers=12,\n", " n_channels=2,\n", ")\n", "\n", "audio_arrays_24khz = [a[:, int(start_s * 24000) : int(end_s * 24000)] for a in audio_arrays_24khz]\n", "\n", "# semantic encode \n", "semantic_codes = encode([audio_arrays_24khz[0].mean(axis=0, keepdim=True)], SAMPLE_RATE, EMBEDDING_RATE)\n", "semantic_codes = np.array(semantic_codes)[0, :, 0]\n", "\n", "print(semantic_codes.shape)\n", "if semantic_codes.shape[0] < 9000: # pad with 4000\n", " semantic_codes = np.pad(semantic_codes, (0, 9000 - semantic_codes.shape[0]), mode=\"constant\", constant_values=4000)\n", "semantic_codes = torch.from_numpy(semantic_codes).long().cuda() \n", "print(semantic_codes.shape)\n", "\n", "\n", "if \"Don't Know Why\" in audio_path:\n", " lyrics = \"[Verse]\\nI waited til I saw the sun\\n I don't know when I didn't come\\nI left you by the house of fun\\nI don't know why\"\n", " tags = [\"jazz\", \"pop\", \"female vocal\"]\n", "else:\n", " lyrics = \"\"\n", " tags = []" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "cfg_scales = [6.0]\n", "seeds = [np.random.randint(0, 2**32 - 1) for n in range(3)]\n", "\n", "sample_size = 36000\n", "\n", "\n", "for seed in seeds:\n", " for n in range(len(cfg_scales)):\n", " print(cfg_scales[n], seed)\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic_and_text(\n", " model,\n", " semantic_codes,\n", " tags,\n", " lyrics,\n", " steps=num_steps,\n", " cfg_scale=cfg_scales[n],\n", " sample_size=sample_size,\n", " sample_rate=48000,\n", " seed=seed,\n", " sampler_type=\"dpmpp-2m-sde\",\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n", " pred_zq /= scale_factor\n", " #Wpred_zq = torch.randn_like(pred_zq)\n", " print(\"pred_zq\", pred_zq.shape)\n", "\n", " pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() \n", " pred_audio /= pred_audio.abs().max().clamp(1e-8)\n", " print(pred_audio.mean())\n", "\n", " # save audio\n", " pred_audio_filepath = os.path.join(out_dir, f\"{seed}-pred.wav\")\n", " torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n", " IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Chunked" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Loop inference" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import math\n", "\n", "def calculate_blocks_from_seconds(total_seconds: float) -> int:\n", " if total_seconds <= 30:\n", " return 1\n", " else:\n", " return math.ceil((total_seconds - 30) / 25) + 1" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# compute number of blocks needed\n", "# we will operate on blocks of 30 seconds\n", "# this translates to 3000 embeddings. \n", "# however, after the first block, we will condition on the output of the previous block\n", "# we will use the last 5 seconds (500 embeddings) of the previous block to condition on\n", "# plus, we will use the last 5 seconds of the semantic codes from previous block (125 tokens)\n", "num_steps = 100\n", "context_dur_s = 5.0\n", "semantic_window_size = 750\n", "semantic_context_size = int(25 * context_dur_s)\n", "semantic_hop_size = int(semantic_window_size - semantic_context_size)\n", "\n", "print(\"semantic_context_size\", semantic_context_size, \"semantic_hop_size\", semantic_hop_size)\n", "\n", "latents_window_size = 3000\n", "latents_context_size = int(100 * context_dur_s)\n", "latents_hop_size = int(latents_window_size - latents_context_size)\n", "\n", "print(\"latents_context_size\", latents_context_size, \"latents_hop_size\", latents_hop_size)\n", "\n", "num_blocks = calculate_blocks_from_seconds(end_s - start_s)\n", "print(num_blocks)\n", "\n", "# init final latent array\n", "latents_output = torch.zeros(1, 128, latents_window_size * num_blocks).cuda()\n", "\n", "\n", "# iterate over blocks with index\n", "for block_idx in range(num_blocks):\n", " # for the first block, we use the start of the conditioning from semantic\n", " # for the rest of the blocks, we use a hop size of \n", " if block_idx == 0:\n", " semantic_start_idx = 0\n", " else:\n", " semantic_start_idx = block_idx * semantic_hop_size\n", " semantic_end_idx = semantic_start_idx + semantic_window_size\n", " print(semantic_start_idx, semantic_end_idx)\n", "\n", " # get the conditioning from the semantic codes\n", " semantic_codes_block = semantic_codes[semantic_start_idx:semantic_end_idx]\n", "\n", " # run the diffusion process on current block\n", " if block_idx == 0:\n", " # we have no context or conditioning from previous block\n", " latents_block = None\n", " else:\n", " latents_start_idx = latents_window_size - latents_context_size\n", " latents_end_idx = latents_start_idx + latents_context_size\n", " print(latents_start_idx, latents_end_idx) \n", " latents_block = upsampled_latents[...,latents_start_idx:latents_end_idx]\n", "\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic_and_text(\n", " model,\n", " semantic_codes_block,\n", " tags,\n", " lyrics,\n", " steps=num_steps,\n", " cfg_scale=cfg_scales[n],\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " seed=seed,\n", " sampler_type=\"dpmpp-3m-sde\",\n", " latent_context=latents_block,\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n", " pred_zq /= scale_factor\n", " #Wpred_zq = torch.randn_like(pred_zq)\n", " print(\"pred_zq\", pred_zq.shape)\n", " if block_idx == 0:\n", " write_start_idx = 0\n", " write_end_idx = latents_window_size\n", " pred_zq_block = pred_zq\n", " else:\n", " pred_zq_block = pred_zq[...,latents_context_size:]\n", " write_start_idx = write_end_idx\n", " write_end_idx = write_start_idx + (latents_window_size - latents_context_size)\n", "\n", " print(\"write_start_idx\", write_start_idx, \"write_end_idx\", write_end_idx)\n", " print(\"pred_zq_block\", pred_zq_block.shape)\n", "\n", " latents_output[...,write_start_idx:write_end_idx] = pred_zq_block\n", "\n", "# decode the latents\n", "with torch.no_grad():\n", " pred_audio = vae_model.decode(latents_output)[0].detach().cpu() \n", " pred_audio /= pred_audio.abs().max().clamp(1e-8)\n", " print(pred_audio.mean())\n", "\n", "# save audio\n", "IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n" ] }, { "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 }