{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# !pip install auraloss torchlibrosa" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n", "\n", "import torch\n", "import funcy\n", "import IPython\n", "import torchaudio\n", "import numpy as np\n", "import pyloudnorm as pyln\n", "\n", "from tqdm import tqdm\n", "from time import perf_counter\n", "from typing import Optional, List\n", "from ear.utils import load_audio, apply_normalization\n", "from ear.system import EarSystem\n", "from suno_utils.audio import Audio\n", "\n", "from suno_utils.utils.s3 import read_from_s3" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def compare_quality(system, audio_a: torch.Tensor, audio_b: torch.Tensor):\n", " \"\"\" Compare the quality of two audio files using the trained model.\n", "\n", " Parameters\n", " ----------\n", " system : EarSystem\n", " the trained model\n", " audio_a : torch.Tensor\n", " audio tensor of shape (2, num_frames)\n", " audio_b : torch.Tensor \n", " audio tensor of shape (2, num_frames)\n", "\n", " Returns\n", " -------\n", " pref : torch.Tensor\n", " preference prediction\n", "\n", " quant : torch.Tensor\n", " quantification prediction\n", "\n", " \"\"\"\n", "\n", " # move audio_a and audio_b to same device as system parameters\n", " audio_a = audio_a.to(system.device)\n", " audio_b = audio_b.to(system.device)\n", "\n", " # first, embed the audio that will be evaluated\n", " with torch.no_grad():\n", " embeds_a = system.embed(audio_a)\n", " embeds_b = system.embed(audio_b)\n", "\n", " # aggregate embeddings over time with a moving mean of frame size\n", " embeds_a = torch.nn.functional.adaptive_avg_pool1d(embeds_a.permute(0, 2, 1), 137).permute(0, 2, 1)\n", " embeds_b = torch.nn.functional.adaptive_avg_pool1d(embeds_b.permute(0, 2, 1), 137).permute(0, 2, 1)\n", "\n", " # concat embeds into singular tensors\n", " embeds = torch.cat((embeds_a, embeds_b), dim=-1)\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 and chunks\n", " pref_preds = pref_preds.mean(dim=1).mean(dim=0)\n", " quant_preds = quant_preds.mean(dim=1).mean(dim=0)\n", " pref = torch.sigmoid(pref_preds)\n", " quant = torch.argmax(quant_preds, dim=0).float()\n", "\n", " return pref, quant\n", "\n", "\n", "def prepare_audio(\n", " audio: Audio,\n", " num_frames: int,\n", " start_s: float = None,\n", " end_s: float = None,\n", "):\n", " sample_rate = audio.sample_rate\n", " audio = torch.from_numpy(audio.array_float)\n", "\n", " if audio.shape[0] != 2:\n", " audio = audio.repeat(2, 1)\n", "\n", " # crop audio based on metadata example\n", " if start_s is not None and end_s is not None:\n", " start_frame = int(start_s * sample_rate)\n", " end_frame = int(end_s * sample_rate)\n", " audio = audio[:, start_frame:end_frame]\n", "\n", " # if the file is long, only take part of it\n", " if audio.shape[-1] > (sample_rate * 120):\n", " audio = audio[:, : sample_rate * 120]\n", "\n", " # downmix and resample decoded audio to 24khz\n", " audio = torchaudio.functional.resample(audio, sample_rate, 24_000)\n", "\n", " # pad by repeating the signal if shorter than window\n", " if audio.shape[-1] < num_frames:\n", " pad_size = num_frames - audio.shape[-1]\n", " audio = torch.nn.functional.pad(audio, (1, pad_size), mode=\"replicate\")\n", "\n", " # chunk into non-overlapping blocks of num_frames\n", " audio_chunks = []\n", " num_chunks = audio.shape[-1] // num_frames\n", " for n in range(num_chunks):\n", " start_idx = n * num_frames\n", " end_idx = start_idx + num_frames\n", " audio_chunks.append(audio[:, start_idx:end_idx])\n", "\n", " # loudness norm\n", " meter = pyln.Meter(24_000)\n", "\n", " for audio_chunk_idx in range(len(audio_chunks)):\n", " x_lufs_db = meter.integrated_loudness(audio.T.numpy())\n", " if x_lufs_db == -float(\"inf\"):\n", " gain_lin = 1.0\n", " else:\n", " delta_lufs_db = -20.0 - x_lufs_db\n", " gain_lin = 10.0 ** (np.clip(delta_lufs_db, a_min=-120, a_max=48.0) / 20.0)\n", " audio_chunks[audio_chunk_idx] *= gain_lin\n", "\n", " # take the last chunk keeping the list\n", " #audio_chunks = audio_chunks[-1:]\n", "\n", " return torch.stack(audio_chunks)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load pretrained ear model\n", "device = \"cuda\"\n", "checkpoint_filepath = \"s3://suno-data/christian/ear/w5p4nhzn-epoch=27.cpkt\"\n", "load_f = funcy.partial(EarSystem.load_from_checkpoint, map_location=\"cpu\")\n", "system = read_from_s3(checkpoint_filepath, read_f=load_f)\n", "system.to(device)\n", "system.eval()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load audio files\n", "NUM_FRAMES = 131072\n", "\n", "audio_a_filepath = \"s3://suno-data/datasets/harvest/genius_hq/audio/Ej4f12q7FRw.webm\"\n", "audio_b_filepath = \"s3://suno-data/datasets/harvest/genius_hq/audio/PkMdMl_Kv8U.webm\"\n", "\n", "audio_a = Audio.from_s3(audio_a_filepath, n_channels=2)\n", "audio_b = Audio.from_s3(audio_b_filepath, n_channels=2)\n", "\n", "prep_audio_a = prepare_audio(audio_a, NUM_FRAMES)\n", "prep_audio_b = prepare_audio(audio_b, NUM_FRAMES)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# compare quality\n", "pref, quant = compare_quality(system, prep_audio_a, prep_audio_b)\n", "print(pref)\n", "\n", "print(\"audio_a is preferred\") if pref < 0.5 else print(\"audio_b is preferred\")\n", "print(\"quantification level (higher is larger difference):\", quant.item())\n", "# Note: quantification level is not well calibrated rn, but higher values indicate larger differences" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import pandas as pd\n", "df = pd.read_pickle(\n", " \"/home/tony/Data/Preference/30b_v2/interesting_clips_v4_t_3_20240901_full.pkl\"\n", ") \n", "print(\"Preference data shape\", df.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "df = df.sort_values(by=[\"request_id\", \"preference\"])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def process_request_id(request_id, negative_id, positive_id):\n", " \"\"\"Take the request id, pos and neg ids, and return the preference and quantification\"\"\"\n", " try:\n", " negative_audio = Audio.from_s3(f\"s3://suno-data-uploads/studio/uploads/{negative_id}.mp3\", n_channels=2)\n", " positive_audio = Audio.from_s3(f\"s3://suno-data-uploads/studio/uploads/{positive_id}.mp3\", n_channels=2)\n", " negative_prep = prepare_audio(negative_audio, NUM_FRAMES)\n", " positive_prep = prepare_audio(positive_audio, NUM_FRAMES)\n", " pref, quant = compare_quality(system, negative_prep, positive_prep)\n", " # print(pref)\n", " # print(\"audio_a is preferred\") if pref < 0.5 else print(\"audio_b is preferred\")\n", " # print(\"quantification level (higher is larger difference):\", quant.item())\n", " return {request_id: (pref.cpu().numpy()[0], quant.cpu().numpy())}\n", " except Exception as e:\n", " print(e)\n", " return {request_id: None}" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "request_id_to_pref = {}\n", "request_jobs = []\n", "unique_request_ids = df[\"request_id\"].unique()\n", "grouped = df.groupby(\"request_id\")\n", "\n", "for request_id in tqdm(unique_request_ids):\n", " group = grouped.get_group(request_id)\n", " assert (group.iloc[0][\"preference\"], group.iloc[1][\"preference\"]) == (0, 1)\n", " negative_id, positive_id = group[\"s3_id\"].tolist()\n", " request_jobs.append((request_id, negative_id, positive_id))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from multiprocessing import Pool\n", "\n", "with Pool(8) as p:\n", " results = p.starmap(process_request_id, request_jobs[:8], chunksize=1)\n", "\n", "request_id_to_pref = {k: v for result in results for k, v in result.items() if v is not None}" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "request_id_to_pref" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "request_jobs[0]" ] }, { "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.14" } }, "nbformat": 4, "nbformat_minor": 2 }