# -*- coding: utf-8 -*- """ Benchmark: how many 30s chunks/second can we load & encode """ from __future__ import annotations import os import sys import time import soxr import torch import random import tempfile import numpy as np from typing import Any, Dict, Iterable, List, Optional, Tuple from contextlib import contextmanager from suno_utils.audio import Audio from suno_utils.utils.opusfile import OpusFile from torch.utils.data import IterableDataset, DataLoader from tokenizers import Tokenizer from tqdm import tqdm # If these are in a local repo, keep your path append — but do it once, at top. sys.path.append("/home/christian/code/neon/sunoDiff") # Replace with your actual audio module imports # from your_audio_lib import Audio, OpusFile from helpers import ( dist_barrier, download_s3_file, get_filename, read_jsonl, write_jsonl, print_with_time_master, respell_random_words_in_text, ) # test some of the local data CODEC_FILEPATH = "s3://suno-data/minz/models/dac_vae_tuned_25hz.pth" SEMANTIC_FILEPATH = "s3://suno-data/georg/models/semantic/mert_25.pt" SEMANTIC_CLUSTERS_FILEPATH = "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy" from suno_utils.tasks.dac_vae_fixed_25hz import ( preload_models as preload_codec_models, decode as codec_decode, encode as codec_encode, ) from suno_utils.tasks.mert_25 import ( preload_models as preload_semantic_models, encode as semantic_encode, ) # ----------------------------- # Text prep # ----------------------------- def augment_text_training(tags: List[str], text: str) -> str: """Combine tags and text for training; tags may be empty.""" tags = tags or [] text = text or "" return (",".join(tags) + " " + text).strip() def prepare_text_inference(tags: List[str], text: str) -> str: """Combine tags and text for inference; tags may be empty.""" tags = tags or [] text = text or "" return (",".join(tags) + " " + text).strip() # ----------------------------- # Tokenizer # ----------------------------- @contextmanager def _download_from_s3_if_needed(maybe_s3_filepath: str): """If path is S3, download to a temp file and yield local path.""" temp_dir = None tmp_filepath = maybe_s3_filepath if maybe_s3_filepath.startswith("s3://"): temp_dir = tempfile.TemporaryDirectory() filename = get_filename(maybe_s3_filepath, keep_ext=True) tmp_filepath = os.path.join(temp_dir.name, filename) download_s3_file(maybe_s3_filepath, tmp_filepath) try: yield tmp_filepath finally: if temp_dir is not None: temp_dir.cleanup() def load_tokenizer( tokenizer_filepath: str = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json", pad_token: str = "[PAD]", extra_specials: Optional[List[str]] = None, ) -> Tokenizer: """Load and minimally configure a Tokenizer.""" extra_specials = extra_specials or ["\n"] with _download_from_s3_if_needed(tokenizer_filepath) as tmp_fp: tok = Tokenizer.from_file(tmp_fp) tok.add_special_tokens(extra_specials) # Ensure PAD exists; fall back to adding if missing pad_id = tok.token_to_id(pad_token) if pad_id is None: tok.add_special_tokens([pad_token]) pad_id = tok.token_to_id(pad_token) # Attach convenient attributes tok.pad_token = pad_token tok.pad_idx = pad_id return tok # ----------------------------- # Audio slicing + meta handling # ----------------------------- SAMPLE_RATE = 48_000 def _resample_to_mert(arr): assert arr.ndim == 2 assert arr.shape[0] == 2 out_arr = soxr.resample(arr.T, 48_000, 24_000).mean(axis=1).astype(np.float32) audio = Audio.from_array_float(out_arr, sample_rate=24000, max_allowed_val=12) return audio def _choose_window( duration_s: float, chunk_s: float, ctx_s: float, use_ctx: bool, ) -> Tuple[float, float, float, float, bool]: """ Choose a random window from the audio file. Returns: full_start, full_end, target_start, target_end, use_ctx """ total_needed = chunk_s + ctx_s has_enough_for_context = duration_s >= total_needed # Only use context if requested AND we have enough audio use_ctx = use_ctx and has_enough_for_context if use_ctx: # Use both target and context: read 60s, context first 30s, target last 30s max_start = max(0.0, duration_s - total_needed) full_start = random.uniform(0.0, max_start) full_end = full_start + total_needed target_start = full_start + ctx_s target_end = full_start + total_needed else: # Use only target: read 30s max_start = max(0.0, duration_s - chunk_s) full_start = random.uniform(0.0, max_start) full_end = full_start + chunk_s target_start = full_start target_end = full_start + chunk_s return full_start, full_end, target_start, target_end, use_ctx def _extract_audio_segments( path: str, duration_s: float, audio_chunk_s: float, audio_ctx_s: float, ) -> Tuple["Audio", Optional["Audio"], float, float]: """Extract target and context audio segments from the file. Returns: audio_target, audio_ctx, target_start, target_end """ def read_opus_window(start_s: float, total_len_s: float) -> "Audio": buf_size = int(SAMPLE_RATE * total_len_s) float_arr = OpusFile(path=path).read( buf_size=buf_size, float_samples=True, from_position=int(SAMPLE_RATE * max(0.0, start_s)), ) return Audio.from_array_float( float_arr, sample_rate=SAMPLE_RATE, max_allowed_val=12 ) # Handle short files if duration_s < audio_chunk_s: window = read_opus_window(0.0, duration_s) audio_target = window.pad_to_length(audio_chunk_s) return audio_target, None, 0.0, duration_s # Choose window and extract segments use_ctx = random.random() < 0.5 full_start, full_end, target_start, target_end, use_ctx = _choose_window( duration_s=duration_s, chunk_s=audio_chunk_s, ctx_s=audio_ctx_s, use_ctx=use_ctx, ) window = read_opus_window(full_start, full_end - full_start) if use_ctx: # When using context: first 30s is context, last 30s is target audio_ctx = window.get_segment(from_s=0.0, to_s=audio_ctx_s) audio_target = window.get_segment( from_s=audio_ctx_s, to_s=audio_ctx_s + audio_chunk_s ).pad_to_length(audio_chunk_s) else: # When not using context: only target audio (first 30s of window) audio_ctx = None audio_target = window.get_segment(from_s=0.0, to_s=audio_chunk_s).pad_to_length( audio_chunk_s ) return audio_target, audio_ctx, target_start, target_end def _extract_lyrics_for_window( text_aligned: List[Tuple[float, float, str]], window_start: float, window_end: float, ) -> str: """ Extract lyrics that fall within the given time window. Args: text_aligned: list of (start_time, end_time, text) tuples window_start: start time of the window window_end: end time of the window Returns: Concatenated text from lyrics fully contained within the window """ if not text_aligned: return "" parts = [ text for (s, e, text) in text_aligned if (s >= window_start) and (e <= window_end) ] return "".join(parts).strip() def _process_text( meta: Dict[str, Any], target_start: float, target_end: float, tokenizer: Tokenizer, cond_text_len: int, is_training: bool, text_drop_prob: float, ) -> torch.Tensor: """Process text data into tokenized tensor.""" # Extract text for the target window tags = meta.get("tags") or [] if "text_aligned" in meta and meta["text_aligned"]: text = _extract_lyrics_for_window( meta["text_aligned"], target_start, target_end ) else: text = "" # Combine tags and text if is_training: full_text = ( "" if random.random() <= text_drop_prob else augment_text_training(tags, text) ) else: full_text = prepare_text_inference(tags, text) # Tokenize and pad/truncate ids = tokenizer.encode(full_text).ids[:cond_text_len] if tokenizer.pad_idx is None: raise ValueError("Tokenizer must define a pad_idx.") if len(ids) < cond_text_len: ids += [tokenizer.pad_idx] * (cond_text_len - len(ids)) return torch.tensor(ids, dtype=torch.long) def load_meta( meta: Dict[str, Any], tokenizer: Tokenizer, cond_text_len: int, audio_chunk_s: float, audio_ctx_s: float, is_training: bool, text_drop_prob: float = 0.1, ): """Load audio and text data from a meta entry. Assumes meta['local_filepath'] is an OPUS file. Guarantees: - audio_target is exactly `audio_chunk_s` seconds (pads with silence if needed) - If the file is shorter than `audio_chunk_s`, audio_ctx is None """ duration_s = float(meta["duration_s"]) path = meta["local_filepath"] # Extract audio segments audio_target, audio_ctx, target_start, target_end = _extract_audio_segments( path=path, duration_s=duration_s, audio_chunk_s=audio_chunk_s, audio_ctx_s=audio_ctx_s, ) # Process text data text_codes = _process_text( meta=meta, target_start=target_start, target_end=target_end, tokenizer=tokenizer, cond_text_len=cond_text_len, is_training=is_training, text_drop_prob=text_drop_prob, ) # Create 24kHz mono version for semantic encoding audio_target_24k = _resample_to_mert(audio_target.array_float) return audio_target, audio_ctx, text_codes, audio_target_24k # ----------------------------- # Dataset & collate # ----------------------------- class DynamicDataset(IterableDataset): def __init__( self, metas: Iterable[Dict[str, Any]], *, audio_chunk_s: float = 30.0, audio_ctx_s: float = 30.0, cond_text_len: int = 1536, text_drop_prob: float = 0.1, is_training: bool = False, tokenizer: Optional[Tokenizer] = None, tokenizer_path: str = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json", ): self.metas = list(metas) self.audio_chunk_s = float(audio_chunk_s) self.audio_ctx_s = float(audio_ctx_s) self.cond_text_len = int(cond_text_len) self.is_training = bool(is_training) self.text_drop_prob = float(text_drop_prob) self.tokenizer = tokenizer or load_tokenizer(tokenizer_path) def __iter__(self): # IterableDataset must not rely on shuffle=True in DataLoader. for meta in self.metas: try: audio_target, audio_ctx, text_codes, audio_target_24k = load_meta( meta=meta, tokenizer=self.tokenizer, cond_text_len=self.cond_text_len, audio_chunk_s=self.audio_chunk_s, audio_ctx_s=self.audio_ctx_s, is_training=self.is_training, text_drop_prob=self.text_drop_prob, ) yield audio_target, audio_ctx, text_codes, audio_target_24k except Exception as e: print(f"[WARN] Error loading meta {meta.get('id', '')}: {e}") continue def collate_fn(batch: List[Tuple[Any, Optional[Any], torch.Tensor, Any]]): audio_target_list = [item[0] for item in batch] audio_ctx_list = [item[1] for item in batch] text_codes_list = [item[2] for item in batch] audio_target_24k_list = [item[3] for item in batch] # If your encoder expects tensors, you can add padding/stacking here instead. return audio_target_list, audio_ctx_list, text_codes_list, audio_target_24k_list def _as_tensor_batch(x): """Convert codec output to a torch tensor (batch-first).""" if isinstance(x, torch.Tensor): return x.detach().cpu() try: import numpy as np if isinstance(x, np.ndarray): return torch.from_numpy(x) except Exception: pass if isinstance(x, (list, tuple)): elems = [_as_tensor_batch(e) for e in x] return torch.stack(elems, dim=0) return torch.as_tensor(x) def benchmark( metas: Iterable[Dict[str, Any]], *, batch_size: int = 2, num_workers: int = 1, num_eval_batches: int = 100, codec_encode_fn=None, audio_chunk_s: float = 30.02, audio_ctx_s: float = 30.02, cond_text_len: int = 1536, is_training: bool = False, text_drop_prob: float = 0.1, n_vae_tokens: int = 750, vae_dim: int = 128, ): """ codec_encode_fn: function(audio_list, normalize_volume=False) -> (B, T, D) Expected output shape: (B, n_vae_tokens, vae_dim) Builds: - vae_target: (B, n_vae_tokens, vae_dim) - audio_ctx_vae: (B, n_vae_tokens, vae_dim) (zeros where ctx=None) - audio_ctx_mask: (B, 1, 1) (1.0 where ctx exists, 0.0 otherwise) """ assert codec_encode_fn is not None, "Please pass codec_encode_fn=your_encoder" dataset = DynamicDataset( metas=metas, audio_chunk_s=audio_chunk_s, audio_ctx_s=audio_ctx_s, cond_text_len=cond_text_len, is_training=is_training, text_drop_prob=text_drop_prob, ) loader = DataLoader( dataset, batch_size=batch_size, shuffle=False, # IterableDataset must not be shuffled here collate_fn=collate_fn, num_workers=num_workers, pin_memory=False, ) seen = 0 t0 = time.time() for i, ( audio_target_list, audio_ctx_list, text_codes_list, audio_target_24k_list, ) in enumerate(tqdm(loader, desc="Processing batches")): # ---------------------------- # Encode audio targets (always present) # ---------------------------- vae_target = _as_tensor_batch( codec_encode_fn(audio_target_list, normalize_volume=False) ) # Encode semantic targets (24kHz mono audio) semantic_codes = _as_tensor_batch(semantic_encode(audio_target_24k_list)) if vae_target.ndim != 3: raise ValueError( f"codec_encode(target) must return (B, T, D); got {vae_target.shape}" ) B, T, D = vae_target.shape if T != n_vae_tokens or D != vae_dim: raise ValueError( f"codec_encode(target) shape mismatch: expected (B, {n_vae_tokens}, {vae_dim}), " f"got (B, {T}, {D})" ) device, dtype = vae_target.device, vae_target.dtype # ---------------------------- # Encode audio context (may be None) # ---------------------------- audio_ctx_vae = torch.zeros( (B, n_vae_tokens, vae_dim), dtype=dtype, device=device ) audio_ctx_mask = torch.zeros((B, n_vae_tokens, 1), dtype=dtype, device=device) ctx_indices = [idx for idx, a in enumerate(audio_ctx_list) if a is not None] if ctx_indices: ctx_batch = [audio_ctx_list[idx] for idx in ctx_indices] encoded_ctx = _as_tensor_batch( codec_encode_fn(ctx_batch, normalize_volume=False) ) if encoded_ctx.ndim != 3: raise ValueError( f"codec_encode(ctx) must return (b, T, D); got {encoded_ctx.shape}" ) b_ctx, T_ctx, D_ctx = encoded_ctx.shape if T_ctx != n_vae_tokens or D_ctx != vae_dim: raise ValueError( f"codec_encode(ctx) shape mismatch: expected (b, {n_vae_tokens}, {vae_dim}), " f"got (b, {T_ctx}, {D_ctx})" ) # Scatter ctx results into full batch tensors for k, global_idx in enumerate(ctx_indices): audio_ctx_vae[global_idx] = encoded_ctx[k] audio_ctx_mask[global_idx] = 1.0 # ---------------------------- # Here you can run downstream model / profiling / etc. # ---------------------------- # Example: dummy use so these aren't optimized out _ = ( vae_target, audio_ctx_vae, audio_ctx_mask, text_codes_list, semantic_codes, ) # special case to decode audio print( vae_target.shape, audio_ctx_vae.shape, audio_ctx_mask.shape, semantic_codes.shape, ) # vae_audio_decoded = codec_decode(vae_target[0]) # audio_ctx_vae_decoded = codec_decode(audio_ctx_vae[0]) # combine the audio combined_vae = torch.cat([vae_target, audio_ctx_vae], dim=1) for n in range(combined_vae.shape[0]): combined_vae_decoded = codec_decode(combined_vae[n]) combined_vae_decoded.write_hq_mp3(f"combined_vae_decoded_{n}.mp3") seen += 1 if seen >= num_eval_batches: break # ---------------------------- # Timing summary # ---------------------------- t1 = time.time() elapsed = t1 - t0 time_per_batch = elapsed / max(1, seen) chunks_per_batch = batch_size chunks_per_second = ( chunks_per_batch / time_per_batch if time_per_batch > 0 else float("inf") ) print(f"Time taken: {elapsed:.3f}s for {seen} batches") print(f"Time per batch: {time_per_batch:.4f}s") print(f"~Chunks/sec (30s targets): {chunks_per_second:.2f}") def main(): print("Loading models...") _ = preload_codec_models(CODEC_FILEPATH) _ = preload_semantic_models(SEMANTIC_FILEPATH, SEMANTIC_CLUSTERS_FILEPATH) print("Loading metas...") filepath = "/app2/suno/data/auk_v0/metas_v8_tr_mini.jsonl" metas = read_jsonl(filepath) print(len(metas)) print("Benchmarking...") benchmark(metas, codec_encode_fn=codec_encode) if __name__ == "__main__": main()