{ "cells": [ { "cell_type": "code", "execution_count": 27, "metadata": {}, "outputs": [], "source": [ "import os\n", "import json\n", "import time\n", "import tqdm\n", "import random\n", "import numpy as np\n", "import webdataset as wds\n", "from suno_utils.audio import Audio\n", "from webdataset.handlers import reraise_exception\n", "from torch.utils.data import DataLoader, IterableDataset\n", "\n", "\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", " # add a bit of noise\n", " noise = np.random.randn(*audio.shape) * 1e-5\n", " audio = audio + noise\n", "\n", " # augmentation\n", " # if is_shuffle:\n", " # max_abs = np.max(np.abs(audio))\n", " # upper_bound = min(2.0, 1.0 / max_abs if max_abs > 0 else 2.0)\n", " # amp_factor = np.random.uniform(0.5, upper_bound) # +- 6dB\n", " # audio = audio * amp_factor\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 IterableWebDataset(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" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# get dataset\n", "dataset = IterableWebDataset(\n", " s3_path = \"s3://webdataset/diffusion_mix_fix_10s\",\n", " shard_start = 0,\n", " shard_end = 100,\n", " split = \"train\",\n", " sample_length = 48000,\n", " num_iterations = 1000,\n", " world_size = 1,\n", " )\n", "\n", "# Create a DataLoader (optional, for batching)\n", "dataloader = DataLoader(dataset, batch_size=80, num_workers=20)\n", "\n", "# Iterate through 100 steps of the dataset\n", "start_time = time.time()\n", "\n", "keys = []\n", "wavs = []\n", "for i, out in tqdm.tqdm(enumerate(dataloader)):\n", "# for _ in tqdm.tqdm(range(100)):\n", " try:\n", " keys.append(out['key'])\n", " wavs.append(out['wav'])\n", " except StopIteration:\n", " print(\"Reached the end of the dataset before 100 iterations.\")\n", " break\n", " if i > 10:\n", " break\n", "\n", "end_time = time.time()\n", "total_time = end_time - start_time\n", "print(f\"Webdataset Time taken for 100 iterations: {total_time:.2f} seconds\")\n", "print(f\"Webdataset Average time per iteration: {total_time / 100:.2f} seconds\")\n", "print(keys[0][:10])" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(keys)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "# noise test, take some silence, add varying levels of noise and measure the energy\n", "for noise_level in [1e-8, 1e-7, 1e-6, 1e-5, 1e-4]:\n", " silence = torch.zeros(48000).half().float() # one second of silence\n", " noise = torch.randn_like(silence) * noise_level\n", " wav = silence + noise\n", " energy = (wav ** 2).mean().item()\n", " dynamic_range = 20 * np.log10(1/noise_level)\n", " print(f\"Noise level: {noise_level}, Energy: {energy}, Dynamic range: {dynamic_range}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import tqdm\n", "import torchaudio\n", "import matplotlib.pyplot as plt\n", "output_dir = \"/home/christian/code/christian/debug/\"\n", "os.makedirs(output_dir, exist_ok=True)\n", "\n", "energies = []\n", "for key, wav in tqdm.tqdm(zip(keys, wavs)):\n", " bs, chs, seqlen = wav.shape\n", " for i in range(bs):\n", " filepath = os.path.join(output_dir, f\"{key[i]}_{i}.wav\")\n", " #print(filepath)\n", " # measure the energy of each example and plot distribution\n", " energy = (wav[i] ** 2).mean().item()\n", " energies.append(energy)\n", "\n", " if energy < 0.0001:\n", " print(key[i])\n", " print(energy)\n", " torchaudio.save(filepath, wav[i], 48000)\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(np.mean(energies), np.std(energies), np.max(energies), np.min(energies))\n", "plt.hist(energies, bins=100)\n", "plt.show()" ] }, { "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 }