import os import sys import time import modal import torch import pathlib import funcy import json import pandas as pd import numpy as np import tempfile import torchaudio import pyloudnorm as pyln from suno_utils.audio import Audio from suno_utils.gpt.generation import GenerationConfig from suno_utils.gpt.engine import Engine from suno_utils.gpt.generation_engine import make_request from suno_utils.worker.settings import s3_client # from suno_utils.worker.modal_base import MODAL_MOUNTS from suno_utils.gpt import chirp_v2_5 as chirp_v2 from suno_utils.utils.text import read_jsonl from suno_utils.utils.s3 import list_s3_dir, read_from_s3 from suno_utils.tasks.hoot import ( encode_filepaths, encode, clean_text, encode_and_align, _get_alignable_tokens, _assign_word_timings, get_aligned_lyrics, ctc_align, load_model_list, get_word_timing_from_audio_and_lyrics, _merge_into_lines, decode_logits, load_model, ) from suno_utils.tasks.shimmerscore import shimmerscore from suno_utils.utils.metrics import get_cer from suno_utils.tasks.ear import load_model from suno_utils.diffusion.generation import ( preload_dit_model, preload_tokenizer, TOKENIZER_FILEPATH, SEMANTIC_MODEL_FILEPATH, SEMANTIC_CLUSTERS_FILEPATH, _retrieve_models, # preload_ear_model, ) from suno_utils.tasks.mert_25 import ( preload_models as preload_semantic_models, encode as encode_semantic, ) from suno_utils.diffusion import generation as diffusion_gen from suno_utils.tasks.upsample_engine import UpsampleEngine, Request, Job # from suno_utils.tasks.upsample_engine_old import UpsampleEngine, Request, Job MOUNT_PATH = "/suno/models" 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/", } ), modal.Secret.from_name("api-callback-token"), modal.Secret.from_name("datadog-metrics"), ] 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", ] ) .dockerfile_commands( [ "COPY --from=datadog/serverless-init:1.2.1 /datadog-init /app/datadog-init", 'ENTRYPOINT ["/app/datadog-init"]', ] ) .pip_install("torch==2.5.1", "torchaudio==2.5.1") .pip_install( "flashinfer-python", index_url="https://flashinfer.ai/whl/cu124/torch2.5/" ) .pip_install_private_repos( "github.com/suno-ai/glockenspiel.git@a5dba4e50#subdirectory=descript-audio-codec&egg=descript-audio-codec", git_user="mcamac", secrets=[modal.Secret.from_name("victor-modal-github-token")], ) .pip_install_private_repos( "github.com/suno-ai/neon.git@1c83548#subdirectory=hoot", git_user="mcamac", secrets=[modal.Secret.from_name("victor-modal-github-token")], ) .pip_install( "boto3", "transformers", "tokenizers", "encodec", "ctc_segmentation", "psutil", "redis", "pydantic", "nnAudio", "rpyc", "biopython>=1.81", # TODO: don't love this depdendency, for hoot "pynvml", # for torch cuda utilization "torchsde", "ninja", "wheel", ) .pip_install_from_pyproject( "/home/christian/code/glockenspiel/suno_utils/pyproject.toml", ) .run_commands( # This is really slow "git clone https://github.com/Dao-AILab/flash-attention.git", "cd flash-attention/hopper && python setup.py install", gpu="h100", ) .apt_install( "libogg0", "libopus0", "opus-tools", ) .pip_install("transformers==4.44.0", "wandb") ) def calculate_stereo_width(waveform): # Split into left and right channels left = waveform[0] right = waveform[1] # Compute mid/side representation mid = (left + right) / 2 side = (left - right) / 2 # Compute RMS energy of mid and side channels mid_energy = torch.sqrt(torch.mean(mid**2)) side_energy = torch.sqrt(torch.mean(side**2)) # Compute stereo width based on mid/side ratio # Normalize to range 0-1 using sigmoid-like function width_ratio = (side_energy / (mid_energy + 1e-8)).item() stereo_width = 2 * (1 / (1 + np.exp(-width_ratio)) - 0.5) return stereo_width def calculate_loudness(waveform, sr): meter = pyln.Meter(sr) lufs_db = meter.integrated_loudness(waveform) return lufs_db def calculate_loudness_factor(waveform, sr): meter = pyln.Meter(sr) normalized_waveform = waveform / np.clip(np.max(np.abs(waveform)), 1e-10, None) lufs_db = meter.integrated_loudness(normalized_waveform) return lufs_db def calculate_possible_clipped_samples(waveform): # Convert to numpy if it's a torch tensor if isinstance(waveform, torch.Tensor): waveform_np = waveform.numpy() else: waveform_np = waveform # Count samples that are at or above the clipping threshold return ( np.sum(np.abs(waveform_np) >= 1.0).item() if isinstance(np.sum(np.abs(waveform_np) >= 1.0), torch.Tensor) else np.sum(np.abs(waveform_np) >= 1.0) ) def calculate_average_spectrum_db(waveform, n_fft=16384, hop_length=8192): # Keep as torch tensor or convert to torch tensor if it's numpy if not isinstance(waveform, torch.Tensor): waveform = torch.from_numpy(waveform) # Calculate the average spectrum using STFT for efficiency n_fft = 2048 # Choose an appropriate FFT size hop_length = n_fft // 4 # Standard hop length # Compute STFT using torch if waveform.dim() > 1: # For stereo, 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, ) # Get magnitude stft_magnitude = torch.abs(stft) stft_results.append(stft_magnitude) # Average across time frames for each channel magnitude_spectrum = torch.stack( [torch.mean(stft, dim=1) for stft in stft_results] ) else: # For mono stft = torch.stft( waveform, n_fft=n_fft, hop_length=hop_length, window=torch.hann_window(n_fft), return_complex=True, ) # Get magnitude stft_magnitude = torch.abs(stft) magnitude_spectrum = torch.mean(stft_magnitude, dim=1) # Convert to dB scale spectrum_db = 20 * torch.log10( magnitude_spectrum + 1e-10 ) # Adding small value to avoid log(0) # Convert to numpy for consistency with the rest of the code return spectrum_db.numpy() def calculate_average_stereo_spectrum(waveform, sr): assert waveform.dim() == 2 and waveform.shape[0] == 2 # split into left and right channels left = waveform[0] right = waveform[1] # compute mid and side channels mid = (left + right) / 2 side = (left - right) / 2 # calculate spectrum for mid and side channels spectrum_mid = calculate_average_spectrum_db(mid, sr) spectrum_side = calculate_average_spectrum_db(side, sr) return spectrum_mid, spectrum_side def calculate_spectrum_evolution(waveform, sr, n_fft=16384, hop_length=8192): # compute spectrum for first 30s waveform_first = waveform[:, : 30 * sr] spectrum_first = calculate_average_spectrum_db(waveform_first, n_fft, hop_length) # compute spectrum for last 30s waveform_last = waveform[:, -30 * sr :] spectrum_last = calculate_average_spectrum_db(waveform_last, n_fft, hop_length) return spectrum_first, spectrum_last def analyze_audio(filepath): audio, sr = torchaudio.load(filepath) if sr != 48000: audio = torchaudio.functional.resample(audio, sr, 48000) # calculate loudness lufs_db = calculate_loudness(audio.permute(1, 0).numpy(), sr) lufs_db_factor = calculate_loudness_factor(audio.permute(1, 0).numpy(), sr) # calculate stereo width stereo_width = calculate_stereo_width(audio) # clipped samples clipped_samples = calculate_possible_clipped_samples(audio) # average spectrum average_spectrum_db = calculate_average_spectrum_db(audio) # average stereo spectrum average_stereo_spectrum_mid, average_stereo_spectrum_side = ( calculate_average_stereo_spectrum(audio, sr) ) # spectrum evolution spectrum_first, spectrum_last = calculate_spectrum_evolution(audio, sr) return { "lufs_db": lufs_db, "lufs_db_factor": lufs_db_factor, "stereo_width": stereo_width, "clipped_samples": clipped_samples, "average_spectrum_db": average_spectrum_db, "average_spectrum_db_first": spectrum_first, "average_spectrum_db_last": spectrum_last, "average_stereo_spectrum_mid": average_stereo_spectrum_mid, "average_stereo_spectrum_side": average_stereo_spectrum_side, } def third_octave_bands(sr, fmin=20.0, fmax=None): """ Compute 1/3-octave band center frequencies and edges. """ if fmax is None: fmax = sr / 2.0 k = np.arange(-30, 30) # wide enough range f_center = 1000.0 * (2.0 ** (k / 3.0)) # ISO 1/3 octave centers f_center = f_center[(f_center >= fmin) & (f_center <= fmax)] f_lower = f_center / (2 ** (1 / 6)) f_upper = f_center * (2 ** (1 / 6)) return f_center, f_lower, f_upper def third_octave_response_db(waveform: torch.Tensor, sr: int): """ Compute 1/3 octave magnitude response in dB from waveform. Args: waveform (torch.Tensor): shape (n_samples,) or (1, n_samples) sr (int): sample rate Returns: freqs (np.ndarray): band center frequencies mags_db (torch.Tensor): band magnitudes in dB """ if waveform.ndim > 1: waveform = waveform.squeeze(0) n = waveform.numel() spec = torch.fft.rfft(waveform) mag = torch.abs(spec) / n freqs = torch.fft.rfftfreq(n, d=1.0 / sr) # Get bands f_center, f_lower, f_upper = third_octave_bands(sr) band_mags = [] for fl, fu in zip(f_lower, f_upper): idx = (freqs >= fl) & (freqs < fu) if idx.any(): band_mags.append(mag[idx].mean()) else: band_mags.append(torch.tensor(0.0)) band_mags = torch.stack(band_mags) # Convert to dB (avoid log(0)) mags_db = 20 * torch.log10(band_mags + 1e-12) return f_center, mags_db def spectrum_mse_from_vae( vae_latents: np.ndarray, mean_spectrum: np.ndarray, sample_rate: int = 48000, ): """ Load VAE latents from local path or s3:// URI, decode to audio, compute 1/3-octave spectrum (mono), normalize to the provided mean_spectrum level, and return the spectrum, normalized spectrum, and MSE vs mean_spectrum. """ def _to_mono_torch(arr: np.ndarray) -> torch.Tensor: t = torch.from_numpy(arr) return t.mean(dim=0) if t.ndim > 1 else t # 1) Load latents and decode to audio audio = codec_decode(vae_latents) # must expose .array_float (np.ndarray) # 2) Third-octave spectrum (mono) mono = _to_mono_torch(audio.array_float) center_freqs, spectrum_t = third_octave_response_db(mono, sample_rate) spectrum = spectrum_t.numpy() mono_first = mono[: 30 * sample_rate] mono_last = mono[-30 * sample_rate :] _, spectrum_t_last = third_octave_response_db(mono_last, sample_rate) _, spectrum_t_first = third_octave_response_db(mono_first, sample_rate) spectrum_delta = (spectrum_t_first - spectrum_t_last).abs() total_delta = spectrum_delta.sum() # 3) Level-normalize to batch mean spectrum if spectrum.shape != mean_spectrum.shape: raise ValueError( f"mean_spectrum shape {mean_spectrum.shape} != spectrum shape {spectrum.shape}" ) mean_of_mean = float(np.mean(mean_spectrum)) spec_mean = float(np.mean(spectrum)) spectrum_normalized = spectrum - (spec_mean - mean_of_mean) # 4) MSE vs mean_spectrum mse = float(np.mean((spectrum_normalized - mean_spectrum) ** 2)) return { "center_freqs": center_freqs, # np.ndarray [n_bands] "spectrum": spectrum, # np.ndarray [n_bands], dB "spectrum_normalized": spectrum_normalized, # np.ndarray [n_bands], dB "spectrum_delta": spectrum_delta.numpy(), # np.ndarray [n_bands], dB "total_delta": total_delta.numpy(), # float "mse": mse, # float } def _reload_models_if_needed(hoot_filepath, hoot_tokenizer_filepath): model_list = load_model_list( checkpoint_filepath=hoot_filepath, tokenizer_filepath=hoot_tokenizer_filepath, n_gpus=None, ) model = model_list[0]["model"] tokenizer = model_list[0]["tokenizer"] return model_list, model, tokenizer class GenerateWorker: def __init__( self, dit_model_filepath: str, output_path: str, mean_spectrum: np.ndarray, ): self.output_path = output_path self.mean_spectrum = mean_spectrum start_time = time.time() print("Start loading models") # load diffusion model num_gpus = torch.cuda.device_count() cuda_device = torch.cuda.current_device() print(f"Found {num_gpus} GPUs. Using GPU {cuda_device}.") codec_filepath = CODEC_FILEPATH _ = preload_codec_models(codec_filepath) # _ = preload_ear_model(ear_model_filepath) print( f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds" ) @staticmethod def download_models(dir_path=MOUNT_PATH): """Download diffusion models.""" print("Start downloading models") _ = diffusion_gen.get_model_if_needed( CODEC_FILEPATH, cache_dir=dir_path, ) _ = diffusion_gen.get_model_if_needed( diffusion_gen.SEMANTIC_MODEL_FILEPATH, cache_dir=dir_path ) _ = diffusion_gen.get_model_if_needed( diffusion_gen.SEMANTIC_CLUSTERS_FILEPATH, cache_dir=dir_path ) _ = diffusion_gen.get_model_if_needed( diffusion_gen.CODEC_FILEPATH, cache_dir=dir_path ) print("Finish downloading models") def generate(self, work_item): global models """Generate audio from a work item.""" # when saving to s3 we will use a structure like this: # self.output_path/ # item_id/ # original_semantic.npz # 0_generated_audio.mp3 # 0_generated_semantic.npz # 0_metadata.json # 1_generated_audio.mp3 # 1_generated_semantic.npz # 1_metadata.json # ... # metadata.json for dirname in work_item: print(f"Processing {dirname}...") # get all the files in the item_id directory search_dir = f"{BASE_S3_DIR}{dirname}/" print(f"Searching in {search_dir}") files = list_s3_dir(search_dir) all_files = [f[0] for f in files] all_files = [os.path.basename(f) for f in all_files] filtered_files = [f for f in all_files if f.endswith(".npz")] filtered_files = [f for f in filtered_files if "upsampled_vae" in f] if len(filtered_files) == 0: print(f"No upsampled vae files found for {dirname}") for f in filtered_files: filename = f.replace("_upsampled_vae.npz", "") s3_filepath = os.path.join( "christian/outputs/v3-base-data-ctx-rs-t3", f"{dirname}", f"{filename}_spectrum.npz", ) # check if the file exists if f"{filename}_spectrum.npz" in all_files: print(f"Spectrum already exists for {dirname}") print() continue vae_latents = read_from_s3( f"{BASE_S3_DIR}{dirname}/{f}", read_f=np.load, )["vae_latents"] print(vae_latents.shape) result = spectrum_mse_from_vae(vae_latents, self.mean_spectrum) with tempfile.TemporaryDirectory() as td: output_filepath = os.path.join( td, f"{filename}_spectrum.npz", ) # Save all items in result as npz (including arrays and scalars) np.savez(output_filepath, **result) s3_client.upload_file(output_filepath, "suno-data", s3_filepath) print(f"Saved spectrum to {s3_filepath}") CODEC_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth" from suno_utils.tasks.dac_vae_fixed_25hz import ( preload_models as preload_codec_models, decode as codec_decode, encode as codec_encode, decode_stream_to_full_audio, ) def download_model_wrapper_d(): # this print is necessary to have modal rerun this when MODEL changes # Modal tracks referenced global variables # Change the name of the function to force a rerun print("Downloading model for history encoder") GenerateWorker.download_models() image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS) APP_NAME = f"batch-metadata-backfill" app = modal.App(APP_NAME, image=image, secrets=SECRETS) N_MAX_REPLICAS = 64 GPU_TYPE = "A10G" # H100, A10G @app.cls( # gpu=modal.gpu.H100(count=1), gpu=GPU_TYPE, cpu=4, secrets=SECRETS, timeout=2 * 60 * 60, container_idle_timeout=240, # mounts=MODAL_MOUNTS, memory=15000, concurrency_limit=N_MAX_REPLICAS, ) class GenerateStub: def __init__(self, dit_ckpt: str, output_path: str, mean_spectrum: np.ndarray): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = GenerateWorker( dit_ckpt, output_path, mean_spectrum, ) @modal.method() def generate(self, work_item: list[dict]): return self.worker.generate(work_item) DIT_MODEL_FILEPATH = None OUTPUT_STR = None MEAN_SPECTRUM = None BASE_S3_DIR = "s3://suno-data/christian/outputs/v3-base-data-ctx-rs-t3/" @app.local_entrypoint() def main(): # find all the dirnames from # s3://suno-data/christian/outputs/v3-base-data-ctx-rs-t3/ dirnames = list_s3_dir( f"{BASE_S3_DIR}", ) dirnames = [os.path.dirname(f[0]) for f in dirnames] dirnames = [f.split("/")[-1] for f in dirnames] dirnames = list(set(dirnames)) print(f"Total dirnames: {len(dirnames)}") spectrum_csv = ( "/home/christian/code/christian/metadata/genius_t6_sampled_10k_spectrums.csv" ) # load the mean spectrum from the csv spectrums_df = pd.read_csv(spectrum_csv) # compute the mean spectrum # now we have a list of spectrums so we want to compute the mean spectrum # keep in mind each spectrum is a tuple of (freqs, mags) spectrums = spectrums_df.to_numpy() mean_spectrum = np.mean(spectrums, axis=0) print("mean_spectrum shape: ", mean_spectrum.shape) # just do the first 128 work_items = dirnames # [:128] print(f"Total work items: {len(work_items)}") ## now remove these from the work items # work_items = [f for f in work_items if f["id"] not in existing_ids] print(f"Total work items remaining: {len(work_items)}") chunksize = 16 # num of prompts per worker worker = GenerateStub( DIT_MODEL_FILEPATH, f"christian/outputs/{OUTPUT_STR}", mean_spectrum ) work_items = list(funcy.chunks(chunksize, work_items)) print(f"Chunksize: {chunksize}, total chunks: {len(work_items)}") # print("Testing inference...") # t0 = time.time() # for work_item in work_items[:1]: # _ = worker.generate.remote(work_item) # print(f"{int(round(time.time()-t0))}s for test") print("Running batch inference...") t0 = time.time() _ = list(worker.generate.map(work_items)) print(round((time.time() - t0) / 60 / 60), "h total for batch generation")