{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n", "import torch\n", "import torchaudio\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "# ---------------- 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", "def apply_channel_imbalance(audio: torch.Tensor, sample_rate: float, imbalance: float):\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", "def apply_highpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float):\n", " return torchaudio.functional.highpass_biquad(audio, sample_rate, cutoff_hz)\n", "\n", "def apply_lowpass(audio: torch.Tensor, sample_rate: float, cutoff_hz: float):\n", " return torchaudio.functional.lowpass_biquad(audio, sample_rate, cutoff_hz)\n", "\n", "def apply_bandpass(audio: torch.Tensor, sample_rate: float, lowcut_hz: float, highcut_hz: float):\n", " return torchaudio.functional.bandpass_biquad(audio, sample_rate, lowcut_hz, highcut_hz)\n", "\n", "def apply_tanh_distortion(audio: torch.Tensor, sample_rate: float, gain_db: float):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " return torch.tanh(audio * gain_lin)\n", "\n", "def apply_clipping_distortion(audio: torch.Tensor, sample_rate: float, gain_db: float):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " return (audio * gain_lin).clamp(-1, 1)\n", "\n", "def apply_white_noise(audio: torch.Tensor, sample_rate: float, gain_db: float):\n", " gain_lin = 10 ** (gain_db / 20.0)\n", " noise = gain_lin * torch.randn(audio.shape)\n", " return noise + audio\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\n", "\n", "def apply_hum(audio: torch.Tensor, sample_rate: float, amplitude: float):\n", " t = torch.arange(audio.shape[-1], device=audio.device) / sample_rate\n", " hum = amplitude * torch.sin(2 * np.pi * 60 * t)\n", " # Add harmonics at 120Hz and 180Hz\n", " hum += (amplitude * 0.5) * torch.sin(2 * np.pi * 120 * t)\n", " hum += (amplitude * 0.25) * torch.sin(2 * np.pi * 180 * t)\n", " return audio + hum.expand_as(audio)\n", "\n", "def apply_comb_filter(audio: torch.Tensor, sample_rate: float, delay_ms: float, gain_db: float):\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", "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", "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", "def apply_reverb(audio: torch.Tensor, 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): # -10-10 dB\n", " \n", " effects = [\n", " [\"reverb\", str(reverberance), str(hf_damping), str(room_scale), \n", " str(stereo_depth), str(pre_delay), str(wet_gain)]\n", " ]\n", " out, _ = torchaudio.sox_effects.apply_effects_tensor(audio, sample_rate, effects)\n", " return out\n" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "corruptions = {\n", " \"stereo_to_mono\": {\n", " \"fn\" : apply_stereo_to_mono,\n", " \"params\" : {}\n", " },\n", " \"channel_imbalance\": {\n", " \"fn\" : apply_channel_imbalance,\n", " \"params\" : {\n", " \"imbalance\" : [-1, -0.9, -0.75, -0.5, 0.5, 0.75, 0.9, 1]\n", " }\n", " },\n", " \"lowpass\": {\n", " \"fn\" : apply_lowpass,\n", " \"params\" : {\n", " \"cutoff_hz\" : [1000, 2000, 4000, 6000, 8000, 10000]\n", " }\n", " },\n", " \"highpass\": {\n", " \"fn\" : apply_highpass,\n", " \"params\" : {\n", " \"cutoff_hz\" : [60, 100, 250, 500, 1000, 2000, 4000]\n", " }\n", " },\n", " \"tanh_distortion\": {\n", " \"fn\" : apply_tanh_distortion,\n", " \"params\" : {\n", " \"gain_db\" : [8, 12, 16, 20, 24, 32]\n", " }\n", " },\n", " \"clipping_distortion\": {\n", " \"fn\" : apply_clipping_distortion,\n", " \"params\" : {\n", " \"gain_db\" : [8, 12, 16, 20, 24, 32]\n", " }\n", " },\n", " \"white_noise\": {\n", " \"fn\" : apply_white_noise,\n", " \"params\" : {\n", " \"gain_db\" : [-52, -48, -42, -36]\n", " }\n", " },\n", " \"hum\": {\n", " \"fn\" : apply_hum,\n", " \"params\" : {\n", " \"amplitude\" : [0.01, 0.02, 0.04, 0.08]\n", " }\n", " },\n", " \"comb_filter\": {\n", " \"fn\" : apply_comb_filter,\n", " \"params\" : {\n", " \"delay_ms\" : [0.1, 0.4, 0.8, 1.6, 3.2, 6.4, 12.8, 25.6],\n", " \"gain_db\" : [3, 6, 12, 24]\n", " }\n", " },\n", " \"reduce_bit_depth\": {\n", " \"fn\" : apply_reduce_bit_depth,\n", " \"params\" : {\n", " \"bits\" : [4, 6, 8, 10, 12]\n", " }\n", " },\n", " \"add_clicks\": {\n", " \"fn\" : apply_add_clicks,\n", " \"params\" : {\n", " \"density\" : [0.00001, 0.0001, 0.001, 0.01]\n", " }\n", " },\n", " \"reverb\": {\n", " \"fn\" : apply_reverb,\n", " \"params\" : {\n", " \"reverberance\" : [0, 10, 25, 50, 100],\n", " \"hf_damping\" : [0, 10, 25, 50, 100],\n", " \"room_scale\" : [0, 10, 25, 50, 100],\n", " \"stereo_depth\" : [0, 10, 25, 50, 100],\n", " \"pre_delay\" : [0, 10, 25, 50, 100],\n", " \"wet_gain\" : [-10, -5, 0, 5, 10]\n", " }\n", " },\n", " \"audio_codec\": {\n", " \"fn\" : apply_audio_codec,\n", " \"params\" : {\n", " \"bit_rate\" : [8_000, 16_000, 24_000, 32_000, 40_000, 48_000, 64_000, 80_000, 96_000, 112_000, 128_000]\n", " }\n", " }\n", "}\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "# create presets (dict of dicts) \n", "num_presets = 100\n", "\n", "presets = {}\n", "for i in range(num_presets):\n", " preset = {}\n", " # the order of the corruptions is random\n", " # the number of corruptions is random\n", " # the parameters of the corruptions are random\n", " rand_perm = np.random.permutation(list(corruptions.keys()))\n", " #num_corruptions = np.random.randint(1, len(corruptions))\n", " num_corruptions = min(1 + int(np.random.exponential(0.9)), len(corruptions))\n", " for j in range(num_corruptions):\n", " corruption_name = rand_perm[j]\n", " # get the corruption function and parameters\n", " corruption = corruptions[corruption_name]\n", "\n", " # add the corruption to the preset\n", " preset[corruption_name] = {\n", " \"fn\" : corruption[\"fn\"].__name__,\n", " \"params\" : {}\n", " }\n", "\n", " # get the parameters\n", " params = corruption[\"params\"]\n", " for param_name, param_range in params.items():\n", " param_value = np.random.choice(param_range)\n", " # if int, convert to int\n", " if isinstance(param_value, np.int64):\n", " param_value = int(param_value)\n", " if isinstance(param_value, np.float64):\n", " param_value = float(param_value)\n", " preset[corruption_name][\"params\"][param_name] = param_value\n", "\n", " presets[i] = preset\n", "\n", "for preset_idx, preset in presets.items():\n", " print(f\"Preset {preset_idx}:\")\n", " for corruption_name, corruption_info in preset.items():\n", " print(f\" {corruption_name}: {corruption_info}\")\n", "\n", "# save the presets to a json file\n", "with open(f\"/home/christian/code/christian/metadata/presets-{num_presets}.json\", \"w\") as f:\n", " json.dump(presets, f, indent=4)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# loop through the presets and make sure they are all unique\n", "# a preset is unique if the set of corruptions and then parameter values are unique\n", "for preset_idx, preset in presets.items():\n", " for other_preset_idx, other_preset in presets.items():\n", " if preset_idx != other_preset_idx:\n", " if preset.keys() == other_preset.keys():\n", " print(f\"Preset {preset_idx} and {other_preset_idx} are the same\")\n", "\n" ] }, { "cell_type": "code", "execution_count": 40, "metadata": {}, "outputs": [], "source": [ "# now we have a function that takes a preset and applies the relevant corruptions to the audio\n", "def apply_preset(audio: torch.Tensor, sr: float, preset: dict):\n", " for corruption_name, corruption_info in preset.items():\n", " #audio = corruption_info[\"fn\"](audio, sr, **corruption_info[\"params\"])\n", " audio = corruptions[corruption_name][\"fn\"](audio, sr, **corruption_info[\"params\"])\n", " return audio\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import IPython\n", "test_audio_filepath = \"/home/christian/audio/reference-audio-wav/02 Dreams.wav\"\n", "test_audio, sr = torchaudio.load(test_audio_filepath)\n", "test_audio = test_audio[:, sr*60:sr*70]\n", "\n", "preset_idx = np.random.randint(0, len(presets))\n", "print(f\"Preset {preset_idx}: {presets[preset_idx]}\")\n", "test_audio_preset = apply_preset(test_audio, sr, presets[preset_idx])\n", "\n", "IPython.display.display(IPython.display.Audio(test_audio.numpy(), rate=sr))\n", "IPython.display.display(IPython.display.Audio(test_audio_preset.numpy(), rate=sr))\n" ] }, { "cell_type": "code", "execution_count": 30, "metadata": {}, "outputs": [], "source": [ "from typing import List\n", "from tqdm import tqdm\n", "import torchaudio\n", "import itertools\n", "class CorruptAudioDataset(torch.utils.data.Dataset):\n", " def __init__(\n", " self,\n", " filepaths: List[str],\n", " presets: List[dict],\n", " sample_rate: int,\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.presets = presets\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", "\n", " self.buffer = []\n", " self._reload_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", " preset_idx = np.random.randint(0, len(self.presets))\n", " buffer_idx = np.random.randint(0, len(self.buffer))\n", " audio = self.buffer[buffer_idx]\n", " corrupted_audio = apply_preset(\n", " audio, self.sample_rate, self.presets[preset_idx]\n", " )\n", "\n", " # ensure nothing is out of range\n", " if audio.abs().max() > 1.0:\n", " audio = audio / audio.abs().max()\n", " if corrupted_audio.abs().max() > 1.0:\n", " corrupted_audio = corrupted_audio / corrupted_audio.abs().max()\n", "\n", " # apply random gain reduction \n", " if np.random.uniform() < 0.5:\n", " gain_reduction_db = np.random.uniform(-10, 0)\n", " audio *= 10 ** (gain_reduction_db / 20.0)\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", " return audio, corrupted_audio, preset_idx\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Model" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [ "import math\n", "import torch\n", "import torch.nn as nn\n", "\n", "class SinusoidalPositionalEncoding(nn.Module):\n", " def __init__(self, hidden_dim):\n", " super().__init__()\n", " position = torch.arange(10000).unsqueeze(1)\n", " div_term = torch.exp(torch.arange(0, hidden_dim, 2) * -(math.log(10000.0) / hidden_dim))\n", " pe = torch.zeros(10000, hidden_dim)\n", " pe[:, 0::2] = torch.sin(position * div_term)\n", " pe[:, 1::2] = torch.cos(position * div_term)\n", " self.register_buffer('pe', pe)\n", "\n", " def forward(self, x):\n", " # x: (batch, num_patches, hidden_dim)\n", " return x + self.pe[:x.size(1)]\n", "\n", "class AudioQualityModel(nn.Module):\n", " def __init__(self, \n", " num_presets: int,\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", " super().__init__()\n", " self.hidden_dim = hidden_dim\n", " self.latent_dim = latent_dim\n", "\n", " # wav2vec2.0 feature encoder\n", " self.conv_layers = nn.Sequential(\n", " nn.Conv1d(2, hidden_dim, kernel_size=10, stride=5),\n", " nn.GroupNorm(num_groups=hidden_dim, num_channels=hidden_dim),\n", " nn.ReLU(),\n", " *[nn.Sequential(\n", " nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, stride=2),\n", " nn.GroupNorm(num_groups=hidden_dim, num_channels=hidden_dim),\n", " nn.ReLU()\n", " ) for _ in range(num_conv_layers)]\n", " )\n", " \n", " # Position embedding\n", " self.pos_embed = SinusoidalPositionalEncoding(hidden_dim)\n", " \n", " # Transformer encoder\n", " encoder_layer = nn.TransformerEncoderLayer(\n", " d_model=hidden_dim,\n", " nhead=num_heads,\n", " dim_feedforward=hidden_dim * 4,\n", " dropout=dropout\n", " )\n", " self.transformer = nn.TransformerEncoder(encoder_layer, num_transformer_layers)\n", " \n", " # Output head\n", " self.mlp_head = nn.Sequential(\n", " nn.LayerNorm(hidden_dim),\n", " nn.Linear(hidden_dim, latent_dim) # Single quality score output\n", " )\n", "\n", " # Projection head for pre-training\n", " self.proj_head = nn.Sequential(\n", " nn.LayerNorm(latent_dim*2),\n", " nn.Linear(latent_dim*2, 512),\n", " nn.GELU(),\n", " nn.Linear(512, 512),\n", " nn.GELU(),\n", " nn.Linear(512, 512),\n", " nn.GELU(),\n", " nn.Linear(512, num_presets)\n", " )\n", "\n", " def get_embeddings(self, x):\n", " # x shape: (batch_size, 2, time)\n", " x = self.conv_layers(x) # (batch, 512, seq)\n", " x = x.transpose(1, 2) # (batch, seq, 512)\n", " \n", " # Add position embeddings\n", " x = self.pos_embed(x) # Maintains (batch, seq, 512)\n", " \n", " # Transformer expects (seq, batch, dim)\n", " x = x.transpose(0, 1) # (seq, batch, 512)\n", " x = self.transformer(x)\n", " x = x.transpose(0, 1) # (batch, seq, 512)\n", " \n", " x = x.mean(dim=1) # (batch, 512)\n", " return self.mlp_head(x)\n", "\n", " def forward(self, input_audio, output_audio):\n", " # Stack on batch dimension\n", " stacked_audio = torch.cat([input_audio, output_audio], dim=0)\n", " # Get embeddings in one pass\n", " stacked_embeds = self.get_embeddings(stacked_audio)\n", " # Split back\n", " input_embed, output_embed = torch.chunk(stacked_embeds, 2, dim=0)\n", " return self.proj_head(torch.cat([input_embed, output_embed], dim=1))\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# first we want to define the corruption \"presets\" and save this as a json file \n", "# we will then want a function that will take one present (as dict) and then apply the relevant corruptions to the audio\n", "import glob\n", "from tqdm import tqdm\n", "\n", "# setup dataset\n", "audio_filepaths = glob.glob(\"/app/suno/data/audio_2ch_48khz_lg/val/genius_hq/*.wav\")\n", "audio_filepaths = audio_filepaths[:100]\n", "dataset = CorruptAudioDataset(audio_filepaths, presets, sample_rate=48000, chunk_size_s=10.0, buffer_size=10_000)\n", "dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=True, num_workers=4, persistent_workers=True)\n", "\n", "# setup model\n", "model = AudioQualityModel(len(presets))\n", "num_params = sum(p.numel() for p in model.parameters())\n", "print(f\"Number of parameters: {num_params/1e6:0.1f}M\")\n", "model.cuda()\n", "\n", "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n", "# pass through the model\n", "pbar = tqdm(dataloader, total=len(dataloader))\n", "for batch in pbar:\n", " optimizer.zero_grad()\n", " \n", " audio, corrupted_audio, preset_idx = batch\n", " audio = audio.cuda()\n", " corrupted_audio = corrupted_audio.cuda()\n", " preset_idx = preset_idx.cuda()\n", " \n", " scores = model(audio, corrupted_audio)\n", "\n", " loss = torch.nn.functional.cross_entropy(scores, preset_idx)\n", " loss.backward()\n", " optimizer.step()\n", "\n", " # compute the accuracy\n", " accuracy = (scores.argmax(dim=1) == preset_idx).float().mean()\n", " pbar.set_postfix({\"loss\": loss.item(), \"accuracy\": accuracy.item()})\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "model = AudioQualityModel()\n", "# count the parameters\n", "num_params = sum(p.numel() for p in model.parameters())\n", "print(f\"Number of parameters: {num_params/1e6:0.1f}M\")\n", "\n", "# count parameters in the transformer encoder, mlp head, and conv layers\n", "num_params_transformer = sum(p.numel() for p in model.transformer.parameters())\n", "num_params_mlp_head = sum(p.numel() for p in model.mlp_head.parameters())\n", "num_params_conv_layers = sum(p.numel() for p in model.conv_layers.parameters())\n", "#print(f\"Number of parameters in pos_embed: {num_params_pos_embed/1e6:0.1f}M\")\n", "print(f\"Number of parameters in transformer encoder: {num_params_transformer/1e6:0.1f}M\")\n", "print(f\"Number of parameters in mlp head: {num_params_mlp_head/1e6:0.1f}M\")\n", "print(f\"Number of parameters in conv layers: {num_params_conv_layers/1e6:0.1f}M\")\n", "\n", "\n", "with torch.no_grad():\n", " x = torch.randn(1, 2, 48000*10)\n", " y = model(x)\n", " print(f\"Output shape: {y.shape}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import IPython\n", "import glob\n", "# test dataset\n", "# setup dataset\n", "audio_filepaths = glob.glob(\"/app/suno/data/audio_2ch_48khz_lg/val/genius_hq/*.wav\")\n", "audio_filepaths = audio_filepaths[:100]\n", "dataset = CorruptAudioDataset(audio_filepaths, presets, num_workers=4, sample_rate=48000, chunk_size_s=10.0, buffer_size=100)\n", "dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=True, num_workers=4, persistent_workers=True)\n", "\n", "input_maxes = []\n", "output_maxes = []\n", "# pass through the model\n", "for i in range(10):\n", " pbar = tqdm(dataloader, total=len(dataloader))\n", " for batch in pbar:\n", " audio, corrupted_audio, preset_idx = batch\n", " input_max = audio.abs().max()\n", " output_max = corrupted_audio.abs().max()\n", " input_maxes.append(input_max)\n", " output_maxes.append(output_max)\n", "\n", " print(f\"Input max: {torch.stack(input_maxes).mean():0.4f} +/- {torch.stack(input_maxes).std():0.4f}\")\n", " print(f\"Output max: {torch.stack(output_maxes).mean():0.4f} +/- {torch.stack(output_maxes).std():0.4f}\")\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 }