{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import glob\n", "import torchaudio\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"5\"\n", "import torch\n", "import pyloudnorm as pyln\n", "from dac.model.discriminator2 import Discriminator as Discriminator_import" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load discriminator\n", "from dac.model.discriminator2 import Discriminator as Discriminator_import\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 Discriminator_import.__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": 5, "metadata": {}, "outputs": [], "source": [ "def preprocess(y):\n", " # Remove DC offset\n", " y = y - y.mean(dim=-1, keepdims=True)\n", " # Peak normalize the volume of input audio\n", " y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9)\n", " return y\n", "\n", "def evaluate_discriminator(model, audio: torch.Tensor):\n", " with torch.no_grad():\n", " y = model(preprocess(audio.unsqueeze(0)))\n", " loss_g = 0\n", " for y_elem in y:\n", " loss_g += torch.mean((1 - y_elem[-1]) ** 2)\n", " return loss_g\n", "\n", "def loudness_normalize(audio: torch.Tensor, taget_loudness: float = -16.0, sr: int = 48000):\n", " meter = pyln.Meter(sr) # create loudness meter\n", " loudness = meter.integrated_loudness(audio.permute(1, 0).numpy())\n", " loudness_delta = taget_loudness - loudness\n", " loudness_delta_ln = 10 ** (loudness_delta / 20)\n", " audio = audio * loudness_delta_ln\n", " return audio" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "x_real, _ = torchaudio.load(\"/home/christian/code/christian/notebooks/outputs/vae-17092024/02 Dreams-begin-input.wav\")\n", "x_real = loudness_normalize(x_real)\n", "# loudness normalize\n", "x_real = x_real.to(device)\n", "\n", "\n", "x_cycled, _ = torchaudio.load(\"/home/christian/code/christian/notebooks/outputs/vae-17092024/02 Dreams-begin-cycled-25hz_128_vae_peaq_kl_0.005.wav\")\n", "x_cycled = loudness_normalize(x_cycled)\n", "x_cycled = x_cycled.to(device)\n", "\n", "#x_cycled = x_real + torch.randn_like(x_real) * 0.6\n", "\n", "real_loss = evaluate_discriminator(model_disc, x_real)\n", "cycled_loss = evaluate_discriminator(model_disc, x_cycled)\n", "\n", "# seems like higher means better quality \n", "print(f\"Real loss: {real_loss}\")\n", "print(f\"Cycled loss: {cycled_loss}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "base_dir = \"outputs/step-test2\"\n", "disct_losses = []\n", "for steps in [2, 4, 8, 16, 32, 64]:\n", " x_a, _ = torchaudio.load(f\"{base_dir}/golden-steps={steps}.mp3\")\n", " #x_a = loudness_normalize(x_a)\n", " x_a = x_a.to(device)\n", " a_loss = evaluate_discriminator(model_disc, x_a)\n", " disct_losses.append(a_loss.item())\n", " print(f\"steps={steps}: {a_loss}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from ear.utils import load_audio, apply_normalization\n", "from ear.system import EarSystem\n", "# load pretrained ear model\n", "ckpt_path = \"/home/christian/code/christian/checkpoints/w5p4nhzn-epoch=27.cpkt\"\n", "if not os.path.isfile(ckpt_path):\n", " os.system(\n", " f\"aws s3 cp s3://suno-data/christian/ear/w5p4nhzn-epoch=27.cpkt /home/christian/code/christian/checkpoints\"\n", " )\n", "system = EarSystem.load_from_checkpoint(ckpt_path)\n", "system.cuda()\n", "system.eval()\n", "\n", "NUM_FRAMES = 131072" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load reference audio used for quality comparision\n", "# ref_dir = \"/app/suno/christian/data/codec_audio/reference-audio-wav-mono-24khz/\"\n", "# ref_filepaths = glob.glob(os.path.join(ref_dir, \"*.input.wav\"))\n", "# ref_filepaths = np.random.choice(ref_filepaths, num_compare)\n", "ref_filepaths = [\n", " \"/home/christian/audio/reference-audio-wav-mono-24khz/02 Dreams.wav\",\n", " \"/home/christian/audio/reference-audio-wav-mono-24khz/01 Mario Takes A Walk.wav\",\n", " \"/home/christian/audio/reference-audio-wav-mono-24khz/02 Freddie Freeloader.wav\",\n", " \"/home/christian/audio/reference-audio-wav-mono-24khz/09 Sounds Like Hallelujah.wav\",\n", " \"/home/christian/audio/reference-audio-wav-mono-24khz/03 Your New Aesthetic.wav\",\n", "]\n", "\n", "ref_audios = [\n", " load_audio(\n", " filepath,\n", " num_frames=NUM_FRAMES,\n", " target_sample_rate=system.hparams.sample_rate,\n", " )\n", " for filepath in ref_filepaths\n", "]\n", "ref_audios = torch.stack(ref_audios)\n", "print(\"ref_audios\", ref_audios.shape)\n", "ref_audio = ref_audios.cuda()\n", "\n", "# first precompute the reference embeddings\n", "ref_embeds = system.embed(ref_audios)\n", "print(\"ref_embeds\", ref_embeds.shape)\n" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [], "source": [ "def evaluate_ear(audio: torch.Tensor, ref_embeds: torch.Tensor):\n", " bs = audio.shape[0]\n", " num_refs = ref_embeds.shape[0]\n", "\n", " # first, embed the audio that will be evaluated\n", " with torch.no_grad():\n", " eval_embeds = system.embed(audio)\n", "\n", " # aggregate the eval_embed\n", " eval_embeds = eval_embeds.mean(dim=1, keepdim=True)\n", " ref_embeds = ref_embeds.mean(dim=1, keepdim=True)\n", "\n", " print(\"eval_embeds\", eval_embeds.shape)\n", " print(\"ref_embeds\", ref_embeds.shape)\n", "\n", " # eval_embeds has shape (bs, embed_dim)\n", " # ref_embeds has shape (num_refs, embed_dim)\n", " # now copy the eval and reference embeds to evaluate against all\n", " ref_embeds = ref_embeds.repeat(bs, 1, 1)\n", " eval_embeds = eval_embeds.repeat(num_refs, 1, 1)\n", "\n", " # concat embeds into singular tensors\n", " embeds = torch.cat((eval_embeds, ref_embeds), dim=-1)\n", " # print(\"embeds\", embeds.shape)\n", "\n", " # no run through the projection to make predictions\n", " with torch.no_grad():\n", " pref_preds = system.pref_classifier(embeds)\n", " quant_preds = system.quant_classifier(embeds)\n", "\n", " # print(pref_preds.shape, quant_preds.shape)\n", "\n", " # get a final score by taking mean across seq of preds\n", " pref_preds = pref_preds.mean(dim=1).squeeze(1)\n", " quant_preds = quant_preds.mean(dim=1).squeeze(1)\n", " pref = torch.sigmoid(pref_preds)\n", " quant = torch.argmax(quant_preds, dim=1).float()\n", "\n", " # aggregate predictions across the reference recordings\n", " prefs = pref.view(bs, -1).mean()\n", " quants = quant.view(bs, -1).mean()\n", " scores = -((prefs * 2) - 1) * (quants + 1)\n", "\n", " return prefs\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "base_dir = \"outputs/step-test2\"\n", "quality_losses = []\n", "for steps in [2, 4, 8, 16, 32, 64]:\n", " x_a, _ = torchaudio.load(f\"{base_dir}/golden-steps={steps}.mp3\")\n", " x_a = loudness_normalize(x_a)\n", " x_a = x_a.to(device)\n", " a_loss = evaluate_ear(x_a.unsqueeze(0), ref_embeds)\n", " quality_losses.append(1 - a_loss.item())\n", " print(f\"steps={steps}: {a_loss}\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# make a plot of losses across steps\n", "import matplotlib.pyplot as plt\n", "\n", "fig, axs = plt.subplots(2, 1, figsize=(5, 4), sharex=True)\n", "\n", "axs[0].plot(disct_losses, label=\"Discriminator\")\n", "axs[1].plot(quality_losses, label=\"Quality\")\n", "axs[1].set_xticks(range(len(disct_losses)), [2, 4, 8, 16, 32, 64])\n", "axs[0].set_ylabel(\"Discriminator\")\n", "axs[1].set_xlabel(\"Steps\")\n", "axs[1].set_ylabel(\"Quality\")\n", "axs[0].grid(c=\"lightgray\")\n", "axs[1].grid(c=\"lightgray\")\n", "plt.show()" ] }, { "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 }