import os import json import torch import torchaudio import pyloudnorm as pyln from tqdm import tqdm from suno_utils.utils.s3 import read_from_s3 from suno_utils.utils.text import read_jsonl, write_jsonl from torchaudio.transforms import MelSpectrogram def compute_stereo_width_simple(waveform, sample_rate): """ Compute stereo width using a simple time-domain approach. Returns a value between 0 (mono) and 1 (maximum stereo spread). The calculation is based on comparing the difference signal (L-R) to the sum signal (L+R). A higher difference relative to the sum indicates more stereo content. """ # Ensure stereo if waveform.size(0) == 1: raise ValueError("Audio file must be stereo (2 channels)") elif waveform.size(0) > 2: waveform = waveform[:2, :] # Take first two channels if more exist # Get left and right channels left = waveform[0] right = waveform[1] # Compute difference and sum signals difference = left - right sum_signal = left + right # Compute RMS (Root Mean Square) energy of both signals diff_energy = torch.sqrt(torch.mean(difference**2)) sum_energy = torch.sqrt(torch.mean(sum_signal**2)) # Compute width as ratio of difference to total energy # Normalize to be between 0 and 1 width = (diff_energy / (sum_energy + 1e-8)).item() width = min(width, 1.0) # Clip to maximum of 1 return width import torch from torchaudio.transforms import MelSpectrogram def analyze_spectral_balance(waveform, sample_rate): """ Analyze the spectral balance of an audio file and determine if it's bassy, mid-focused, or bright, while also computing the spectral centroid. Returns: - character: 'bassy', 'mid_focused', or 'bright' - bass_ratio, mid_ratio, high_ratio: Energy ratios - spectral_centroid: The average spectral centroid of the waveform """ # Convert to mono if stereo if waveform.size(0) > 1: waveform = torch.mean(waveform, dim=0, keepdim=True) # Create mel spectrogram mel_spec = MelSpectrogram( sample_rate=sample_rate, n_fft=2048, hop_length=512, n_mels=128, f_min=20, f_max=20000, )(waveform) # Convert to dB scale mel_spec_db = torch.log10(mel_spec + 1e-9) # Calculate average energy in each frequency band bass_energy = torch.mean(mel_spec_db[:, :40]).item() # ~20-250 Hz mid_energy = torch.mean(mel_spec_db[:, 40:80]).item() # ~250-4000 Hz high_energy = torch.mean(mel_spec_db[:, 80:]).item() # ~4000-20000 Hz # Calculate relative ratios total_energy = bass_energy + mid_energy + high_energy bass_ratio = bass_energy / total_energy mid_ratio = mid_energy / total_energy high_ratio = high_energy / total_energy # Determine dominant characteristic if bass_ratio > max(mid_ratio, high_ratio): character = "bassy" elif mid_ratio > max(bass_ratio, high_ratio): character = "mid_focused" else: character = "bright" # Spectral Centroid computation # Frequency bins for the mel spectrogram mel_frequencies = torch.linspace(20, 20000, 128) spectral_centroid = torch.sum( mel_frequencies * torch.mean(mel_spec, dim=-1) ) / torch.sum(torch.mean(mel_spec, dim=-1)) spectral_centroid = spectral_centroid.item() return character, bass_ratio, mid_ratio, high_ratio, spectral_centroid def analyze_loudness_factor(waveform, sample_rate): """ Analyze loudness factor of an audio file and provide descriptive characteristics. Loudness factor is the LUFS measurement after peak normalization. """ # peak normalize normalized_audio = waveform / torch.max(torch.abs(waveform)).clamp(min=1e-9) # Measure LUFS meter = pyln.Meter(sample_rate) loudness_factor = meter.integrated_loudness(normalized_audio.permute(1, 0).numpy()) # Categorize and select descriptors if loudness_factor < -16: category = "dynamic" elif loudness_factor < -10: category = "moderate" else: category = "compressed" return category, loudness_factor def analyze_clipping(audio_tensor, threshold=0.99, sample_rate=44100): """ Vectorized analysis of audio clipping artifacts after peak normalization. Parameters: audio_tensor: torch.Tensor Audio samples (shape: [channels, samples] or [samples]) threshold: float Threshold for considering a sample clipped (0.0 to 1.0) sample_rate: int Audio sample rate in Hz, used to normalize clip percentage Returns: tuple (total_clipped_samples, clips_per_second) """ if not isinstance(audio_tensor, torch.Tensor): audio_tensor = torch.tensor(audio_tensor) if audio_tensor.dim() == 1: audio_tensor = audio_tensor.unsqueeze(0) # Peak normalize original_peak = audio_tensor.abs().max().clamp(min=1e-9) audio_tensor = audio_tensor / original_peak # Find clipped samples across all channels all_clips = audio_tensor.abs() >= threshold # Calculate total clipped samples (max across channels) total_clips = all_clips.sum(dim=-1).max().item() # Calculate clips per second duration_seconds = audio_tensor.shape[-1] / sample_rate clips_per_second = total_clips / duration_seconds return total_clips, clips_per_second def download_audio(s3_filepath: str, example_id: str, tmp_dir: str): filename = os.path.basename(s3_filepath) out_filepath = os.path.join(tmp_dir, f"{example_id}-{filename}") # only download the file if its not already downloaded if not os.path.isfile(out_filepath): os.system(f"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1") return out_filepath class AudioMetadataDataset(torch.utils.data.Dataset): def __init__( self, metas, existing_ids: list = None, tmp_dir: str = "/mnt/localdisk/tmp/cjs", ): # we need to filter metas to only use ones with s3_filepath # but, we have to be careful to maintain the index of the original metadata filtered_metas = [] for meta_idx, meta in enumerate(metas): if "audio_filepath" in meta: meta["orig_idx"] = meta_idx filtered_metas.append(meta) elif "s3_filepath" in meta: meta["orig_idx"] = meta_idx meta["audio_filepath"] = meta["s3_filepath"] filtered_metas.append(meta) num_filtered = len(filtered_metas) num_original = len(metas) percent_remaining = (num_filtered / num_original) * 100 print( f"{num_filtered}/{num_original} ({percent_remaining:0.2f}%) examples have s3_filepath." ) self.metas = filtered_metas self.tmp_dir = tmp_dir self.existing_ids = existing_ids os.makedirs(self.tmp_dir, exist_ok=True) def __len__(self): return len(self.metas) def __getitem__(self, idx): meta = self.metas[idx] if meta["id"] in self.existing_ids: return None filepath = download_audio(meta["audio_filepath"], meta["id"], self.tmp_dir) try: audio, sample_rate = torchaudio.load(filepath) # crop audio to max of 4min audio = audio[:, : int(4 * 60 * sample_rate)] # ensure audio is stereo if audio.size(0) == 1: audio = audio.repeat(2, 1) elif audio.size(0) > 2: audio = audio[:2, :] except Exception as e: print(f"Error: {e}") print(f"Failed to process: {filepath}") return None # delete audio file os.remove(filepath) # do the analysis here with torch.no_grad(): category, loudness_factor = analyze_loudness_factor(audio, sample_rate) spectral_character, bass_ratio, mid_ratio, high_ratio, spectral_centroid = ( analyze_spectral_balance(audio, sample_rate) ) stereo_width = compute_stereo_width_simple(audio, sample_rate) total_clips, clips_per_second = analyze_clipping( audio, sample_rate=sample_rate ) new_meta = { "id": meta["id"], "features": { "loudness_factor": f"{loudness_factor:0.3f}", "spectral_character": spectral_character, "spectral_centroid": f"{spectral_centroid:0.3f}", "bass_ratio": f"{bass_ratio:0.3f}", "mid_ratio": f"{mid_ratio:0.3f}", "high_ratio": f"{high_ratio:0.3f}", "stereo_width": f"{stereo_width:0.3f}", "total_clips": total_clips, "clips_per_second": f"{clips_per_second:0.3f}", }, } return new_meta def collate_fn(batch): return batch if __name__ == "__main__": batch_size = 200 num_workers = 200 # base_metas_filepath = ( # "/home/christian/code/christian/metadata/genius_hq_metas.jsonl" # ) # out_metas_filepath = ( # "/home/christian/code/christian/metadata/genius_hq_metas_audio_production.jsonl" # ) base_metas_filepath = ( "/home/christian/code/christian/metadata/youtube_music_metas.jsonl" ) out_metas_filepath = "/home/christian/code/christian/metadata/youtube_music_metas_audio_production.jsonl" # base_metas_filepath = ( # "/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl" # ) # out_metas_filepath = "/home/christian/code/christian/metadata/genius_hq_metas_filtered_audio_production.jsonl" if os.path.isfile(out_metas_filepath): print(f"Output file already exists: {out_metas_filepath}") # load the existing file existing_metas = read_jsonl(out_metas_filepath) print(f"Loaded {len(existing_metas)} existing metas") existing_ids = [meta["id"] for meta in existing_metas] else: existing_ids = [] base_metas = read_jsonl(base_metas_filepath) results = {} print(len(base_metas)) # create a meta dataset meta_dataset = AudioMetadataDataset(base_metas, existing_ids, tmp_dir="/tmp/cjs") meta_dataloader = torch.utils.data.DataLoader( meta_dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=collate_fn, ) for new_metas in tqdm(meta_dataloader): new_metas = [meta for meta in new_metas if meta is not None] write_jsonl(new_metas, out_metas_filepath, do_append=True)