{ "cells": [ { "cell_type": "code", "execution_count": 94, "metadata": {}, "outputs": [], "source": [ "import os\n", "import json\n", "import glob\n", "import math\n", "import time\n", "import wandb\n", "import torch\n", "import auraloss\n", "import random\n", "import itertools\n", "import torchaudio\n", "import numpy as np\n", "import torch.nn as nn\n", "import matplotlib.pyplot as plt\n", "import torch.distributed as dist\n", "\n", "from tqdm import tqdm\n", "from typing import Tuple, Dict, List, Set, Any\n", "from torch.utils.data import DistributedSampler\n", "from torch.nn.parallel import DistributedDataParallel\n", "from torch.optim.lr_scheduler import LinearLR, ChainedScheduler\n", "\n", "\n", "# ---------------- functional corruptions with parameters ----------------\n", "def apply_stereo_to_mono(audio: torch.Tensor, sample_rate: float):\n", " return audio.mean(dim=0, keepdims=True).repeat(2, 1)\n", "\n", "\n", "def apply_channel_imbalance(\n", " audio: torch.Tensor, sample_rate: float, imbalance: float = 0.0\n", "):\n", " if not -1 <= imbalance <= 1 or audio.shape[-2] != 2:\n", " raise ValueError(\"Invalid input\")\n", " out = audio.clone()\n", " l_gain, r_gain = (1.0 - imbalance, 1.0) if imbalance > 0 else (1.0, 1.0 + imbalance)\n", " out[0, :], out[1, :] = out[0, :] * l_gain, out[1, :] * r_gain\n", " return out\n", "\n", "\n", "def apply_highpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float = 1000.0):\n", " return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_hz)\n", "\n", "\n", "def apply_lowpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float = 1000.0):\n", " return torchaudio.functional.lowpass_biquad(audio, sample_rate, cutoff_hz)\n", "\n", "\n", "def apply_bandpass(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " central_freq: float = 1000.0,\n", " bandwidth: float = 0.707,\n", "):\n", " return torchaudio.functional.bandpass_biquad(\n", " audio, sample_rate, central_freq, bandwidth\n", " )\n", "\n", "\n", "def apply_tanh_distortion(\n", " audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0\n", "):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " return torch.tanh(audio * gain_lin)\n", "\n", "\n", "def apply_clipping_distortion(\n", " audio: torch.Tensor, sample_rate: float, gain_db: float = 0.0\n", "):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " return (audio * gain_lin).clamp(-1, 1)\n", "\n", "\n", "def apply_noise(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " gain_db: float = 0.0,\n", " noise_type: str = \"white\",\n", "):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " noise = torch.randn_like(audio)\n", "\n", " if noise_type == \"white\":\n", " return audio + gain_lin * noise\n", " elif noise_type == \"pink\":\n", " b = torch.tensor([0.049922035, -0.095993537, 0.050612699, -0.004408786])\n", " a = torch.tensor([1, -2.494956002, 2.017265875, -0.522189400])\n", " noise = torchaudio.functional.filtfilt(noise, a, b)\n", " noise /= noise.abs().max()\n", " return audio + gain_lin * noise\n", " else:\n", " raise ValueError(f\"Invalid noise type: {noise_type}\")\n", "\n", "\n", "def apply_dc_offset(\n", " audio: torch.Tensor, sample_rate: float, offset: float = 0.0, mode: str = \"constant\"\n", "):\n", " if mode == \"constant\":\n", " return audio + offset\n", " elif mode == \"ramp\":\n", " ramp = torch.linspace(0, offset, audio.shape[-1])\n", " return audio + ramp\n", " else:\n", " raise ValueError(f\"Invalid mode: {mode}\")\n", "\n", "\n", "def apply_flip_polarity(audio: torch.Tensor, sample_rate: float):\n", " audio = audio.clone()\n", " channel = torch.randint(0, audio.shape[0], (1,)).item()\n", " audio[channel] *= -1\n", " return audio\n", "\n", "\n", "def apply_audio_codec(\n", " audio: torch.Tensor, sample_rate: float, bit_rate: int = 16000, n_passes: int = 1\n", "):\n", " for _ in range(n_passes):\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"mp3\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " audio = effector.apply(audio.T, sample_rate).T\n", " return audio\n", "\n", "\n", "def apply_hum(\n", " audio: torch.Tensor, sample_rate: float, amplitude: float = 0.0, freq: float = 0.0\n", "):\n", " t = torch.arange(audio.shape[-1], device=audio.device) / sample_rate\n", " hum = amplitude * torch.sin(2 * np.pi * freq * t)\n", " # Add harmonics at 2x\n", " hum += (amplitude * 0.5) * torch.sin(2 * np.pi * 2 * freq * t)\n", " return audio + hum.expand_as(audio)\n", "\n", "\n", "def apply_comb_filter(\n", " audio: torch.Tensor, sample_rate: float, delay_ms: float = 0.0, gain_db: float = 0.0\n", "):\n", " delay_samples = int(delay_ms * sample_rate / 1000)\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " delayed = torch.roll(audio, shifts=delay_samples, dims=-1)\n", " return audio + gain_lin * delayed\n", "\n", "\n", "def apply_reduce_bit_depth(audio: torch.Tensor, sample_rate: float, bits: int = 8):\n", " steps = 2**bits\n", " return (audio.clamp(-1, 1) * 0.5 + 0.5) * (steps - 1) // 1 / (steps - 1) * 2 - 1\n", "\n", "\n", "def apply_add_clicks(audio: torch.Tensor, sample_rate: float, density: float = 0.001):\n", " mask = torch.rand_like(audio) < density\n", " clicks = (torch.rand_like(audio) * 2 - 1) * mask\n", " return audio + clicks\n", "\n", "\n", "def apply_stereo_width(audio: torch.Tensor, sample_rate: float, width: float = 1.0):\n", " left, right = audio[0], audio[1]\n", " mid = (left + right) * 0.5\n", " side = (left - right) * 0.5\n", " side = (\n", " side * width\n", " ) # when width < 1, side is narrower, when width > 1, side is wider\n", " return torch.stack([mid + side, mid - side])\n", "\n", "\n", "def apply_spectral_mask(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " threshold: float = -60,\n", " ratio: float = 0.5,\n", " n_fft: int = 2048,\n", "):\n", " window = torch.hann_window(n_fft).to(audio.device)\n", " spec = torch.stft(audio, n_fft, n_fft // 4, window=window, return_complex=True)\n", " mask = torch.where(20 * torch.log10(torch.abs(spec) + 1e-8) < threshold, ratio, 1.0)\n", " return torch.istft(\n", " spec * mask, n_fft, n_fft // 4, window=window, length=audio.shape[-1]\n", " )\n", "\n", "\n", "def apply_time_stretch(audio: torch.Tensor, sample_rate: float, rate: float = 1.0):\n", " effects = [\n", " [\"tempo\", str(rate)],\n", " ]\n", " return torchaudio.sox_effects.apply_effects_tensor(audio, sample_rate, effects)[0]\n", "\n", "\n", "def apply_wow_flutter(\n", " audio: torch.Tensor, sample_rate: int = 48000, rate: float = 5.0, depth: float = 0.1\n", "):\n", " t = torch.arange(audio.shape[-1], device=audio.device) / sample_rate\n", " mod = depth * torch.sin(2 * torch.pi * rate * t)\n", "\n", " # Convert modulation to sample offsets\n", " offsets = (mod * sample_rate).long()\n", "\n", " # Apply time-varying delay\n", " output = torch.zeros_like(audio)\n", " for i in range(audio.shape[-1]):\n", " idx = max(0, min(i + offsets[i].item(), audio.shape[-1] - 1))\n", " output[..., i] = audio[..., idx]\n", " return output\n", "\n", "\n", "def apply_wow_flutter_fast(\n", " audio: torch.Tensor, sample_rate: int = 48000, rate: float = 5.0, depth: float = 0.1\n", "):\n", " t = torch.arange(audio.shape[-1], device=audio.device) / sample_rate\n", " mod = depth * torch.sin(2 * torch.pi * rate * t)\n", "\n", " indices = torch.arange(audio.shape[-1], device=audio.device)\n", " indices = indices + (mod * sample_rate).long()\n", " indices = torch.clamp(indices, 0, audio.shape[-1] - 1)\n", "\n", " while indices.dim() < audio.dim():\n", " indices = indices.unsqueeze(0)\n", " indices = indices.expand_as(audio)\n", "\n", " output = torch.gather(audio, -1, indices)\n", " return output\n", "\n", "\n", "def apply_stereo_fold(audio, sample_rate):\n", " mono = audio.mean(dim=0, keepdim=True)\n", " # Add phase issues\n", " return torch.cat([mono, -mono], dim=0)\n", "\n", "\n", "def apply_ring_modulation(audio, sample_rate, freq=440, mix=0.2):\n", " samples = audio.shape[-1]\n", " t = torch.linspace(0, samples / sample_rate, samples, device=audio.device)\n", "\n", " freq = freq + 10 * torch.sin(2 * torch.pi * 0.5 * t)\n", " phase = 2 * torch.pi * freq * t\n", " carrier = torch.sin(phase).view(1, -1) # Changed from (1,1,-1) to (1,-1)\n", "\n", " modulated = audio * carrier\n", " return (1 - mix) * audio + mix * modulated\n", "\n", "\n", "def apply_white_noise_burst(\n", " audio,\n", " sample_rate,\n", " noise_level=0.1,\n", " min_burst_length=500,\n", " max_burst_length=8000,\n", " p_burst=0.01,\n", "):\n", " # Use shortest burst length to determine number of segments\n", " num_segments = audio.shape[-1] // min_burst_length\n", "\n", " # Generate random burst lengths and levels\n", " burst_lengths = torch.randint(\n", " min_burst_length, max_burst_length, (num_segments,), device=audio.device\n", " )\n", " burst_levels = noise_level * (0.5 + torch.rand(num_segments, device=audio.device))\n", " burst_mask = (\n", " torch.rand(num_segments, device=audio.device) < p_burst\n", " ).bool() # Changed to bool\n", "\n", " # Create index tensor for the full audio length\n", " indices = torch.arange(audio.shape[-1], device=audio.device)\n", "\n", " # Create cumulative positions\n", " positions = torch.cumsum(burst_lengths, dim=0)\n", " starts = torch.cat([torch.tensor([0], device=audio.device), positions[:-1]])\n", "\n", " # Create mask using broadcasting\n", " mask = torch.zeros(audio.shape[-1], device=audio.device)\n", " valid_mask = (indices.unsqueeze(0) >= starts.unsqueeze(1)) & (\n", " indices.unsqueeze(0) < positions.unsqueeze(1)\n", " )\n", " valid_mask = valid_mask & burst_mask.unsqueeze(1) # Now both are boolean\n", "\n", " # Convert boolean mask to burst levels\n", " mask = (valid_mask.float() * burst_levels.unsqueeze(1)).max(dim=0)[0]\n", "\n", " # Expand mask to match audio dimensions\n", " mask = mask.view(1, -1).expand_as(audio)\n", "\n", " # Apply noise\n", " noise = torch.randn_like(audio)\n", " return audio + noise * mask\n", "\n", "\n", "def apply_quantize_zero(audio, sample_rate, threshold=0.001):\n", " # first normalize to -1, 1\n", " audio = audio / audio.abs().max().clamp(1e-8)\n", " audio_quantized = torch.where(torch.abs(audio) < threshold, 0, audio)\n", " return audio_quantized\n", "\n", "\n", "def apply_phase_randomize(\n", " audio: torch.Tensor, sample_rate: float, block_size: int = 2048, mix: float = 0.75\n", "):\n", " window = torch.hann_window(block_size, device=audio.device)\n", " # Process each channel\n", " output = []\n", " for channel in audio:\n", " stft = torch.stft(channel, block_size, window=window, return_complex=True)\n", " mag = stft.abs()\n", " random_phase = torch.exp(2j * torch.pi * torch.rand_like(stft))\n", " channel_out = torch.istft(\n", " mag * random_phase, block_size, window=window, length=channel.shape[-1]\n", " )\n", " output.append(channel_out)\n", " return (1 - mix) * audio + mix * torch.stack(output)\n", "\n", "def apply_freq_boost(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " freq_hz: float = 1000.0,\n", " gain_db: float = 10.0,\n", "):\n", " # Design peaking filter with Q=0.707\n", " w0 = 2 * math.pi * freq_hz / sample_rate\n", " alpha = math.sin(w0) / (2 * 0.707)\n", " A = 10 ** (gain_db / 40.0)\n", "\n", " b0 = 1 + alpha * A\n", " b1 = -2 * math.cos(w0)\n", " b2 = 1 - alpha * A\n", " a0 = 1 + alpha / A\n", " a1 = -2 * math.cos(w0)\n", " a2 = 1 - alpha / A\n", "\n", " # Normalize coefficients by a0\n", " b0 = b0 / a0\n", " b1 = b1 / a0\n", " b2 = b2 / a0\n", " a1 = a1 / a0\n", " a2 = a2 / a0\n", " a0 = 1.0\n", "\n", " return torchaudio.functional.biquad(audio, b0, b1, b2, a0, a1, a2)\n", "\n", "\n", "def apply_reverb(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " reverberance: int = 50, # 0-100\n", " hf_damping: int = 50, # 0-100\n", " room_scale: int = 100, # 0-100\n", " stereo_depth: int = 100, # 0-100\n", " pre_delay: float = 0, # 0-200ms\n", " wet_gain: float = 0,\n", "): # -10-10 dB\n", "\n", " effects = [\n", " [\n", " \"reverb\",\n", " str(reverberance),\n", " str(hf_damping),\n", " str(room_scale),\n", " str(stereo_depth),\n", " str(pre_delay),\n", " str(wet_gain),\n", " ]\n", " ]\n", " out, _ = torchaudio.sox_effects.apply_effects_tensor(audio, sample_rate, effects)\n", " return out\n", "\n", "\n", "def apply_audio_codec_advanced(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " bit_rate: int = 16000,\n", " n_passes: int = 1,\n", " codec_type: str = \"mp3\",\n", "):\n", " for _ in range(n_passes):\n", " if codec_type == \"mp3\":\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"mp3\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " elif codec_type == \"ogg-vorbis\":\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"ogg\",\n", " encoder=\"vorbis\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " elif codec_type == \"opus\":\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"ogg\",\n", " encoder=\"opus\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " else:\n", " raise ValueError(f\"Invalid codec type: {codec_type}\")\n", " audio = effector.apply(audio.T, sample_rate).T\n", " return audio\n", "\n", "\n", "def apply_audio_codec(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " bit_rate: int = 16000,\n", " n_passes: int = 1,\n", " format_str: str = \"mp3\",\n", "):\n", " for _ in range(n_passes):\n", " effector = torchaudio.io.AudioEffector(\n", " format=format_str,\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " audio = effector.apply(audio.T, sample_rate).T\n", " return audio\n", "\n", "# Separate function mapping\n", "\n", "# Separate function mapping\n", "corruption_functions = {\n", " \"stereo_to_mono\": apply_stereo_to_mono,\n", " \"channel_imbalance\": apply_channel_imbalance,\n", " \"lowpass\": apply_lowpass,\n", " \"bandpass\": apply_bandpass,\n", " \"highpass\": apply_highpass,\n", " \"tanh_distortion\": apply_tanh_distortion,\n", " \"clipping_distortion\": apply_clipping_distortion,\n", " \"noise\": apply_noise,\n", " \"hum\": apply_hum,\n", " \"comb_filter\": apply_comb_filter,\n", " \"reduce_bit_depth\": apply_reduce_bit_depth,\n", " \"add_clicks\": apply_add_clicks,\n", " \"reverb\": apply_reverb,\n", " \"audio_codec\": apply_audio_codec,\n", " \"dc_offset\": apply_dc_offset,\n", " \"flip_polarity\": apply_flip_polarity,\n", " \"stereo_width\": apply_stereo_width,\n", " \"spectral_mask\": apply_spectral_mask,\n", " \"time_stretch\": apply_time_stretch,\n", " \"wow_flutter\": apply_wow_flutter_fast,\n", " \"stereo_fold\": apply_stereo_fold,\n", " \"ring_modulation\": apply_ring_modulation,\n", " \"white_noise_burst\": apply_white_noise_burst,\n", " \"quantize_zero\": apply_quantize_zero,\n", " \"phase_randomize\": apply_phase_randomize,\n", " \"freq_boost\": apply_freq_boost,\n", " \"audio_codec_advanced\": apply_audio_codec_advanced,\n", "}\n", "\n", "\n", "# now we have a function that takes a preset and applies the relevant corruptions to the audio\n", "# Modified apply_preset function to use both config and functions\n", "# now we have a function that takes a preset and applies the relevant corruptions to the audio\n", "# Modified apply_preset function to use both config and functions\n", "def apply_preset(audio: torch.Tensor, sr: float, preset: dict, functions: dict):\n", " chs, seq_len = audio.shape\n", " for corruption_name, corruption_info in preset.items():\n", " audio = functions[corruption_name](\n", " audio.clone(), sr, **corruption_info[\"params\"]\n", " )\n", " audio = torch.clamp(audio, -1, 1) # clip to -1, 1\n", "\n", " # repeat pad to original length\n", " if audio.shape[-1] < seq_len:\n", " audio = audio.repeat(1, seq_len)\n", "\n", " # crop to original length\n", " audio = audio[..., :seq_len]\n", "\n", " return audio\n", "\n", "\n", "def create_label_encoder(corruptions_dict: dict) -> Dict[str, int]:\n", " \"\"\"\n", " Create a mapping from all possible corruption parameter combinations to indices.\n", " For corruptions with multiple parameters, creates labels for all combinations.\n", " \"\"\"\n", " label_to_idx = {}\n", " idx = 0\n", "\n", " for corruption_name, corruption_info in corruptions_dict.items():\n", " params = corruption_info[\"params\"]\n", "\n", " # If corruption has no parameters\n", " if not params:\n", " label = f\"{corruption_name}\"\n", " label_to_idx[label] = idx\n", " idx += 1\n", " continue\n", "\n", " # Get all parameter names and their possible values\n", " param_names = list(params.keys())\n", " param_values = [params[name] for name in param_names]\n", "\n", " # Generate all possible combinations of parameter values\n", " for values in itertools.product(*param_values):\n", " # Create parameter string\n", " param_str = \",\".join(\n", " f\"{name}={value}\" for name, value in zip(param_names, values)\n", " )\n", " label = f\"{corruption_name}:{param_str}\"\n", " label_to_idx[label] = idx\n", " idx += 1\n", "\n", " return label_to_idx\n", "\n", "\n", "def sample_n_corruptions(max_corruptions: int = 5, p: float = 0.5):\n", " probs = np.array([(1 - p) ** i * p for i in range(max_corruptions)])\n", " probs = probs / probs.sum()\n", " return np.random.choice(np.arange(max_corruptions), p=probs) + 1\n", "\n", "\n", "# Similarly, we need to update how we generate labels in the preset function\n", "def generate_random_preset(\n", " corruptions_dict: dict,\n", " max_corruptions: int = 10,\n", " no_corruption_probability: float = 0.01,\n", "):\n", " \"\"\"\n", " Generate a random preset and its corresponding labels.\n", " Handles corruptions with multiple parameters.\n", " \"\"\"\n", " if random.random() < no_corruption_probability:\n", " return {}, set()\n", "\n", " max_corruptions = len(corruptions_dict)\n", "\n", " # sample n_corruptions from exponential distribution\n", " n_corruptions = sample_n_corruptions(max_corruptions, p=0.5)\n", " if n_corruptions == 0:\n", " return {}, set()\n", "\n", " if len(corruptions_dict) == 1:\n", " selected_corruptions = [list(corruptions_dict.keys())[0]]\n", " else:\n", " selected_corruptions = random.sample(\n", " list(corruptions_dict.keys()), n_corruptions\n", " )\n", "\n", " preset = {}\n", " labels = set()\n", "\n", " for corruption_name in selected_corruptions:\n", " corruption_info = corruptions_dict[corruption_name]\n", " params = {}\n", "\n", " # If no parameters, just add the corruption name\n", " if not corruption_info[\"params\"]:\n", " preset[corruption_name] = {\"params\": {}}\n", " labels.add(corruption_name)\n", " continue\n", "\n", " # Generate parameters and create combined label\n", " param_strs = []\n", " for param_name, param_values in corruption_info[\"params\"].items():\n", " param_value = random.choice(param_values)\n", " params[param_name] = param_value\n", " param_strs.append(f\"{param_name}={param_value}\")\n", "\n", " preset[corruption_name] = {\"params\": params}\n", " # Create single label with all parameters\n", " label = f\"{corruption_name}:{','.join(param_strs)}\"\n", " labels.add(label)\n", "\n", " return preset, labels\n", "\n", "\n", "def labels_to_tensor(\n", " labels: Set[str], label_encoder: Dict[str, int], device: str = \"cpu\"\n", ") -> torch.Tensor:\n", " \"\"\"\n", " Convert a set of labels to a binary tensor.\n", " \"\"\"\n", " output = torch.zeros(len(label_encoder), dtype=torch.float32, device=device)\n", " for label in labels:\n", " if label in label_encoder:\n", " output[label_encoder[label]] = 1.0\n", " return output\n", "\n", "def tensor_to_corruptions(label_tensor: torch.Tensor, label_encoder: Dict[str, int]) -> List[str]:\n", " return [list(label_encoder.keys())[i] for i in torch.nonzero(label_tensor).flatten().tolist()]\n", "\n", "\n", "class CorruptAudioDataset(torch.utils.data.Dataset):\n", " def __init__(\n", " self,\n", " manifest_path: str,\n", " label_encoder: Dict[str, int],\n", " corruptions: Dict[str, Dict[str, Any]],\n", " sample_rate: int,\n", " max_corruptions: int = 10,\n", " no_corruption_probability: float = 0.5,\n", " num_workers: int = 1,\n", " chunk_size_s: float = 5.0,\n", " buffer_size: int = 50_000,\n", " max_chunks_per_file: int = 100,\n", " num_versions: int = 1,\n", " random_crop: bool = False,\n", " ):\n", " self.manifest_path = manifest_path\n", " self.sample_rate = sample_rate\n", " self.chunk_size_s = chunk_size_s\n", " self.buffer_size = buffer_size\n", " self.chunk_size_samples = int(chunk_size_s * sample_rate)\n", " self.num_workers = num_workers\n", " self.max_corruptions = max_corruptions\n", " self.no_corruption_probability = no_corruption_probability\n", " self.max_chunks_per_file = max_chunks_per_file\n", " self.num_versions = num_versions\n", " self.random_crop = random_crop\n", "\n", " self.preprocess_chunk_size_samples = self.chunk_size_samples\n", " # if random_crop, then we adjust chunks to be larger\n", " if self.random_crop:\n", " self.preprocess_chunk_size_samples *= 1.25\n", " self.preprocess_chunk_size_samples = int(self.preprocess_chunk_size_samples)\n", "\n", " # assert self.num_versions > 1, \"num_versions must be greater than 1\"\n", "\n", " self.label_encoder = label_encoder\n", " self.num_labels = len(self.label_encoder)\n", " self.corruptions = corruptions\n", " self.items_since_last_reload = buffer_size # force a reload\n", " self.buffer = []\n", "\n", " self.loss_fn = auraloss.freq.MelSTFTLoss(\n", " sample_rate=self.sample_rate,\n", " fft_size=2048,\n", " win_length=2048,\n", " hop_size=1024,\n", " n_mels=64,\n", " )\n", "\n", " # self.loss_fn = auraloss.time.SISDRLoss()\n", "\n", " # load manifest\n", " with open(manifest_path, \"r\") as f:\n", " self.filepaths = [line.strip() for line in f.readlines()]\n", " print(f\"Loaded {len(self.filepaths)} filepaths from {manifest_path}\")\n", "\n", " def __len__(self):\n", " return self.buffer_size * self.num_workers\n", "\n", " def _reload_buffer(self):\n", " self.buffer = []\n", " rand_idxs = torch.randperm(len(self.filepaths))\n", "\n", " # max rand_idxs repeat endlessly\n", " rand_idxs = itertools.cycle(rand_idxs)\n", " # pbar = tqdm(rand_idxs, total=len(self.filepaths), desc=\"Loading audio buffer\")\n", " for idx in rand_idxs:\n", " if len(self.buffer) >= self.buffer_size:\n", " break\n", "\n", " try:\n", " filepath = self.filepaths[idx]\n", " audio, sr = torchaudio.load(filepath)\n", "\n", " if sr != self.sample_rate:\n", " audio = torchaudio.functional.resample(audio, sr, self.sample_rate)\n", "\n", " # Pad if needed to ensure consistent chunk size\n", " if audio.shape[-1] < self.preprocess_chunk_size_samples:\n", " continue\n", "\n", " # Split into chunks\n", " chunks = audio.unfold(\n", " -1,\n", " self.preprocess_chunk_size_samples,\n", " self.preprocess_chunk_size_samples,\n", " )\n", " chunks = chunks.chunk(chunks.shape[1], dim=1)\n", "\n", " # Filter chunks by minimum length\n", " valid_chunks = [\n", " chunk.squeeze(1)\n", " for chunk in chunks\n", " if chunk.shape[-1] >= self.preprocess_chunk_size_samples\n", " ]\n", "\n", " # filter out chunks of silence\n", " valid_chunks = [\n", " chunk for chunk in valid_chunks if (chunk.abs() ** 2).mean() > 0.001\n", " ]\n", "\n", " # limit to max_chunks_per_file\n", " valid_chunks = valid_chunks[: self.max_chunks_per_file]\n", "\n", " self.buffer.extend(valid_chunks)\n", "\n", " # pbar.set_postfix({\"buffer_size\": len(self.buffer)})\n", "\n", " except Exception as e:\n", " print(f\"Error loading {filepath}: {e}\")\n", " continue\n", " self.items_since_last_reload = 0\n", "\n", " def __getitem__(self, _):\n", " # reload buffer if needed\n", " if self.items_since_last_reload >= len(self.buffer):\n", " self._reload_buffer()\n", "\n", " # sample a random audio from the buffer\n", " buffer_idx = np.random.randint(0, len(self.buffer))\n", " audio = self.buffer[buffer_idx].clone()\n", "\n", " corrupted_audios = []\n", " label_tensors = []\n", " mse_losses = []\n", " for version_idx in range(self.num_versions):\n", "\n", " preset, labels = generate_random_preset(\n", " self.corruptions,\n", " self.max_corruptions,\n", " self.no_corruption_probability,\n", " )\n", "\n", " if False:\n", " corrupted_audio = (\n", " apply_preset(\n", " audio,\n", " self.sample_rate,\n", " preset,\n", " corruption_functions,\n", " )\n", " if preset\n", " else audio\n", " )\n", " # always apply the same corruption if version_idx > 1\n", " if version_idx >= 1:\n", " corrupted_audio = torch.tanh(audio * 12.0)\n", " else:\n", " corrupted_audio = audio.clone()\n", " # corrupted_audio = corrupted_audio / corrupted_audio.abs().max().clamp(1e-8)\n", "\n", " # random crop to chunk size\n", " if self.random_crop:\n", " start_idx = np.random.randint(\n", " 0, corrupted_audio.shape[-1] - self.chunk_size_samples\n", " )\n", " end_idx = start_idx + self.chunk_size_samples\n", " corrupted_audio_crop = corrupted_audio[..., start_idx:end_idx]\n", " audio_crop = audio[..., start_idx:end_idx]\n", " else:\n", " corrupted_audio_crop = corrupted_audio\n", " audio_crop = audio\n", "\n", " # lets clip anything out of range here\n", " corrupted_audio_crop = torch.clamp(corrupted_audio_crop, -1, 1)\n", " # label_tensor = labels_to_tensor(labels, self.label_encoder)\n", " self.items_since_last_reload += 1\n", "\n", " # label_tensors.append(label_tensor)\n", "\n", " # measure the mse between the corrupted audio and the original audio\n", " # audio_crop_norm = audio_crop / audio_crop.abs().max().clamp(1e-8)\n", " # corrupted_audio_crop_norm = (\n", " # corrupted_audio_crop / corrupted_audio_crop.abs().max().clamp(1e-8)\n", " # )\n", " # mse = (corrupted_audio_norm - audio_norm).pow(2).mean()\n", " mse = self.loss_fn(\n", " audio_crop.unsqueeze(0), corrupted_audio_crop.unsqueeze(0)\n", " )\n", " # clamp the loss from 0 to 20.0\n", " # mse = torch.log(mse + 1e-3)\n", " #mse = torch.clamp(mse, 0.0, 10.0) / 10.0\n", " mse_losses.append(mse)\n", "\n", " if np.random.uniform() < 0.5:\n", " # peak normalize\n", " corrupted_audio_crop = (\n", " corrupted_audio_crop / corrupted_audio_crop.abs().max().clamp(1e-8)\n", " )\n", " audio_crop = audio_crop / audio_crop.abs().max().clamp(1e-8)\n", "\n", " if np.random.uniform() < 0.5:\n", " gain_reduction_db = np.random.uniform(-12, 0)\n", " corrupted_audio_crop *= 10 ** (gain_reduction_db / 20.0)\n", "\n", " corrupted_audios.append(corrupted_audio_crop)\n", "\n", " if self.num_versions == 1:\n", " return corrupted_audios[0], mse_losses[0]\n", "\n", " # create a new label for which version has lower mse\n", " mse_losses = torch.tensor(mse_losses)\n", " label_tensor = mse_losses.argmin().float()\n", "\n", " return torch.stack(corrupted_audios), label_tensor" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "from torch.utils.data import DistributedSampler\n", "\n", "local_rank = 0\n", "\n", "run_config = {\n", " \"training\": {\n", " \"max_steps\": 1_000_000,\n", " \"run_name\": \"genius-hq\",\n", " \"project_name\": \"ear-v2\",\n", " \"lr\": 2e-5,\n", " \"grad_clip_norm\": 10.0,\n", " },\n", " \"model\": {\n", " \"hidden_dim\": 1024,\n", " \"latent_dim\": 128,\n", " \"num_heads\": 8,\n", " \"num_conv_layers\": 6,\n", " \"num_transformer_layers\": 12,\n", " \"dropout\": 0.1,\n", " \"compare\": True,\n", " },\n", " \"dataset\": {\n", " \"corruptions_config\": \"/home/christian/code/christian/metadata/corruptions_config_lpf.json\",\n", " \"train_manifest\": \"/app/suno/data/audio_2ch_48khz_lg/ear_train_filtered_v2_with_gens.csv\",\n", " \"val_manifest\": \"/app/suno/data/audio_2ch_48khz_lg/ear_val.csv\",\n", " \"max_corruptions\": 1,\n", " \"no_corruption_probability\": 0.2,\n", " \"batch_size\": 1,\n", " \"num_workers\": 8,\n", " \"chunk_size_s\": 5.0,\n", " \"buffer_size\": 100,\n", " \"sample_rate\": 48_000,\n", " \"random_crop\": True,\n", " },\n", "}\n", "\n", "corruptions = json.load(open(run_config[\"dataset\"][\"corruptions_config\"]))\n", "label_encoder = create_label_encoder(corruptions)\n", "\n", "# setup dataset\n", "train_filepaths = run_config[\"dataset\"][\"train_manifest\"]\n", "train_dataset = CorruptAudioDataset(\n", " train_filepaths,\n", " label_encoder,\n", " corruptions,\n", " sample_rate=run_config[\"dataset\"][\"sample_rate\"],\n", " num_workers=run_config[\"dataset\"][\"num_workers\"],\n", " chunk_size_s=run_config[\"dataset\"][\"chunk_size_s\"],\n", " buffer_size=run_config[\"dataset\"][\"buffer_size\"],\n", " max_corruptions=run_config[\"dataset\"][\"max_corruptions\"],\n", " no_corruption_probability=run_config[\"dataset\"][\"no_corruption_probability\"],\n", " num_versions=2,\n", " random_crop=run_config[\"dataset\"][\"random_crop\"],\n", ")\n", "\n", "#train_sampler = DistributedSampler(train_dataset, rank=local_rank, shuffle=True)\n", "train_loader = torch.utils.data.DataLoader(\n", " train_dataset,\n", " batch_size=run_config[\"dataset\"][\"batch_size\"],\n", " num_workers=run_config[\"dataset\"][\"num_workers\"],\n", " persistent_workers=True, # this is necessary for the buffer to work\n", ")" ] }, { "cell_type": "code", "execution_count": 96, "metadata": {}, "outputs": [], "source": [ "\n", "import IPython\n", "import matplotlib.pyplot as plt\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "peaks = []\n", "energy = []\n", "labels = []\n", "input_audios = []\n", "corrupted_audios = []\n", "num_corruptions = []\n", "\n", "for bidx, batch in enumerate(tqdm(train_loader)):\n", " audio, label = batch\n", " input_audio = audio[:,0,...]\n", " corrupted_audio = audio[:,1,...]\n", " #print(label, audio.shape)\n", " peaks.append(input_audio.abs().max())\n", " energy.append((input_audio ** 2).mean())\n", " labels.append(label)\n", " input_audios.append(input_audio)\n", " corrupted_audios.append(corrupted_audio)\n", " num_corruptions.append(label.sum())\n", " if bidx > 100:\n", " break\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "labels_tensor = torch.stack(labels).view(-1)\n", "print(labels_tensor.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# make a histogram of the labels\n", "plt.hist(labels_tensor.cpu(), bins=100)\n", "plt.show()\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "for n in [20]:\n", " idx = n\n", " bidx = 0\n", " input_audio = input_audios[idx][bidx]\n", " corrupted_audio = corrupted_audios[idx][bidx]\n", " print(labels[idx][bidx])\n", " IPython.display.display(IPython.display.Audio(input_audio[0].cpu().numpy(), rate=48000, normalize=True))\n", " IPython.display.display(IPython.display.Audio(corrupted_audio[0].cpu().numpy(), rate=48000, normalize=True))\n", "\n", "# 0 means A has lower loss (better fidelity)\n", "# 1 means B has lower loss (better fidelity)\n" ] }, { "cell_type": "code", "execution_count": 27, "metadata": {}, "outputs": [], "source": [ "from torch.nn import functional as F\n", "\n", "def bradley_terry_loss(\n", " r_i: torch.Tensor, r_j: torch.Tensor, labels: torch.Tensor\n", ") -> torch.Tensor:\n", " \"\"\"\n", " Compute Bradley-Terry loss for paired comparisons.\n", "\n", " Args:\n", " r_i: Logits/scores for first options in pairs, shape (batch_size,)\n", " r_j: Logits/scores for second options in pairs, shape (batch_size,)\n", " labels: Binary tensor indicating whether first option (1) or second option (0)\n", " was preferred, shape (batch_size,)\n", "\n", " Returns:\n", " Mean loss value as a torch.Tensor\n", " \"\"\"\n", " idx = 0 # batch idx\n", " print(r_i[idx], r_j[idx], labels[idx])\n", " # Compute negative log likelihood using logsigmoid for numerical stability\n", " loss = -((1-labels) * F.logsigmoid(r_i - r_j) + labels * F.logsigmoid(r_j - r_i))\n", " return loss.mean()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "r_i = torch.tensor([1, 2, 3, 3]).float()\n", "r_j = torch.tensor([4, 4, 4, 4]).float()\n", "labels = torch.tensor([1, 0, 1, 0]).float()\n", "\n", "loss = bradley_terry_loss(r_i, r_j, labels)\n", "print(loss)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "peaks = np.array(peaks)\n", "energy = np.array(energy)\n", "print(peaks.mean(), peaks.std(), peaks.min(), peaks.max())\n", "print(energy.mean(), energy.std(), energy.min(), energy.max())\n", "plt.hist(peaks, bins=100)\n", "plt.show()\n", "plt.hist(energy, bins=np.logspace(np.log10(energy.min()), np.log10(energy.max()), 100))\n", "plt.xscale(\"log\")\n", "plt.show()\n", "\n", "plt.hist(num_corruptions, bins=100)\n", "plt.show()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# find the examples with top 5 lowest energy\n", "top_5_low_energy_idxs = np.argsort(energy)[:5]\n", "for idx in top_5_low_energy_idxs:\n", " # convert label to string\n", " label_str = labels[idx]\n", " label = tensor_to_corruptions(label_str, label_encoder)\n", " print(label)\n", " audio = audios[idx]\n", " print((audio ** 2).mean())\n", " IPython.display.display(IPython.display.Audio(audio[0].cpu().numpy(), rate=48000, normalize=False))\n", "\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "scores = torch.randn(200) * 0.000\n", "all_preds = torch.sigmoid(scores) > 0.5\n", "all_preds = all_preds.numpy()\n", "all_labels = torch.randint(0, 2, (200,)).numpy()\n", "\n", "accuracy = np.mean((all_preds == all_labels).astype(float))\n", "\n", "print(accuracy)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "preprocess_chunk_size_samples = int(5.0 * 48000) + 2 * 48000\n", "audio = torch.randn(2, int(60.0 * 48000))\n", "\n", "# Split into chunks\n", "chunks = audio.unfold(\n", " -1,\n", " preprocess_chunk_size_samples,\n", " preprocess_chunk_size_samples,\n", ")\n", "print(chunks.shape)\n", "chunks = chunks.chunk(chunks.shape[1], dim=1)\n", "for chunk in chunks:\n", " print(chunk.shape)\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 }