{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"4\"\n", "\n", "import torch\n", "\n", "import torch.nn as nn\n", "import math" ] }, { "cell_type": "code", "execution_count": 2, "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" ] }, { "cell_type": "code", "execution_count": 102, "metadata": {}, "outputs": [], "source": [ "import torchaudio\n", "import numpy as np\n", "\n", "def corrupt_audio(\n", " audio,\n", " sample_rate,\n", " highpass_prob=0.1,\n", " lowpass_prob=0.1,\n", " noise_prob=0.1,\n", " tanh_prob=0.1,\n", " clip_prob=0.1,\n", " preemphasis_prob=0.05,\n", " deemphasis_prob=0.05,\n", " bass_boost_prob=0.1,\n", " treble_boost_prob=0.1,\n", " mp3_prob=0.90,\n", "):\n", " if np.random.uniform() < highpass_prob: # highpass\n", " # sample on a log scale\n", " freq_hz = 10 ** np.random.uniform(np.log10(20), np.log10(20000))\n", " print(f\"Highpass: {freq_hz}\")\n", " audio = torchaudio.functional.highpass_biquad(audio, sample_rate, freq_hz)\n", " if np.random.uniform() < lowpass_prob: # lowpass\n", " freq_hz = 10 ** np.random.uniform(np.log10(20), np.log10(20000))\n", " print(f\"Lowpass: {freq_hz}\")\n", " audio = torchaudio.functional.lowpass_biquad(audio, sample_rate, freq_hz)\n", " if np.random.uniform() < bass_boost_prob: # bass boost\n", " gain_db = np.random.uniform(6, 12)\n", " freq_hz = np.random.uniform(20, 240)\n", " audio = torchaudio.functional.bass_biquad(audio, sample_rate, gain_db, freq_hz)\n", " if np.random.uniform() < treble_boost_prob: # treble boost\n", " gain_db = np.random.uniform(6, 12)\n", " freq_hz = np.random.uniform(1000, 10000)\n", " audio = torchaudio.functional.treble_biquad(audio, sample_rate, gain_db, freq_hz)\n", " if np.random.uniform() < noise_prob: # noise\n", " noise_gain_db = np.random.uniform(-48, -24)\n", " audio = audio + torch.randn_like(audio) * 10 ** (noise_gain_db / 20.0)\n", " if np.random.uniform() < tanh_prob: # tanh\n", " gain_db = np.random.uniform(12, 24)\n", " audio = torch.tanh(audio * 10 ** (gain_db / 20.0))\n", " if np.random.uniform() < clip_prob: # clip\n", " gain_db = np.random.uniform(12, 24)\n", " audio = torch.clamp(audio * 10 ** (gain_db / 20.0), -1, 1)\n", " if np.random.uniform() < preemphasis_prob: # preemphasis\n", " coeff = np.random.uniform(0.75, 1.0)\n", " audio = torchaudio.functional.preemphasis(audio, coeff)\n", " if np.random.uniform() < deemphasis_prob: # deemphasis\n", " coeff = np.random.uniform(0.75, 1.0)\n", " audio = torchaudio.functional.deemphasis(audio, coeff)\n", " if np.random.uniform() < mp3_prob: # mp3\n", " bit_rate = np.random.choice(\n", " [\n", " 8000,\n", " 16000,\n", " 24000,\n", " 32000,\n", " 48000,\n", " 64000,\n", " 96000,\n", " 112000,\n", " 128000,\n", " ]\n", " )\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 torch.clamp(audio, -1.0, 1.0)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "import IPython\n", "\n", "audio, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n", "# crop to 30s\n", "audio = audio[:, :30*sr]\n", "\n", "corrupted_audio = corrupt_audio(audio, sr)\n", "#corrupted_audio = torchaudio.functional.deemphasis(audio, 0.96)\n", "print(corrupted_audio.abs().max())\n", "\n", "IPython.display.display(IPython.display.Audio(corrupted_audio, rate=sr))\n", "#IPython.display.display(IPython.display.Audio(audio, rate=sr))\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#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-03-12_00-13-20_s1753/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", "\n" ] }, { "cell_type": "code", "execution_count": 44, "metadata": {}, "outputs": [], "source": [ "\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" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "audio_filepaths = [\n", " \"/home/christian/code/christian/notebooks/audio/dreams-no-ctx.mp3\",\n", " \"/home/christian/code/christian/notebooks/audio/dreams-with-ctx-2.mp3\"\n", " ]\n", "\n", "results = []\n", "for audio_filepath in tqdm(audio_filepaths):\n", " x, sr = torchaudio.load(audio_filepath)\n", "\n", " if x.shape[1] < 60*sr:\n", " print(f\"Skipping {audio_filepath} because it's too short\")\n", " continue\n", "\n", " # crop to 30s and use the last 30s\n", " # first 30s\n", " x_first = x[:, :30*sr]\n", "\n", " #x_corrupt = x_first.clone()\n", " #x_corrupt = apply_audio_codec(x_corrupt, sr, 8000, 1)\n", "\n", " # lower score is better, \n", " # run inference\n", " with torch.no_grad():\n", " first_scores, first_mean_score = model.get_score(x_first.cuda())\n", " #corrupt_scores, corrupt_mean_score = model.get_score(x_corrupt.cuda())\n", "\n", " print(audio_filepath, first_mean_score)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# run infernce on a given audio file\n", "import torchaudio\n", "import IPython\n", "import numpy as np\n", "import glob\n", "from tqdm import tqdm\n", "#x, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n", "\n", "#audio_filepaths = glob.glob(\"/home/christian/audio/reference-audio-wav/*.wav\")\n", "\n", "#audio_filepaths = [\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"]\n", "audio_filepaths = glob.glob(\"/home/christian/code/christian/notebooks/audio/ctx2/*.mp3\")\n", "#audio_filepaths = glob.glob(\"/home/christian/code/christian/notebooks/audio/s3/*.mp3\")\n", "print(len(audio_filepaths))\n", "\n", "# for the current model, higher is worse\n", "\n", "results = []\n", "for audio_filepath in tqdm(audio_filepaths):\n", " x, sr = torchaudio.load(audio_filepath)\n", "\n", " if x.shape[1] < 60*sr:\n", " print(f\"Skipping {audio_filepath} because it's too short\")\n", " continue\n", "\n", " # crop to 30s and use the last 30s\n", " # first 30s\n", " x_first = x[:, :30*sr]\n", " x_last = x[:, -30*sr:]\n", "\n", " # lower score is better, \n", " # run inference\n", " with torch.no_grad():\n", " first_scores, first_mean_score = model.get_score(x_first.cuda())\n", " last_scores, last_mean_score = model.get_score(x_last.cuda())\n", "\n", " # delta mean score\n", " delta_mean_score = last_mean_score - first_mean_score\n", "\n", " # print the top 5 predictions\n", " # store results\n", " results.append({\n", " \"filepath\": audio_filepath,\n", " \"first_scores\": first_scores,\n", " \"first_mean_score\": first_mean_score,\n", " \"last_scores\": last_scores,\n", " \"last_mean_score\": last_mean_score,\n", " \"delta_mean_score\": delta_mean_score,\n", " \"first_audio\": x_first.numpy(),\n", " \"last_audio\": x_last.numpy(),\n", " })\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [], "source": [ "# sort results by preds\n", "results = sorted(results, key=lambda x: x[\"last_mean_score\"], reverse=True)\n", "\n", "\n", "\n", "#print(results[-1][\"filepath\"])\n", "#print(results[-1][\"first_mean_score\"])\n", "#IPython.display.display(IPython.display.Audio(results[-1][\"first_audio\"], rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# listen to the top 1 and bottom 1\n", "idx = 0\n", "print(results[idx][\"filepath\"])\n", "print(results[idx][\"first_mean_score\"], results[idx][\"last_mean_score\"], results[idx][\"delta_mean_score\"])\n", "IPython.display.display(IPython.display.Audio(results[idx][\"first_audio\"], rate=48000))\n", "IPython.display.display(IPython.display.Audio(results[idx][\"last_audio\"], rate=48000))" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [], "source": [ "# parallelize this\n", "from tqdm import tqdm\n", "import torchaudio\n", "import glob\n", "import numpy as np\n", "import torch\n", "from torch.utils.data import Dataset, DataLoader\n", "\n", "class AudioDataset(Dataset):\n", " def __init__(self, filepaths, start_s=0.0, duration_s=60.0, target_sr=48000, chunk_duration_s=5.0):\n", " self.filepaths = filepaths\n", " self.start_s = start_s\n", " self.duration_s = duration_s\n", " self.target_sr = target_sr\n", " self.chunk_duration_s = chunk_duration_s\n", " \n", " def __len__(self):\n", " return len(self.filepaths)\n", " \n", " def __getitem__(self, idx):\n", " filepath = self.filepaths[idx]\n", " try:\n", " audio, sr = torchaudio.load(filepath)\n", " if sr != self.target_sr:\n", " audio = torchaudio.functional.resample(audio, sr, self.target_sr)\n", " \n", " # Crop to specified duration\n", " audio_duration_s = audio.shape[1] / self.target_sr\n", " start_s = np.random.uniform(0, audio_duration_s - self.duration_s)\n", " start_idx = int(start_s * self.target_sr)\n", " end_idx = int((start_s + self.duration_s) * self.target_sr)\n", " audio = audio[:, start_idx:end_idx]\n", "\n", " if audio.shape[1] < self.target_sr * self.duration_s:\n", " return None\n", "\n", " # now fold in chunks\n", " chunk_size = int(self.target_sr * self.chunk_duration_s)\n", " audio_chunks = audio.unfold(1, chunk_size, chunk_size)\n", " audio_chunks = audio_chunks.reshape(-1, 2, chunk_size)\n", "\n", " # Normalize\n", " #audio = audio / audio.abs().max().clamp(1e-6)\n", " \n", " return {\n", " \"audio\": audio,\n", " \"audio_chunks\": audio_chunks,\n", " \"filepath\": filepath\n", " }\n", " except Exception as e:\n", " print(f\"Error loading {filepath}: {str(e)}\")\n", " return None\n", "\n", "def collate_fn(batch):\n", " # Filter out None values from failed loads\n", " batch = [b for b in batch if b is not None]\n", " if not batch:\n", " return None\n", " \n", " return {\n", " 'audio': torch.stack([item['audio'] for item in batch]),\n", " 'audio_chunks': torch.stack([item['audio_chunks'] for item in batch]),\n", " 'filepath': [item['filepath'] for item in batch]\n", " }\n", "\n", "idx_to_label = {idx: label for label, idx in label_encoder.items()}\n", "\n", "# Setup\n", "root_dir = \"/app/suno/data/audio_2ch_48khz_lg/train/youtube_music/\"\n", "filepaths = glob.glob(os.path.join(root_dir, \"*.wav\"))\n", "filepaths = filepaths[:10_000]\n", "\n", "# Create dataset and dataloader\n", "dataset = AudioDataset(filepaths)\n", "dataloader = DataLoader(\n", " dataset,\n", " batch_size=1, # Adjust based on your GPU memory\n", " num_workers=16, # Adjust based on your CPU cores\n", " collate_fn=collate_fn,\n", " shuffle=False,\n", " pin_memory=True\n", ")\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "results = []\n", "\n", "with torch.no_grad():\n", " for batch in tqdm(dataloader):\n", " if batch is None:\n", " continue\n", " \n", " # Move batch to GPU\n", " audio = batch['audio']\n", " audio_chunks = batch['audio_chunks'].cuda()\n", " filepaths = batch['filepath']\n", " \n", " # Get predictions\n", " preds = model(audio_chunks.squeeze(0))\n", " #probs = torch.sigmoid(preds).cpu().numpy()\n", " \n", " # Process each item in the batch\n", " for i in range(len(filepaths)):\n", " #label_probs = {idx_to_label[j]: prob for j, prob in enumerate(probs[i])}\n", " \n", " results.append({\n", " \"filepath\": filepaths[i],\n", " \"preds\": preds.mean().item(),\n", " #\"label_predictions\": label_probs,\n", " \"audio\": audio[i].cpu().numpy(),\n", " })" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# histogram of preds\n", "import matplotlib.pyplot as plt\n", "all_preds = [result[\"preds\"] for result in results]\n", "plt.hist(all_preds, bins=100)\n", "plt.show()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# sort results by preds\n", "import random\n", "results = sorted(results, key=lambda x: x[\"preds\"], reverse=True)\n", "# print the top 10 results\n", "num_to_print = 5\n", "num_printed = 0\n", "\n", "# randomly sample 5 results with preds above 0.5\n", "#results_sublist = [result for result in results if result[\"preds\"] > -0.1]\n", "results_sublist = results\n", "print(len(results_sublist))\n", "#results_sublist = random.sample(results_sublist, 5)\n", "\n", "for result in results_sublist:\n", "\n", " # measure the energy of the audio\n", " energy = np.mean(result[\"audio\"]**2)\n", " #print(f\"energy: {energy}\")\n", " if energy < 1e-6:\n", " print(\"energy is too low, skipping\")\n", " continue\n", "\n", " print(result[\"filepath\"])\n", " print(result[\"preds\"] * 60.0)\n", " #IPython.display.display(IPython.display.Audio(result[\"audio\"], rate=48000))\n", " num_printed += 1\n", " if num_printed >= num_to_print:\n", " break\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# parallelize this\n", "import glob\n", "import torch\n", "import random\n", "import torchaudio\n", "import numpy as np\n", "from tqdm import tqdm\n", "from torch.utils.data import Dataset, DataLoader\n", "from suno_utils.utils.text import read_jsonl\n", "\n", "# Setup\n", "metas_filepath = \"/home/christian/code/christian/metadata/genius_hq_metas.jsonl\"\n", "metas = read_jsonl(metas_filepath)\n", "print(len(metas))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\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\n", "\n", "# metas audio dataset\n", "class MetasAudioDataset(Dataset):\n", " def __init__(self, metas, duration_s=5.0, target_sr=48000, min_duration_s=60.0, num_chunks=2):\n", " self.metas = metas\n", " self.duration_s = duration_s\n", " self.target_sr = target_sr\n", " self.min_duration_s = min_duration_s\n", " self.num_chunks = num_chunks\n", " \n", " def __len__(self):\n", " return len(self.metas)\n", " \n", " def __getitem__(self, idx):\n", " meta = self.metas[idx]\n", "\n", " # get the s3 filepath \n", " if \"audio_filepath\" in meta:\n", " s3_filepath = meta[\"audio_filepath\"]\n", " elif \"s3_filepath\" in meta:\n", " s3_filepath = meta[\"s3_filepath\"]\n", " else:\n", " return None\n", "\n", " # download the audio\n", " filepath = download_audio(s3_filepath, meta[\"id\"], \"/mnt/localdisk/tmp/cjs\")\n", "\n", " try:\n", " audio, sr = torchaudio.load(filepath)\n", " if sr != self.target_sr:\n", " audio = torchaudio.functional.resample(audio, sr, self.target_sr)\n", " \n", " # ensure stereo\n", " if audio.shape[0] == 1:\n", " audio = audio.repeat(2, 1)\n", " elif audio.shape[0] > 2:\n", " audio = audio[:2, :] \n", "\n", " audio_duration_s = audio.shape[1] / self.target_sr\n", " if audio_duration_s < self.min_duration_s:\n", " return None\n", "\n", " # chunk into duration_s chunks\n", " chunk_size = int(self.target_sr * self.duration_s)\n", " chunks = torch.split(audio, chunk_size, dim=1)\n", " chunks = [chunk / chunk.abs().max().clamp(1e-6) for chunk in chunks]\n", " # drop last chunk if its less than duration_s\n", " if len(chunks[-1]) < chunk_size:\n", " chunks = chunks[:-1]\n", " # randomly sample N chunks\n", " audio = torch.stack(random.sample(chunks, self.num_chunks), dim=0)\n", " \n", " return {\n", " 'audio': audio,\n", " 'filepath': filepath,\n", " 'id': meta[\"id\"]\n", " }\n", " except Exception as e:\n", " print(f\"Error loading {filepath}: {str(e)}\")\n", " return None\n", "\n", "def collate_fn(batch):\n", " # Filter out None values from failed loads\n", " batch = [b for b in batch if b is not None]\n", " if not batch:\n", " return None\n", " \n", " return {\n", " 'id': [item['id'] for item in batch],\n", " 'audio': torch.stack([item['audio'] for item in batch]),\n", " 'filepath': [item['filepath'] for item in batch]\n", " }\n", "\n", "idx_to_label = {idx: label for label, idx in label_encoder.items()}\n", "\n", "# Create dataset and dataloader\n", "dataset = MetasAudioDataset(metas[:10_000])\n", "dataloader = DataLoader(\n", " dataset,\n", " batch_size=4, # effective batch is 8*5=40\n", " num_workers=16, # Adjust based on your CPU cores\n", " collate_fn=collate_fn,\n", " shuffle=False,\n", " pin_memory=True\n", ")\n", "\n", "results = []\n", "model = model.cuda()\n", "model.eval()\n", "\n", "# output jsonl file\n", "output_filepath = \"/home/christian/code/christian/metadata/genius_hq_metas_audio_corruptions.jsonl\"\n", "# open file handle for appending\n", "with open(output_filepath, \"a\") as f:\n", "\n", " # first read in the existing file\n", " existing_ids = set()\n", " with open(output_filepath, \"r\") as f:\n", " for line in f:\n", " data = json.loads(line)\n", " existing_ids.add(data[\"id\"])\n", "\n", " with torch.no_grad():\n", " for batch in tqdm(dataloader):\n", " if batch is None:\n", " continue\n", " \n", " # Move batch to GPU\n", " audio = batch['audio'].cuda()\n", " filepaths = batch['filepath']\n", "\n", " bs, n_chunks, n_channels, n_samples = audio.shape\n", " \n", " # move chunk dim to batch dim (perform in parallel)\n", " audio = audio.view(-1, n_channels, n_samples)\n", "\n", " # Get predictions\n", " preds = model(audio)\n", " preds = preds.view(bs, n_chunks, -1)\n", " probs = torch.sigmoid(preds).cpu()\n", " probs = probs.mean(dim=1)\n", " print(probs.shape)\n", " # enforce correct shape \n", " probs = probs.view(bs, -1)\n", " print(probs.shape)\n", " # Process each item in the batch\n", " for i in range(len(filepaths)):\n", " if batch[\"id\"][i] in existing_ids:\n", " continue\n", "\n", " label_probs = {idx_to_label[j]: float(prob) for j, prob in enumerate(probs[i])}\n", " \n", " # get the logits\n", " logits = probs[i].squeeze().tolist()\n", " # make json serializable\n", " logits = [float(logit) for logit in logits]\n", "\n", " result = {\n", " \"id\": batch[\"id\"][i],\n", " }\n", "\n", " for label, prob in label_probs.items():\n", " result[label] = prob\n", "\n", " # write results to file\n", " f.write(json.dumps(result) + \"\\n\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# load corruption results\n", "corruption_results = read_jsonl(output_filepath)\n", "print(len(corruption_results))\n", "\n", "# convert to dataframe\n", "import pandas as pd\n", "df = pd.DataFrame(corruption_results)\n", "\n", "columns = [col for col in df.columns if col != \"id\"]\n", "\n", "for col in columns:\n", " # compute upper 95% threshold\n", " threshold = df[col].quantile(0.99)\n", " # find top 10 rows for this column\n", " top_rows = df.sort_values(by=col, ascending=False).head(5)\n", " print(top_rows[[\"id\", col]])\n", " \n", "#for idx, row in df.iterrows():\n", "# for label, prob in row.items():\n", "# if label == \"id\":\n", "# continue\n", "# if prob > 0.5:\n", "# print(f\"{row['id']:<40} : {label:<40} : {prob:.3f}\")\n" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "# metas map\n", "metas_map = {meta[\"id\"]: meta for meta in metas}" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# download the the audio file \n", "meta_id = \"16ee543e-e719-49bc-ba95-9963f0f6b0e1\"\n", "s3_filepath = metas_map[meta_id][\"audio_filepath\"]\n", "filepath = download_audio(s3_filepath, meta_id, \"/mnt/localdisk/tmp/cjs\")\n", "print(filepath)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from tqdm import tqdm\n", "import torchaudio\n", "import glob\n", "import numpy as np\n", "\n", "# apply one corruption to a random sample\n", "root_dir = \"/app/suno/data/audio_2ch_48khz_lg/train/spot_genres/\"\n", "filepaths = glob.glob(os.path.join(root_dir, \"*.wav\"))\n", "filepaths = filepaths[:5000]\n", "\n", "idx_to_label = {idx: label for label, idx in label_encoder.items()}\n", "\n", "results = []\n", "\n", "for filepath in tqdm(filepaths):\n", " audio, sr = torchaudio.load(filepath)\n", " if sr != 48000:\n", " audio = torchaudio.functional.resample(audio, sr, 48000)\n", "\n", " # crop to 5 seconds\n", " start_s = 60.0\n", " end_s = start_s + 5.0\n", " audio = audio[:, int(start_s*48000):int(end_s*48000)]\n", "\n", " if audio.shape[1] < 48000*5.0:\n", " continue\n", "\n", " audio /= audio.abs().max().clamp(1e-6)\n", " audio = audio.cuda()\n", "\n", " with torch.no_grad():\n", " preds = model(audio.unsqueeze(0))\n", "\n", " # Get probabilities\n", " probs = torch.sigmoid(preds).cpu().numpy()[0]\n", " \n", " # Create a dictionary mapping class names to their probabilities\n", " label_probs = {idx_to_label[i]: prob for i, prob in enumerate(probs)}\n", " \n", " results.append({\n", " \"filepath\": filepath,\n", " \"raw_preds\": probs, # Keep the raw predictions if needed\n", " \"label_predictions\": label_probs, # Add the labeled predictions\n", " \"audio\": audio.cpu().numpy(),\n", " })" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(len(results))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# shuffle results\n", "import IPython\n", "import matplotlib.pyplot as plt\n", "#random.shuffle(results)\n", "\n", "corrupt_str = \"highpass:cutoff_hz=250\"\n", "\n", "for idx in range(len(results)):\n", " corruptions = results[idx][\"label_predictions\"]\n", "\n", " # compute codec score\n", " codec_score = 0\n", " for corruption, prob in corruptions.items():\n", " if corrupt_str in corruption:\n", " codec_score += prob\n", "\n", " for corruption, prob in corruptions.items():\n", " # check for nan in prob\n", " if np.isnan(prob):\n", " print(corruption, prob)\n", "\n", " results[idx][\"codec_score\"] = codec_score\n", " results[idx][\"total_score\"] = sum(corruptions.values())\n", " results[idx][\"quality_score\"] = compute_quality_score(corruptions)\n", " #IPython.display.display(IPython.display.Audio(result[\"audio\"], rate=48000))\n", "\n", "codec_scores = [result[\"codec_score\"] for result in results]\n", "total_scores = [result[\"total_score\"] for result in results]\n", "quality_scores = [result[\"quality_score\"] for result in results]\n", "print(len(quality_scores), quality_scores)\n", "plt.hist(quality_scores, bins=100)\n", "plt.show()\n", "\n", "# find results with the top 10 codec scores\n", "top_results = sorted(results, key=lambda x: x[\"quality_score\"], reverse=False)[10:20]\n", "\n", "# print the top results and their scores\n", "for result in top_results:\n", " print(result[\"filepath\"])\n", " print(result[\"quality_score\"])\n", " #print(result[\"codec_score\"])\n", " #for label, prob in result[\"label_predictions\"].items():\n", " # if corrupt_str in label:\n", " # print(f\"{label:<40} : {prob:.3f}\")\n", " IPython.display.display(IPython.display.Audio(result[\"audio\"], rate=48000))\n" ] }, { "cell_type": "code", "execution_count": 125, "metadata": {}, "outputs": [], "source": [ "def compute_quality_score(predictions_dict):\n", " weights = {\n", " 'noise': 1.5,\n", " 'clipping_distortion': 1.4,\n", " 'white_noise_burst': 1.3,\n", " 'audio_codec': 4.0,\n", " 'tanh_distortion': 1.2,\n", " 'bandpass': 1.2,\n", " 'lowpass': 1.1,\n", " 'highpass': 1.1,\n", " 'hum': 1.1,\n", " 'reverb': 0.9,\n", " 'comb_filter': 0.9,\n", " 'wow_flutter': 0.8,\n", " 'ring_modulation': 0.8,\n", " 'channel_imbalance': 0.7,\n", " 'stereo_width': 0.7,\n", " 'phase_randomize': 0.6,\n", " 'spectral_mask': 0.6\n", " }\n", " \n", " severity_scales = {\n", " 'lowpass': lambda params: 1 + (8000 - float(params['cutoff_hz'])) / 8000,\n", " 'highpass': lambda params: 1 + float(params['cutoff_hz']) / 4000,\n", " 'audio_codec': lambda params: 1 + (128000 - float(params['bit_rate'])) / 128000,\n", " 'clipping_distortion': lambda params: 1 + float(params['gain_db']) / 20,\n", " 'tanh_distortion': lambda params: 1 + float(params['gain_db']) / 20,\n", " 'hum': lambda params: 1 + float(params['amplitude']) / 0.5,\n", " 'wow_flutter': lambda params: 1 + float(params['depth']) / 0.005\n", " }\n", " \n", " def extract_params(corruption_str):\n", " if ':' not in corruption_str:\n", " return {}\n", " param_str = corruption_str.split(':', 1)[1]\n", " params = {}\n", " for param in param_str.split(','):\n", " if '=' in param:\n", " key, value = param.split('=')\n", " params[key.strip()] = value.strip()\n", " return params\n", " \n", " def transform_prediction(p, severity_mult=1.0):\n", " if p < 1e-6: return 0\n", " return -10 * (p ** 0.3) * severity_mult\n", " \n", " score = 100\n", " \n", " for corruption_type, weight in weights.items():\n", " relevant_items = [(k, v) for k, v in predictions_dict.items() if k.startswith(corruption_type)]\n", " if relevant_items:\n", " try:\n", " max_pred = max(pred for _, pred in relevant_items if pred is not None and pred == pred)\n", " max_pred_key = next(k for k, v in relevant_items if v == max_pred)\n", " \n", " # Apply severity scaling if available\n", " severity_mult = 1.0\n", " if corruption_type in severity_scales:\n", " params = extract_params(max_pred_key)\n", " try:\n", " severity_mult = severity_scales[corruption_type](params)\n", " except (KeyError, ValueError):\n", " pass\n", " \n", " penalty = transform_prediction(max_pred, severity_mult) * weight\n", " score += penalty\n", " except ValueError:\n", " continue\n", " \n", " return max(0, min(100, score))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# pick two random results\n", "import random\n", "random_results = random.sample(results, 2)\n", "for result in random_results:\n", " print(result[\"filepath\"])\n", " print(result[\"total_score\"])\n", " corruptions = result[\"label_predictions\"]\n", "\n", " for label, prob in corruptions.items():\n", " if prob > 0.3:\n", " print(f\"{label:<40} : {prob:.3f}\")\n", "\n", " quality_score = compute_quality_score(corruptions)\n", " print(f\"quality score: {quality_score}\")\n", " IPython.display.display(IPython.display.Audio(result[\"audio\"], rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "quality_score = compute_quality_score(corruptions)\n", "print(f\"quality score: {quality_score}\")\n", "\n", "for label, prob in corruptions.items():\n", " print(f\"{label:<40} : {prob:.3f}\")\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# shuffle results\n", "import IPython\n", "import matplotlib.pyplot as plt\n", "#random.shuffle(results)\n", "\n", "threshold = 0.3 \n", "\n", "for idx in range(len(results)):\n", " corruptions = results[idx][\"label_predictions\"]\n", " results[idx][\"total_score\"] = sum(prob > threshold for prob in corruptions.values())\n", " detected_corruptions = []\n", "\n", " for label, prob in corruptions.items():\n", " if prob > threshold:\n", " detected_corruptions.append(label)\n", "\n", " results[idx][\"detected_corruptions\"] = detected_corruptions\n", "\n", "total_scores = [result[\"total_score\"] for result in results]\n", "plt.hist(total_scores, bins=10)\n", "plt.show()\n", "\n", "# find results with the top 10 codec scores\n", "top_results = sorted(results, key=lambda x: x[\"total_score\"], reverse=True)[10:20]\n", "\n", "# print the top results and their scores\n", "for result in top_results:\n", " print(result[\"filepath\"])\n", " print(result[\"total_score\"])\n", " print(result[\"detected_corruptions\"])\n", " IPython.display.display(IPython.display.Audio(result[\"audio\"], rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(len(results))\n", "import matplotlib.pyplot as plt\n", "plt.hist([result[\"preds\"][0] for result in results], bins=100)\n", "plt.show()" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# create a score for each result for codec\n", "codec_scores = []\n", "for result in results:\n", " score = 0\n", " for corruption, prob in result[\"corruptions\"].items():\n", " if \"lowpass\" in corruption:\n", " score += prob\n", " codec_scores.append(score)\n", "\n", "nonzero_scores = [score for score in codec_scores if score > 0]\n", "plt.hist(nonzero_scores, bins=100)\n", "plt.show()\n", "\n", "# get the indices of the top 10 scores\n", "top_indices = sorted(range(len(codec_scores)), key=lambda i: codec_scores[i], reverse=True)[:10]\n", "\n", "# get the corresponding results\n", "top_results = [results[i] for i in top_indices]\n", "top_scores = [codec_scores[i] for i in top_indices]\n", "\n", "# print the top results and their scores\n", "for result, score in zip(top_results[:5], top_scores[:5]):\n", " print(result[\"filepath\"])\n", " print(result[\"preds\"][0])\n", " print(result[\"corruptions\"])\n", " print(score)\n", " IPython.display.display(IPython.display.Audio(result[\"audio\"], rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# sort results by preds\n", "results.sort(key=lambda x: x[\"preds\"][0], reverse=True)\n", "\n", "# print the first 10 results\n", "for result in results[0:5]:\n", " print(result[\"filepath\"])\n", " print(result[\"preds\"][0])\n", " print(result[\"corruptions\"])\n", " IPython.display.display(IPython.display.Audio(result[\"audio\"], rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# pick a clip at random\n", "random_result = random.choice(results)\n", "print(random_result[\"filepath\"])\n", "print(random_result[\"preds\"][0])\n", "for corruption, prob in random_result[\"corruptions\"].items():\n", " print(f\"{corruption:<40} : {prob:.3f}\")\n", "IPython.display.display(IPython.display.Audio(random_result[\"audio\"], rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# test on dpo data\n", "import pandas as pd\n", "import torchaudio\n", "import os\n", "\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\n", "\n", "\n", "BASE_S3_DIR = \"s3://suno-data-uploads/studio/uploads\"\n", "\n", "# load pkl file\n", "#DATA_PKL_PATH = (\n", "# \"/home/tony/Data/Preference/up_v3/interesting_clips_up_u_3_20241216_full.pkl\"\n", "#)\n", "DATA_PKL_PATH = \"/home/tony/Data/Preference/up_v3/interesting_clips_up_u_3_20250105_full.pkl\"\n", "\n", "df = pd.read_pickle(DATA_PKL_PATH)\n", "print(f\"df shape: {df.shape}\")\n", "\n", "df_indices = list(range(0, len(df), 2))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import IPython\n", "from tqdm import tqdm\n", "agree = []\n", "pbar = tqdm(df_indices[0:100])\n", "for idx in pbar:\n", " negative_row = df.iloc[idx]\n", " positive_row = df.iloc[idx+1]\n", "\n", " preds_list = []\n", " \n", " for row in [negative_row, positive_row]:\n", " s3_filepath = f\"{BASE_S3_DIR}/{row['id_x']}.mp3\"\n", " filepath = download_audio(s3_filepath, row['id_x'], \"/mnt/localdisk/tmp/cjs\")\n", " try:\n", " audio, sample_rate = torchaudio.load(filepath)\n", "\n", " if audio.shape[1] < 48000*60.0:\n", " continue \n", " # crop audio to max of 30sec\n", " start_s = 0.0\n", " end_s = start_s + 60.0\n", " audio = audio[:, int(start_s*sample_rate):int(end_s*sample_rate)]\n", " # convert into 5s chunks\n", " chunk_size = int(48000*5.0)\n", " num_chunks = audio.shape[1] // chunk_size\n", " chunks = torch.chunk(audio, num_chunks, dim=1)\n", " chunks = [chunk / chunk.abs().max().clamp(1e-6) for chunk in chunks]\n", " audio = torch.stack(chunks, dim=0)\n", " # ensure audio is stereo\n", " if audio.size(0) == 1:\n", " audio = audio.repeat(2, 1)\n", " elif audio.size(0) > 2:\n", " audio = audio[:2, :]\n", " except Exception as e:\n", " print(f\"Error: {e}\")\n", " print(f\"Failed to process: {filepath}\")\n", " continue\n", "\n", " # run inference on audio\n", " audio = audio.cuda()\n", " with torch.no_grad():\n", " preds = model(audio)\n", "\n", " clip_score = torch.sigmoid(preds).sum().item()\n", " preds_list.append(clip_score)\n", "\n", " #print(f\"neg: {preds_list[0]:.3f} pos: {preds_list[1]:.3f}\")\n", " if len(preds_list) < 2:\n", " continue\n", "\n", " if preds_list[0] > preds_list[1]:\n", " agree.append(1)\n", " else:\n", " agree.append(0)\n", "\n", " pbar.set_description(f\"agree: {sum(agree)}/{len(agree)} ({sum(agree)/len(agree)*100:.3f}%)\")\n", "\n", " #IPython.display.display(IPython.display.Audio(audio, rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import math\n", "\n", "def sigma_to_t(sigma):\n", " return sigma.atan() / math.pi * 2\n", "\n", "sigma = torch.tensor(1.0)\n", "t = sigma_to_t(sigma)\n", "print(t)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# embed 500 samples\n", "# take 10 of these clips and apply strong mp3 compression\n", "# embed everything \n", "# get a prototype embed by averaging the 10 compressed clips\n", "# then look at 10 nearest neighbors in the original space" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "import torchaudio\n", "\n", "def apply_lowpass_filter(audio, sample_rate, cutoff_freq):\n", " return torchaudio.functional.lowpass_biquad(audio, sample_rate, cutoff_freq)\n", "\n", "def apply_highpass_filter(audio, sample_rate, cutoff_freq):\n", " return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_freq)\n", "\n", "def apply_audio_codec(audio: torch.Tensor, sample_rate: float, bit_rate: int):\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"mp3\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " return effector.apply(audio.T, sample_rate).T" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "def embed_files(filepaths, batch_size=16, apply_corruption=False):\n", " audios = []\n", " embeds = []\n", " # chunk filepaths into batches\n", " for i in tqdm(range(0, len(filepaths), batch_size)):\n", " batch_filepaths = filepaths[i:i+batch_size]\n", " audio_batch = []\n", " for filepath in batch_filepaths:\n", " audio, sr = torchaudio.load(filepath)\n", " if sr != 48000:\n", " audio = torchaudio.functional.resample(audio, sr, 48000)\n", "\n", " if audio.shape[1] < 48000*65.0:\n", " continue\n", " \n", " #crop to 5 seconds\n", " start_s = 60.0\n", " end_s = start_s + 5.0\n", " audio = audio[:, int(start_s*48000):int(end_s*48000)]\n", "\n", " # apply optional corruption\n", " if apply_corruption:\n", " random_choice = random.random()\n", " if random_choice > 0.66:\n", " audio = apply_lowpass_filter(audio, 48000, 5000)\n", " elif random_choice > 0.33:\n", " audio = apply_highpass_filter(audio, 48000, 500)\n", " else:\n", " audio = apply_audio_codec(audio, 48000, 16000)\n", "\n", " audios.append(audio)\n", " audio_batch.append(audio)\n", "\n", " # batchify\n", " audio_batch = torch.stack(audio_batch, dim=0).cuda()\n", " with torch.no_grad():\n", " embeds_batch = model.get_embeddings(audio_batch).cpu()\n", " # split embeds to list of tensors\n", " embeds.extend(torch.split(embeds_batch, batch_size))\n", "\n", " return audios, embeds\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import glob\n", "from tqdm import tqdm\n", "import os\n", "import random\n", "\n", "\n", "root_dir = \"/app/suno/data/audio_2ch_48khz_lg/train/genius_hq\"\n", "filepaths = glob.glob(os.path.join(root_dir, \"*.wav\"))\n", "print(len(filepaths))\n", "main_filepaths = filepaths[:500]\n", "proto_filepaths = filepaths[1000:1100]\n", "print(len(proto_filepaths))\n", "\n", "main_embeds = []\n", "main_audios = []\n", "audios, embeds = embed_files(main_filepaths)\n", "main_embeds = embeds\n", "main_audios = audios\n", "\n", "main_embeds = torch.cat(main_embeds, dim=0)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# embed the proto files\n", "proto_embeds = []\n", "proto_audios = []\n", "audios, embeds = embed_files(proto_filepaths, apply_corruption=True)\n", "proto_embeds = embeds\n", "proto_audios = audios\n", "\n", "proto_embeds = torch.cat(proto_embeds, dim=0)\n", "\n", "# build the query by averaging the proto embeddings\n", "query = proto_embeds.mean(dim=0)\n", "print(query.shape)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import IPython\n", "# listen to the first proto audio\n", "IPython.display.display(IPython.display.Audio(proto_audios[1].cpu(), rate=48000))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# nearest neighbors in the original space\n", "import IPython\n", "from sklearn.neighbors import NearestNeighbors\n", "\n", "# Assuming your data:\n", "# embeds: shape (n_examples, n_features)\n", "# query: shape (1, n_features) or (n_features,)\n", "\n", "# Initialize and fit\n", "nn = NearestNeighbors(n_neighbors=5) # Adjust number of neighbors as needed\n", "nn.fit(main_embeds)\n", "\n", "# Find nearest neighbors\n", "distances, indices = nn.kneighbors(query.reshape(1, -1))\n", "print(distances)\n", "print(indices)\n", "\n", "for idx in indices[0]:\n", " IPython.display.display(IPython.display.Audio(main_audios[idx].cpu(), rate=48000))\n" ] }, { "cell_type": "code", "execution_count": 89, "metadata": {}, "outputs": [], "source": [ "# add one query recording by applying strong mp3 compression\n", "query_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "query_audio, sr = torchaudio.load(query_filepath)\n", "if sr != 48000:\n", " query_audio = torchaudio.functional.resample(query_audio, sr, 48000)\n", "start_s = 60.0\n", "end_s = start_s + 5.0\n", "query_audio = query_audio[:, int(start_s*48000):int(end_s*48000)]\n", "query_audio = apply_audio_codec(query_audio, 48000, 16000)\n", "query_audio = query_audio.cuda()\n", "with torch.no_grad():\n", " query_embed = model.get_embeddings(query_audio.unsqueeze(0))\n", "embeds = torch.cat([embeds, query_embed.cpu()], dim=0)\n", "labels.append(\"g\")\n", "audios.append(query_audio)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import IPython\n", "# listen to query audio\n", "IPython.display.display(IPython.display.Audio(query_audio.cpu(), rate=48000))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# use umap to reduce dimensionality\n", "import umap\n", "reducer = umap.UMAP()\n", "embeds_2d = reducer.fit_transform(embeds)\n", "print(embeds_2d.shape)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# now create a scatter plot\n", "import matplotlib.pyplot as plt\n", "plt.scatter(embeds_2d[:, 0], embeds_2d[:, 1], c=labels, alpha=0.5)\n", "plt.show()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(len([label for label in labels if label == \"b\"]))\n", "print(len([label for label in labels if label == \"r\"]))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from sklearn.cluster import KMeans\n", "from sklearn.preprocessing import StandardScaler\n", "\n", "# Initialize kmeans\n", "kmeans = KMeans(n_clusters=2, random_state=1, verbose=2, tol=1e-6)\n", "\n", "# Fit the model and get cluster assignments in one step\n", "scaler = StandardScaler()\n", "embeds_2d_scaled = scaler.fit_transform(embeds_2d) \n", "pred_labels = kmeans.fit_predict(embeds_2d_scaled) # X is your data of shape (examples, features)\n", "\n", "# Get cluster centers\n", "centers = kmeans.cluster_centers_\n", "\n", "# Get cluster assignments for new data\n", "#new_labels = kmeans.predict(embeds_2d)\n", "\n", "# add the centroids to the plot\n", "# first convert to 2d\n", "plt.scatter(embeds_2d_scaled[:, 0], embeds_2d_scaled[:, 1], c=labels, alpha=0.5)\n", "plt.scatter(centers[:, 0], centers[:, 1], c=\"black\", marker=\"x\", label=\"Centroids\")\n", "plt.legend()\n", "plt.show()\n" ] }, { "cell_type": "code", "execution_count": 274, "metadata": {}, "outputs": [], "source": [ "import numpy as np\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(audio: torch.Tensor, sample_rate: float, imbalance: float = 0.0):\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, sample_rate: float, central_freq: float=1000.0, 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(audio: torch.Tensor, sample_rate: float, gain_db: float=0.0):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " return torch.tanh(audio * gain_lin)\n", "\n", "\n", "def apply_clipping_distortion(audio: torch.Tensor, sample_rate: float, gain_db: float=0.0):\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, sample_rate: float, gain_db: float=0.0, 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(audio: torch.Tensor, sample_rate: float, amplitude: float=0.0, freq: float=0.0):\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 = side * width # 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(audio: torch.Tensor, sample_rate: float, threshold: float = -60, ratio: float = 0.5, n_fft: int = 2048):\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(spec * mask, n_fft, n_fft//4, window=window, length=audio.shape[-1])\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(audio: torch.Tensor, sample_rate: int = 48000, rate: float = 5.0, depth: float = 0.1):\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", "def apply_wow_flutter_fast(audio: torch.Tensor, sample_rate: int = 48000, rate: float = 5.0, depth: float = 0.1):\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", "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", "\n", "# --- new corruptions from dec 18 2024 ---\n", "\n", "def apply_compression(\n", " waveform: torch.Tensor,\n", " sample_rate: int,\n", " attack_time: float = 10.0,\n", " release_time: float = 100.0,\n", " threshold: float = -20.0,\n", " ratio: float = 2.0,\n", " knee_width: float = 1.0,\n", " gain: float = 0.0\n", "): \n", " # Format the transfer characteristic points\n", " transfer = f\"{knee_width}:{threshold},{ratio}\"\n", " \n", " effects = [\n", " \"compand\",\n", " str(attack_time/1000), # Attack time in seconds\n", " str(release_time/1000), # Release time in seconds\n", " transfer\n", " ]\n", " \n", " processed_waveform, new_sample_rate = torchaudio.sox_effects.apply_effects_tensor(\n", " waveform,\n", " sample_rate,\n", " [effects]\n", " )\n", " \n", " return processed_waveform, new_sample_rate\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", "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(audio, sample_rate, noise_level=0.1, min_burst_length=500, max_burst_length=8000, p_burst=0.01):\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(min_burst_length, max_burst_length, (num_segments,), device=audio.device)\n", " burst_levels = noise_level * (0.5 + torch.rand(num_segments, device=audio.device))\n", " burst_mask = (torch.rand(num_segments, device=audio.device) < p_burst).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)) & (indices.unsqueeze(0) < positions.unsqueeze(1))\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", "def apply_quantize_zero(audio, threshold=0.001):\n", " return torch.where(torch.abs(audio) < threshold, 0, audio)\n", "\n", "def apply_phase_randomize(audio: torch.Tensor, sample_rate: float, block_size: int = 2048, mix: float = 0.75):\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(mag * random_phase, block_size, window=window, length=channel.shape[-1])\n", " output.append(channel_out)\n", " return (1 - mix) * audio + mix * torch.stack(output)\n" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "import torch\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", "def apply_stereo_to_mono(audio: torch.Tensor, sample_rate: float):\n", " return audio.mean(dim=0, keepdims=True).repeat(2, 1)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# test some corruptions\n", "# add one query recording by applying strong mp3 compression\n", "import IPython.display\n", "import torchaudio\n", "\n", "#query_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "#query_filepath = \"/home/christian/audio/reference-audio-wav/02 Take Five.wav\"\n", "query_filepath = \"/home/christian/audio/reference-audio-wav/02 The Metronome.wav\"\n", "query_audio, sr = torchaudio.load(query_filepath)\n", "if sr != 48000:\n", " query_audio = torchaudio.functional.resample(query_audio, sr, 48000)\n", "start_s = 60.0\n", "end_s = start_s + 10.0\n", "query_audio = query_audio[:, int(start_s*48000):int(end_s*48000)]\n", "\n", "#preset, labels = generate_random_preset(corruptions, no_corruption_probability=0.05)\n", "#for idx, label in enumerate(labels):\n", "# print(idx, label)\n", "\n", "\n", " \n", "#corrupted_audio = apply_preset(query_audio, 48000, preset, corruptions, corruption_functions)\n", "#corrupted_audio = apply_quantize_zero(query_audio, threshold=0.025) # 0.025, 0.05, 0.1, 0.2\n", "#corrupted_audio = apply_phase_randomize(query_audio, 48000)\n", "corrupted_audio = apply_stereo_width(query_audio, 48000, width=4.0)\n", "print(corrupted_audio.shape)\n", "\n", "\n", "IPython.display.display(IPython.display.Audio(query_audio.cpu(), rate=48000))\n", "IPython.display.display(IPython.display.Audio(corrupted_audio.cpu(), rate=48000))\n" ] }, { "cell_type": "code", "execution_count": 268, "metadata": {}, "outputs": [], "source": [ "def sample_n_corruptions(max_corruptions: int = 5, p: float = 0.5, batch_size: int = 1):\n", " # Create geometric distribution weights\n", " probs = torch.tensor([(1-p)**i * p for i in range(max_corruptions)], device='cuda')\n", " probs = probs / probs.sum() # Normalize to sum to 1\n", " \n", " # Sample from this distribution\n", " return torch.multinomial(probs.expand(batch_size, -1), 1).squeeze(-1) + 1\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", " 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.333, batch_size=1)\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", "def apply_preset(\n", " audio: torch.Tensor, sr: float, preset: dict, config: dict, functions: dict\n", "):\n", " chs, seq_len = audio.shape\n", " for corruption_name, corruption_info in preset.items():\n", " audio = functions[corruption_name](audio, sr, **corruption_info[\"params\"])\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", "# 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", "}\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import time\n", "# lets profile the speed of each corruption\n", "query_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "query_audio, sr = torchaudio.load(query_filepath)\n", "if sr != 48000:\n", " query_audio = torchaudio.functional.resample(query_audio, sr, 48000)\n", "start_s = 60.0\n", "end_s = start_s + 10.0\n", "query_audio = query_audio[:, int(start_s*48000):int(end_s*48000)]\n", "\n", "for corruption_name, corruption_function in corruption_functions.items():\n", " timings = []\n", " for n in range(20):\n", " start_time = time.time()\n", " corrupted_audio = corruption_function(query_audio, 48000)\n", " end_time = time.time()\n", " timings.append(end_time - start_time)\n", " print(f\"{corruption_name}: {np.mean(timings)*1000:.2f} ms (std: {np.std(timings)*1000:.2f} ms)\")\n", "\n", " " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "p = 0.5\n", "probs = torch.tensor([(1-p)**i * p for i in range(100)], device='cuda')\n", "probs = probs / probs.sum() # Normalize to sum to 1\n", "print(probs * 100)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "\n", "def apply_shimmer(\n", " audio: torch.Tensor, \n", " sample_rate: float,\n", " ring_freq: float = 6000,\n", " threshold_mult: float = 2.0,\n", " decay_length: int = 1000,\n", " mix: float = 0.8\n", "):\n", " # Sum to mono for trigger detection\n", " mono = torch.mean(audio, dim=0) if audio.dim() == 2 else audio\n", " \n", " # Detect transients\n", " diff = torch.abs(torch.diff(mono))\n", " threshold = torch.mean(diff) + threshold_mult * torch.std(diff)\n", " triggers = (diff > threshold).float()\n", " \n", " # Generate ringing\n", " ring = torchaudio.functional.bandpass_biquad(\n", " triggers[None, None, :], # Add batch and channel dims\n", " sample_rate,\n", " central_freq=ring_freq,\n", " Q=50.0\n", " )\n", " \n", " # Apply decay\n", " decay = torch.exp(-torch.linspace(0, 5, decay_length)).to(audio.device)\n", " ring = torch.nn.functional.conv1d(\n", " ring,\n", " decay.view(1, 1, -1),\n", " padding='same'\n", " )\n", " \n", " # Expand back to stereo if needed\n", " if audio.dim() == 2:\n", " ring = ring.squeeze(0).repeat(2, 1)\n", " else:\n", " ring = ring.squeeze(0).squeeze(0)\n", " \n", " # apply reverb to the ring\n", " ring = apply_reverb(ring, 48000, reverberance=100, hf_damping=100, room_scale=100, stereo_depth=100, pre_delay=0, wet_gain=0)\n", "\n", " # Pad to match original length and mix\n", " ring = torch.nn.functional.pad(ring, (0, 1))\n", " return audio + (mix * ring)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "# Define how to process each sample\n", "def process_sample(sample, sample_length, is_shuffle):\n", " json_data = json.loads(sample['json'])\n", "\n", " # make audio_bytes to be even for stereo support\n", " if len(sample[\"audio_bytes\"]) / json_data[\"audio_byte_width\"] % 2 != 0:\n", " sample[\"audio_bytes\"] = sample[\"audio_bytes\"][:-json_data[\"audio_byte_width\"]]\n", " \n", " # audio\n", " audio = Audio(\n", " sample['audio_bytes'], \n", " sample_rate=json_data['audio_sample_rate'],\n", " byte_width=json_data[\"audio_byte_width\"], \n", " n_channels=2,\n", " ).normalize_volume().array_float\n", "\n", " # get random 1s chunk\n", " start_idx = np.random.randint(0, audio.shape[1] - sample_length) if is_shuffle else 0\n", " audio = audio[:, start_idx:start_idx + sample_length]\n", "\n", " # Parse the JSON data\n", " return {\n", " 'wav': audio,\n", " 'key': json_data['__key__'],\n", " }\n", "\n", "def retry(attempts=3):\n", " def handler(exn):\n", " if attempts > 0:\n", " print(f\"Retrying after error: {exn}\")\n", " return attempts - 1 # Reduce the number of attempts\n", " else:\n", " reraise_exception(exn) # Raise the error if retry attempts are exhausted\n", " return handler\n", "\n", "\n", "class BufferedWebDataset(IterableDataset):\n", " def __init__(\n", " self, \n", " s3_path,\n", " shard_start,\n", " shard_end,\n", " split, \n", " sample_length,\n", " num_iterations,\n", " world_size,\n", " ):\n", " assert split in [\"train\", \"valid\"]\n", " is_shuffle = (split == \"train\")\n", " urls = [os.path.join(s3_path, split, f\"shard_{i:06d}.tar\") for i in range(shard_start, shard_end)]\n", " if is_shuffle:\n", " random.shuffle(urls)\n", " self.urls = [f'pipe:aws s3 cp {url} --profile oracle --endpoint-url https://lrkg2trbk8ge.compat.objectstorage.us-chicago-1.oraclecloud.com -' for url in urls]\n", " self.num_iterations = num_iterations\n", " self.world_size = world_size\n", " self.sample_length = sample_length\n", " self.is_shuffle = is_shuffle\n", "\n", " def __iter__(self):\n", " for url in self.urls:\n", " # get dataset\n", " if self.world_size > 1:\n", " dataset = wds.WebDataset(url, nodesplitter=wds.split_by_worker)\n", " else:\n", " dataset = wds.WebDataset(url)\n", " if self.is_shuffle:\n", " dataset = dataset.shuffle(200)\n", " dataset = dataset.map(lambda sample: process_sample(sample, self.sample_length, self.is_shuffle), handler=retry(attempts=5))\n", "\n", " # iterate\n", " for sample in dataset:\n", " yield sample\n", "\n", " def __len__(self):\n", " return self.num_iterations\n", "\n", "\n", "class CorruptAudioDataset(torch.utils.data.Dataset):\n", " def __init__(\n", " self,\n", " filepaths: List[str],\n", " label_encoder: Dict[str, int],\n", " corruptions: Dict[str, Dict[str, Any]],\n", " sample_rate: int,\n", " min_corruptions: int = 0,\n", " max_corruptions: int = 4,\n", " no_corruption_probability: float = 0.1,\n", " num_workers: int = 1,\n", " chunk_size_s: float = 10.0,\n", " buffer_size: int = 50_000,\n", " ):\n", " self.filepaths = filepaths\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.min_corruptions = min_corruptions\n", " self.max_corruptions = max_corruptions\n", " self.no_corruption_probability = no_corruption_probability\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", " 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", " print(\"Reloading buffer...\")\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 pbar:\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.chunk_size_samples:\n", " continue\n", "\n", " # Split into chunks\n", " chunks = audio.unfold(\n", " -1, self.chunk_size_samples, self.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.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", " 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", " if self.items_since_last_reload >= len(self.buffer):\n", " self._reload_buffer()\n", "\n", " # get a random preset and apply it to the audio\n", " buffer_idx = np.random.randint(0, len(self.buffer))\n", " audio = self.buffer[buffer_idx].clone()\n", "\n", " preset, labels = generate_random_preset(\n", " self.corruptions,\n", " self.min_corruptions,\n", " self.max_corruptions,\n", " self.no_corruption_probability,\n", " )\n", "\n", " if preset: # If there are corruptions to apply\n", " corrupted_audio = apply_preset(\n", " audio,\n", " self.sample_rate,\n", " preset,\n", " self.corruptions,\n", " corruption_functions,\n", " )\n", " else:\n", " corrupted_audio = audio # Use original audio if no corruptions\n", "\n", " label_tensor = labels_to_tensor(labels, self.label_encoder)\n", "\n", " # ensure nothing is out of range\n", " if corrupted_audio.abs().max() > 1.0:\n", " corrupted_audio = corrupted_audio / corrupted_audio.abs().max()\n", "\n", " if np.random.uniform() < 0.5:\n", " gain_reduction_db = np.random.uniform(-10, 0)\n", " corrupted_audio *= 10 ** (gain_reduction_db / 20.0)\n", "\n", " self.items_since_last_reload += 1\n", "\n", " return corrupted_audio, label_tensor" ] } ], "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 }