{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import glob\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"\n", "import glob\n", "import torch\n", "import IPython\n", "import numpy as np\n", "import funcy\n", "\n", "import torchaudio\n", "\n", "from tqdm import tqdm\n", "from dac.model.dac4 import DAC\n", "from dac.model.discriminator2 import Discriminator as Discriminator_import\n", "from dac.nn import loss as loss_import\n", "from dac.utils.accelerator import Accelerator\n", "from dac.utils import load_model\n", "\n", "from suno_utils.utils.s3 import read_from_s3\n", "from suno_utils.models.dac.model.sac_vae import SAC\n", "\n", "import pyloudnorm as pyln\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "models = {}" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio_filepaths = glob.glob(\"/home/christian/audio/reference-audio-wav/*.wav\")\n", "print(len(audio_filepaths))\n", "\n", "import sys\n", "sys.path.insert(0, \"/home/christian/code/neon/suno-codec-training/sac/models\")\n", "\n", "#from sac_vae import SunoCodecVAE\n", "from dac_vae import DACVAE\n", "\n", "\n", "\n", "model_checkpoints = {\n", " #\"dac_1e-4_16d_750hz\": \"/home/christian/logs/dac_1e-4_16d_750hz/last-v3.ckpt\",\n", " #\"dac_1e-4_16d_375hz\" : \"/home/christian/logs/dac_1e-4_16d_375hz/last.ckpt\",\n", " #\"dac_1e-4_8d_375hz\" : \"/home/christian/logs/dac_1e-4_8d_375hz/last.ckpt\",\n", " \"dac_1e-4_128d_25hz\" : \"/home/christian/logs/dac_1e-4_128d_25hz/last.ckpt\",\n", " #\"dac_1e-4_128d_25hz_noise\" : \"/home/christian/logs/dac_1e-4_128d_25hz_noise/last.ckpt\",\n", " #\"dac_1e-4_128_25hz_noise_no_norm\" : \"/home/christian/logs/dac_1e-4_128_25hz_noise_no_norm/last.ckpt\",\n", "}\n", "\n", "for model_name, checkpoint_filepath in model_checkpoints.items():\n", " S = torch.load(checkpoint_filepath)\n", "\n", " # init model\n", " codec = DACVAE(**S[\"hyperparameters\"][\"kwargs\"])\n", "\n", " # load model\n", " codec_ckpt = {k[6:]: v for k, v in S[\"state_dict\"].items() if k.split(\".\")[0] == \"codec\"}\n", " codec.load_state_dict(codec_ckpt)\n", " codec = codec.cuda()\n", " codec.eval()\n", "\n", " models[model_name] = codec\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch.nn as nn\n", "import typing\n", "from typing import List\n", "from collections import namedtuple\n", "from torchaudio.transforms import MelSpectrogram\n", "\n", "# Multi-scale melspectrogram loss\n", "STFTParams = namedtuple(\n", " \"STFTParams\",\n", " [\"window_length\", \"hop_length\", \"window_type\", \"match_stride\", \"padding_type\"],\n", ")\n", "\"\"\"\n", "STFTParams object is a container that holds STFT parameters - window_length,\n", "hop_length, and window_type. Not all parameters need to be specified. Ones that\n", "are not specified will be inferred by the AudioSignal parameters.\n", "\n", "Parameters\n", "----------\n", "window_length : int, optional\n", " Window length of STFT, by default ``0.032 * self.sample_rate``.\n", "hop_length : int, optional\n", " Hop length of STFT, by default ``window_length // 4``.\n", "window_type : str, optional\n", " Type of window to use, by default ``sqrt\\_hann``.\n", "match_stride : bool, optional\n", " Whether to match the stride of convolutional layers, by default False\n", "padding_type : str, optional\n", " Type of padding to use, by default 'reflect'\n", "\"\"\"\n", "STFTParams.__new__.__defaults__ = (None, None, None, None, None)\n", "\n", "class MelSpectrogramLoss(nn.Module):\n", " \"\"\"Compute distance between mel spectrograms. Can be used\n", " in a multi-scale way.\n", "\n", " Parameters\n", " ----------\n", " n_mels : List[int]\n", " Number of mels per STFT, by default [150, 80],\n", " window_lengths : List[int], optional\n", " Length of each window of each STFT, by default [2048, 512]\n", " loss_fn : typing.Callable, optional\n", " How to compare each loss, by default nn.L1Loss()\n", " clamp_eps : float, optional\n", " Clamp on the log magnitude, below, by default 1e-5\n", " mag_weight : float, optional\n", " Weight of raw magnitude portion of loss, by default 1.0\n", " log_weight : float, optional\n", " Weight of log magnitude portion of loss, by default 1.0\n", " pow : float, optional\n", " Power to raise magnitude to before taking log, by default 2.0\n", " weight : float, optional\n", " Weight of this loss, by default 1.0\n", " match_stride : bool, optional\n", " Whether to match the stride of convolutional layers, by default False\n", "\n", " Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py\n", " \"\"\"\n", "\n", " def __init__(\n", " self,\n", " n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],\n", " window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],\n", " loss_fn: typing.Callable = nn.L1Loss(),\n", " clamp_eps: float = 1e-5,\n", " mag_weight: float = 1.0,\n", " log_weight: float = 1.0,\n", " pow: float = 1.0,\n", " weight: float = 1.0,\n", " match_stride: bool = False,\n", " mel_fmin: List[float] = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],\n", " mel_fmax: List[float] = [None, None, None, None, None, None, None],\n", " window_type: str = None,\n", " ):\n", " super().__init__()\n", " stft_params = [\n", " STFTParams(\n", " window_length=w,\n", " hop_length=w // 4,\n", " match_stride=match_stride,\n", " window_type=window_type,\n", " )\n", " for w in window_lengths\n", " ]\n", " self.n_mels = n_mels\n", " self.loss_fn = loss_fn\n", " self.clamp_eps = clamp_eps\n", " self.log_weight = log_weight\n", " self.mag_weight = mag_weight\n", " self.weight = weight\n", " self.mel_fmin = mel_fmin\n", " self.mel_fmax = mel_fmax\n", " self.pow = pow\n", " self.melspecs = nn.ModuleList([\n", " MelSpectrogram(\n", " sample_rate=48000,\n", " n_mels=n_mel,\n", " f_min=fmin,\n", " f_max=fmax,\n", " n_fft=s.window_length,\n", " win_length=s.window_length,\n", " hop_length=s.hop_length,\n", " power=pow,\n", " )\n", " for n_mel, fmin, fmax, s in zip(n_mels, mel_fmin, mel_fmax, stft_params)\n", " ])\n", " self.eps = 1e-10\n", "\n", " def forward(self, x, y):\n", " \"\"\"Computes mel loss between an estimate and a reference\n", " signal.\n", "\n", " Parameters\n", " ----------\n", " x : torch.tensor\n", " Estimate signal\n", " y : torch.tensor\n", " Reference signal\n", "\n", " Returns\n", " -------\n", " torch.Tensor\n", " Mel loss.\n", " \"\"\"\n", " loss = 0.0\n", " for mel_spec in self.melspecs:\n", " x_mels = mel_spec(x)\n", " y_mels = mel_spec(y)\n", "\n", " loss += self.loss_fn(torch.log10(x_mels + self.eps), torch.log10(y_mels + self.eps))\n", " return loss\n", " \n", "\n", "mel_loss = MelSpectrogramLoss()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load audio and encode then decode\n", "audio_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "x, sr = torchaudio.load(audio_filepath)\n", "\n", "if sr != 48000:\n", " x = torchaudio.functional.resample(x, sr, 48000)\n", "\n", "start_idx = 0\n", "end_idx = start_idx + int(48000*10.0)\n", "x = x[:,start_idx:end_idx]\n", "x = x.to(\"cuda:0\")\n", "\n", "x = torch.zeros_like(x)\n", "length = x.shape[-1]\n", "\n", "with torch.no_grad():\n", " x = models[\"dac_1e-4_128d_25hz\"].preprocess(x.unsqueeze(0), 48000)\n", " print(x.shape)\n", " results = models[\"dac_1e-4_128d_25hz\"].encode(x)\n", " z = results[\"z\"].half().float()\n", " print(z.shape)\n", " audio_reconstructed = models[\"dac_1e-4_128d_25hz\"].decode(z)\n", " print(audio_reconstructed)\n", " print(audio_reconstructed.shape)\n", " audio_reconstructed = audio_reconstructed[..., :length]\n", "\n", " mel_loss_value = mel_loss(x.cpu(), audio_reconstructed.cpu())\n", " print(mel_loss_value)\n", "\n", "IPython.display.Audio(audio_reconstructed.cpu().squeeze().numpy(), rate=48000)\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "torch.log10(torch.tensor(1e-10))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "x = x.cpu()\n", "y = audio_reconstructed.cpu()\n", "\n", "loss_fn = nn.L1Loss()\n", "eps = 1e-10\n", "\n", "loss = 0.0\n", "for mel_spec in mel_loss.melspecs:\n", " x_mels = mel_spec(x)\n", " y_mels = mel_spec(y)\n", "\n", " print(x_mels.shape)\n", " print(y_mels.shape)\n", "\n", " fig, axs = plt.subplots(2, 1)\n", " axs[0].pcolormesh(torch.log10(x_mels[0, 0].squeeze() + eps).numpy())\n", " axs[1].pcolormesh(torch.log10(y_mels[0, 0].squeeze() + eps).numpy())\n", " #print(torch.log10(x_mels[0, 0].squeeze()).numpy().shape)\n", " print(y_mels[0, 0].squeeze().numpy())\n", " print()\n", " print(y_mels[0, 0].squeeze().numpy() + eps)\n", " print()\n", " print(torch.log10(y_mels[0, 0].squeeze() + eps).numpy())\n", " plt.show()\n", "\n", " loss_value = loss_fn(torch.log10(x_mels + eps), torch.log10(y_mels + eps))\n", " loss += loss_value\n", " print(loss_value)\n", "\n", "print(loss)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "centroid_path = \"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\"\n", "centroids = read_from_s3(centroid_path, read_f=np.load)\n", "print(centroids.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio_filepaths = glob.glob(\"/home/christian/audio/reference-audio-wav/*.wav\")\n", "print(len(audio_filepaths))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# filter args\n", "device=\"cuda:0\"\n", "ckpt_path = \"/app/suno/christian/checkpoints/dac/mw_vae_peaq_128_fix/mw_vae_peaq_128_fix.pth\"\n", "sd = torch.load(ckpt_path)\n", "sd[\"metadata\"][\"kwargs\"] = {\n", " k: v for k, v in sd[\"metadata\"][\"kwargs\"].items() if k in DAC.__init__.__code__.co_varnames\n", "}\n", "model_25hz = DAC(**sd[\"metadata\"][\"kwargs\"])\n", "model_25hz.load_state_dict(sd[\"state_dict\"])\n", "model_25hz.eval()\n", "model_25hz.to(device)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load disriminator\n", "\n", "device=\"cuda:0\"\n", "model_name = \"100hz_128_vae_peaq_kl_0.005\"\n", "ckpt_path = f\"/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/discriminator/weights.pth\"\n", "\n", "if not os.path.isfile(ckpt_path):\n", " raise ValueError(f\"Checkpoint not found: {ckpt_path}\") \n", "\n", "print(f\"Loading model {model_name} from {ckpt_path}\")\n", "sd = torch.load(ckpt_path)\n", "sd[\"metadata\"][\"kwargs\"] = {\n", " k: v for k, v in sd[\"metadata\"][\"kwargs\"].items() if k in DAC.__init__.__code__.co_varnames\n", "}\n", "model_disc = Discriminator_import(**sd[\"metadata\"][\"kwargs\"])\n", "model_disc.load_state_dict(sd[\"state_dict\"])\n", "model_disc.eval()\n", "model_disc.to(device)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from dac.model.dac4 import DAC\n", "from suno_utils.utils.s3 import read_from_s3\n", "\n", "models = {}\n", "# 25hz_128_vae_peaq_kl_0.005\n", "model_names = [\"25hz_128_vae_kl_less_freq\"] \n", "\n", "# load models\n", "#\"/app/suno/christian/checkpoints/vae/25hz_64_vae_peaq_kl_0.005/best/dac/weights.pth\"\n", "\n", "for model_name in model_names:\n", " # filter args\n", " device=\"cuda:0\"\n", " ckpt_path = f\"/app/suno/christian/checkpoints/vae/{model_name}/best/dac/weights.pth\"\n", "\n", " if not os.path.isfile(ckpt_path):\n", " raise ValueError(f\"Checkpoint not found: {ckpt_path}\") \n", "\n", " print(f\"Loading model {model_name} from {ckpt_path}\")\n", " sd = torch.load(ckpt_path)\n", " sd[\"metadata\"][\"kwargs\"] = {\n", " k: v for k, v in sd[\"metadata\"][\"kwargs\"].items() if k in DAC.__init__.__code__.co_varnames\n", " }\n", " #if \"vae_noise_kl_0.005\" in model_name:\n", " # sd[\"metadata\"][\"kwargs\"][\"noise_enhance\"] = True\n", " print(sd[\"metadata\"][\"kwargs\"])\n", " model_vae = DAC(**sd[\"metadata\"][\"kwargs\"])\n", " model_vae.load_state_dict(sd[\"state_dict\"])\n", " model_vae.eval()\n", " model_vae.to(device)\n", " models[model_name] = model_vae" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from suno_utils.tasks.sac_vae_25hz_lowdim import preload_models\n", "from suno_utils.tasks.sac_vae_25hz_lowdim import encode, decode\n", "\n", "_ = preload_models(checkpoint_filepath = \"s3://suno-data/minz/models/sac_vae_25hz_64d.pth\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load baseline 100Hz VAE model\n", "device = \"cuda:0\"\n", "\n", "checkpoint_filepath = \"s3://suno-data/christian/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", "model_25hz = DAC(**sd[\"metadata\"][\"kwargs\"])\n", "model_25hz.load_state_dict(sd[\"state_dict\"])\n", "model_25hz.eval()\n", "model_25hz.to(device)\n", "\n", "models[\"25hz_vae_peaq_kl_0.005\"] = model_25hz\n", "\n", "# load prod DAC\n", "\n", "checkpoint_filepath = \"s3://suno-data/georg/models/codec/dac_2c_25x12.pt\"\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", "model_25hz = DAC(**sd[\"metadata\"][\"kwargs\"])\n", "model_25hz.load_state_dict(sd[\"state_dict\"])\n", "model_25hz.eval()\n", "model_25hz.to(device)\n", "\n", "models[\"dac_2c_25x12\"] = model_25hz" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#checkpoint_filepath = \"s3://suno-data/minz/models/sac_vae_25hz.pth\"\n", "checkpoint_filepath = \"s3://suno-data/minz/models/sac_vae_25hz_64d.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 SAC.__init__.__code__.co_varnames\n", "}\n", "model_minz_25hz = SAC(**sd[\"metadata\"][\"kwargs\"])\n", "model_minz_25hz.load_state_dict(sd[\"state_dict\"])\n", "model_minz_25hz.eval()\n", "model_minz_25hz.to(device)\n", "\n", "models[\"sac_vae_25hz\"] = model_minz_25hz" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(models.keys())" ] }, { "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 Take Five.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", "]\n", "\n", "audio_paths = glob.glob(os.path.join(\"/home/christian/audio/reference-audio-wav/*.wav\"))\n", "#audio_paths = glob.glob(os.path.join(\"/home/christian/audio/50_genre_songs/*.mp3\"))\n", "\n", "file_ext = \".wav\"\n", "#outdir = \"outputs/codec-50-genres-20241126\"\n", "outdir = \"outputs/reference-audio-20250121\"\n", "os.makedirs(outdir, exist_ok=True)\n", "\n", "meter = pyln.Meter(48000)\n", "\n", "# code cycle examples\n", "\n", "for audio_filepath in tqdm(audio_paths):\n", " print(audio_filepath)\n", " audio, sr = torchaudio.load(audio_filepath)\n", "\n", " for float_16 in [True]:\n", "\n", " # normalize the whole tracks\n", " target_lufs = -16.0\n", " in_lufs = meter.integrated_loudness(audio.cpu().permute(1, 0).numpy().squeeze())\n", " delta_lufs = target_lufs - in_lufs\n", " audio *= 10 ** (delta_lufs / 20.0)\n", " print(audio.abs().max())\n", "\n", " for section in [\"middle\"]:\n", "\n", " if section == \"middle\":\n", " start_idx = audio.shape[-1] // 2\n", " elif section == \"begin\":\n", " start_idx = 524288\n", " else:\n", " start_idx = 0\n", "\n", " end_idx = start_idx + 524288\n", " \n", " audio_section = audio[:, start_idx:end_idx]\n", " \n", " audio_48k = torchaudio.functional.resample(audio_section, sr, 48000)\n", " #audio_12k = torchaudio.functional.resample(audio, sr, 12000)\n", "\n", "\n", " audio_48k = audio_48k.to(\"cuda:0\")\n", " #audio_48k *= 10 ** (-6.0/20.0) # 6dB of headroom\n", "\n", " filename = os.path.basename(audio_filepath).replace(file_ext, \"\")\n", " input_filepath = os.path.join(outdir, f\"{filename}-{section}-input.wav\")\n", " #cycled_25hz_filepath = os.path.join(outdir, f\"{filename}-{section}-25hz-cycled.wav\")\n", " #cycled_100hz_mean_filepath = os.path.join(outdir, f\"{filename}-{section}-100hz-mean-cycled.wav\")\n", "\n", " torchaudio.save(input_filepath, audio_48k.cpu().squeeze(), 48000)\n", " \n", " for model_name, model in models.items():\n", " with torch.no_grad():\n", " outputs = model.encode(audio_48k.unsqueeze(0))\n", "\n", " if \"dac_2c_25x12\" in model_name:\n", " codes = outputs[\"codes\"]\n", " z, _, _ = model.quantizer.from_codes(codes)\n", " else:\n", " z = outputs[\"z\"]\n", " \n", " # decode\n", " if float_16:\n", " z = z.half()\n", " z = z.float()\n", "\n", " audio_cycled = model.decode(z)\n", "\n", " #mean = outputs[\"mean\"]\n", " #scale = outputs[\"scale\"]\n", " #print(\"z\", \"min\", z.min(), \"max\", z.max(), \"mean\", z.mean(), \"std\", z.std())\n", " float_type = \"float16\" if float_16 else \"float32\"\n", " audio_cycled_filepath = os.path.join(outdir, f\"{filename}-{section}-{float_type}-cycled-{model_name}.wav\")\n", " torchaudio.save(audio_cycled_filepath, audio_cycled.cpu().squeeze(), 48000)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "torch.nn.functional.softplus(torch.tensor(1.0)) + 1e-4" ] }, { "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 }