{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import nnAudio.features as feat\n", "import torchaudio.transforms\n", "\n", "from suno_utils.utils.s3 import read_from_s3\n", "\n", "N_BINS = 240\n", "N_BANDS = 8\n", "\n", "\n", "def apply_log_filter(stft_output, filter_matrix):\n", " \"\"\"\n", " Apply the logarithmic filter matrix to the Short-Time Fourier Transform (STFT) output.\n", "\n", " This function applies a precomputed logarithmic filter matrix to the STFT output of an audio signal\n", " to reduce its dimensionality and to capture the energy in logarithmically spaced frequency bands.\n", "\n", " Parameters\n", " ----------\n", " stft_output : torch.Tensor\n", " A tensor representing the STFT output with shape (batch_size, num_bins, num_frames), where\n", " num_bins is the number of frequency bins and num_frames is the number of time frames.\n", " filter_matrix : torch.Tensor\n", " A tensor representing the logarithmic filter matrix with shape (num_bands, num_bins), where\n", " num_bands is the number of logarithmically spaced frequency bands.\n", "\n", " Returns\n", " -------\n", " torch.Tensor\n", " A tensor representing the filtered STFT output with shape (batch_size, num_bands, num_frames).\n", " Each band contains the aggregated energy from the corresponding set of frequency bins.\n", " \"\"\"\n", " stft_output_transposed = stft_output.transpose(1, 2)\n", " filtered_output_transposed = torch.matmul(stft_output_transposed, filter_matrix.T)\n", " filtered_output = filtered_output_transposed.transpose(1, 2)\n", " return filtered_output\n", "\n", "\n", "def evaluate_bpm(model: torch.nn.Module, eval_audio: torch.Tensor, device: str):\n", " \"\"\"\n", " Args:\n", " system (torch.nn.Module):\n", " eval_audio (torch.Tensor): Audio to evaluate with shape (bs, n_harmonics=6, n_bins, n_bands)\n", "\n", " \"\"\"\n", " with torch.no_grad():\n", " eval_audio = eval_audio.to(device)\n", " outputs = model(eval_audio)\n", " probs = torch.softmax(outputs, dim=1)\n", " confs, preds = torch.max(probs, 1)\n", " return torch.tensor([class_to_bpm(pred) for pred in preds.tolist()]), torch.tensor(\n", " confs.tolist()\n", " )\n", "\n", "\n", "def class_to_bpm(class_index, min_bpm=30, max_bpm=286, num_classes=256):\n", " \"\"\"Map a class index back to a BPM value (to the center of the class interval).\"\"\"\n", " class_width = (max_bpm - min_bpm) / num_classes\n", " bpm = min_bpm + class_width * (class_index)\n", " return bpm\n", "\n", "\n", "def compute_hcqm(y, stft_spec, band_filter, cqt_specs):\n", " \"\"\"\n", " Compute the Harmonic Constant-Q Modulation (HCQM) for an input signal.\n", "\n", " As described by Foroughmand & Peeters in\n", " \"Deep-Rhythm for Tempo Estimation and Rhythm Pattern Recognition\", 2019\n", "\n", " Parameters:\n", " - y (Tensor): The input signal tensor of shape (batch_size, num_samples).\n", " - stft_spec (STFT object): An object to compute the Short-Time Fourier Transform (STFT).\n", " - band_filter (Tensor): A filter matrix of shape (num_bands, num_bins) to apply to the STFT.\n", " - cqt_specs (list of CQT objects): A list of Constant-Q Transform (CQT) objects for different harmonics / bands\n", "\n", " Returns:\n", " - hcqm (Tensor): The computed HCQM of shape (batch_size, N_BINS, N_BANDS, N_HARMONICS), where 6 corresponds to the number of different harmonics analyzed.\n", " \"\"\"\n", " stft = stft_spec(y)\n", " stft_bands = apply_log_filter(stft, band_filter)\n", " stft_bands_flat = stft_bands.reshape(\n", " stft.size(0) * stft_bands.size(1), stft_bands.size(2)\n", " )\n", " osf_flat = onset_strength(y=stft_bands_flat)\n", " hcqm = torch.zeros((stft.size(0) * N_BANDS, N_BINS, 6))\n", " for h, spec in enumerate(cqt_specs):\n", " hcqm[:, :, h] = spec(osf_flat).mean(-1)\n", " hcqm = hcqm.reshape(stft_bands.size(0), N_BINS, N_BANDS, 6)\n", " return hcqm\n", "\n", "\n", "def create_log_filter(num_bins, num_bands):\n", " log_bins = (\n", " np.logspace(np.log10(1), np.log10(num_bins), num=num_bands + 1, base=10.0) - 1\n", " )\n", " log_bins = np.unique(np.round(log_bins).astype(int))\n", " filter_matrix = torch.zeros(num_bands, num_bins)\n", " for i in range(num_bands):\n", " if i < num_bands - 1:\n", " start_bin, end_bin = log_bins[i], log_bins[i + 1]\n", " else:\n", " start_bin, end_bin = log_bins[i], num_bins\n", " filter_matrix[i, start_bin:end_bin] = 1 / (end_bin - start_bin)\n", "\n", " return filter_matrix\n", "\n", "\n", "def load_tempo_model(model_path: str):\n", " model = DeepRhythmModel()\n", " ckpt = read_from_s3(model_path, read_f=torch.load)\n", " model.load_state_dict(ckpt)\n", " model.cuda()\n", " model.eval()\n", " return model\n", "\n", "\n", "def make_kernels(len_audio=22050 * 8, sr=22050):\n", " n_fft = 2048\n", " hop = 512\n", " n_fft_bins = int(1 + n_fft / 2)\n", " band_filter = create_log_filter(n_fft_bins, N_BANDS)\n", " stft_spec = feat.stft.STFT(\n", " sr=sr, n_fft=n_fft, hop_length=hop, output_format=\"Magnitude\", verbose=False\n", " )\n", " cqt_specs = []\n", " for h in [1 / 2, 1, 2, 3, 4, 5]:\n", " # Convert from BPM to Hz\n", " fmin = (32.7 * h) / 60\n", " sr_cqt = len_audio // (hop * 8)\n", " fmax = sr_cqt / 2\n", " num_octaves = np.log2(fmax / fmin)\n", " bins_per_octave = N_BINS / num_octaves\n", " cqt_spec = feat.cqt.CQT(\n", " sr=sr_cqt,\n", " hop_length=len_audio // hop,\n", " n_bins=N_BINS,\n", " bins_per_octave=bins_per_octave,\n", " fmin=fmin,\n", " output_format=\"Magnitude\",\n", " verbose=False,\n", " pad_mode=\"constant\",\n", " )\n", " cqt_specs.append(cqt_spec)\n", " return stft_spec, band_filter, cqt_specs\n", "\n", "\n", "def onset_strength(\n", " y=None,\n", " n_fft=2048,\n", " hop_length=512,\n", " lag=1,\n", " ref=None,\n", " detrend=False,\n", " center=True,\n", " aggregate=None,\n", "):\n", " \"\"\"\n", " Compute the onset strength of an audio signal or a spectrogram.\n", "\n", " The onset strength is a measure of the increase in energy of an audio signal.\n", "\n", " Parameters\n", " ----------\n", " y : torch.Tensor, optional\n", " The raw audio waveform, expected to be a 2D tensor of shape (batch_size, time_samples).\n", " If provided, it will be used to compute the spectrogram internally. Default is None.\n", " n_fft : int, optional\n", " The number of FFT components. Default is 2048.\n", " hop_length : int, optional\n", " The number of samples between successive frames. Default is 512.\n", " lag : int, optional\n", " The lag between frames for computing the difference in energy. Default is 1.\n", " ref : torch.Tensor, optional\n", " The reference spectrogram to which the energy difference is computed. If None, the\n", " spectrogram provided by `S` or computed from `y` is used as the reference. Default is None.\n", " detrend : bool, optional\n", " If True, remove the mean from the onset envelope. Default is False.\n", " center : bool, optional\n", " If True, pad the time dimension of the onset envelope so that frames are centered around\n", " their timestamps. Default is True.\n", " aggregate : callable, optional\n", " A function to aggregate the channels dimension (e.g., torch.mean, torch.sum). If None,\n", " the mean is used. Default is None.\n", "\n", " Returns\n", " -------\n", " torch.Tensor\n", " The onset strength envelope, a 2D tensor of shape (batch_size, time_frames).\n", "\n", " \"\"\"\n", " # Ensure y is reshaped to (batch, channels, time) if it's not already\n", " if y is not None and y.dim() == 2:\n", " y = y.unsqueeze(1)\n", "\n", " S = torchaudio.transforms.AmplitudeToDB(top_db=80)(y)\n", " ref = S\n", "\n", " # Compute difference to reference, spaced by lag\n", " onset_env = S[..., lag:] - ref[..., :-lag]\n", " onset_env = torch.clamp(onset_env, min=0.0) # Discard negatives\n", "\n", " if aggregate is None:\n", " aggregate = torch.mean\n", " if callable(aggregate):\n", " onset_env = aggregate(onset_env, dim=-2)\n", "\n", " # Padding and detrending\n", " pad_width = lag\n", " if center:\n", " pad_width += n_fft // (2 * hop_length)\n", " onset_env = F.pad(onset_env, (pad_width, 0), \"constant\", 0)\n", "\n", " if detrend:\n", " onset_env -= onset_env.mean(dim=-1, keepdim=True)\n", "\n", " if center:\n", " onset_env = onset_env[..., : S.shape[-1]]\n", " return onset_env\n", "\n", "\n", "def prepare_audio(\n", " audio_filepath: str,\n", " start_s: float = None,\n", " end_s: float = None,\n", "):\n", " # audio = Audio.from_s3(s3_filepath)\n", " # sample_rate = audio.sample_rate\n", " # audio = torch.from_numpy(audio.array_float)\n", " audio, sample_rate = torchaudio.load(audio_filepath)\n", " audio = audio.mean(dim=0)\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", " # downmix and resample decoded audio to appropriate sr, also set num_frames\n", " audio = torchaudio.functional.resample(audio, sample_rate, 22_050)\n", " num_frames = int(8 * 22_050)\n", " # clip 8 seconds (but lagged to avoid intros)\n", " if (\n", " audio.shape[-1] > num_frames and audio.shape[-1] < 15 * 22_050 + num_frames\n", " ): # if longer than 8 but shorter than 23 seconds\n", " audio = audio[audio.shape[-1] - num_frames :] # take last 8 seconds\n", " elif audio.shape[-1] > num_frames: # if \"normal\" (> 23 seconds)\n", " start_ix = 15 * 22_050 # take 00:15 to 00:23\n", " audio = audio[start_ix : start_ix + num_frames]\n", " elif (\n", " audio.shape[-1] < num_frames\n", " ): # pad by repeating the signal if shorter than window\n", " pad_size = num_frames - audio.shape[-1]\n", " audio = torch.tensor(\n", " np.pad(audio.detach().cpu().numpy(), (0, pad_size), \"wrap\")\n", " )\n", " audio = preprocess_tempo_audio(audio)\n", "\n", " return audio\n", "\n", "\n", "def preprocess_tempo_audio(audio: torch.Tensor):\n", " stft_spec, band_filter, cqt_specs = make_kernels()\n", " input = torch.unsqueeze(audio, 0)\n", " preprocessed_audio = compute_hcqm(input, stft_spec, band_filter, cqt_specs).permute(\n", " 0, 3, 1, 2\n", " )\n", " return preprocessed_audio\n", "\n", "\n", "class DeepRhythmModel(nn.Module):\n", " def __init__(self, num_classes=256):\n", " super(DeepRhythmModel, self).__init__()\n", " # input shape is (6, 240, 8)\n", " self.num_classes = num_classes\n", " self.conv1 = nn.Conv2d(\n", " in_channels=6, out_channels=128, kernel_size=(4, 6), padding=\"same\"\n", " )\n", " self.bn1 = nn.BatchNorm2d(128)\n", " self.conv2 = nn.Conv2d(\n", " in_channels=128, out_channels=64, kernel_size=(4, 6), padding=\"same\"\n", " )\n", " self.bn2 = nn.BatchNorm2d(64)\n", " self.conv3 = nn.Conv2d(\n", " in_channels=64, out_channels=64, kernel_size=(4, 6), padding=\"same\"\n", " )\n", " self.bn3 = nn.BatchNorm2d(64)\n", " self.conv4 = nn.Conv2d(\n", " in_channels=64, out_channels=32, kernel_size=(4, 6), padding=\"same\"\n", " )\n", " self.bn4 = nn.BatchNorm2d(32)\n", " self.conv5 = nn.Conv2d(in_channels=32, out_channels=8, kernel_size=(120, 6))\n", " self.bn5 = nn.BatchNorm2d(8)\n", " self.fc1 = nn.Linear(2904, 256)\n", " self.elu = nn.ELU()\n", " self.dropout = nn.Dropout(0.5)\n", " self.fc2 = nn.Linear(256, num_classes)\n", " self._initialize_weights()\n", "\n", " def forward(self, x):\n", " x = F.relu(self.bn1(self.conv1(x)))\n", " x = F.relu(self.bn2(self.conv2(x)))\n", " x = F.relu(self.bn3(self.conv3(x)))\n", " x = F.relu(self.bn4(self.conv4(x)))\n", " x = F.relu(self.bn5(self.conv5(x)))\n", " x = x.reshape(x.size(0), -1)\n", " x = self.dropout(self.elu(self.fc1(x)))\n", " x = self.fc2(x)\n", " return x\n", "\n", " def _initialize_weights(self):\n", " for m in self.modules():\n", " if isinstance(m, nn.Conv2d):\n", " nn.init.kaiming_normal_(m.weight, mode=\"fan_out\", nonlinearity=\"relu\")\n", " elif isinstance(m, nn.BatchNorm2d):\n", " nn.init.constant_(m.weight, 1)\n", " nn.init.constant_(m.bias, 0)\n", " elif isinstance(m, nn.Linear):\n", " nn.init.xavier_normal_(m.weight)\n", " nn.init.constant_(m.bias, 0)\n", "\n", "import os\n", "\n", "def download_audio(s3_filepath: str, example_id: str, tmp_dir: str):\n", " filename = os.path.basename(s3_filepath)\n", " out_filepath = os.path.join(tmp_dir, f\"{example_id}-{filename}\")\n", " # only download the file if its not already downloaded\n", " if not os.path.isfile(out_filepath):\n", " os.system(f\"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1\")\n", " return out_filepath" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load pretrained tempo model\n", "tempo_model_path = \"s3://suno-data/christian/deeprhythm-0.5.pth\"\n", "tempo_model = load_tempo_model(tempo_model_path)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n", "filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "\n", "prep_audio = prepare_audio(filepath)\n", "evaluate_bpm(tempo_model, prep_audio, \"cuda\")\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 }