import os import time import modal import torch import funcy import json import numpy as np import tempfile import torchaudio import pyloudnorm as pyln import logging from typing import Dict, Any, Optional, Tuple import librosa import soundfile as sf from suno_utils.utils.text import read_jsonl from suno_utils.utils.s3 import list_s3_dir, read_from_s3 from suno_utils.worker.settings import s3_client # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Modal setup aws_secret = modal.Secret.from_name("studio-aws") SECRETS = [ aws_secret, modal.Secret.from_dict( { "SUNO_ASSETS_PATH": "/suno/models/assets", "XDG_CACHE_HOME": "/suno/models/", } ), ] base_image = ( modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.10") .apt_install("curl", "ffmpeg", "sox", "unzip", "libsox-fmt-mp3", "zlib1g-dev", "git", "clang") .run_commands( [ 'curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"', "unzip -q awscliv2.zip", "./aws/install", ] ) .pip_install("torch==2.5.1", "torchaudio==2.5.1") .pip_install( "boto3", "psutil", "numpy", "librosa", "soundfile", "pyloudnorm", "funcy", "tqdm" ) .add_local_python_source("suno_utils", copy=False) ) # Audio analysis functions (same as your original code but cleaned up) def calculate_stereo_width(waveform): """Calculate stereo width from waveform. Returns None for mono.""" if waveform.shape[0] == 1: return None left, right = waveform[0], waveform[1] mid = (left + right) / 2 side = (left - right) / 2 mid_energy = torch.sqrt(torch.mean(mid**2)) side_energy = torch.sqrt(torch.mean(side**2)) width_ratio = (side_energy / (mid_energy + 1e-8)).item() return 2 * (1 / (1 + np.exp(-width_ratio)) - 0.5) def calculate_loudness_safe(waveform, sr, normalize=False): """Calculate LUFS loudness safely. Returns None if audio too short.""" try: meter = pyln.Meter(sr) if normalize: waveform = waveform / np.clip(np.max(np.abs(waveform)), 1e-10, None) return meter.integrated_loudness(waveform) except ValueError as e: if "greater than the block size" in str(e) or "length" in str(e).lower(): return None raise e def calculate_alternative_loudness(waveform): """Calculate RMS and peak loudness for short samples.""" waveform_np = waveform.numpy() if isinstance(waveform, torch.Tensor) else waveform rms = np.sqrt(np.mean(waveform_np**2)) peak = np.max(np.abs(waveform_np)) rms_db = 20 * np.log10(rms + 1e-10) peak_db = 20 * np.log10(peak + 1e-10) return rms_db, peak_db def calculate_possible_clipped_samples(waveform): """Count samples that might be clipped.""" waveform_np = waveform.numpy() if isinstance(waveform, torch.Tensor) else waveform clipped = np.sum(np.abs(waveform_np) >= 1.0) return clipped.item() if isinstance(clipped, torch.Tensor) else clipped def calculate_average_spectrum_db(waveform, n_fft=2048, hop_length=None): """Calculate average spectrum in dB. Handles very short audio gracefully.""" if hop_length is None: hop_length = n_fft // 4 if not isinstance(waveform, torch.Tensor): waveform = torch.from_numpy(waveform) audio_length = waveform.shape[-1] if audio_length < n_fft: n_fft = max(2 ** int(np.log2(audio_length)), 32) hop_length = n_fft // 4 if audio_length < n_fft: return None try: if waveform.dim() > 1: # Multi-channel: compute STFT for each channel stft_results = [] for channel in range(waveform.shape[0]): stft = torch.stft( waveform[channel], n_fft=n_fft, hop_length=hop_length, window=torch.hann_window(n_fft), return_complex=True ) stft_results.append(torch.abs(stft)) magnitude_spectrum = torch.stack([torch.mean(stft, dim=1) for stft in stft_results]) else: # Mono stft = torch.stft( waveform, n_fft=n_fft, hop_length=hop_length, window=torch.hann_window(n_fft), return_complex=True ) magnitude_spectrum = torch.mean(torch.abs(stft), dim=1) spectrum_db = 20 * torch.log10(magnitude_spectrum + 1e-10) return spectrum_db.numpy() except Exception as e: logger.debug(f"STFT failed for audio length {audio_length}, n_fft {n_fft}: {e}") return None def calculate_average_stereo_spectrum(waveform, sr): """Calculate mid/side spectrum for stereo audio.""" if waveform.shape[0] == 1: spectrum = calculate_average_spectrum_db(waveform) return (spectrum, np.zeros_like(spectrum)) if spectrum is not None else (None, None) left, right = waveform[0], waveform[1] mid = (left + right) / 2 side = (left - right) / 2 spectrum_mid = calculate_average_spectrum_db(mid) spectrum_side = calculate_average_spectrum_db(side) if spectrum_mid is None or spectrum_side is None: return None, None return spectrum_mid, spectrum_side def calculate_spectrum_evolution(waveform, sr, n_fft=2048, hop_length=None): """Calculate spectrum evolution (first 30s vs last 30s).""" if hop_length is None: hop_length = n_fft // 4 audio_length = waveform.shape[-1] duration_samples = 30 * sr if audio_length < n_fft: return None, None if audio_length < duration_samples: # Short audio: use entire clip for both spectrum = calculate_average_spectrum_db(waveform, n_fft, hop_length) return (spectrum.copy(), spectrum) if spectrum is not None else (None, None) # Long audio: compare first and last 30s spectrum_first = calculate_average_spectrum_db(waveform[..., :duration_samples], n_fft, hop_length) spectrum_last = calculate_average_spectrum_db(waveform[..., -duration_samples:], n_fft, hop_length) if spectrum_first is None or spectrum_last is None: return None, None return spectrum_first, spectrum_last def load_audio_robust(filepath: str) -> Tuple[Optional[torch.Tensor], Optional[int], str]: """Load audio file with multiple fallback methods.""" methods = [ ("torchaudio", lambda: torchaudio.load(filepath)), ("librosa", lambda: librosa.load(filepath, sr=None, mono=False)), ("soundfile", lambda: sf.read(filepath)) ] for method_name, load_func in methods: try: if method_name == "torchaudio": audio, sr = load_func() return audio, sr, method_name elif method_name == "librosa": audio, sr = load_func() if audio.ndim == 1: audio = torch.from_numpy(audio).unsqueeze(0) else: audio = torch.from_numpy(audio) if audio.shape[0] > audio.shape[1]: audio = audio.T return audio.float(), int(sr), method_name elif method_name == "soundfile": audio, sr = load_func() audio = torch.from_numpy(audio.T if audio.ndim > 1 else audio.reshape(1, -1)) return audio.float(), int(sr), method_name except Exception as e: logger.debug(f"Failed to load {filepath} with {method_name}: {e}") continue return None, None, "Failed to load with all methods: torchaudio, librosa, soundfile" def validate_audio(audio: torch.Tensor, sr: int) -> Tuple[bool, str]: """Validate loaded audio for basic sanity checks.""" try: if audio.numel() == 0: return False, "Audio file is empty" if sr < 1000 or sr > 192000: return False, f"Unusual sample rate: {sr} Hz" duration = audio.shape[-1] / sr if duration < 0.001: return False, f"Audio too short: {duration:.3f} seconds" if duration > 3600: return False, f"Audio too long: {duration:.1f} seconds" if torch.isnan(audio).any(): return False, "Audio contains NaN values" if torch.isinf(audio).any(): return False, "Audio contains infinite values" if torch.max(torch.abs(audio)) < 1e-10: return False, "Audio is completely silent" return True, "Valid" except Exception as e: return False, f"Validation error: {e}" def convert_multichannel_audio(audio: torch.Tensor) -> torch.Tensor: """Convert multi-channel audio (>5 channels) to stereo for LUFS compatibility.""" num_channels = audio.shape[0] if num_channels > 5: logger.debug(f"Converting {num_channels}-channel audio to stereo") return audio[:2] return audio def create_failed_result(row_id: str, error_message: str, load_method: str = None) -> Dict[str, Any]: """Create a standardized failed result dictionary.""" return { "id": row_id, "lufs_db": None, "lufs_db_factor": None, "rms_loudness_db": None, "peak_loudness_db": None, "stereo_width": None, "clipped_samples": None, "average_spectrum_db": None, "average_spectrum_db_first": None, "average_spectrum_db_last": None, "average_stereo_spectrum_mid": None, "average_stereo_spectrum_side": None, "processing_success": False, "error_message": error_message, "load_method": load_method, "duration_seconds": None } def to_list_if_array(value): """Convert numpy array to list, otherwise return as-is.""" return value.tolist() if isinstance(value, np.ndarray) else value def make_json_serializable(obj): """Convert numpy types to JSON serializable types.""" if isinstance(obj, np.integer): return int(obj) elif isinstance(obj, np.floating): return float(obj) elif isinstance(obj, np.ndarray): return obj.tolist() elif isinstance(obj, dict): return {key: make_json_serializable(value) for key, value in obj.items()} elif isinstance(obj, list): return [make_json_serializable(item) for item in obj] else: return obj def analyze_single_audio(row_id: str, s3_filepath: str) -> Dict[str, Any]: """Analyze a single audio file and return statistics.""" try: # Load audio from S3 audio, sr, load_method = read_from_s3(s3_filepath, read_f=load_audio_robust) if audio is None: return create_failed_result(row_id, f"Failed to load audio: {load_method}") # Validate audio is_valid, validation_message = validate_audio(audio, sr) if not is_valid: return create_failed_result(row_id, f"Invalid audio: {validation_message}", load_method) duration_seconds = audio.shape[-1] / sr # Handle multi-channel audio for LUFS compatibility audio_for_lufs = convert_multichannel_audio(audio) # Resample to 48kHz if needed if sr != 48000: audio = torchaudio.functional.resample(audio, sr, 48000) audio_for_lufs = torchaudio.functional.resample(audio_for_lufs, sr, 48000) sr = 48000 # Convert to numpy for loudness calculations (channels, samples) -> (samples, channels) audio_np = audio_for_lufs.permute(1, 0).numpy() # Calculate loudness metrics lufs_db = calculate_loudness_safe(audio_np, sr) lufs_db_factor = calculate_loudness_safe(audio_np, sr, normalize=True) if lufs_db is None: rms_loudness, peak_loudness = calculate_alternative_loudness(audio_np) else: rms_loudness = peak_loudness = None # Calculate other metrics stereo_width = calculate_stereo_width(audio) clipped_samples = calculate_possible_clipped_samples(audio) average_spectrum_db = calculate_average_spectrum_db(audio) # Calculate stereo spectrum average_stereo_spectrum_mid, average_stereo_spectrum_side = calculate_average_stereo_spectrum(audio, sr) # Calculate spectrum evolution spectrum_first, spectrum_last = calculate_spectrum_evolution(audio, sr) return { "id": row_id, "lufs_db": float(lufs_db) if lufs_db is not None else None, "lufs_db_factor": float(lufs_db_factor) if lufs_db_factor is not None else None, "rms_loudness_db": float(rms_loudness) if rms_loudness is not None else None, "peak_loudness_db": float(peak_loudness) if peak_loudness is not None else None, "stereo_width": float(stereo_width) if stereo_width is not None else None, "clipped_samples": int(clipped_samples) if clipped_samples is not None else None, "average_spectrum_db": to_list_if_array(average_spectrum_db), "average_spectrum_db_first": to_list_if_array(spectrum_first), "average_spectrum_db_last": to_list_if_array(spectrum_last), "average_stereo_spectrum_mid": to_list_if_array(average_stereo_spectrum_mid), "average_stereo_spectrum_side": to_list_if_array(average_stereo_spectrum_side), "processing_success": True, "error_message": None, "load_method": load_method, "duration_seconds": float(duration_seconds) } except Exception as e: logger.error(f"Error processing {s3_filepath} (ID: {row_id}): {e}") return create_failed_result(row_id, str(e)) class AudioAnalysisWorker: def __init__(self, output_path: str): self.output_path = output_path logger.info("AudioAnalysisWorker initialized") def analyze_batch(self, work_items: list[dict]) -> list[dict]: """Analyze a batch of audio files and return results.""" results = [] batch_size = len(work_items) logger.info(f"Processing batch of {batch_size} files") for i, item in enumerate(work_items): item_id = str(item["id"]) s3_filepath = item["s3_filepath"] # Log progress every 10 files or for small batches #if i % 10 == 0 or batch_size <= 20: # logger.info(f"Processing {i+1}/{batch_size}: {item_id}") # Analyze the audio file result = analyze_single_audio(item_id, s3_filepath) results.append(result) # Save individual result to S3 immediately try: with tempfile.TemporaryDirectory() as td: result_path = os.path.join(td, f"{item_id}_analysis.json") # Make sure result is JSON serializable json_safe_result = make_json_serializable(result) with open(result_path, 'w') as f: json.dump(json_safe_result, f) s3_result_path = os.path.join( self.output_path, f"{item_id}_analysis.json" ) s3_client.upload_file( result_path, "suno-data", s3_result_path, ExtraArgs={"ContentType": "application/json"} ) except Exception as e: logger.error(f"Failed to save result for {item_id}: {e}") logger.error(f"Result data types: {[(k, type(v)) for k, v in result.items()]}") # Log batch statistics success_count = sum(1 for r in results if r['processing_success']) failure_count = batch_size - success_count success_rate = success_count / batch_size * 100 logger.info(f"Batch complete: {success_count}/{batch_size} successful ({success_rate:.1f}%)") if failure_count > 0: failed_ids = [r['id'] for r in results if not r['processing_success']] logger.warning(f"Failed files: {failed_ids[:5]}{'...' if len(failed_ids) > 5 else ''}") return results # Modal app setup app = modal.App("audio-analysis-worker", image=base_image, secrets=SECRETS) BATCH_SIZE = 500 # Process 50 files per worker N_MAX_REPLICAS = 100 # Maximum number of concurrent workers @app.cls( cpu=4, memory=8000, # 8GB memory secrets=SECRETS, timeout=30 * 60, # 30 minute timeout per batch container_idle_timeout=60, # 1 minute idle timeout concurrency_limit=N_MAX_REPLICAS, ) class AudioAnalysisStub: def __init__(self, output_path: str): self.worker = AudioAnalysisWorker(output_path) @modal.method() def analyze_batch(self, work_items: list[dict]) -> list[dict]: return self.worker.analyze_batch(work_items) @app.local_entrypoint() def main( input_jsonl_path: str, output_path: str, max_files: Optional[int] = None, resume: bool = True ): """ Main entry point for audio analysis. Args: input_jsonl_path: Path to JSONL file containing audio metadata output_path: S3 path for storing results (e.g., "christian/audio_analysis/batch_001") max_files: Optional limit on number of files to process resume: Whether to skip files that already have results """ print(f"Loading metadata from: {input_jsonl_path}") # Load work items if input_jsonl_path.startswith("s3://"): work_items = read_from_s3(input_jsonl_path, read_f=read_jsonl) else: work_items = read_jsonl(input_jsonl_path) print(f"Loaded {len(work_items)} total work items") # Apply max_files limit if specified if max_files is not None: work_items = work_items[:max_files] print(f"Limited to {len(work_items)} files") # Resume functionality: check which files already have results if resume: print("Checking for existing results...") try: existing_files = list_s3_dir(f"s3://suno-data/{output_path}/") existing_files = [f[0] for f in existing_files if f[0].endswith("_analysis.json")] existing_ids = set() for filepath in existing_files: # Extract ID from filename: {id}_analysis.json filename = os.path.basename(filepath) if filename.endswith("_analysis.json"): file_id = filename[:-len("_analysis.json")] existing_ids.add(file_id) original_count = len(work_items) work_items = [item for item in work_items if str(item["id"]) not in existing_ids] skipped_count = original_count - len(work_items) print(f"Found {len(existing_ids)} existing results") print(f"Skipped {skipped_count} files, {len(work_items)} remaining to process") except Exception as e: print(f"Warning: Could not check existing results: {e}") print("Proceeding without resume...") if len(work_items) == 0: print("No files to process!") return # Create batches for parallel processing work_batches = list(funcy.chunks(BATCH_SIZE, work_items)) print(f"Created {len(work_batches)} batches (batch size: {BATCH_SIZE})") # Initialize worker worker = AudioAnalysisStub(output_path) print("Starting batch processing...") start_time = time.time() # Process all batches batch_results = list(worker.analyze_batch.map(work_batches)) # Flatten results all_results = [] for batch_result in batch_results: all_results.extend(batch_result) processing_time = time.time() - start_time # Save consolidated results print("Saving consolidated results...") with tempfile.TemporaryDirectory() as td: consolidated_path = os.path.join(td, "consolidated_results.jsonl") with open(consolidated_path, 'w') as f: for result in all_results: json_safe_result = make_json_serializable(result) f.write(json.dumps(json_safe_result) + '\n') s3_consolidated_path = os.path.join(output_path, "consolidated_results.jsonl") s3_client.upload_file( consolidated_path, "suno-data", s3_consolidated_path, ExtraArgs={"ContentType": "application/jsonl"} ) # Final statistics total_files = len(all_results) successful_files = sum(1 for r in all_results if r['processing_success']) failed_files = total_files - successful_files success_rate = successful_files / total_files * 100 if total_files > 0 else 0 print(f"\n=== ANALYSIS COMPLETE ===") print(f"Total files processed: {total_files}") print(f"Successful: {successful_files}") print(f"Failed: {failed_files}") print(f"Success rate: {success_rate:.2f}%") print(f"Processing time: {processing_time/60:.1f} minutes") print(f"Results saved to: s3://suno-data/{output_path}/") if failed_files > 0: failed_results = [r for r in all_results if not r['processing_success']] error_counts = {} for result in failed_results: error = result['error_message'] error_counts[error] = error_counts.get(error, 0) + 1 print(f"\nMost common errors:") for error, count in sorted(error_counts.items(), key=lambda x: x[1], reverse=True)[:5]: print(f" {count}: {error}")