{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/christian/miniconda3/envs/suno_env/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n", "/home/christian/code/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": [ "# autoreload\n", "%load_ext autoreload\n", "%autoreload 2\n", "\n", "import os\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\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,\n", " upsample_diffusion_from_codes,\n", " denoising_diffusion_from_codes,\n", " upsample_diffusion_from_semantic,\n", " upsample_diffusion_from_discrete,\n", " upsample_diffusion_from_uncond,\n", " upsample_diffusion_from_lyrics,\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.utils.s3 import read_from_s3\n", "from suno_utils.models.dac.nn.quantize_2 import ResidualVectorQuantize\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", ")\n", "\n", "\n", " \n", "\n", "def decode_vq(model, codes, n_quantizers):\n", " z_q = 0\n", " for i, quantizer in enumerate(model.quantizer.quantizers[:n_quantizers]):\n", " _z_q = quantizer.embed_code(codes[:, :, i]).transpose(1, 2)\n", " _z_q = quantizer.out_proj(_z_q)\n", " z_q += _z_q.transpose(1, 2)\n", "\n", " return z_q" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ckpt_path = \"/home/christian/code/neon/stable-audio-tools/checkpoints/diffusion_semantic+text_100hz_200m_epoch=0-step=140000.ckpt\"\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", "if \"25hz\" in ckpt_path:\n", " model_type = \"semantic\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_semantic_48khz_lg_vae.json\"\n", " )\n", "\n", " # load VAE\n", " vae_model = load_vae_model(\n", " #checkpoint_filepath=\"s3://suno-data/christian/mw_vae_peaq_128_fix.pth\",\n", " checkpoint_filepath=\"s3://suno-data/georg/models/codec/dac_vae_128.pth\",\n", " device=\"cuda\"\n", " )\n", "\n", "elif \"100hz\" in ckpt_path:\n", " device = \"cuda:0\"\n", " # checkpoint_filepath = \"/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth\"\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", " if \"uncond\" in ckpt_path:\n", " model_type = \"uncond\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_uncond_48khz_lg_vae.json\"\n", " )\n", " elif \"lyrics\" in ckpt_path:\n", " model_type = \"lyrics\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_lyrics_48khz_lg_vae.json\"\n", " )\n", " elif \"semantic+text_100hz_200m_scale=2.5\" in ckpt_path:\n", " model_type = \"semantic+text_100hz_200m\"\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 \"semantic+text_100hz_1b_scale=2.5\" in ckpt_path:\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", " else:\n", " model_type = \"discrete\"\n", " model_config_path = os.path.join(\n", " config_dir, \"stable_audio_2_0_discrete_48khz_lg_vae.json\"\n", " )\n", "\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", "print(model_config)\n", "# model_config[\"model\"][\"diffusion\"][\"config\"][\"use_checkpointing\"] = True\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", "\n", "\n", "#from suno_utils.tasks.dac_vae_peaq import (\n", "# preload_models as preload_vae_models,\n", "# encode as vae_encode,\n", "# decode as vae_decode,\n", "#)\n", "#_ = preload_codec_models(\"s3://suno-data/georg/models/codec/dac_2c_25x12.pt\")\n", "#_ = preload_vae_models(\"s3://suno-data/georg/models/codec/dac_vae_128.pth\"))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Unconditional" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_steps = 250\n", "loudnorm = True\n", "n_tokens_memmap = 1000\n", "\n", "\n", "out_dir = f\"outputs/val-{ckpt_name}\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "for n in range(250):\n", "\n", " seed = np.random.randint(0, 2**32 - 1)\n", " with torch.no_grad():\n", " # randomly sample number of steps\n", " num_steps = np.random.choice([10, 50, 100, 250, 500, 1000])\n", " num_steps = 250\n", " upsampled_latents = upsample_diffusion_from_uncond(\n", " model,\n", " steps=num_steps,\n", " cfg_scale=1.0,\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " seed=seed,\n", " sampler_type=\"dpmpp-3m-sde\",\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\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", " #target_audio_filepath = os.path.join(out_dir, f\"{rand_idx}-target.wav\")\n", " pred_audio_filepath = os.path.join(out_dir, f\"{seed}-num-steps={num_steps}-pred.wav\")\n", "\n", " #torchaudio.save(target_audio_filepath, target_audio.cpu().squeeze(), 48000)\n", " torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n", "\n", " #IPython.display.display(IPython.display.Audio(data=target_audio.cpu().squeeze().numpy(), rate=48000))\n", " #IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Refine" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_steps = 250\n", "loudnorm = True\n", "n_tokens_memmap = 1000\n", "\n", "# load audio file\n", "audio, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n", "audio_48khz = torchaudio.functional.resample(audio, sr, 48000)\n", "\n", "start_frame = audio_48khz.shape[1] // 2\n", "end_frame = start_frame + int(48000 * 10)\n", "audio_48khz = audio_48khz[:, start_frame:end_frame]\n", "\n", "# embed with VAE\n", "with torch.no_grad():\n", " vae_latents = vae_model.encode(audio_48khz[None].to(\"cuda\"))[\"z\"]\n", "print(vae_latents.shape)\n", "# set to init noise\n", "init_audio = vae_latents\n", "\n", "init_noise = torch.randn([1, 128, vae_latents.shape[-1]], device=device) * 0.00001\n", "\n", "out_dir = f\"outputs/val-{ckpt_name}-refine\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "for n in range(3):\n", "\n", " seed = np.random.randint(0, 2**32 - 1)\n", " with torch.no_grad():\n", " # randomly sample number of steps\n", " num_steps = 2\n", " upsampled_latents = upsample_diffusion_from_uncond(\n", " model,\n", " steps=num_steps,\n", " cfg_scale=1.0,\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " seed=seed,\n", " init_audio=init_audio,\n", " init_noise=init_noise,\n", " sampler_type=\"dpmpp-3m-sde\",\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\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", " #target_audio_filepath = os.path.join(out_dir, f\"{rand_idx}-target.wav\")\n", " pred_audio_filepath = os.path.join(out_dir, f\"{seed}-num-steps={num_steps}-pred.wav\")\n", "\n", " #torchaudio.save(target_audio_filepath, target_audio.cpu().squeeze(), 48000)\n", " torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n", "\n", " IPython.display.display(IPython.display.Audio(data=audio_48khz.cpu().squeeze().numpy(), rate=48000))\n", " IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Lyric conditioned" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Val set" ] }, { "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_v3/data_vae_val.bin\"\n", "semantic_memmap_path = \"/app/suno/data/chirp_v4/vae_v3/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(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))\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Raw audio" ] }, { "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", "\n", "audio_path = \"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n", "\n", "audio_arrays_48khz = load_audio_mp(\n", " [audio_path],\n", " target_sample_rate=48000,\n", " max_duration_s=30.0,\n", " normalize_volume=True,\n", " num_workers=12,\n", " n_channels=2,\n", ")\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", " max_duration_s=30.05,\n", " normalize_volume=True,\n", " num_workers=12,\n", " n_channels=2,\n", ")\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", "semantic_codes = torch.from_numpy(semantic_codes).long().cuda() \n", "print(semantic_codes.shape)\n", "\n", "\n", "tags = [\"jazz\", \"pop\", \"female vocal\"]\n", "dreams_lyrics = \"[Verse 1] Now here you go again, you say you want your freedom, well who am I to keep you down?\"\n", "dont_know_why_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" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## From npz" ] }, { "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", "gen_id = \"ee467d00-5813-4a74-9792-c9ae4a09d344\"\n", "s3_filepath = f\"s3://suno-data-uploads/studio/uploads/{gen_id}.npz\"\n", "data = read_from_s3(s3_filepath, read_f=np.load)\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", "semantic_codes = codes[:, 0].astype(np.uint16)\n", "semantic_codes = torch.from_numpy(semantic_codes).long().cuda()\n", "print(semantic_codes.shape)\n", "\n", "tags = []\n", "lyrics = \"\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Inference" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_steps = 250\n", "loudnorm = True\n", "n_tokens_memmap = 3000\n", "\n", "out_dir = f\"outputs/val-{ckpt_name}\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "#tags = [\"pop\"]\n", "\n", "lyrics_a = \" Almost heaven, west virginia, blue ridge mountains, shenandoah river\" \n", "lyrics_b = \" country roads, take me home to the place I belong, west virginia, mountain mama\"\n", "lyrics_c = \" So this is why So this is what makes life divine\"\n", "lyrics_d = \" hello darkness my old friend\"\n", "lyrics_e = \" Everywhere I'm lookin' now I'm surrounded by your embrace Baby, I can see your halo\"\n", "lyrics_train = \" Come Holy Spirit, trapeze away again, the Lord of sin is free.\"\n", "lyrics_empty = \"\"\n", "#lyrics_train = \" here are scarlet flags? They were starting to mind me Where are planting lights on all of you?\"\n", "lyrics_val = \"Hello world! Try to sing this? Is this in the validation set? I'm gonna love you. Ooh, I'm gonna love you. Ready? Yeah\"\n", "lyrics_val = \"[Chorus]\\nSteady bragging 'bout them bodies like that shit impress me\\nLet me see you make it out and go and get a check in\\nThese niggas tryna ride the wave like a fuckin' jet-ski\\nHeard through the grapevine you tried to disrespect me\\nAll that hating ain't got a nigga paid yet\\nBroke nigga, ain't flew inside a plane yet\\nAll that hating ain't got a nigga paid yet\\nBroke nigga, ain't flew inside a plane yet, uh\\n[Verse 1]\\nHell nah, I ain't worried about no frivolous movement\\nGot the chopper right here if a nigga get stupid\\nAnd you know I'ma use it 'cause it's more than just music\\n'Cause I got the cash money, I'm the '03 Juvie\"\n", "\n", "#tags = []\n", "#lyrics = \"\"\n", "\n", "semantic_codes_pad = (torch.ones(750) * 4000).long().cuda()\n", "#tags = []\n", "#lyrics = lyrics_b\n", "\n", "#print(semantic_codes[0])\n", "#print(tags)\n", "#print(lyrics)\n", "\n", "\n", "cfg_scales = [1.0, 2.0, 3.0, 6.0]\n", "cfg_scales = [1.0, 1.1, 1.3]\n", "\n", "for n in range(len(cfg_scales)):\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=48000,\n", " seed=seed,\n", " sampler_type=\"dpmpp-2m-sde\",\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\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": [ "# GPT + Diffusion" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%load_ext autoreload\n", "%autoreload 2\n", "\n", "# setup GPT\n", "from suno_utils.gpt.chirp_v2_5 import (\n", " GenerationConfig,\n", " decode_stream,\n", " preload_models,\n", " prep_gconf,\n", " codec_decode_stream_to_full_audio,\n", " generate,\n", ")\n", "\n", "from suno_utils.gpt.generation import clean_models, _load_model\n", "\n", "def load_gpt_model(\n", " ckpt_path=None,\n", " tokenizer_path=None,\n", " use_gpu=True,\n", " force_reload=False,\n", " use_tp=False,\n", "):\n", " if torch.cuda.device_count() == 0 or not use_gpu:\n", " device = \"cpu\"\n", " else:\n", " device = \"cuda\"\n", " model_key = \"main_model\"\n", " \n", " if ckpt_path is None:\n", " raise ValueError(\"model not initialized, need checkpoint path. maybe run `preload_models`?\")\n", " clean_models(model_key=model_key)\n", " model, tokenizer = _load_model(ckpt_path, tokenizer_path, device, use_tp=use_tp)\n", "\n", " return model\n", "\n", "\n", "#gpt_ckpt_path = \"/app/suno/data/dpo/models/model_13b_full.pt\" # model before fine-tuning\n", "#gpt_ckpt_path = \"/app/suno/checkpoints/2024-07-12_13-28-28/last_ckpt.pt\" # 2b trained only on semantic\n", "#gpt_ckpt_path = \"/app/suno/checkpoints/2024-07-13_20-42-16/last_ckpt.pt\"\n", "#gpt_ckpt_path = \"/app/suno/checkpoints/2024-07-14_21-24-26/last_ckpt.pt\"\n", "gpt_ckpt_path = \"/app/suno/checkpoints/2024-07-16_18-46-08/last_ckpt.pt\"\n", "tokenizer_path = \"/app/suno/data/dpo/models/tokenizer_60k.json\"\n", "\n", "preload_models(\n", " gpt_ckpt_path=gpt_ckpt_path,\n", " tokenizer_path=tokenizer_path,\n", " load_gpt=True,\n", " load_semantic=False,\n", " load_codec_device=\"cuda\",\n", ")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "\n", "text_tags_list = [\n", " # \"indie-pop, energetic, rock, psychedelic\",\n", " # \"jazz, female vocal, smooth\",\n", " # \"r&b, soul, funk\",\n", " \"pop, high energy, female vocal\",\n", " \"rock, high energy, male vocal\",\n", " \"indie rock, female vocal, psychedelic\",\n", " # \"moody, blues, soulful\",\n", " # \"classical, cinematic, female vocal\",\n", " # \"metal, hardcore, dark\",\n", " \"male vocal, bluegrass, guitar\",\n", " # \"female vocal, voice, singing, guitar, folk, jazz, r&b\",\n", "]\n", "\n", "text_a = \"\"\"\n", "Almost Heaven, West Virginia\n", "Blue Ridge Mountains, Shenandoah River\n", "Life is old there, older than the trees\n", "Younger than the mountains, growing like a breeze\n", "\n", "Country roads, take me home\n", "To the place I belong\n", "West Virginia, mountain mama\n", "Take me home, country roads\n", "\n", "[outro]\n", "\"\"\"\n", "\n", "text_b = \"\"\"\n", "[Verse 1]\n", "If I had the superpower to reverse \n", "Time back to a spring century before \n", "When everyone was hustling to preserve \n", "Unless I wielded powers to reverse \n", "\n", "[Instrumental Solo]\n", "\n", "\"\"\"\n", "text_empty = \"\"\n", "\n", "random_seed = 20\n", "#text_tags = \"\"\n", "text_tags = \"bluegrass, female vocal\"\n", "#text_tags = \"indie rock, female vocal, psychedelic, acoustic\"\n", "\n", "general_config = dict(\n", " #cfg_coef=2.5, \n", " #min_eos_p=0.1,\n", " #eos_pad_duration_s=0,\n", " #cfg_coef_tags=1.0,\n", " #cfg_coef_tags_max_steps=None, # collect the data for now\n", " #n_repeat_tags=3,\n", " #use_whisper=False,\n", " text_start_control_tags=\"{start}\",\n", " # text_end_control_tags=\"{end}\",\n", " random_seed=random_seed,\n", " n_batch=1,\n", ")\n", "\n", "gconfig = GenerationConfig(\n", " text=text_a,\n", " text_tags=text_tags,\n", " max_gen_duration_s=60,\n", " temp_semantic=0.9,\n", " #top_k_semantic=None,\n", " #top_p_semantic=None,\n", " **general_config,\n", ")\n", "\n", "semantic_codes = generate(gconfig)\n", "semantic_codes = torch.from_numpy(semantic_codes[0]).squeeze()\n", "semantic_codes = semantic_codes.cuda()\n", "\n", "# evaluate on val example\n", "print(semantic_codes.shape)\n", "n_chunks = semantic_codes.shape[0] // 250\n", "pred_zqs = []\n", "\n", "num_steps = 250\n", "loudnorm = True\n", "\n", "for chunk_idx in range(n_chunks): \n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic(\n", " model,\n", " semantic_codes[chunk_idx * 250 : (chunk_idx + 1) * 250],\n", " steps=num_steps,\n", " cfg_scale=1.0,\n", " sample_size=250,\n", " sample_rate=48000,\n", " )\n", " pred_zq = upsampled_latents\n", " print(\"pred_zq\", pred_zq.shape)\n", " pred_zqs.append(pred_zq)\n", "\n", "pred_zq = torch.cat(pred_zqs, dim=-1)\n", "print(pred_zq.shape)\n", "pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() \n", "\n", "#IPython.display.display(IPython.display.Audio(data=target_audio.cpu().squeeze().numpy(), rate=48000))\n", "IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Diffusion on real semantic" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "N_TOKENS_MEMMAP = 250\n", "VAE_DIM = 128\n", "use_val = True\n", "\n", "if use_val:\n", " semantic_memmap_path = \"/app/suno/data/chirp_v4/vae/data_semantic_val.bin\"\n", " vae_memmap_path = \"/app/suno/data/chirp_v4/vae/data_vae_val.bin\"\n", "else:\n", " semantic_memmap_path = \"/app/suno/data/chirp_v4/vae/data_semantic_tr.bin\"\n", " vae_memmap_path = \"/app/suno/data/chirp_v4/vae/data_vae_tr.bin\"\n", "\n", "semantic_data = np.memmap(semantic_memmap_path, dtype=np.uint16, mode=\"r\")\n", "semantic_data = semantic_data.reshape(-1, N_TOKENS_MEMMAP)\n", "print(\"semantic_data\", semantic_data.shape)\n", "\n", "vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode=\"r\")\n", "vae_data = vae_data.reshape(-1, N_TOKENS_MEMMAP, VAE_DIM)\n", "vae_data = vae_data\n", "print(\"vae_data\", vae_data.shape)\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "n_tokens_memmap = 1000\n", "vae_dim = 128\n", "discrete_memmap_path = \"/app/suno/christian/data/vae/vae_discrete_val.bin\"\n", "vae_memmap_path = \"/app/suno/christian/data/vae/vae_val.bin\"\n", "\n", "vae_memmap_path = \"/app/suno/christian/data/suno_diffusion_tiktok_covers/vae_val.bin\"\n", "\n", "discrete_data = np.memmap(discrete_memmap_path, dtype=np.uint16, mode=\"r\")\n", "discrete_data = discrete_data.reshape(-1, n_tokens_memmap)\n", "print(\"discrete_data\", discrete_data.shape)\n", "\n", "vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode=\"r\")\n", "vae_data = vae_data.reshape(-1, vae_dim, n_tokens_memmap)\n", "print(\"vae_data\", vae_data.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# decode random\n", "rand_idx = np.random.randint(0, vae_data.shape[0])\n", "vae_latents = torch.from_numpy(vae_data[rand_idx, ...].copy()).cuda()\n", "print(\"vae_latents\", vae_latents)\n", "\n", "target_audio = vae_model.decode(vae_latents.unsqueeze(0))[0].detach().cpu() \n", "target_audio /= target_audio.abs().max()\n", "IPython.display.display(IPython.display.Audio(data=target_audio.cpu().squeeze().numpy(), rate=48000))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "num_steps = 250\n", "loudnorm = True\n", "\n", "out_dir = \"outputs/val-diffusion-uncond-vae-100hz-1b-epoch-2-step-70000\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "for n in range(10):\n", " rand_idx = np.random.randint(0, discrete_data.shape[0])\n", "\n", " # evaluate on val example\n", " #semantic_codes = torch.from_numpy(semantic_data[rand_idx, ...].copy()).long().cuda()\n", " discrete_codes = torch.from_numpy(discrete_data[rand_idx, ...].copy()).long().cuda()\n", " vae_latents = torch.from_numpy(vae_data[rand_idx, ...].copy()).cuda()\n", "\n", " print(discrete_codes.shape, vae_latents.shape)\n", "\n", " target_audio = vae_model.decode(vae_latents.unsqueeze(0))[0].detach().cpu() \n", " target_audio /= target_audio.abs().max()\n", "\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_discrete(\n", " model,\n", " discrete_codes,\n", " steps=num_steps,\n", " cfg_scale=1.0,\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\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()\n", "\n", " # save audio\n", " target_audio_filepath = os.path.join(out_dir, f\"{rand_idx}-target.wav\")\n", " pred_audio_filepath = os.path.join(out_dir, f\"{rand_idx}-pred.wav\")\n", "\n", " torchaudio.save(target_audio_filepath, target_audio.cpu().squeeze(), 48000)\n", " torchaudio.save(pred_audio_filepath, pred_audio.cpu().squeeze(), 48000)\n", "\n", " IPython.display.display(IPython.display.Audio(data=target_audio.cpu().squeeze().numpy(), rate=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": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load npz file with numpy\n", "gen_npz = np.load(\"/app/suno/christian/data/npz/a5e2198a-f352-4abb-9a24-7f81b143ded3.npz\")\n", "print(gen_npz.files)\n", "data = gen_npz[\"v3.0_raw\"]\n", "print(data.shape)\n", "print(semantic_codes.shape)\n", "semantic_codes = torch.from_numpy(semantic_codes).cuda()\n", "\n", "num_frames = semantic_codes.shape[0] // 250\n", "\n", "pred_zqs = []\n", "for n in range(num_frames):\n", "\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic(\n", " model,\n", " semantic_codes[n*250:(n+1)*250],\n", " steps=1000,\n", " cfg_scale=1.0,\n", " sample_size=semantic_codes.shape[0],\n", " sample_rate=48000,\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n", " print(\"pred_zq\", pred_zq.shape)\n", " pred_zqs.append(pred_zq)\n", "\n", " pred_zq = torch.cat(pred_zqs, dim=-1)\n", " print(pred_zq.shape)\n", " #pred_audio = vae_model.decode(pred_zq)[0].detach().cpu() \n", " pred_audio = vae_decode(pred_zq.squeeze(0).permute(1,0).detach())\n", " pred_audio = torch.from_numpy(pred_audio.array_float)\n", " pred_audio /= pred_audio.abs().max()\n", "\n", "IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Diffusion on File" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import faiss\n", "centroid_path = \"s3://suno-data/christian/vae_100hz_32768.npy\"\n", "centroids = read_from_s3(centroid_path, read_f=np.load)\n", "print(centroids.shape)\n", "d = 128\n", "\n", "# create faiss kmeans index using centroids\n", "index_cpu = faiss.IndexFlatL2(d)\n", "index_cpu.add(centroids)\n", "\n", "# to encode, pass through the encoder and then quantize via index" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio_paths = [ \n", "# \"/home/christian/audio/reference-audio-wav/09 Sounds Like Hallelujah.wav\",\n", "#\"/home/christian/audio/reference-audio-wav/02 Freddie Freeloader.wav\",\n", "#\"/home/christian/audio/reference-audio-wav/01 No Son Of Mine.wav\",\n", "#\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\",\n", "# \"/home/christian/audio/reference-audio-wav/01 J.S. Bach Suite No.1, S.1007, G major - I. Prelude.wav\",\n", "#\"/home/christian/audio/reference-audio-wav/04 Fuckwithmeyouknowigotit.wav\",\n", "#\"/home/christian/audio/reference-audio-wav/03 Always Be.wav\",\n", "#\"/home/christian/audio/reference-audio-wav/Norah Jones - Don't Know Why [1LH4vnrM-Vs].wav\"\n", "#\"/home/christian/audio/reference-audio-wav/Speak For Me [omeNvD8IddM].wav\" \n", "#\"/home/christian/audio/reference-audio-wav/Laufey - From The Start (Official Music Video) [lSD_L-xic9o].wav\"\n", "\"/home/christian/audio/reference-audio-wav/08 Get Lucky.wav\"\n", "]\n", "\n", "from tqdm import tqdm\n", "#audio_paths = glob.glob(os.path.join(\"/home/christian/audio/reference-audio-wav/*.wav\"))\n", "file_ext = \".wav\"\n", "\n", "out_dir = \"outputs/diffusion-22072024\"\n", "os.makedirs(out_dir, exist_ok=True)\n", "\n", "for audio_path in tqdm(audio_paths):\n", " print(audio_path) \n", " audio_input, sr = torchaudio.load(audio_path)\n", "\n", " if audio_input.shape[0] == 1:\n", " audio_input = audio_input.repeat(2, 1)\n", "\n", " if sr != 48000:\n", " audio_input = torchaudio.functional.resample(audio_input, sr, 48000)\n", "\n", " start_frame = audio_input.shape[-1] // 2\n", " #start_frame = 0\n", " end_frame = start_frame + int(10 * 48000)# + 1024\n", "\n", " audio_input = audio_input[:, start_frame:end_frame]\n", "\n", " audio_input_24khz = torchaudio.functional.resample(audio_input, 48000, 24000)\n", " audio_input_24khz = audio_input_24khz.mean(dim=0, keepdim=True)\n", " audio_input_24khz = apply_normalization(audio_input_24khz, 24_000, target_loudness_lufs_db=-16.0)\n", "\n", " #audio_input_24khz *= 10 ** (-6.0/20.0)\n", "\n", " if \"25hz\" in ckpt_path:\n", " #semantic_codes = semantic_encode([audio_input_24khz])\n", " #semantic_codes = torch.from_numpy(semantic_codes[0][:,0]).long().cuda()\n", " semantic_codes = semantic_encode_files([audio_path])[0]\n", " print(\"semantic_codes\", semantic_codes.shape)\n", " start_code = semantic_codes.shape[0] // 2\n", " #start_code = 0\n", " end_code = start_code + 250\n", " semantic_codes = torch.from_numpy(semantic_codes[start_code:end_code, 0]).long().cuda()\n", " print(\"semantic_codes\", semantic_codes.shape)\n", "\n", " target_cycled_audio = vae_model(audio_input.unsqueeze(0).cuda())[\"audio\"].detach().cpu() \n", " #target_cycled_audio = vae_decode(audio_input_24khz.permute(1,0).cuda())\n", " #target_cycled_audio = torch.from_numpy(target_cycled_audio.array_float)\n", " target_cycled_audio /= target_cycled_audio.abs().max()\n", "\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_semantic(\n", " model,\n", " semantic_codes,\n", " steps=1000,\n", " cfg_scale=1.0,\n", " sample_size=semantic_codes.shape[0],\n", " sample_rate=48000,\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\n", " print(\"pred_zq\", pred_zq.shape)\n", "\n", " elif \"100hz\" in ckpt_path:\n", " with torch.no_grad():\n", " target_cycled_audio = vae_model(audio_input.unsqueeze(0).cuda())[\"audio\"].detach().cpu() \n", " #target_cycled_audio = vae_decode(audio_input_24khz.permute(1,0).cuda())\n", " #target_cycled_audio = torch.from_numpy(target_cycled_audio.array_float)\n", " target_cycled_audio /= target_cycled_audio.abs().max()\n", "\n", " with torch.no_grad():\n", " vae_latents = vae_model.encode(audio_input.unsqueeze(0).cuda())[\"z\"].detach().cpu()\n", " vae_latents = vae_latents.squeeze(0).permute(1,0)\n", " \n", " print(vae_latents.shape)\n", " discrete_codes = index_cpu.search(, 1)[1]\n", " discrete_codes = torch.from_numpy(discrete_codes).long().cuda().squeeze()\n", " print(discrete_codes.shape)\n", "\n", " with torch.no_grad():\n", " upsampled_latents = upsample_diffusion_from_discrete(\n", " model,\n", " discrete_codes,\n", " steps=250,\n", " cfg_scale=1.0,\n", " sample_size=n_tokens_memmap,\n", " sample_rate=48000,\n", " )\n", " pred_zq = upsampled_latents#.squeeze()#.permute(1, 0)\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()\n", "\n", " input_filepath = os.path.join(out_dir, os.path.basename(audio_path).replace(\".wav\", \"_input.wav\"))\n", " target_filepath = os.path.join(out_dir, os.path.basename(audio_path).replace(\".wav\", \"_target.wav\"))\n", " pred_filepath = os.path.join(out_dir, os.path.basename(audio_path).replace(\".wav\", \"_pred.wav\"))\n", "\n", " torchaudio.save(input_filepath, audio_input.cpu().squeeze(), 48000)\n", " torchaudio.save(target_filepath, target_cycled_audio.cpu().squeeze(), 48000)\n", " torchaudio.save(pred_filepath, pred_audio.cpu().squeeze(), 48000)\n", "\n", " #IPython.display.display(IPython.display.Audio(data=audio_input.cpu().squeeze().numpy(), rate=48000))\n", " #IPython.display.display(IPython.display.Audio(data=target_cycled_audio.cpu().squeeze().numpy(), rate=48000))\n", " #IPython.display.display(IPython.display.Audio(data=pred_audio.cpu().squeeze().numpy(), rate=48000))\n", "\n", "\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 }