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 random 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.tasks.ear_v3 import load_checkpoint from typing import Optional, Dict, Any import torch.fft as fft from suno_utils.diffusion.generation import ( preload_dit_model, preload_tokenizer, TOKENIZER_FILEPATH, SEMANTIC_MODEL_FILEPATH, SEMANTIC_CLUSTERS_FILEPATH, _retrieve_models, ) 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 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 combine_width_octave( width_deltas, # array-like of stereo width deltas (can be signed) octave_totals, # array-like of total octave deltas (>=0) method="robust", # "robust" or "minmax" ): w = np.abs(np.asarray(width_deltas, dtype=float)) o = np.asarray(octave_totals, dtype=float) def robust_norm(x): med = np.median(x) q1, q3 = np.percentile(x, 25), np.percentile(x, 75) iqr = max(q3 - q1, 1e-12) return (x - med) / iqr def minmax_norm(x): xmin, xmax = float(np.min(x)), float(np.max(x)) return (x - xmin) / (max(xmax - xmin, 1e-12)) if method == "robust": z_w, z_o = robust_norm(w), robust_norm(o) combo = 0.5 * (z_w + z_o) # map to 0–100 for readability score = 100 * minmax_norm(combo) elif method == "minmax": n_w, n_o = minmax_norm(w), minmax_norm(o) score = 100 * (0.5 * (n_w + n_o)) else: raise ValueError("method must be 'robust' or 'minmax'") return score # higher = more change; use (100 - score) for stability 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 = 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 stereo_width(waveform: torch.Tensor): """ Compute the stereo width of a waveform. """ # can you implement this? # Assume waveform shape is (2, seq_len) if waveform.ndim != 2 or waveform.shape[0] != 2: raise ValueError( "waveform must have shape (2, seq_len) for stereo width calculation" ) left = waveform[0] right = waveform[1] # Compute correlation coefficient between L and R left = left - left.mean() right = right - right.mean() numerator = (left * right).mean() denominator = torch.sqrt((left**2).mean() * (right**2).mean()) + 1e-12 corr = numerator / denominator # Stereo width: 0 = mono, 1 = fully wide (L and R uncorrelated), -1 = fully out of phase width = torch.sqrt(1 - corr**2) return width.item() # Assumes third_octave_bands, third_octave_response_db, stereo_width # are defined exactly as in your snippet above. def analyze_audio_first_last( audio: torch.Tensor, sample_rate: int, segment_duration: int = 60, ) -> Dict[str, Any]: """ Compute 1/3-octave response and stereo width for the first and last segments of a single audio tensor, and return their deltas. Args: audio: Tensor of shape (channels, samples) or (samples,) sample_rate: sample rate in Hz segment_duration: segment length in seconds for the first/last comparison Returns: { 'center_freqs': np.ndarray, 'third_octave_first': Tensor[dB], 'third_octave_last': Tensor[dB], 'delta_response': Tensor[dB], # first - last 'total_delta': Tensor[scalar], # L1 magnitude of delta_response 'stereo_width_first': float or None, 'stereo_width_last': float or None, 'stereo_width_delta': float or None, # first - last } """ if not isinstance(audio, torch.Tensor): raise TypeError("audio must be a torch.Tensor") if audio.numel() == 0: raise ValueError("audio is empty") # Determine segment length in samples (clip to available length) seg_len = min(audio.shape[-1], segment_duration * sample_rate) # Helper: mono mixdown for 1/3-octave analysis mono = audio.mean(dim=0) if audio.ndim > 1 else audio first_mono = mono[:seg_len] last_mono = mono[-seg_len:] # 1/3-octave responses (dB) center_freqs_first, oct_first = third_octave_response_db(first_mono, sample_rate) center_freqs_last, oct_last = third_octave_response_db(last_mono, sample_rate) # Centers should match; keep the first as canonical if ( len(center_freqs_first) != len(center_freqs_last) or (center_freqs_first != center_freqs_last).any() ): raise RuntimeError( "Mismatched third-octave centers between first and last segments." ) delta = oct_first - oct_last total_delta = delta.abs().sum() # Stereo width (if stereo input) def maybe_width(x: torch.Tensor) -> Optional[float]: if x.ndim == 2 and x.shape[0] == 2 and x.shape[1] > 0: return stereo_width(x) return None first_full = audio[:, :seg_len] if audio.ndim == 2 else audio last_full = audio[:, -seg_len:] if audio.ndim == 2 else audio width_first = maybe_width(first_full) width_last = maybe_width(last_full) width_delta = None if (width_first is not None) and (width_last is not None): width_delta = float(width_first - width_last) return { "center_freqs": center_freqs_first, # np.ndarray "third_octave_first": oct_first, # Tensor[dB] "third_octave_last": oct_last, # Tensor[dB] "delta_response": delta, # Tensor[dB] "total_delta": total_delta, # Tensor[scalar] "stereo_width_first": width_first, # float or None "stereo_width_last": width_last, # float or None "stereo_width_delta": width_delta, # float or None } def best_worst_by_robust_combo( *features, weights=None, directions=None, return_scores=True ): """ Robust-rank items given multiple feature lists. Parameters ---------- *features : array-like One or more equal-length sequences (one value per item per feature). weights : array-like or None Optional weights per feature (same length as number of features). Defaults to equal weights. directions : array-like of {+1, -1} or None Optional sign per feature. +1 means higher is better for that feature, -1 means lower is better. Defaults to +1 for all. return_scores : bool If True, also return the final combined scores array. Returns ------- best_idx : int worst_idx : int (scores) : np.ndarray, only if return_scores=True """ if len(features) == 0: raise ValueError("Provide at least one feature.") X = [np.asarray(f, dtype=float) for f in features] n = len(X[0]) if any(len(f) != n for f in X): lens = [len(f) for f in X] raise ValueError(f"All features must have the same length. Got lengths: {lens}") m = len(X) # number of features # defaults if weights is None: weights = np.ones(m, dtype=float) else: weights = np.asarray(weights, dtype=float) if len(weights) != m: raise ValueError("weights must match number of features") if directions is None: directions = np.ones(m, dtype=float) else: directions = np.asarray(directions, dtype=float) if len(directions) != m: raise ValueError("directions must match number of features") if not np.all(np.isin(directions, [+1, -1])): raise ValueError("directions must be +1 or -1") # normalize weights to sum to 1 wsum = weights.sum() if wsum <= 0: raise ValueError("Sum of weights must be > 0") weights = weights / wsum def robust_norm(x): med = np.median(x) q1, q3 = np.percentile(x, 25), np.percentile(x, 75) iqr = max(q3 - q1, 1e-12) return (x - med) / iqr # robust-normalize each feature, apply direction (+1/-1), then weight and sum scores = np.zeros(n, dtype=float) for j, x in enumerate(X): z = robust_norm(x) * directions[j] scores += weights[j] * z # best = largest score; worst = smallest score best_idx = int(np.argmax(scores)) worst_idx = int(np.argmin(scores)) if return_scores: return best_idx, worst_idx, scores return best_idx, worst_idx def best_pair_maxmin_bruteforce( clips, metrics=("cer", "shimmer_score", "ear_v3_score"), better_is=("lower", "lower", "higher"), tau=None, # None, scalar, or list/tuple per metric (positive margins) tiebreak="sum", # "sum" or "weighted" weights=None, # weights for "weighted" tiebreak ): """ Find (pos, neg) maximizing the minimum per-feature margin, with direction-aware metrics. We transform each feature so that "higher is better" by multiplying with sign: sign_j = +1 if higher-is-better, -1 if lower-is-better For a pair (p, n), margin vector in transformed space: diff' = (sign * S[p]) - (sign * S[n]) Objective: maximize min(diff'). Thresholds tau (if any) apply to diff' (diff' >= tau). Returns dict with pos/neg indices & clips, per-metric deltas (original sign), objective value, tiebreak value, and a bottleneck report. """ # Build score matrix (N, d) S = np.array([[clip[m] for m in metrics] for clip in clips], dtype=float) N, d = S.shape # Direction: map to signs so that higher is better for all dir_map = {"higher": 1.0, "lower": -1.0} signs = np.array([dir_map[b] for b in better_is], dtype=float) S_prime = S * signs # transformed space # Normalize thresholds into transformed space (positive margins) tau_vec = None if tau is None else np.broadcast_to(np.array(tau, dtype=float), (d,)) # Tiebreak setup if tiebreak not in ("sum", "weighted"): raise ValueError("tiebreak must be 'sum' or 'weighted'") if tiebreak == "weighted": if weights is None: weights = np.ones(d, dtype=float) w = np.array(weights, dtype=float) if w.shape != (d,): raise ValueError("weights must have length equal to number of metrics") best_primary = None best_tie = None best_pair = None best_payload = None for p in range(N): for n in range(N): if p == n: continue diff_prime = S_prime[p] - S_prime[n] # margins in transformed space if (tau_vec is not None) and not np.all(diff_prime >= tau_vec): continue # Primary objective: maximize the minimum margin min_margin = float(np.min(diff_prime)) # Tiebreaker tie_val = ( float(np.sum(diff_prime)) if tiebreak == "sum" else float(np.dot(diff_prime, w)) ) cand_better = ( (best_primary is None) or (min_margin > best_primary) or ( min_margin == best_primary and (best_tie is None or tie_val > best_tie) ) ) if cand_better: best_primary, best_tie = min_margin, tie_val best_pair = (p, n) # For interpretability: raw deltas in original space/signs delta_raw = S[p] - S[n] # Bottleneck index = where diff' is minimal bottleneck_idx = int(np.argmin(diff_prime)) best_payload = { "delta_per_metric": { m: float(delta_raw[i]) for i, m in enumerate(metrics) }, "diff_prime": diff_prime.copy(), "bottleneck_metric": metrics[bottleneck_idx], "bottleneck_margin": float(diff_prime[bottleneck_idx]), "bottleneck_index": bottleneck_idx, } if best_pair is None: return None p, n = best_pair return { "pos_idx": p, "neg_idx": n, "pos_clip": clips[p], "neg_clip": clips[n], "delta": best_payload["delta_per_metric"], # original units/signs "objective_min_margin": best_primary, # in transformed space "tiebreak_value": best_tie, "metrics": list(metrics), "better_is": list(better_is), "thresholds_applied": (tau is not None), "bottleneck": { "metric": best_payload["bottleneck_metric"], "margin_transformed": best_payload["bottleneck_margin"], "metric_index": best_payload["bottleneck_index"], }, } 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 _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, test_type: str, ): self.output_path = output_path self.test_type = test_type 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}.") tokenizer_filepath = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json" semantic_model_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt" semantic_clusters_filepath = ( "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy" ) codec_filepath = CODEC_FILEPATH _ = diffusion_gen.preload_dit_model( dit_model_filepath=dit_model_filepath, use_ema_if_exists=True, compile=False, weights_precision=torch.bfloat16, ) _ = preload_tokenizer(tokenizer_filepath) _ = preload_semantic_models(semantic_model_filepath, semantic_clusters_filepath) _ = preload_codec_models(codec_filepath) self.diffusion_engine = UpsampleEngine(min_chunk_size=30 * 25) # load hoot model print("Loading hoot model") hoot_filepath = "s3://suno-data/christian/checkpoints/hoot_v3.pt" hoot_tokenizer_filepath = ( "s3://suno-data/christian/checkpoints/hoot_v3_tokenizer_v3.pt" ) self.hoot_model_list, self.hoot_model, self.hoot_tokenizer = ( _reload_models_if_needed(hoot_filepath, hoot_tokenizer_filepath) ) # load ear model ear_model_filepath = "s3://suno-data/christian/checkpoints/ear/ear_v2_s3080.pt" self.ear_model = load_model(ear_model_filepath, compile=True) self.ear_model_v3 = load_checkpoint(EAR_MODEL_V3_FILEPATH, device=cuda_device) print( f"Finish loading models. Took {round(time.time() - start_time, 2)} seconds" ) @staticmethod def download_models(dit_model_filepath, 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 ) _ = diffusion_gen.get_model_if_needed(dit_model_filepath, cache_dir=dir_path) # _ = chirp_v2._get_model_if_needed(gpt_ckpt, cache_dir=dir_path) # _ = chirp_v2._get_model_if_needed(chirp_v2.TOKENIZER_PATH, 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 item in work_item: item_id = item["id"] # if we have s3_filepath, this is original audio # so we will need to encode vae latents and semantic codes # the s3_filepath will be the positive example and the upsampled audio will be the negative if "s3_filepath" in item: s3_filepath = item["s3_filepath"] # read audio from s3 and then semantic encode audio = Audio.from_s3(s3_filepath, n_channels=2) try: audio_normed = audio.normalize_loudness(target_integrated_lufs=-14) except Exception as e: print(f"Error normalizing loudness: {e}") continue if SAVE_CYCLED_VAE: # encode the vae latents vae_latents_cycled = codec_encode(audio_normed) with tempfile.TemporaryDirectory() as td: vae_latents_path = os.path.join( td, f"{item_id}_original_vae.npz" ) np.savez( vae_latents_path, vae_latents=vae_latents_cycled.astype(np.float16), ) s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_original_vae.npz", ) s3_client.upload_file( vae_latents_path, "suno-data", s3_filepath, ExtraArgs={"ContentType": "application/octet-stream"}, ) if CYCLE_ONLY: continue # also encode semantic codes codes = encode_semantic( audio.convert(sample_rate=24_000, byte_width=2, n_channels=1) ).astype(np.int64) if SAVE_CODES: with tempfile.TemporaryDirectory() as td: codes_path = os.path.join(td, f"{item_id}_semantic.npz") np.savez(codes_path, codes=codes) s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_semantic.npz", ) s3_client.upload_file( codes_path, "suno-data", s3_filepath, ExtraArgs={"ContentType": "application/octet-stream"}, ) # if we dont have the s3_filepath, this is a generated audio # so we will pull the semantic codes from s3 else: # load the semantic codes from s3 s3_filepath = f"s3://suno-data-uploads/studio/uploads/{item_id}.npz" try: data = read_from_s3(s3_filepath, read_f=np.load) if "v3.0_raw" in data: codes = data["v3.0_raw"] elif "v3.5_raw" in data: codes = data["v3.5_raw"] elif "v4.0_raw" in data: codes = data["v4.0_raw"] elif "v4.5_raw" in data: codes = data["v4.5_raw"] elif "v5.0_raw" in data: codes = data["v5.0_raw"] else: raise ValueError(f"No codes found for {item_id}") except Exception as e: print(f"Error loading {s3_filepath}: {e}") continue semantic_codes = torch.from_numpy(codes[:, 0]).long() # .cuda() # if tags is a list, join it into a string if isinstance(item["tags"], list): tags_str = ", ".join(item["tags"][:5]) else: tags_str = item["tags"] if tags_str is None: tags_str = "" lyrics = item["text"] # Calculate original audio CER if we have original audio original_cer = None if "s3_filepath" in item and lyrics != "": print("Calculating original audio CER...") with tempfile.TemporaryDirectory() as td: original_audio_path = os.path.join(td, f"{item_id}_original.mp3") audio.write_hq_mp3(original_audio_path) s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_original.mp3", ) # copy this to s3 s3_client.upload_file( original_audio_path, "suno-data", s3_filepath, ExtraArgs={"ContentType": "application/octet-stream"}, ) out = encode_filepaths([original_audio_path], return_logits=True) basic_cleaned_lyrics = lyrics decoded_preds = self.hoot_tokenizer.decode_logits( out[0], prior_text=basic_cleaned_lyrics ) true_text_norm = clean_text(basic_cleaned_lyrics) original_cer = round(get_cer(true_text_norm, decoded_preds), 3) print(f"Original audio CER: {original_cer}") if DIFFUSION_SEED is None: diffusion_seed = 0 for char in item_id: diffusion_seed = (diffusion_seed * 31 + ord(char)) % 1000000 print(f"Generated seed from '{item_id}': {diffusion_seed}") else: diffusion_seed = DIFFUSION_SEED print(f"Using seed from config: {diffusion_seed}") diffusion_steps = DIFFUSION_STEPS noise_ctx_level = np.random.uniform(0.5, 1.0) noise_ctx_pad_len = NOISE_CTX_PAD_LEN rho = RHO sigma_min = SIGMA_MIN sigma_max = SIGMA_MAX min_steps = 8 max_steps = 24 # because of distill model we pin these parameters if "distill" in DIT_MODEL_FILEPATH: text_cfg_coef = 1.0 else: text_cfg_coef = np.random.uniform(1.0, 3.0) # 50% of the time we will run a trajectory to an intermediate point # 50% of the time we will run a trajectory to the end use_trajectory = False if use_trajectory: print("Running trajectory to an intermediate point") # so, first we will run a trajectory to an intermediate point # then we will use the last 30s of the trajectory to run two alternative n+1 trajectories # we will then save the two alternative trajectories and the original trajectory (conditioning only) # this will give us 3 VAEs sequences technically, but for simplicity we will just copy the conditioning # to the start of both alternative trajectories. diffusion_steps = np.random.randint(min_steps, max_steps + 1) init_cfg = diffusion_gen.DiffusionGenerationConfig( steps=diffusion_steps, lyrics=lyrics, tags=tags_str, text_cfg_coef=text_cfg_coef, ctx_cfg_coef=CTX_CFG_SCALE, codec_scale_factor=CODEC_SCALE_FACTOR, scale_ctx_vector=SCALE_CTX_VECTOR, noise_ctx_level=noise_ctx_level, noise_ctx_pad_len=noise_ctx_pad_len, rho=rho, sigma_min=sigma_min, sigma_max=sigma_max, seed=diffusion_seed, objective=OBJECTIVE, sampler_type=SAMPLER_TYPE, ) # we will crop the semantic codes to a random length leaving at least 750 tokens at the end crop_len = random.randint(750, semantic_codes.shape[0] - 750) init_semantic_codes = semantic_codes[:crop_len] # convert crop_len to a time in seconds start_s = crop_len / 25 print( f"Cropping semantic codes to {crop_len} tokens, {semantic_codes.shape[0]} total" ) cropped_semantic_codes = semantic_codes[crop_len:][ :750 ] # just one chunk request = Request( id="dummy", generation_config=init_cfg, tokens=init_semantic_codes, input_tokens_finished=True, ) result = self.diffusion_engine.run_request(request) history_vae_latents = torch.concat(result.vae_latents) # save out the history vae latents with tempfile.TemporaryDirectory() as td: history_vae_latents_path = os.path.join( td, f"{item_id}_history_vae.npz" ) np.savez( history_vae_latents_path, vae_latents=history_vae_latents.cpu() .numpy() .astype(np.float16), ) s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_history_vae.npz", ) s3_client.upload_file( history_vae_latents_path, "suno-data", s3_filepath, ExtraArgs={"ContentType": "application/octet-stream"}, ) elif False: # pick a random starting index and crop to one chunk if semantic_codes.shape[0] <= 750: print( f"Warning: semantic_codes too short ({semantic_codes.shape[0]} tokens), skipping" ) continue start_idx = np.random.randint(0, semantic_codes.shape[0] - 750) cropped_semantic_codes = semantic_codes[start_idx : start_idx + 750] print( f"Cropping semantic codes to {cropped_semantic_codes.shape[0]} tokens, {semantic_codes.shape[0]} total" ) start_s = start_idx / 25 print( f"Cropping semantic codes to {start_s} seconds, {semantic_codes.shape[0]} total" ) else: cropped_semantic_codes = semantic_codes start_s = 0 with tempfile.TemporaryDirectory() as td: # copy semantic codes to s3 also semantic_codes_path = os.path.join(td, f"{item_id}_semantic.npz") np.savez( semantic_codes_path, semantic_codes=cropped_semantic_codes.cpu().numpy(), ) s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_semantic.npz", ) s3_client.upload_file( semantic_codes_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "application/octet-stream", }, ) # now we will run the n+1 trajectories and find the best pair generated_clips = [] # Store all generated clips # num_upsamples = np.random.randint(4, MAX_UPSAMPLES + 1) num_upsamples = 10 for n in range(MAX_UPSAMPLES): # we will crop the semantic codes to a random length leaving at least 750 tokens at the end diffusion_steps = np.random.randint(min_steps, max_steps + 1) diffusion_seed = np.random.randint(0, 1000000) noise_ctx_level = np.random.uniform(0.5, 0.75) if "distill" in DIT_MODEL_FILEPATH: text_cfg_coef = 1.0 else: text_cfg_coef = np.random.uniform(1.0, 3.0) generation_history_latents = ( history_vae_latents[:-750] if use_trajectory else None ) # check generation_history_latents length is 750 if generation_history_latents is not None: if generation_history_latents.shape[0] != 750: generation_history_latents = None gen_cfg = diffusion_gen.DiffusionGenerationConfig( steps=diffusion_steps, lyrics=lyrics, tags=tags_str, text_cfg_coef=text_cfg_coef, ctx_cfg_coef=CTX_CFG_SCALE, codec_scale_factor=CODEC_SCALE_FACTOR, scale_ctx_vector=SCALE_CTX_VECTOR, noise_ctx_level=noise_ctx_level, noise_ctx_pad_len=noise_ctx_pad_len, rho=rho, sigma_min=sigma_min, sigma_max=sigma_max, seed=diffusion_seed, objective=OBJECTIVE, sampler_type=SAMPLER_TYPE, # we will use the last 30s of the trajectory to run the n+1 trajectory generation_history_latents=generation_history_latents, ) request = Request( id="dummy", generation_config=gen_cfg, tokens=cropped_semantic_codes, input_tokens_finished=True, ) result = self.diffusion_engine.run_request(request) output_vae_latents = torch.concat(result.vae_latents) upsampled_audio = decode_stream_to_full_audio(output_vae_latents) metadata = { # "original_audio": item["s3_filepath"], "filename": f"{item_id}_{CHECKPOINT_NAME}_{n}.mp3", "id": item_id, "text": lyrics, "tags": tags_str, "start_s": start_s, "diffusion": { "steps": int(gen_cfg.steps), "seed": int(gen_cfg.seed), "text_cfg_coef": float(gen_cfg.text_cfg_coef), "noise_ctx_level": float(gen_cfg.noise_ctx_level), "codec_scale_factor": float(gen_cfg.codec_scale_factor), "scale_ctx_vector": gen_cfg.scale_ctx_vector, "sampler_type": gen_cfg.sampler_type, }, } # we want to save out # audio file of the final audio # npz of the estimated semantics # metadata json with the original prompt and tags # copy some stuff # we want to copy the original audio and the npz of the original semantics with tempfile.TemporaryDirectory() as td: # save audio to s3 upsampled_audio_path = os.path.join(td, f"{item_id}_{n}.mp3") upsampled_audio.write_hq_mp3(upsampled_audio_path) # run hoot evaluation (CER) # check if the lyrics are not empty if lyrics != "" and RUN_HOOT: out = encode_filepaths( [upsampled_audio_path], return_logits=True ) # get lyrics from basic_cleaned_lyrics = lyrics decoded_preds = self.hoot_tokenizer.decode_logits( out[0], prior_text=basic_cleaned_lyrics ) true_text_norm = clean_text(basic_cleaned_lyrics) cer_val = round(get_cer(true_text_norm, decoded_preds), 3) metadata["hoot_cer"] = cer_val # add cer to metadata print("hoot cer: ", metadata["hoot_cer"]) else: metadata["hoot_cer"] = None cer_val = None # run ear evaluation (quality score) # check if the length of the audio is greater than 5 seconds if upsampled_audio.duration_s > 5.0: quality_score = self.ear_model.get_score(upsampled_audio_path) metadata["ear_score"] = ( quality_score # add quality score to metadata ) print("ear score: ", metadata["ear_score"]) else: metadata["ear_score"] = None # run shimmer score evaluation shimmer_score = shimmerscore(upsampled_audio_path) metadata["shimmer_score"] = shimmer_score # lets crop the batches to multiple of 750s vae_latents_batched = output_vae_latents.unsqueeze(0).float().cuda() vae_latents_batched = vae_latents_batched[ :, : 750 * (vae_latents_batched.shape[1] // 750), : ] vae_latents_batched = ( vae_latents_batched.unfold(1, 750, 750) .squeeze(0) .permute(0, 2, 1) ) with torch.no_grad(): ear_v3_scores = self.ear_model_v3( vae_latents_batched * CODEC_SCALE_FACTOR ) ear_v3_scores = ear_v3_scores.cpu().numpy() # extract the scores, first chunk, last chunk, and mean of chunks first_chunk_scores = ear_v3_scores[0] last_chunk_scores = ear_v3_scores[-1] mean_chunk_scores = ear_v3_scores.mean() metadata["ear_v3_score"] = mean_chunk_scores metadata["ear_v3_score_first"] = first_chunk_scores metadata["ear_v3_score_last"] = last_chunk_scores metadata["ear_v3_score_list"] = ear_v3_scores.tolist() # compute the similarity score in audio between the two clips # run the other evals audio_analysis = analyze_audio(upsampled_audio_path) # add all the keys in audio_analysis to metadata for key, value in audio_analysis.items(): metadata[key] = value # compute the audio analysis results = analyze_audio_first_last( torch.from_numpy(upsampled_audio.array_float), 48000 ) print(f"stereo_width_delta: {results['stereo_width_delta']}") print(f"total_delta (octave): {results['total_delta']}") # merge everything in results to metadata metadata.update(results) # Store this clip for later comparison clip_data = { "n": n, "cer": cer_val, "shimmer_score": shimmer_score, # "ear_v3_score": ear_v3_score.item(), "metadata": metadata.copy(), "upsampled_audio": upsampled_audio, "output_vae_latents": output_vae_latents, "cropped_semantic_codes": cropped_semantic_codes, "gen_cfg": gen_cfg, "first_last_analysis": results, } generated_clips.append(clip_data) # pick the best clip pair via octave and stereo width stereo_width_deltas = [] octave_deltas = [] ear_v3_scores = [] shimmer_scores = [] for n, generated_clip_data in enumerate(generated_clips): stereo_width_deltas.append( generated_clip_data["metadata"]["stereo_width_delta"] ) octave_deltas.append(generated_clip_data["metadata"]["total_delta"]) ear_v3_scores.append( generated_clip_data["metadata"]["ear_v3_score_first"] ) shimmer_scores.append(generated_clip_data["metadata"]["shimmer_score"]) directions = [-1, -1, 1, -1] weights = [1.0, 1.5, 1.0, 0.5] print("Ranking...") best_idx, worst_idx, scores = best_worst_by_robust_combo( np.abs(stereo_width_deltas), octave_deltas, ear_v3_scores, shimmer_scores, weights=weights, directions=directions, ) print("scores: ", scores) # print the stereo_width_delta, total_delta, and ear_v3_score_first for the best and worse print("best_idx: ", best_idx) print(generated_clips[best_idx]["metadata"]["stereo_width_delta"]) print(generated_clips[best_idx]["metadata"]["total_delta"]) print(generated_clips[best_idx]["metadata"]["ear_v3_score_first"]) print(generated_clips[best_idx]["metadata"]["shimmer_score"]) print("worst_idx: ", worst_idx) print(generated_clips[worst_idx]["metadata"]["stereo_width_delta"]) print(generated_clips[worst_idx]["metadata"]["total_delta"]) print(generated_clips[worst_idx]["metadata"]["ear_v3_score_first"]) print(generated_clips[worst_idx]["metadata"]["shimmer_score"]) print("--------------------------------") # this makes neg 0 index, and positive 1 index always for n, item_idx in enumerate([worst_idx, best_idx]): clip_data = generated_clips[item_idx] upsampled_audio = clip_data["upsampled_audio"] output_vae_latents = clip_data["output_vae_latents"] metadata = clip_data["metadata"] # save the upsampled audio with tempfile.TemporaryDirectory() as td: if SAVE_MP3: upsampled_audio_path = os.path.join(td, f"{item_id}_{n}.mp3") upsampled_audio.write_hq_mp3(upsampled_audio_path) s3_filepath = os.path.join( self.output_path, f"{item_id}", (f"{item_id}_{CHECKPOINT_NAME}_{n}.mp3"), ) print(f"Uploading clip {item_idx} to {s3_filepath}") s3_client.upload_file( upsampled_audio_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "audio/mpeg", }, ) # save the vae latents vae_latents_path = os.path.join( td, (f"{item_id}_{CHECKPOINT_NAME}_{n}_upsampled_vae.npz"), ) np.savez( vae_latents_path, vae_latents=output_vae_latents.cpu().numpy().astype(np.float16), ) s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_{CHECKPOINT_NAME}_{n}_upsampled_vae.npz", ) s3_client.upload_file( vae_latents_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "application/octet-stream", }, ) # save the metadata metadata_path = os.path.join(td, f"{item_id}_{n}_metadata.npz") np.savez(metadata_path, **metadata) # move the metadata to s3 s3_filepath = os.path.join( self.output_path, f"{item_id}", f"{item_id}_{CHECKPOINT_NAME}_{n}__metadata.npz", ) s3_client.upload_file( metadata_path, "suno-data", s3_filepath, ExtraArgs={ "ContentType": "application/octet-stream", }, ) # prod diffusion model # DIT_MODEL_FILEPATH = "s3://suno-data/tony/tmp/diff/dit_v3_dpo_t10_3k_5e6_b100_t25.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/georg/tmp/2b_prefix_ft.pt" # 2b v45 # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2_600_000.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/dpo/v1_t6_1E6_beta100_n8_bt4.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_2E6_beta100_n8_bt4.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_5E6_beta100_n16_bt4.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_5E6_beta100_n8_bt4_30k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t10_5E6_beta100_n8_bt4_9k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v45_2b_shared_ctx_s3784.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_t6_5E6_beta100_n16_bt4_9k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_t11_5E6_beta100_n16_bt4_9k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_t6_5E6_beta100_n16_bt4_6k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t12_5E6_beta100_n16_bt4_9k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t13_5E6_beta100_n16_bt4_9k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t13_5E6_beta100_n16_bt4_9k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t15_5E6_beta100_n16_bt4_9k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_10k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_20k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_base_200k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_v1_t6_5E6_beta100_n16_bt4_30k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_100k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_5e5_150k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_base_600k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_t3_5e5_200k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_50k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_v1_t17_5E6_beta100_n16_bt4_9k_repro_main_3k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t17_5E6_beta100_n16_bt4_9k_repro_mai_ema.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_infill_v1_t17_5E6_beta100_n16_bt4_9k_repro_main_3k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_100k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_5E6_beta1000_n16_bt4_9k_repro_main_6k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_5E6_beta1000_n16_bt4_9k_repro_main_3k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_5E6_beta1000_n16_bt4_2k_repro_main.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_5E6_beta1000_n16_bt4_2k_repro_main.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3+t5_40k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_t4_5e5_180k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_ft_ear_t4_5e5_270k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_2E6_beta100_n16_bt4_1k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3_190k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_3e5_t6_100k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3_210k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_2E6_beta100_n16_bt4_1p5k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t17_2E6_beta5000_n8_bt4_1p5k_fix.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_base_ear_sft_t3_230k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_2E6_beta5000_n8_bt4_3k_fix.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_5E6_beta100_n8_bt4_3k_acc16_fix_3k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t6_5E6_beta100_n8_bt4_3k_acc16_fix_last.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_3e5_t6_250k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18_1E6_beta100_n32_bt4_3k_acc4_fix_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t6_dpo_t18_1E6_beta100_n32_bt4_1k_acc8_fix_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t6_dpo_t18_1E6_beta100_n32_bt4_1k_acc8_fix_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18_1E6_beta100_n32_bt4_1k_acc8_fix_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_genius_t6_sampled_10k_5E6_beta100_n12_bt4_3k_acc4_fix_6k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t18_dpo_t18_5E6_beta100_n16_bt4_1k_acc8_fix_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/dit_v6_dpo_t11_9k_5e6_b100.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v1_t6_5E6_beta100_n16_bt4.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18_1E6_beta100_n32_bt4_1k_acc8_fix_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_sft_t6_dpo_t6_5E6_beta100_n32_bt4_1k_acc4_fix_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v2_t3_1E6_beta100_n16_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_5e5_t6_c0_2_30k.pt" # sft model # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/diff_v2_2b_2mil_ft_infill_20250421_v2.pt" # current prod model # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2mil_ft_8k_infill_apr21_t1_18_cs0.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18+syn_5E6_beta100_n16_bt2_3k_acc_4_6k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18+syn_5E6_beta100_n16_bt2_1k_acc_8_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_1e5_t7_70k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t18+syn_5E6_beta100_n16_bt2_1k_acc4_sft.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_syn_corr_5E6_beta100_n16_bt2_3k_acc4_sft_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_ear_sft_1e5_t7_380k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n32_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n32_bt4_acc4_3k_3k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_5E6_beta100_n16_bt2_acc8_1k_sft_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_5E6_beta100_n16_bt2_acc8_1k_sft_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_syn_5E6_beta100_n16_bt2_acc8_1k_sft_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_syn_5E6_beta100_n16_bt2_acc8_1k_sft_0_1_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_5E6_beta100_n32_bt4_acc4_3k_1k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n16_bt2_acc8_1k_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_1E6_beta100_n32_bt4_acc4_3k_1k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t2_1E6_beta100_n16_bt4_acc4_1k_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t2_1E6_beta100_n16_bt4_acc4_1k_1k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_1E6_beta100_n4_bt2_acc8_1k_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_neg_0_5_1E6_beta100_n4_bt2_acc8_1k_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_1E6_beta100_n8_bt2_acc8_2k_16k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_apr21_t1_18_cs0_t1_ear_1E6_beta100_n4_bt2_acc8_4k_32k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v3_t1_ear_1E6_beta100_n4_bt2_acc8_4k_32k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t6_1E6_beta100_n16_bt4_acc4_1k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t7_1E6_beta100_n16_bt4_acc4_1k_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_og_d3_t7_1E6_beta100_n16_bt4_acc4_1k_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_og_d3_t9_1E6_beta100_n16_bt4_acc4_1k_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_ear_t1_1E6_beta100_n4_bt2_acc8_4k_32k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t9_2E6_beta100_n16_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t19_1E6_beta100_n16_bt4_acc4_3k_12k.pt" # # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t10_1E6_beta100_n16_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_750k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t12_1E6_beta100_n16_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t13_1E6_beta100_n16_bt4_acc4_3k_last.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t15_1E6_beta100_n16_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t16_1E6_beta100_n16_bt4_acc4_3k_10k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t17_1E6_beta100_n16_bt4_acc4_3k_last.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t17_1E6_beta100_n16_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t19_1E6_beta100_n8_bt4_acc8_3k_24k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t18_1E6_beta100_n8_bt4_acc8_3k_24k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t20_1E6_beta100_n16_bt16_acc4_6k_24k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t21_1E6_beta100_n16_bt16_acc4_4k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_25m.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d3_t10_1E6_beta100_n16_bt16_acc8_3k_24k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_v1_t19_5E6_beta100_n16_bt4_acc4_3k_12k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d4_2_5E6_beta100_n16_bt2_acc4_4k_16k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_4b_1m.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_5m.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_4b_1_25m.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d4_3_5E6_beta100_n16_bt2_acc4_4k_16k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v2_infill_d4_t6_freeze_5E6_beta100_n16_bt2_acc4_4k_8k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_v45_infill_5e5_sft_log_snr_20k.pt" DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_resume_1_75m.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_25hz_v45_infill_shared_flow_4b_resume_2m.pt" # DIT_MODEL_FILEPATH = ( # "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_v45_distill_adv_100k.pt" # ) # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v45_2b_step_2mil_ft_8k_infill_apr21_d3_v10.pt" # DIT_MODEL_FILEPATH = ( # "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_v45_distill_adv2.pt" # ) # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_25hz_v45_infill_shared_flow_5e5_sft_50k.pt" # DIT_MODEL_FILEPATH = ( # "s3://suno-data/christian/checkpoints/diffusion/8n_2b_adv_bs2_N5_5e6_001_50k.pt" # ) # DIT_MODEL_FILEPATH = ( # "s3://suno-data/christian/checkpoints/diffusion/4n_2b_adv_bs2_N5_5e6_acc4_t1.pt" # ) # /app/suno/modal/models/tony/tmp/diff/v45_2b_step_2mil_ft_8k_infill_apr21_t1_18_cs0.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/8n_2b_flow_distill_bs1_N5_c5e5_g5e6_alt_dmd_cfg_1p5_cosine_nu1_0p1_residual_70k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/16n_2b_flow_distill_bs1_N5_c5e5_g1e6_alt_dmd_cfg_2_residual_sft_220k.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_s3177_lm_t1_0_7_cut_history_4x_1E6_beta100_n8_bt2_acc4_2k_last.pt" # DIT_MODEL_FILEPATH = "s3://suno-data/christian/checkpoints/diffusion/v3_flow_distill_v1_t18_1E6_beta100_n4_bt2_acc2_4k_last.pt" # DIT_MODEL_FILEPATH = ( # "s3://suno-data/christian/checkpoints/diffusion/4n_25hz_2b_flow_5e5_sft_t8_500k.pt" # ) CHECKPOINT_NAME = DIT_MODEL_FILEPATH.split("/")[-1].replace(".pt", "") # OUTPUT_STR = "farm-speaker-desk" # OUTPUT_STR = "pencil-helmet-window" # OUTPUT_STR = "walking-flowers" # OUTPUT_STR = "v3-bootstrap-data-t4" OUTPUT_STR = "v3-base-data-ctx-rs-t5" # OUTPUT_STR = "v3-base-data-ctx-discogs-subset-t0" # OUTPUT_STR = "v2-infill-data-v1" # OUTPUT_STR = "v2-diff-step-test" # OUTPUT_STR = "real-audio-test" TEST_TYPE = "standard" # standard, steps SAVE_CODES = True SAVE_VAE_LATENTS = True SAVE_CYCLED_VAE = True SAVE_MP3 = False CYCLE_ONLY = False GPU_TYPE = "A10G" # H100, A10G OBJECTIVE = "rectified_flow" if "flow" in DIT_MODEL_FILEPATH else "v" SAMPLER_TYPE = "pingpong" if "distill" in DIT_MODEL_FILEPATH else "dpmpp" EAR_MODEL_V3_FILEPATH = "s3://suno-data/christian/checkpoints/ear/ear_v3_s8263.pt" # extra params RHO = 1.0 SIGMA_MIN = 0.5 SIGMA_MAX = 50.0 DIFFUSION_SEED = None # 42 is default RANDOMIZE_DIFFUSION_PARAMS = True USE_ALTERNATE_CFG = False CTX_CFG_SCALE = 1.0 CTX_STEP_PERCENT = 1.0 MAX_UPSAMPLES = 20 # number of attempts to find good pair N_MAX_REPLICAS = 128 OVERWRITE = False RUN_HOOT = False # IGNORE_HISTORY_EVERY = 2 # use this to add a suffix to the checkpoint name EXTRA_NAME_SUFFIX = "" CHECKPOINT_NAME = CHECKPOINT_NAME + EXTRA_NAME_SUFFIX if "dit_v6_dpo_t11_9k_5e6_b100" in DIT_MODEL_FILEPATH: CODEC_SCALE_FACTOR = 2.5 SCALE_CTX_VECTOR = False NOISE_CTX_LEVEL = 0.0 NOISE_CTX_PAD_LEN = 0 DIFFUSION_STEPS = 10 SEMANTIC_SKIP_FACTOR = 1 CODEC_FILEPATH = "s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth" from suno_utils.tasks.dac_vae_100hz_peaq import ( preload_models as preload_codec_models, decode as codec_decode, encode as codec_encode, decode_stream_to_full_audio, ) else: CODEC_SCALE_FACTOR = 0.4 SCALE_CTX_VECTOR = True NOISE_CTX_LEVEL = 0.75 NOISE_CTX_PAD_LEN = 0 DIFFUSION_STEPS = 2 SEMANTIC_SKIP_FACTOR = 1 NOISE_SCHEDULE = "polyexponential" # NOISE_SCHEDULE = "cosine" # CHUNK_SIZE_SCHEDULE = [10 * 25, 20 * 25, 30 * 25] # 10s, 20s, 30s chunks # CHUNK_SIZE_SCHEDULE = [30 * 25, 30 * 25, 30 * 25] # 30s chunks 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(DIT_MODEL_FILEPATH) image = base_image.run_function(download_model_wrapper_d, secrets=SECRETS) APP_NAME = f"batch-generate-rs-gpu" app = modal.App(APP_NAME, image=image, secrets=SECRETS) @app.cls( # gpu=modal.gpu.H100(count=1), gpu=modal.gpu.A10G(count=1) if GPU_TYPE == "A10G" else modal.gpu.H100(count=1), 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, test_type: str): import torch num_gpus = torch.cuda.device_count() print(f"Found {num_gpus} GPUs.") self.worker = GenerateWorker( dit_ckpt, output_path, test_type, ) @modal.method() def generate(self, work_item: list[dict]): return self.worker.generate(work_item) @app.local_entrypoint() def main(): # checkpoints in s3://suno-data/christian/checkpoints # prompts in s3://suno-data/christian/prompts # outputs in s3://suno-data/christian/outputs # load the positive prompts # work_items = read_from_s3( # "s3://suno-data/christian/sft/pos_interesting_clips_up_u_1_20241201_full.jsonl", # read_f=read_jsonl, # ) # work_items = read_jsonl( # "/home/christian/code/christian/metadata/discogs_subset_sampled_metas.jsonl" # ) # work_items = read_jsonl( # "/home/christian/code/christian/metadata/ear/genius_t6_sampled_10k.jsonl" # ) work_items = read_jsonl( "/home/christian/code/christian/metadata/sft/interesting_clips_bluejay_t1_20250811_public_only_lang_balanced_30k.jsonl" ) # just do the first 1 work_items = work_items # [:10000] print(f"Total work items: {len(work_items)}") if not OVERWRITE: # first check for existing ids in the output path existing_ids = list_s3_dir( f"s3://suno-data/christian/outputs/{OUTPUT_STR}/", ) existing_ids = [os.path.dirname(f[0]).split("/")[-1] for f in existing_ids] existing_ids = list(set(existing_ids)) print(f"Total existing ids: {len(existing_ids)}") ## 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 = 4 # num of prompts per worker worker = GenerateStub( DIT_MODEL_FILEPATH, f"christian/outputs/{OUTPUT_STR}", test_type=TEST_TYPE ) 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")