{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import torch\n", "import funcy\n", "import IPython\n", "import numpy as np\n", "\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\"\n", "\n", "from suno_utils.utils.text import ( \n", " write_jsonl,\n", " read_jsonl,\n", " write_json,\n", " read_json,\n", " normalize_whitespace,\n", ")\n", "\n", "from dac.model.dac4 import DAC\n", "from suno_utils.utils.s3 import read_from_s3\n", "\n", "import torchaudio\n", "import matplotlib.pyplot as plt" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [], "source": [ "import torch\n", "\n", "def zero_crossing_rate(audio, frame_size, hop_length):\n", " \"\"\"\n", " Compute the zero crossing rate of the given audio signal.\n", " \n", " Args:\n", " audio (torch.Tensor): 1D tensor containing the audio samples.\n", " frame_size (int): The size of each frame.\n", " hop_length (int): The hop length between frames.\n", " \n", " Returns:\n", " torch.Tensor: A tensor containing the zero crossing rate for each frame.\n", " \"\"\"\n", " # Number of frames to process\n", " num_frames = 1 + (len(audio) - frame_size) // hop_length\n", " \n", " # Initialize tensor to store zero crossing rate\n", " zcr = torch.zeros(num_frames)\n", " \n", " # Iterate over frames\n", " for i in range(num_frames):\n", " start = i * hop_length\n", " end = start + frame_size\n", " frame = audio[start:end]\n", " \n", " # Compute zero crossings\n", " # We compare signs of adjacent samples; a crossing occurs where the product is negative\n", " crossings = torch.where(frame[:-1] * frame[1:] < 0, 1.0, 0.0)\n", " zcr[i] = torch.mean(crossings)\n", " \n", " return zcr" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn.functional as F\n", "\n", "def spectral_flatness(audio, frame_size, hop_length, eps=1e-10):\n", " \"\"\"\n", " Compute the spectral flatness of the given audio signal.\n", " \n", " Args:\n", " audio (torch.Tensor): 1D tensor containing the audio samples.\n", " frame_size (int): The size of each frame for STFT.\n", " hop_length (int): The hop length for STFT.\n", " eps (float): A small number to avoid division by zero.\n", " \n", " Returns:\n", " torch.Tensor: A tensor containing the spectral flatness of each frame.\n", " \"\"\"\n", " # Compute the Short-Time Fourier Transform (STFT)\n", " stft = torch.stft(audio, n_fft=frame_size, hop_length=hop_length, window=torch.hann_window(frame_size), return_complex=True)\n", " \n", " # Compute the magnitude of the complex numbers in STFT\n", " magnitude = torch.abs(stft)\n", " print(magnitude.shape)\n", " \n", " # Compute the geometric mean\n", " geometric_mean = torch.exp(torch.mean(torch.log(magnitude + eps), dim=0))\n", " \n", " # Compute the arithmetic mean\n", " arithmetic_mean = torch.mean(magnitude, dim=0)\n", " \n", " # Compute the spectral flatness\n", " flatness = geometric_mean / (arithmetic_mean + eps)\n", " \n", " return flatness\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sys\n", "sys.path.insert(0, \"/home/christian/code/christian/scripts\")\n", "\n", "from train_ear import AudioQualityModel, create_label_encoder, CorruptAudioDataset\n", "\n", "#model_filepath = \"/app/suno/christian/checkpoints/ear-v2/2024-12-20_11-26-16_s1551/last_ckpt.pt\"\n", "#model_filepath = \"/app/suno/christian/checkpoints/ear-v2/2025-01-06_14-26-49_s7229/last_ckpt.pt\" # ft\n", "#model_filepath = \"/app/suno/christian/checkpoints/ear-v2/2025-02-10_14-44-30_s3169/last_ckpt.pt\"\n", "#model_filepath = \"/app/suno/christian/checkpoints/ear-v2/2025-02-10_17-07-31_s5412/last_ckpt.pt\"\n", "#model_filepath = \"/app/suno/christian/checkpoints/ear-v2/2025-02-11_22-03-40_s9362/last_ckpt.pt\" # finetune with fewer corruptions\n", "#model_filepath = \"/app/suno/christian/checkpoints/ear-v2/2025-02-14_17-09-21_s8910/last_ckpt.pt\"\n", "model_filepath = \"/app/suno/christian/checkpoints/ear-v2/2025-02-24_15-31-21_s8392/last_ckpt.pt\"\n", "\n", "ckpt = torch.load(model_filepath)\n", "model = AudioQualityModel(**ckpt[\"run_config\"][\"model\"])\n", "state_dict = ckpt[\"model\"]\n", "new_state_dict = {}\n", "for key, value in state_dict.items():\n", " new_key = key.replace(\"module.\", \"\")\n", " new_state_dict[new_key] = value\n", "model.load_state_dict(new_state_dict)\n", "model.eval()\n", "model.cuda()\n", "\n", "# also load corruptions config\n", "#corruptions = ckpt[\"corruptions\"]\n", "#print(corruptions)\n", "#label_encoder = create_label_encoder(corruptions)#\n", "#print(len(label_encoder))\n", "#print(label_encoder)\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "def compute_spectral_centroid(audio_tensor, sample_rate, n_fft=4096, hop_size=2048):\n", " # Split into frames\n", " frame_length = n_fft\n", " hop_length = hop_size\n", " frames = audio_tensor.unfold(1, frame_length, hop_length)\n", " \n", " # Apply Hann window\n", " window = torch.hann_window(frame_length, device=audio_tensor.device)\n", " frames = frames.squeeze(0) * window.unsqueeze(0)\n", " \n", " # Compute FFT for each frame\n", " spectrum = torch.fft.rfft(frames) # [num_frames, n_fft//2 + 1]\n", " freqs = torch.fft.rfftfreq(n_fft, d=1/sample_rate) # [n_fft//2 + 1]\n", " \n", " # Compute magnitudes for each frame\n", " magnitudes = torch.abs(spectrum) # [num_frames, n_fft//2 + 1]\n", " \n", " # Compute centroid for each frame\n", " numerator = torch.sum(freqs.view(1, -1) * magnitudes, dim=1) # Sum over frequencies for each frame\n", " denominator = torch.sum(magnitudes, dim=1)\n", "\n", " # Compute mean centroid across all frames\n", " centroid = numerator / (denominator + 1e-8)\n", "\n", " return centroid\n" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "from suno_utils.tasks.dac_vae_fixed_25hz import encode, decode, preload_models\n", "_ = preload_models(checkpoint_filepath=\"s3://suno-data/minz/models/dac_vae_tuned_25hz.pth\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# lets manually encode a waveform\n", "x, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Take Five.wav\")\n", "x = torchaudio.functional.resample(x, sr, 48000)\n", "\n", "# now lets encode it\n", "latents = encode(x)\n", "print(latents.shape)\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# lets chunk the latents in 750 token chunks (30s)\n", "latents_chunks = torch.from_numpy(latents).permute(1, 0).unfold(1, 750, 750)\n", "print(latents_chunks.shape)\n", "\n", "# now lets measure the statistics of each chunk\n", "# get the mean, and std of each chunk\n", "mean_latents = latents_chunks.mean(dim=(0, 2))\n", "std_latents = latents_chunks.std(dim=(0, 2))\n", "\n", "# now lets plot the statistics\n", "plt.figure(figsize=(10, 5))\n", "#plt.plot(mean_latents.numpy(), label=\"mean\")\n", "plt.plot(std_latents.numpy(), label=\"std\")\n", "plt.legend()\n", "plt.show()\n", "\n", "print(mean_latents.shape, std_latents.shape)\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "\n", "filepaths = [\n", " #\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\",\n", " #\"/home/christian/audio/reference-audio-wav/02 Take Five.wav\",\n", " #\"/home/christian/code/christian/notebooks/audio/decay-1.mp3\", # not a lot\n", " #\"/home/christian/code/christian/notebooks/audio/decay-2.mp3\", # not a lot\n", " #\"/home/christian/code/christian/notebooks/audio/decay-3.mp3\",\n", " #\"/home/christian/code/christian/notebooks/audio/decay-4.mp3\",\n", " #\"/home/christian/code/christian/notebooks/audio/decay-5.mp3\",\n", " #\"/home/christian/code/christian/notebooks/audio/decay-6.mp3\",\n", " #\"/home/christian/code/christian/notebooks/audio/a-with-ctx.mp3\",\n", " #\"/home/christian/code/christian/notebooks/audio/a-without-ctx.mp3\"\n", " \"/home/christian/code/christian/notebooks/audio/dreams-no-ctx.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-2.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-3.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-4.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-5.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-6.mp3\", \n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-7.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-8.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-9.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-10.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-11.mp3\",\n", "]\n", "\n", "plt.figure(figsize=(10, 5))\n", "\n", "slope_results = {}\n", "for filepath in filepaths:\n", " x, sr = torchaudio.load(filepath)\n", " x = torchaudio.functional.resample(x, sr, 48000)\n", " #x = x.cuda()\n", "\n", " n_fft = 4096\n", " hop_size = n_fft // 2\n", "\n", "\n", " # using a window size of \n", " # given the n_fft and hop_size, compute the duration of the window in seconds\n", " window_duration_s = n_fft / sr\n", " num_frames_per_30s = 30 / window_duration_s\n", "\n", " scores = compute_spectral_centroid(x.mean(dim=0, keepdim=True), 48000, n_fft, hop_size)\n", "\n", " #scores = scores.squeeze().cpu().numpy()\n", " #scores = - torch.stack(scores).squeeze().cpu().numpy()\n", " #normalized_scores = scores - scores[0]\n", "\n", " #scores = scores[100:-100]\n", " mean_score = scores.mean()\n", "\n", " # now lets plot the spectral centroid\n", " #plt.plot(scores, label=os.path.basename(filepath))\n", " \n", "\n", " # using a window size of of 300\n", " window_size = 300\n", " # compute the mean of the scores in a window of size window_size\n", " window_scores = scores.unfold(0, window_size, window_size).mean(dim=1)\n", " #plt.plot(window_scores, label=os.path.basename(filepath))\n", "\n", " # now fit a linear regression to the window scores\n", " from sklearn.linear_model import LinearRegression\n", " import numpy as np\n", "\n", " # Create a matrix of features (window positions)\n", " X = np.arange(len(window_scores)).reshape(-1, 1)\n", " y = window_scores\n", "\n", " # Create and fit the linear regression model\n", " model = LinearRegression()\n", " model.fit(X, y)\n", "\n", " # Get the slope (rate of change)\n", " slope = model.coef_[0]\n", "\n", " # now plot the linear regression line\n", " plt.plot(X, model.predict(X), label=f\"{os.path.basename(filepath)}: {slope:.2f}\")\n", " slope_results[os.path.basename(filepath)] = slope\n", "\n", "plt.legend()\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# sort the slope results by slope\n", "sorted_slope_results = sorted(slope_results.items(), key=lambda x: x[1])\n", "\n", "# print the sorted slope results\n", "for filepath, slope in sorted_slope_results:\n", " print(f\"{filepath}: {slope:.2f}\")\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 }