{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "6c8c28ff", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "import torch.fft as fft\n" ] }, { "cell_type": "code", "execution_count": null, "id": "995de6b4", "metadata": {}, "outputs": [], "source": [ "def third_octave_bands(sr, fmin=20.0, fmax=None):\n", " \"\"\"\n", " Compute 1/3-octave band center frequencies and edges.\n", " \"\"\"\n", " if fmax is None:\n", " fmax = sr / 2.0\n", "\n", " k = np.arange(-30, 30) # wide enough range\n", " f_center = 1000.0 * (2.0 ** (k / 3.0)) # ISO 1/3 octave centers\n", " f_center = f_center[(f_center >= fmin) & (f_center <= fmax)]\n", "\n", " f_lower = f_center / (2 ** (1 / 6))\n", " f_upper = f_center * (2 ** (1 / 6))\n", " return f_center, f_lower, f_upper\n", "\n", "\n", "def third_octave_response_db(waveform: torch.Tensor, sr: int):\n", " \"\"\"\n", " Compute 1/3 octave magnitude response in dB from waveform.\n", "\n", " Args:\n", " waveform (torch.Tensor): shape (n_samples,) or (1, n_samples)\n", " sr (int): sample rate\n", "\n", " Returns:\n", " freqs (np.ndarray): band center frequencies\n", " mags_db (torch.Tensor): band magnitudes in dB\n", " \"\"\"\n", " if waveform.ndim > 1:\n", " waveform = waveform.squeeze(0)\n", "\n", " n = waveform.numel()\n", " spec = fft.rfft(waveform)\n", " mag = torch.abs(spec) / n\n", " freqs = torch.fft.rfftfreq(n, d=1.0 / sr)\n", "\n", " # Get bands\n", " f_center, f_lower, f_upper = third_octave_bands(sr)\n", " band_mags = []\n", " for fl, fu in zip(f_lower, f_upper):\n", " idx = (freqs >= fl) & (freqs < fu)\n", " if idx.any():\n", " band_mags.append(mag[idx].mean())\n", " else:\n", " band_mags.append(torch.tensor(0.0))\n", "\n", " band_mags = torch.stack(band_mags)\n", "\n", " # Convert to dB (avoid log(0))\n", " mags_db = 20 * torch.log10(band_mags + 1e-12)\n", "\n", " return f_center, mags_db" ] }, { "cell_type": "code", "execution_count": null, "id": "38d7a9df", "metadata": {}, "outputs": [], "source": [ "from suno_utils.audio import Audio\n", "from suno_utils.utils.text import read_jsonl\n", "\n", "work_items = read_jsonl(\n", " \"/home/christian/code/christian/metadata/ear/genius_t6_sampled_10k.jsonl\"\n", ")\n", "print(len(work_items))" ] }, { "cell_type": "code", "execution_count": null, "id": "2aeffecf", "metadata": {}, "outputs": [], "source": [ "item" ] }, { "cell_type": "code", "execution_count": null, "id": "4366c15e", "metadata": {}, "outputs": [], "source": [ "from tqdm import tqdm\n", "import matplotlib.pyplot as plt\n", "\n", "def get_spectrum(item):\n", " try:\n", " audio = Audio.from_s3(item[\"s3_filepath\"], n_channels=2)\n", " audio_tensor = torch.from_numpy(audio.array_float)\n", " return third_octave_response_db(audio_tensor.mean(dim=0), audio.sample_rate)\n", " except Exception as e:\n", " print(f\"Error getting spectrum for {item['s3_filepath']}: {e}\")\n", " return None\n", "\n", "from joblib import Parallel, delayed\n", "\n", "# Use joblib to parallelize get_spectrum over work_items\n", "spectrums = Parallel(n_jobs=-1)(\n", " delayed(get_spectrum)(item) for item in tqdm(work_items)\n", ")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "feb9bdf8", "metadata": {}, "outputs": [], "source": [ "# save all the spectrums to a file\n", "# we can put this into a csv file with the freqs as the columns and the mags for each freq as rows\n", "\n", "import pandas as pd\n", "\n", "# Create a DataFrame with freqs as columns and spectrums as rows\n", "columns = [f\"{int(np.ceil(f))}\" for f in spectrums[0][0]]\n", "spectrum_df = pd.DataFrame([[v.item() for v in s[1]] for s in spectrums if s is not None], columns=columns)\n", "\n", "# save the dataframe to a csv file\n", "spectrum_df.to_csv(\"/home/christian/code/christian/metadata/genius_t6_sampled_10k_spectrums.csv\", index=False)\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "b2a1dd58", "metadata": {}, "outputs": [], "source": [ "spectrum_df.describe()" ] }, { "cell_type": "code", "execution_count": null, "id": "c3851d11", "metadata": {}, "outputs": [], "source": [ "print(len(spectrums))" ] }, { "cell_type": "code", "execution_count": null, "id": "c4096dca", "metadata": {}, "outputs": [], "source": [ "import pandas as pd\n", "spectrum_df = pd.read_csv(\"/home/christian/code/christian/metadata/genius_t6_sampled_10k_spectrums.csv\")\n", "# convert the data from csv to a list of tuples\n", "spectrums = spectrum_df.to_numpy()\n", "# now we have a list of spectrums so we want to compute the mean spectrum" ] }, { "cell_type": "code", "execution_count": null, "id": "f9380384", "metadata": {}, "outputs": [], "source": [ "spectrums.shape" ] }, { "cell_type": "code", "execution_count": null, "id": "d30fca96", "metadata": {}, "outputs": [], "source": [ "# now we have a list of spectrums so we want to compute the mean spectrum\n", "# keep in mind each spectrum is a tuple of (freqs, mags)\n", "spectrum_freqs = [s[0] for s in spectrums if s is not None]\n", "spectrum_mags = [s[1] for s in spectrums if s is not None]\n", "mean_spectrum = np.mean(spectrum_mags, axis=0)\n", "mean_spectrum_freqs = spectrum_freqs[0]\n" ] }, { "cell_type": "code", "execution_count": null, "id": "20f24e0e", "metadata": {}, "outputs": [], "source": [ "fig, ax = plt.subplots(figsize=(10, 4))\n", "plt.semilogx(mean_spectrum_freqs, mean_spectrum)\n", "# set the x axis to be the octave bands\n", "plt.xticks(mean_spectrum_freqs, [f\"{int(np.ceil(f))}\" for f in mean_spectrum_freqs], rotation=45)\n", "plt.xlabel(\"Frequency (Hz)\")\n", "plt.ylabel(\"Magnitude (dB)\")\n", "plt.title(\"Mean Spectrum of 10k Samples\")\n", "plt.grid(c=\"lightgray\")\n", "plt.tight_layout()\n", "plt.show()\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "f3f88074", "metadata": {}, "outputs": [], "source": [ "import os\n", "# once we have a mean spectrum, we want to compare two different spectrums from different diffusion outputs \n", "# against the mean, which becomes our \"reference\" spectrum\n", "# goal is to be close to the reference spectrum, so that we don't get overly bright or dark outputs (muffled)\n", "# so this can be computed as \n", "\n", "model_name = \"4n_25hz_2b_flow_5e5_sft_t8_500k\"\n", "base_dir = \"/app2/suno/data/christian/outputs/v3-base-data-ctx-rs-t3/\"\n", "\n", "dirnames = os.listdir(base_dir)\n", "# filter to only include dirs\n", "dirnames = [d for d in dirnames if os.path.isdir(os.path.join(base_dir, d))]\n", "print(len(dirnames))\n" ] }, { "cell_type": "code", "execution_count": null, "id": "463892bf", "metadata": {}, "outputs": [], "source": [ "# test some of the local data\n", "CODEC_FILEPATH = \"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\"\n", "\n", "from suno_utils.tasks.dac_vae_fixed_25hz import (\n", " preload_models as preload_codec_models,\n", " decode as codec_decode,\n", " encode as codec_encode,\n", " decode_stream_to_full_audio,\n", ")\n", "\n", "_ = preload_codec_models(CODEC_FILEPATH)\n" ] }, { "cell_type": "code", "execution_count": null, "id": "190e9d8e", "metadata": {}, "outputs": [], "source": [ "def process_dir(base_dir, dirname):\n", " results = {}\n", " semantic_codes_filepath = os.path.join(base_dir, dirname, f\"{dirname}_semantic.npz\")\n", " for n in range(2):\n", " upsampled_vae_filepath = os.path.join(base_dir, dirname, f\"{dirname}_{model_name}_{n}_upsampled_vae.npz\")\n", " metadata_filepath = os.path.join(base_dir, dirname, f\"{dirname}_{model_name}_{n}__metadata.npz\")\n", " try:\n", " with np.load(metadata_filepath, allow_pickle=True) as data:\n", " metadata_npz = dict(data)\n", " except FileNotFoundError:\n", " metadata_npz = None\n", "\n", " if metadata_npz is None:\n", " continue\n", " metadata_dict = {key: metadata_npz[key].tolist() for key in metadata_npz.keys()}\n", "\n", " results[n] = {\n", " \"upsampled_vae_filepath\": upsampled_vae_filepath,\n", " \"metadata\": metadata_dict,\n", " }\n", " return results, semantic_codes_filepath" ] }, { "cell_type": "code", "execution_count": null, "id": "90dc0e1f", "metadata": {}, "outputs": [], "source": [ "dirname = dirnames[130]\n", "results, semantic_codes_filepath = process_dir(base_dir, dirname)\n", "\n", "pos_idx = 1\n", "neg_idx = 0\n", "\n", "pos_vae_latents_filepath = results[pos_idx][\"upsampled_vae_filepath\"]\n", "neg_vae_latents_filepath = results[neg_idx][\"upsampled_vae_filepath\"]\n", "\n", "# load the vae latents\n", "pos_vae_latents = np.load(pos_vae_latents_filepath)[\"vae_latents\"]\n", "neg_vae_latents = np.load(neg_vae_latents_filepath)[\"vae_latents\"]\n", "\n", "# decode the vae latents\n", "pos_audio = codec_decode(pos_vae_latents)\n", "neg_audio = codec_decode(neg_vae_latents)\n", "\n", "# get the spectrums\n", "pos_spectrum_freqs, pos_spectrum = third_octave_response_db(torch.from_numpy(pos_audio.array_float).mean(dim=0), 48000)\n", "neg_spectrum_freqs, neg_spectrum = third_octave_response_db(torch.from_numpy(neg_audio.array_float).mean(dim=0), 48000)\n", "\n", "pos_spectrum = pos_spectrum.numpy()\n", "neg_spectrum = neg_spectrum.numpy()\n", "\n", "pos_mean_energy = np.mean(pos_spectrum)\n", "neg_mean_energy = np.mean(neg_spectrum)\n", "\n", "\n", "pos_mean_energy_delta = pos_mean_energy - np.mean(mean_spectrum)\n", "neg_mean_energy_delta = neg_mean_energy - np.mean(mean_spectrum)\n", "\n", "pos_spectrum_normalized = pos_spectrum - pos_mean_energy_delta\n", "neg_spectrum_normalized = neg_spectrum - neg_mean_energy_delta\n", "\n", "# now measure the distance (MSE) between the pos_spectrum and the mean_spectrum\n", "pos_mse = np.mean((pos_spectrum_normalized - mean_spectrum) ** 2)\n", "neg_mse = np.mean((neg_spectrum_normalized - mean_spectrum) ** 2)\n", "\n", "print(f\"pos mse: {pos_mse}\")\n", "pos_audio.play()\n", "\n", "print(f\"neg mse: {neg_mse}\")\n", "neg_audio.play()\n", "\n", "# create a plot of the pos_spectrum and neg_spectrum, against the mean spectrum \n", "fig, ax = plt.subplots(figsize=(10, 4))\n", "plt.semilogx(mean_spectrum_freqs, mean_spectrum, label=\"mean spectrum\", linewidth=2)\n", "plt.semilogx(pos_spectrum_freqs, pos_spectrum_normalized, label=\"pos spectrum\", linewidth=2)\n", "plt.semilogx(neg_spectrum_freqs, neg_spectrum_normalized, label=\"neg spectrum\", linewidth=2)\n", "plt.xlabel(\"Frequency (Hz)\")\n", "plt.ylabel(\"Magnitude (dB)\")\n", "plt.grid(c=\"lightgray\")\n", "plt.title(\"Spectrum Comparison\")\n", "plt.xticks(mean_spectrum_freqs, [f\"{int(np.ceil(f))}\" for f in mean_spectrum_freqs], rotation=45)\n", "plt.legend()\n", "plt.show()\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "042c9263", "metadata": {}, "outputs": [], "source": [ "results_dict = {}\n", "from tqdm import tqdm\n", "for dirname in tqdm(dirnames[:100]):\n", " results, semantic_codes_filepath = process_dir(base_dir, dirname)\n", "\n", " pos_idx = 1\n", " neg_idx = 0\n", "\n", " try:\n", " pos_vae_latents_filepath = results[pos_idx][\"upsampled_vae_filepath\"]\n", " neg_vae_latents_filepath = results[neg_idx][\"upsampled_vae_filepath\"]\n", " except Exception as e:\n", " print(f\"Error processing {dirname}: {e}\")\n", " continue\n", "\n", " # load the vae latents\n", " pos_vae_latents = np.load(pos_vae_latents_filepath)[\"vae_latents\"]\n", " neg_vae_latents = np.load(neg_vae_latents_filepath)[\"vae_latents\"]\n", "\n", " # decode the vae latents\n", " pos_audio = codec_decode(pos_vae_latents)\n", " neg_audio = codec_decode(neg_vae_latents)\n", "\n", " # get the spectrums\n", " pos_spectrum_freqs, pos_spectrum = third_octave_response_db(torch.from_numpy(pos_audio.array_float).mean(dim=0), 48000)\n", " neg_spectrum_freqs, neg_spectrum = third_octave_response_db(torch.from_numpy(neg_audio.array_float).mean(dim=0), 48000)\n", "\n", " pos_spectrum = pos_spectrum.numpy()\n", " neg_spectrum = neg_spectrum.numpy()\n", "\n", " pos_mean_energy = np.mean(pos_spectrum)\n", " neg_mean_energy = np.mean(neg_spectrum)\n", "\n", " pos_mean_energy_delta = pos_mean_energy - np.mean(mean_spectrum)\n", " neg_mean_energy_delta = neg_mean_energy - np.mean(mean_spectrum)\n", "\n", " pos_spectrum_normalized = pos_spectrum - pos_mean_energy_delta\n", " neg_spectrum_normalized = neg_spectrum - neg_mean_energy_delta\n", "\n", " # now measure the distance (MSE) between the pos_spectrum and the mean_spectrum\n", " pos_mse = np.mean((pos_spectrum_normalized - mean_spectrum) ** 2)\n", " neg_mse = np.mean((neg_spectrum_normalized - mean_spectrum) ** 2)\n", "\n", " results_dict[dirname] = {\n", " \"pos_mse\": pos_mse,\n", " \"neg_mse\": neg_mse,\n", " \"pos_spectrum\": pos_spectrum,\n", " \"neg_spectrum\": neg_spectrum,\n", " \"pos_spectrum_normalized\": pos_spectrum_normalized,\n", " \"neg_spectrum_normalized\": neg_spectrum_normalized,\n", " \"pos_audio\": pos_audio,\n", " \"neg_audio\": neg_audio,\n", " }\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "id": "60925b57", "metadata": {}, "outputs": [], "source": [ "# plot distribution of the pos_mse and neg_mse\n", "pos_mse_list = [v[\"pos_mse\"] for v in results_dict.values()]\n", "neg_mse_list = [v[\"neg_mse\"] for v in results_dict.values()]\n", "\n", "fig, ax = plt.subplots(figsize=(10, 4))\n", "ax.hist(pos_mse_list, bins=25, color='lightgreen', edgecolor='black', label=\"pos mse\")\n", "ax.hist(neg_mse_list, bins=25, color='lightblue', edgecolor='black', label=\"neg mse\")\n", "plt.xlabel(\"MSE\")\n", "plt.ylabel(\"Count\")\n", "plt.legend()\n", "plt.show()\n" ] }, { "cell_type": "code", "execution_count": null, "id": "afc3f905", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "86cb1f75", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "e1eec087", "metadata": {}, "outputs": [], "source": [ "from suno_utils.utils.s3 import read_from_s3\n", "\n", "def spectrum_mse_from_vae(\n", " s3_filepath: str,\n", " mean_spectrum: np.ndarray,\n", " sample_rate: int = 48000,\n", "):\n", " \"\"\"\n", " Load VAE latents from local path or s3:// URI, decode to audio, compute\n", " 1/3-octave spectrum (mono), normalize to the provided mean_spectrum level,\n", " and return the spectrum, normalized spectrum, and MSE vs mean_spectrum.\n", " \"\"\"\n", "\n", " vae_latents = read_from_s3(s3_filepath, read_f=np.load)[\"vae_latents\"]\n", "\n", " def _to_mono_torch(arr: np.ndarray) -> torch.Tensor:\n", " t = torch.from_numpy(arr)\n", " return t.mean(dim=0) if t.ndim > 1 else t\n", "\n", " # 1) Load latents and decode to audio\n", " audio = codec_decode(vae_latents) # must expose .array_float (np.ndarray)\n", "\n", " # 2) Third-octave spectrum (mono)\n", " mono = _to_mono_torch(audio.array_float)\n", " center_freqs, spectrum_t = third_octave_response_db(mono, sample_rate)\n", " spectrum = spectrum_t.numpy()\n", "\n", " # 3) Level-normalize to batch mean spectrum\n", " if spectrum.shape != mean_spectrum.shape:\n", " raise ValueError(\n", " f\"mean_spectrum shape {mean_spectrum.shape} != spectrum shape {spectrum.shape}\"\n", " )\n", "\n", " mean_of_mean = float(np.mean(mean_spectrum))\n", " spec_mean = float(np.mean(spectrum))\n", " spectrum_normalized = spectrum - (spec_mean - mean_of_mean)\n", "\n", " # 4) MSE vs mean_spectrum\n", " mse = float(np.mean((spectrum_normalized - mean_spectrum) ** 2))\n", "\n", " return {\n", " \"center_freqs\": center_freqs, # np.ndarray [n_bands]\n", " \"spectrum\": spectrum, # np.ndarray [n_bands], dB\n", " \"spectrum_normalized\": spectrum_normalized, # np.ndarray [n_bands], dB\n", " \"mse\": mse, # float\n", " \"source\": s3_filepath,\n", " }" ] }, { "cell_type": "code", "execution_count": null, "id": "9739c291", "metadata": {}, "outputs": [], "source": [ "print()" ] }, { "cell_type": "code", "execution_count": null, "id": "e1511056", "metadata": {}, "outputs": [], "source": [ "s3_filepath = f\"s3://suno-data/christian/outputs/v3-base-data-ctx-rs-t3/{dirnames[0]}/{dirnames[0]}_{model_name}_0_upsampled_vae.npz\"\n", "\n", "results = spectrum_mse_from_vae(s3_filepath, mean_spectrum)\n", "\n", "results[\"center_freqs\"]\n", "results[\"spectrum\"]\n", "results[\"spectrum_normalized\"]\n", "results[\"mse\"]\n", "results[\"source\"]" ] }, { "cell_type": "code", "execution_count": null, "id": "03c6c513", "metadata": {}, "outputs": [], "source": [ "results[\"spectrum\"]\n", "results[\"mse\"]" ] }, { "cell_type": "code", "execution_count": null, "id": "96b35d34", "metadata": {}, "outputs": [], "source": [ "import os\n", "import numpy as np\n", "from suno_utils.utils.s3 import list_s3_dir, read_from_s3\n", "from suno_utils.worker.settings import s3_client\n", "\n", "\n", "BASE_S3_DIR = \"s3://suno-data/christian/outputs/v3-base-data-ctx-rs-t3/\"\n", "\n", "dirname = \"d0523db9-9502-4962-aa08-fbe3f613a4fb\"\n", "search_dir = f\"{BASE_S3_DIR}{dirname}/\"\n", "\n", "print(f\"Searching in {search_dir}\")\n", "files = list_s3_dir(\n", " search_dir,\n", ")\n", "all_files = [f[0] for f in files]\n", "all_files = [os.path.basename(f) for f in all_files]\n", "filtered_files = [f for f in all_files if f.endswith(\".npz\")]\n", "filtered_files = [f for f in filtered_files if \"upsampled_vae\" in f]\n", "if len(filtered_files) == 0:\n", " print(f\"No upsampled vae files found for {dirname}\")\n", "\n", "print(all_files)\n", "\n", "for f in filtered_files:\n", " vae_latents = read_from_s3(\n", " f\"{BASE_S3_DIR}{dirname}/{f}\",\n", " read_f=np.load,\n", " )[\"vae_latents\"]\n", " print(vae_latents.shape)\n", "\n", " filename = f.replace(\"_upsampled_vae.npz\", \"\")\n", "\n", " s3_filepath = os.path.join(\n", " \"christian/outputs/v3-base-data-ctx-rs-t3\",\n", " f\"{dirname}\",\n", " f\"{filename}_spectrum.npz\",\n", " )\n", " # check if the file exists\n", " if f\"{filename}_spectrum.npz\" in all_files:\n", " print(f\"Spectrum already exists for {dirname}\")\n", " print()\n", " continue\n", "\n", " result = spectrum_mse_from_vae(vae_latents, mean_spectrum)\n", " print(result)\n", " with tempfile.TemporaryDirectory() as td:\n", " output_filepath = os.path.join(\n", " td,\n", " f\"{filename}_spectrum.npz\",\n", " )\n", " # Save all items in result as npz (including arrays and scalars)\n", " np.savez(output_filepath, **result)\n", "\n", " s3_client.upload_file(output_filepath, \"suno-data\", s3_filepath)\n", " print(f\"Saved spectrum to {s3_filepath}\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "0b632825", "metadata": {}, "outputs": [], "source": [ "BASE_S3_DIR = \"s3://suno-data/christian/outputs/v3-base-data-ctx-rs-t3\"\n", "dirname = \"001b8b57-3534-4213-9129-9bb46cf767e7\"\n", "\n", "# get all the files in the item_id directory\n", "files = list_s3_dir(\n", " f\"{BASE_S3_DIR}/{dirname}/\",\n", ")\n", "files = [f[0] for f in files]\n", "files = [os.path.basename(f) for f in files]\n", "files = [f for f in files if f.endswith(\".npz\")]\n", "files = [f for f in files if \"upsampled_vae\" in f]\n", "if len(files) == 0:\n", " print(f\"No upsampled vae files found for {dirname}\")\n", " \n" ] }, { "cell_type": "code", "execution_count": null, "id": "8038eba7", "metadata": {}, "outputs": [], "source": [ "files = list_s3_dir(\n", " f\"{BASE_S3_DIR}/{dirname}/\",\n", ")\n", "files = [f[0] for f in files]\n", "files = [os.path.basename(f) for f in files]\n", "files = [f for f in files if f.endswith(\".npz\")]\n", "files = [f for f in files if \"upsampled_vae\" in f]\n", "print(files)\n", "if len(files) == 0:\n", " print(f\"No upsampled vae files found for {dirname}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "8dafc9e9", "metadata": {}, "outputs": [], "source": [ "dirnames[0]" ] }, { "cell_type": "code", "execution_count": null, "id": "6f106711", "metadata": {}, "outputs": [], "source": [ "dirnames = [f.split(\"/\")[-1] for f in dirnames]\n" ] }, { "cell_type": "code", "execution_count": null, "id": "cfc7b640", "metadata": {}, "outputs": [], "source": [ "\n" ] } ], "metadata": { "kernelspec": { "display_name": "suno_diff", "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.12.9" } }, "nbformat": 4, "nbformat_minor": 5 }