{ "cells": [ { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "import os\n", "import pandas as pd\n" ] }, { "cell_type": "code", "execution_count": 104, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torchaudio\n", "\n", "def apply_audio_codec_advanced(\n", " audio: torch.Tensor,\n", " sample_rate: float,\n", " bit_rate: int = 16000,\n", " n_passes: int = 1,\n", " codec_type: str = \"mp3\",\n", "):\n", " for _ in range(n_passes):\n", " if codec_type == \"mp3\":\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"mp3\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " elif codec_type == \"ogg-vorbis\":\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"ogg\",\n", " encoder=\"vorbis\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " elif codec_type == \"opus\":\n", " effector = torchaudio.io.AudioEffector(\n", " format=\"ogg\",\n", " encoder=\"opus\",\n", " codec_config=torchaudio.io.CodecConfig(bit_rate=bit_rate),\n", " )\n", " else:\n", " raise ValueError(f\"Invalid codec type: {codec_type}\")\n", " audio = effector.apply(audio.T, sample_rate).T\n", " return audio" ] }, { "cell_type": "code", "execution_count": 105, "metadata": {}, "outputs": [], "source": [ "x, sr = torchaudio.load(\"/home/christian/audio/reference-audio-wav/02 Dreams.wav\")\n", "x = x[:480000]\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "x_codec = apply_audio_codec(x, sr, bit_rate=16000, n_passes=1, codec_type=\"aac\")\n" ] }, { "cell_type": "code", "execution_count": 98, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import numpy as np\n", "import subprocess\n", "import tempfile\n", "import os\n", "import random\n", "from typing import Tuple\n", "\n", "def apply_random_codec(audio_tensor: torch.Tensor, sample_rate: int) -> Tuple[torch.Tensor, dict]:\n", " \"\"\"\n", " Apply a random codec with random settings to an input audio tensor using FFmpeg.\n", " \n", " Args:\n", " audio_tensor (torch.Tensor): Input audio tensor of shape (channels, samples)\n", " sample_rate (int): Sample rate of the audio in Hz\n", " \n", " Returns:\n", " Tuple[torch.Tensor, dict]: Processed audio tensor and dictionary of applied settings\n", " \"\"\"\n", " # Codec settings with proper file extensions and bitrate ranges\n", " CODEC_SETTINGS = {\n", " 'libmp3lame': {\n", " 'ext': 'mp3',\n", " 'bitrate_range': (8, 192), # kbps\n", " },\n", " 'aac': {\n", " 'ext': 'm4a',\n", " 'bitrate_range': (8, 192), # kbps\n", " },\n", " 'libvorbis': {\n", " 'ext': 'ogg',\n", " 'bitrate_range': (45, 192), # kbps\n", " },\n", " 'libopus': {\n", " 'ext': 'opus',\n", " 'bitrate_range': (32, 192), # kbps\n", " 'application': {'voip': 2048, 'audio': 2049, 'lowdelay': 2051},\n", " 'vbr_modes': {'off': 0, 'on': 1, 'constrained': 2},\n", " 'frame_duration': [2.5, 5, 10, 20, 40, 60] # ms\n", " },\n", " 'ac3': {\n", " 'ext': 'ac3',\n", " 'bitrate_range': (32, 192), # kbps\n", " },\n", " 'wmav2': {\n", " 'ext': 'wma',\n", " 'bitrate_range': (8, 192), # kbps\n", " },\n", " 'libopencore_amrnb': {\n", " 'ext': 'amr',\n", " 'bitrate_modes': [4.75, 5.15, 5.9, 6.7, 7.4, 7.95, 10.2, 12.2], # kbps\n", " 'sample_rate': 8000, # Fixed sample rate\n", " 'requires_resample': True\n", " }\n", " }\n", " \n", " # Select random codec\n", " codec = random.choice(list(CODEC_SETTINGS.keys()))\n", " settings = CODEC_SETTINGS[codec]\n", " \n", " # Generate random settings\n", " codec_params = {}\n", " \n", " # Handle bitrate selection\n", " if 'bitrate_modes' in settings:\n", " codec_params['bitrate'] = random.choice(settings['bitrate_modes'])\n", " elif 'bitrate_range' in settings:\n", " codec_params['bitrate'] = random.randint(*settings['bitrate_range'])\n", " else:\n", " raise ValueError(f\"Codec {codec} has no bitrate settings defined\")\n", " \n", " # Add codec-specific parameters\n", " if codec == 'libopus':\n", " # Select random options\n", " app_type = random.choice(list(settings['application'].keys()))\n", " vbr_mode = random.choice(list(settings['vbr_modes'].keys()))\n", " frame_dur = random.choice(settings['frame_duration'])\n", " \n", " codec_params.update({\n", " 'application': app_type,\n", " 'application_value': settings['application'][app_type],\n", " 'vbr': vbr_mode,\n", " 'vbr_value': settings['vbr_modes'][vbr_mode],\n", " 'frame_duration': frame_dur\n", " })\n", " if codec == 'libopus':\n", " codec_params['application'] = random.choice(settings['application'])\n", " \n", " # Validate number of channels\n", " num_channels = audio_tensor.shape[0]\n", " if num_channels not in [1, 2]:\n", " raise ValueError(f\"Expected mono or stereo audio (1 or 2 channels), got {num_channels} channels\")\n", " \n", " # Some codecs only support mono\n", " if codec == 'libopencore_amrnb' and num_channels > 1:\n", " print(\"Warning: Converting to mono for AMR codec\")\n", " audio_tensor = audio_tensor.mean(dim=0, keepdim=True)\n", " num_channels = 1\n", " \n", " # Get output sample rate\n", " output_sample_rate = settings.get('sample_rate', sample_rate)\n", " \n", " # Create temporary files with proper extensions\n", " with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_in, \\\n", " tempfile.NamedTemporaryFile(suffix=f'.{settings[\"ext\"]}', delete=False) as temp_encoded, \\\n", " tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_out:\n", " \n", " try:\n", " # Convert tensor to numpy and save as WAV\n", " audio_numpy = audio_tensor.numpy()\n", " import soundfile as sf\n", " sf.write(temp_in.name, audio_numpy.T, sample_rate, format='WAV', subtype='FLOAT')\n", " \n", " # Build FFmpeg command for encoding\n", " encode_cmd = [\n", " 'ffmpeg',\n", " '-v', 'warning',\n", " '-i', temp_in.name\n", " ]\n", " \n", " # Add resampling if needed\n", " if settings.get('requires_resample', False):\n", " encode_cmd.extend([\n", " '-ar', str(settings['sample_rate']),\n", " '-af', f'aresample=osr={settings[\"sample_rate\"]}'\n", " ])\n", " else:\n", " encode_cmd.extend(['-ar', str(output_sample_rate)])\n", " \n", " encode_cmd.extend([\n", " '-c:a', codec,\n", " '-ac', str(num_channels)\n", " ])\n", " \n", " # Handle bitrate parameter format based on codec\n", " if codec == 'libopencore_amrnb':\n", " encode_cmd.extend(['-ab', f'{codec_params[\"bitrate\"]}']) # AMR uses plain number without k\n", " else:\n", " encode_cmd.extend(['-b:a', f'{codec_params[\"bitrate\"]}k'\n", " ])\n", " \n", " # Add codec-specific parameters\n", " if codec == 'libopus':\n", " encode_cmd.extend([\n", " '-application', str(codec_params['application_value']),\n", " '-vbr', str(codec_params['vbr_value']),\n", " '-frame_duration', str(codec_params['frame_duration'])\n", " ])\n", " elif codec == 'libvorbis' and codec_params['bitrate'] < 64:\n", " encode_cmd.extend(['-q:a', '0'])\n", " \n", " encode_cmd.extend([\n", " '-y',\n", " temp_encoded.name\n", " ])\n", " \n", " # Encode\n", " result = subprocess.run(\n", " encode_cmd,\n", " check=True,\n", " capture_output=True,\n", " text=True\n", " )\n", " \n", " # Decode back to WAV\n", " decode_cmd = [\n", " 'ffmpeg',\n", " '-v', 'warning',\n", " '-i', temp_encoded.name,\n", " '-c:a', 'pcm_f32le',\n", " '-ar', str(sample_rate), # Convert back to original sample rate\n", " '-ac', str(num_channels),\n", " '-y',\n", " temp_out.name\n", " ]\n", " \n", " # Decode\n", " result = subprocess.run(\n", " decode_cmd,\n", " check=True,\n", " capture_output=True,\n", " text=True\n", " )\n", " \n", " # Read processed audio back into tensor\n", " processed_audio, _ = sf.read(temp_out.name)\n", " processed_tensor = torch.from_numpy(processed_audio.T)\n", " \n", " # Validate output tensor shape matches input\n", " if processed_tensor.shape[0] != num_channels:\n", " raise RuntimeError(f\"Channel count mismatch: expected {num_channels}, got {processed_tensor.shape[0]}\")\n", " \n", " except subprocess.CalledProcessError as e:\n", " print(f\"FFmpeg error:\\nCommand: {' '.join(e.cmd)}\\nOutput: {e.stdout}\\nError: {e.stderr}\")\n", " raise\n", " except Exception as e:\n", " print(f\"Error processing audio: {str(e)}\")\n", " raise\n", " finally:\n", " # Clean up temporary files\n", " for temp_file in [temp_in.name, temp_encoded.name, temp_out.name]:\n", " try:\n", " if os.path.exists(temp_file):\n", " os.unlink(temp_file)\n", " except OSError as e:\n", " print(f\"Warning: Could not delete temporary file {temp_file}: {e}\")\n", " \n", " # Return processed audio and settings used\n", " settings_used = {\n", " 'codec': codec,\n", " 'parameters': codec_params,\n", " 'extension': settings['ext'],\n", " 'sample_rate': output_sample_rate,\n", " 'channels': num_channels\n", " }\n", " \n", " return processed_tensor, settings_used" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "\n", "x_crop = x[:,:480000]\n", "x_codec, settings = apply_random_codec(x_crop, sr)\n", "print(settings)\n", "\n", "# crop output to same length as inpu\n", "crop_size = x_codec.shape[1] - x_crop.shape[1]\n", "x_codec = x_codec[:,crop_size:].float()\n", "\n", "\n", "import IPython\n", "IPython.display.display(IPython.display.Audio(x_codec.numpy(), rate=sr))\n", "\n", "import auraloss\n", "si_sdr_loss = auraloss.time.SISDRLoss()\n", "mel_loss = auraloss.freq.MelSTFTLoss(sample_rate=44100)\n", "loss = mel_loss(x_codec.unsqueeze(0), x_crop.unsqueeze(0))\n", "print(loss)\n", "\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "pkl_filepath = \"/home/tony/Data/Preference/up_v1/interesting_clips_up_u_1_20241201_full.pkl\"\n", "\n", "df = pd.read_pickle(pkl_filepath)\n", "print(len(df))\n", "df.head()\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# let's download all the mp3 files with positive preference\n", "# to get all the rows with positive preference \n", "positive_df = df[df[\"preference\"] == True]\n", "print(len(positive_df))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# iterate and get list of s3 paths, then we will download with multiprocessing\n", "s3_ids = positive_df[\"id_x\"].tolist()\n", "print(len(s3_ids))\n", "\n", "\n", "\n", "\n", "\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# first download audios in parallel to local from s3\n", "base_dir = \"/app/suno/christian/data\"\n", "out_dir = \"interesting_clips_up_u_1_20241201_positive\"\n", "os.makedirs(os.path.join(base_dir, out_dir), exist_ok=True)\n", "\n", "def download_audio(song_id):\n", " mp3_filepath = f\"s3://suno-data-uploads/studio/uploads/{song_id}.mp3\"\n", " out_filepath = os.path.join(base_dir, out_dir, f\"{song_id}.mp3\")\n", " # surpress output\n", " if not os.path.exists(out_filepath):\n", " os.system(f\"aws s3 cp {mp3_filepath} {out_filepath} > /dev/null 2>&1\")\n", "\n", "# use joblib for parallel downloads with progress bar\n", "from joblib import Parallel, delayed\n", "from tqdm import tqdm\n", "\n", "# use 32 processes for parallel downloads with progress tracking\n", "Parallel(n_jobs=96, backend=\"loky\")(\n", " delayed(download_audio)(song_id) \n", " for song_id in tqdm(s3_ids, desc=\"Downloading audio files\")\n", ")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import glob\n", "local_filepaths = glob.glob(os.path.join(base_dir, out_dir, \"*.mp3\"))\n", "print(len(local_filepaths))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# now we need to add these audio files to the manifest, we can just append to the existing manifest that we have\n", "manifest_filepath = \"/app/suno/data/audio_2ch_48khz_lg/ear_train_filtered_v2.csv\"\n", "\n", "with open(manifest_filepath, \"r\") as f:\n", " lines = f.readlines()\n", "\n", "print(len(lines))\n", "\n", "# add the new audio files to the manifest\n", "for filepath in tqdm(local_filepaths, desc=\"Adding audio files to manifest\"):\n", " lines.append(f\"{filepath}\\n\")\n", "\n", "# write the updated manifest\n", "out_manifest_filepath = \"/app/suno/data/audio_2ch_48khz_lg/ear_train_filtered_v2_with_gens.csv\"\n", "with open(out_manifest_filepath, \"w\") as f:\n", " f.writelines(lines)\n", "\n", "print(f\"Updated manifest saved to {out_manifest_filepath}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "print(len(lines))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "lines[600_000]" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# check that all files in the manifest exist\n", "for line in tqdm(lines[500_000:]):\n", " if not os.path.exists(line.strip()):\n", " print(f\"File does not exist: {line.strip()}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import random" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "vae_embeds = torch.randn(128, 750)\n", "is_training = True\n", "infill_ctx_len = 750\n", "scale_vae_ctx = True\n", "vae_scale_factor = 1.0\n", "\n", "infill_ctx_vae_embeds = torch.zeros_like(vae_embeds)\n", "infill_ctx_mask = torch.zeros(vae_embeds.shape[1]).bool()\n", "\n", "if is_training and 0.05 <= 0.1:\n", " # 10% chance to infill\n", " infill_ctx_vae_embeds = vae_embeds.clone()\n", " infill_ctx_mask = torch.ones(vae_embeds.shape[1]).bool()\n", "\n", " # Randomly choose one of three cases\n", " case = random.random()\n", " print(case)\n", "\n", " # Determine random mask size up to infill_ctx_len\n", " # don't mask more than 90% of the infill_ctx_len\n", " mask_size = random.randint(1, int(infill_ctx_len * 0.9))\n", " print(mask_size)\n", "\n", " if case < 0.2: # Left side masking\n", " start_idx = 0\n", " elif case > 0.2 and case < 0.8: # Center masking\n", " # Ensure we have enough space on both sides\n", " available_start_positions = vae_embeds.shape[1] - mask_size\n", " if available_start_positions > 0:\n", " start_idx = random.randint(1, available_start_positions - 1)\n", " else:\n", " start_idx = 0\n", " else: # Right side masking\n", " start_idx = vae_embeds.shape[1] - mask_size\n", "\n", " # Apply the masking\n", " infill_ctx_vae_embeds[..., start_idx : start_idx + mask_size] = 0\n", " infill_ctx_mask[start_idx : start_idx + mask_size] = False\n", "\n", " if scale_vae_ctx: # scale the context vector before input to model\n", " infill_ctx_vae_embeds = infill_ctx_vae_embeds * vae_scale_factor\n", "\n", "#print(infill_ctx_vae_embeds)\n", "#print(infill_ctx_mask)\n", "\n", "# make a plot of the mask\n", "import matplotlib.pyplot as plt\n", "plt.plot(infill_ctx_mask.numpy())\n", "plt.show()\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "suno_gpt45", "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.14" } }, "nbformat": 4, "nbformat_minor": 2 }