{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# !ls -lah /home/christian/code/neon/stable-audio-tools/harmonai_train/wkd8lkiv/checkpoints/" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# !PYTHONPATH=\"/home/georg/code/neon/stable-audio-tools/:/home/georg/code/neon/suno-codec/:$PYTHONPATH\" \\\n", "# CUDA_VISIBLE_DEVICES=7 python /home/georg/code/neon/stable-audio-tools/unwrap_model.py \\\n", "# --model-config /home/georg/code/neon/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_semantic+text_48khz_lg_scale=2.5.json \\\n", "# --ckpt-path /home/christian/code/neon/stable-audio-tools/harmonai_train/wkd8lkiv/checkpoints/last.ckpt \\\n", "# --name /home/georg/notebooks/gpu_nb/tmp/1b" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# !ls -lah /home/christian/code/neon/stable-audio-tools/harmonai_train/wvskbbkz/checkpoints/" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# !PYTHONPATH=\"/home/georg/code/neon/stable-audio-tools/:/home/georg/code/neon/suno-codec/:$PYTHONPATH\" \\\n", "# CUDA_VISIBLE_DEVICES=7 python /home/georg/code/neon/stable-audio-tools/unwrap_model.py \\\n", "# --model-config /home/georg/code/neon/stable-audio-tools/stable_audio_tools/configs/model_configs/txt2audio/stable_audio_2_0_semantic+text_48khz_vlg_scale=2.5_fsdp_qk_norm.json \\\n", "# --ckpt-path /home/christian/code/neon/stable-audio-tools/harmonai_train/wvskbbkz/checkpoints/last.ckpt \\\n", "# --name /home/georg/notebooks/gpu_nb/tmp/2b" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sys\n", "sys.path.insert(0, \"/home/georg/code/neon/stable-audio-tools/\")\n", "sys.path.insert(0, \"/home/georg/code/neon/suno-codec/\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "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", ")\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.audio import Audio\n", "from suno_utils.utils.s3 import read_from_s3\n", "from suno_utils.utils.text import normalize_whitespace\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": [ "# torch 2.4.0.dev20240419+cu121" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# !ls -lah /home/christian/code/neon/stable-audio-tools/checkpoints" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# ckpt_path = \"/home/georg/notebooks/gpu_nb/tmp/1b.ckpt\"\n", "ckpt_path = \"/home/georg/notebooks/gpu_nb/tmp/2b.ckpt\"\n", "\n", "main_dir = \"/home/georg/code/neon/stable-audio-tools/\"\n", "config_dir = os.path.join(main_dir, \"stable_audio_tools/configs/model_configs/txt2audio\")\n", "\n", "ckpt_name = os.path.basename(ckpt_path).replace(\".ckpt\", \"\")\n", "\n", "# load VAE model\n", "device = \"cuda:0\"\n", "checkpoint_filepath = \"s3://suno-data/christian/100hz_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 \"200m\" in ckpt_path:\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_sm_scale=2.5.json\"\n", " )\n", "elif \"1b\" in ckpt_path:\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_lg_scale=2.5.json\"\n", " )\n", "elif \"2b\" in ckpt_path:\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic+text_48khz_vlg_scale=2.5_fsdp_qk_norm.json\"\n", " )\n", "else:\n", " raise NotImplementedError(\"not a known config\")\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}\")" ] }, { "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\n", "# from suno_utils.tasks.mert_25 import SAMPLE_RATE as, EMBEDDING_RATE" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "scrolled": true }, "outputs": [], "source": [ "def load_audio(audio_fp, offset_s=0, sample_rate=48_000, add_fragment=False, n_channels=2):\n", " audio = Audio.from_file(audio_fp, sample_rate=sample_rate, n_channels=n_channels)\n", " end_s = offset_s + 30\n", " if add_fragment:\n", " end_s += 0.01\n", " audio_arr = audio.get_segment(from_s=offset_s, to_s=end_s)\n", " return audio_arr\n", "\n", "def load_audio_arr(audio_fp, offset_s=0, sample_rate=48_000, add_fragment=False, n_channels=2):\n", " audio = load_audio(\n", " audio_fp, offset_s=offset_s, sample_rate=sample_rate, add_fragment=add_fragment, n_channels=n_channels\n", " )\n", " audio_arr = audio.array_float\n", " if n_channels == 1:\n", " audio_arr = audio_arr[None]\n", " return audio_arr\n", "\n", "# audio_fp = \"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n", "# audio_fp = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "# audio_fp = \"/home/georg/notebooks/samples/walking_down.mp3\"\n", "# audio = load_audio(audio_fp, offset_s=0, sample_rate=24_000, add_fragment=True, n_channels=1)\n", "# audio_fp = \"/home/georg/notebooks/samples/turkey_time.mp3\"\n", "# audio = load_audio(audio_fp, offset_s=0, sample_rate=24_000, add_fragment=True, n_channels=1)\n", "## https://suno.com/song/4d9e5af0-ba0d-46fb-8226-b2aba89a93de\n", "# audio_fp = \"/home/georg/notebooks/samples/button.mp3\"\n", "# audio = load_audio(audio_fp, offset_s=0, sample_rate=24_000, add_fragment=True, n_channels=1)\n", "# audio_fp = \"/home/georg/notebooks/samples/frans.mp3\"\n", "# audio = load_audio(audio_fp, offset_s=0, sample_rate=24_000, add_fragment=True, n_channels=1)\n", "audio_fp = \"/home/georg/notebooks/samples/pandas.mp3\"\n", "audio = load_audio(audio_fp, offset_s=5, sample_rate=24_000, add_fragment=True, n_channels=1)\n", "# audio_fp = \"/home/georg/notebooks/samples/skibidi.mp3\"\n", "# audio = load_audio(audio_fp, offset_s=110, sample_rate=24_000, add_fragment=True, n_channels=1)\n", "\n", "# audio_fp = \"/home/georg/notebooks/samples/halo.wav\"\n", "# audio = load_audio(audio_fp, offset_s=20, sample_rate=24_000, add_fragment=True, n_channels=1)\n", "\n", "# semantic encode \n", "semantic_codes = semantic_encode(audio)\n", "semantic_codes = semantic_codes.astype(np.int64)[:, 0]\n", "semantic_codes = torch.from_numpy(semantic_codes).long().cuda() \n", "print(semantic_codes.shape)\n", "\n", "audio.play()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "if \"Don't Know Why\" in audio_fp:\n", " lyrics = \"\"\"\n", "[Verse]\n", "I waited til I saw the sun\n", "I don't know when I didn't come\n", "I left you by the house of fun\n", "I don't know why\n", " \"\"\".strip()\n", " tags = [\"jazz\", \"pop\", \"female vocal\"]\n", "elif \"Dreams\" in audio_fp:\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 \"halo\" in audio_fp.lower():\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 \"walking\" in audio_fp.lower():\n", " lyrics = \"\"\"\n", "Walking down the street\n", "Feeling so alive\n", "Got my head in the clouds\n", "Got a gleam in my eye\n", "Every step that i take\n", "Is like a brand new start\n", "No matter where I'm going\n", "I'll always find my part\n", "\n", "Oooo\n", "\n", "Life is like a high\n", " \"\"\".strip()\n", " tags = [\"Pop\"]\n", "elif \"turkey\" in audio_fp.lower():\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\n", " \"\"\".strip()\n", " tags = [\"Pop\"]\n", "elif \"button\" in audio_fp.lower():\n", " lyrics = \"\"\"\n", "Why is Marlon's camera lookin' like a Youtube kids thumbnail.\n", "Marlon's brick is microscale\n", "My drip is not for sale.\n", "I'll run it up and send you straight to hell. (bitch)\n", "\n", "I don't walk, I drive Ferrari.. (bitch)\n", "You don't talk, cause yo' ass is sorry.. (aa)\n", "I shoot my AK-47.. (ya)\n", "boy, you look barely even seven.. (seven)\n", "Suck my cockroach bitch..\n", "Damn bro? You were good, you made me spit quick.. (he sucked my stick)\n", "Don't matter, I wouldn't even let you bootylick you fish stick. (aa)\n", " \"\"\".strip()\n", " tags = [\"Trap\"]\n", "elif \"frans\" in audio_fp.lower():\n", " lyrics = \"\"\"\n", "[Intro]\n", "Övertid (övertid)\n", "Extratid (extra tid)\n", "Mertid (mertid)\n", "Ingen tid (ingen tid)\n", "\n", "[Verse]\n", "Genomgång på morgonen\n", "Humöret är på topp\n", "Men sen när arbetet är på gång\n", "Allting får ett stopp (stopp)\n", " \"\"\".strip()\n", " tags = [\"Funl\", \"Soul\"]\n", "elif \"panda\" in audio_fp.lower():\n", " lyrics = \"\"\"\n", "[Verse]\n", "Es ist Zeit, wild zu sein, \n", "wir sind bereit\n", "Mit dem Panda Sonnenschutz\n", "wir sind bereit\n", "Die Sonne scheint stark,\n", "aber wir sind stärker\n", "\n", "[Chorus]\n", "Panda Sonnenschutz\n", "Panda Sonnenschutz\n", " \"\"\".strip()\n", " tags = [\"EDM\"]\n", "elif \"skibidi\" in audio_fp.lower():\n", " lyrics = \"\"\"\n", "skibidi skibidi skibidi skibidi skibidi skibidi skibidi skibidi skibidi skibidi skibidi\n", " \"\"\"\n", " tags = [\"Choir\", \"Jazz\"]\n", "else:\n", " lyrics = \"\"\n", " tags = []\n", "print(\"TAGS:\", \", \".join(tags)[:30])\n", "print(\"LYRICS:\", normalize_whitespace(lyrics)[:30])" ] }, { "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 = 1000 #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", "sweep_cfg = False\n", "\n", "if sweep_cfg:\n", " cfg_scales = [1.0, 1.5, 3.0, 6.0]\n", "else:\n", " cfg_scales = [2.0] * 2\n", "seed = np.random.randint(0, 2**32 - 1)\n", "\n", "for n in range(len(cfg_scales)):\n", " if not sweep_cfg:\n", " seed = np.random.randint(0, 2**32 - 1)\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=n_tokens_memmap,\n", " sample_rate=48_000,\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", "# break" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Playground" ] }, { "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_v4/data_vae_val.bin\"\n", "semantic_memmap_path = \"/app/suno/data/chirp_v4/vae_v4/data_semantic_val.bin\"\n", "\n", "vae_n_tokens = 3000\n", "semantic_n_tokens = 750\n", "vae_dim = 128\n", "\n", "# load memmap and get semantic\n", "vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode=\"r\")\n", "vae_data = vae_data.reshape(-1, vae_n_tokens, vae_dim)\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_v3/metas_val.jsonl\")\n", "\n", "print(len(metas), vae_data.shape, semantic_data.shape)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "rand_idx = np.random.randint(len(metas))\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", "print(rand_idx)\n", "print(tags)\n", "print(lyrics) \n", "print(semantic_codes.shape, vae_latents.shape)\n", "\n", "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))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.14" } }, "nbformat": 4, "nbformat_minor": 4 }