# -*- coding: utf-8 -*- """ Benchmark: how many 30s chunks/second can we load & encode """ from __future__ import annotations import os import re import time import soxr import math 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 from helpers import ( download_s3_file, get_filename, read_jsonl, print_with_time_master, ) # ----------------------------- # Text prep # ----------------------------- def _space_repl(m): s = m.group() n_newline = s.count("\n") if n_newline >= 2: return "\n\n" elif n_newline == 1: return "\n" return " " def _simplify_whitespace(text, retain_newlines=True): """simplify while respecting up to 2 newlines""" if retain_newlines: text = re.sub(r"\s+", _space_repl, text).strip() else: text = re.sub(r"\s+", " ", text).strip() return text CASE_AUGMENT_FUNCS = [ str.upper, str.lower, str.capitalize, str.title, ] def _augment_tag(s): # case augment if random.random() >= 0.8: s = random.choice(CASE_AUGMENT_FUNCS)(s) # other misc formatting if random.random() >= 0.5: s = s.replace("-", " ").strip() return s def _clean_tag(s, retain_newlines=False): s = s.replace("[", " ").replace("]", " ") return _simplify_whitespace(s, retain_newlines=retain_newlines) MAX_TAG_LEN = 500 MAX_TOT_TAGS_LEN = 1000 def augment_text_training(tags: List[str], lyrics: str): # Clean tags and remove empty ones tags = [clean_tag for tag in tags if len(clean_tag := _clean_tag(tag, retain_newlines=False)) > 0] # Randomly shuffle and truncate tags 50% of the time random.shuffle(tags) if random.random() < 0.5: if len(tags) > 0: tags = tags[: random.randint(1, len(tags))] # Augment each tag tags = [_augment_tag(tag) for tag in tags] # Choose a random character to join tags merge_char = random.choice([", ", " ", "; ", ",", ";"]) # Join tags, truncate if necessary tags_str = f"{merge_char.join([tag[:MAX_TAG_LEN] for tag in tags])[:MAX_TOT_TAGS_LEN]}".strip() # Simplify whitespace in lyrics lyrics = _simplify_whitespace(lyrics, retain_newlines=True) # 5% chance to lowercase lyrics if random.random() < 0.05: lyrics = lyrics.lower() # 5% chance to remove all newlines from lyrics if random.random() < 0.05: lyrics = re.sub(r"\n+", " ", lyrics) lyrics = lyrics.strip() # Choose a random character to join tags and lyrics merge_char = random.choice(["\n\n", "\n", " "]) # Combine tags and lyrics text = "" if len(tags_str) > 0: text += f"[{tags_str}]" if len(lyrics) > 0: if len(text) > 0: merge_char = random.choice(["\n\n", "\n", " "]) text += merge_char text += lyrics return text def prepare_text_inference(tags: List[str], lyrics: str): tags = [clean_tag for tag in tags if len(clean_tag := _clean_tag(tag, retain_newlines=False)) > 0] tags_str = f"{', '.join([tag[:MAX_TAG_LEN] for tag in tags])[:MAX_TOT_TAGS_LEN]}".strip() lyrics = _simplify_whitespace(lyrics, retain_newlines=True) # Combine tags and lyrics text = "" if len(tags_str) > 0: text += f"[{tags_str}]" if len(lyrics) > 0: if len(text) > 0: text += "\n\n" text += lyrics return text # ----------------------------- # 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 # Padding added to audio segments before encoding (cropped after encoding) CROP_PADDING_S = 1.0 # seconds of padding on each side def _normalize_volume(arr, sample_rate, target_db=-16, safe_peak_norm=True, silence_thresh=0.01): assert isinstance(arr, torch.Tensor) assert len(arr.shape) == 2 if arr.abs().max() < silence_thresh or arr.shape[-1] < sample_rate / 2: return arr import pyloudnorm as pyln try: m = pyln.Meter(sample_rate) # create BS.1770 meter lufs_db = m.integrated_loudness(arr.T.cpu().numpy()) # expects (T, ch) gain_factor = np.log(10) / 20 # convert db to amplitude gain = target_db - lufs_db gain = np.exp(gain * gain_factor) assert np.isfinite(gain) and gain > 0 norm_arr = arr * gain if safe_peak_norm and np.abs(norm_arr).max() > 1: norm_arr = norm_arr / np.abs(norm_arr).max() except: return arr return norm_arr def _resample_to_mert(arr, in_sample_rate=48_000, out_sample_rate=24_000): assert arr.ndim == 2 assert arr.shape[0] == 2 out_arr = soxr.resample(arr.T, in_sample_rate, out_sample_rate).mean(axis=1).astype(np.float32) audio = Audio.from_array_float(out_arr, sample_rate=out_sample_rate, 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. With padding, we need extra space: - Context needs CROP_PADDING_S on left and right - Target needs CROP_PADDING_S on left and right - The padding overlaps between ctx and target, so we need: chunk_s + ctx_s + 2 * CROP_PADDING_S total Returns: full_start, full_end, target_start, target_end, use_ctx """ total_needed = chunk_s + ctx_s # With padding, we need extra space: 2 * CROP_PADDING_S (one for left of ctx, one for right of target) total_needed_with_padding = total_needed + 2.0 * CROP_PADDING_S has_enough_for_context = duration_s >= total_needed_with_padding # Only use context if requested AND we have enough audio (including padding) use_ctx = use_ctx and has_enough_for_context if use_ctx: # Use both target and context: context first, then target # We need: full_start >= CROP_PADDING_S (room for padding on left) # We need: target_end + CROP_PADDING_S <= duration_s (room for padding on right) # target_end = full_start + total_needed, so: full_start <= duration_s - total_needed - CROP_PADDING_S max_start = duration_s - total_needed - CROP_PADDING_S # Ensure we have at least CROP_PADDING_S on the left min_start = CROP_PADDING_S if max_start < min_start: # Not enough room, fall back to no context use_ctx = False max_start = max(0.0, duration_s - chunk_s - 2.0 * CROP_PADDING_S) min_start = CROP_PADDING_S if duration_s >= chunk_s + 2.0 * CROP_PADDING_S else 0.0 full_start = random.uniform(min_start, max(max_start, min_start)) target_start = full_start target_end = full_start + chunk_s else: full_start = random.uniform(min_start, max_start) full_end = full_start + total_needed target_start = full_start + ctx_s target_end = full_start + total_needed else: # Use only target: need CROP_PADDING_S on each side total_needed_with_padding = chunk_s + 2.0 * CROP_PADDING_S has_enough = duration_s >= total_needed_with_padding if not has_enough: # If file is too short even with padding, just use what we have full_start = 0.0 target_start = 0.0 else: max_start = duration_s - chunk_s - CROP_PADDING_S min_start = CROP_PADDING_S full_start = random.uniform(min_start, max(max_start, min_start)) target_start = full_start full_end = full_start + chunk_s target_end = full_start + chunk_s return full_start, full_end, target_start, target_end, use_ctx def _remove_silence_frames( audio: np.ndarray, sr: int, *, win_ms: float = 25.0, # frame length (ms) hop_ms: float = 10.0, # frame hop (ms) energy_threshold: float = 1e-4, # threshold on mean squared amplitude per frame pad_ms: float = 20.0, # context to keep around voiced (handled as frame dilation) min_speech_ms: float = 50.0, # drop super-short voiced runs min_silence_ms: float = 30.0, # fill short gaps normalize: bool = False, ): """ Fast CPU silence removal using short-time energy on framed audio (strided view). Works entirely at frame resolution for speed; only concatenation touches segments. Returns concatenated voiced audio, a sample-level mask, and sample segments. """ if audio.ndim != 1: raise ValueError("Pass mono 1D audio.") if audio.size == 0: return audio, np.zeros(0, dtype=bool), np.zeros((0, 2), dtype=np.int64) x = audio.astype(np.float32, copy=False) N = x.size # ---- Framing via strided view (no copies) ---- win = max(1, int(round(sr * win_ms / 1000.0))) hop = max(1, int(round(sr * hop_ms / 1000.0))) if N < win: # pad to at least one full frame pad_right = win - N else: # ensure last frame starts on a hop grid and fully fits n_hops = int(np.ceil((N - win) / hop)) if (N - win) > 0 else 0 total = win + n_hops * hop pad_right = max(0, total - N) if pad_right: x_pad = np.pad(x, (0, pad_right), mode="constant") else: x_pad = x Np = x_pad.size n_frames = 1 + (Np - win) // hop # Build strided 2D view: shape (n_frames, win) itemsize = x_pad.strides[0] frames = np.lib.stride_tricks.as_strided( x_pad, shape=(n_frames, win), strides=(hop * itemsize, itemsize), writeable=False, ) # ---- Frame energy (vectorized) ---- # mean squared amplitude per frame frame_energy = (frames * frames).mean(axis=1) # ---- Threshold -> initial voiced mask (frames) ---- v = frame_energy >= float(energy_threshold) # ---- Morphology on frames (vectorized) ---- # dilation radius from pad_ms pad_frames = int(np.round(pad_ms / hop_ms)) if pad_frames > 0: # binary dilation via 1D convolution on int mask kernel = np.ones(2 * pad_frames + 1, dtype=np.int16) v = np.convolve(v.astype(np.int16), kernel, mode="same") > 0 # Remove very short voiced runs (min_speech) min_speech_frames = int(np.ceil(min_speech_ms / hop_ms)) if min_speech_frames > 1: changes = np.diff(v.astype(np.int8), prepend=0, append=0) starts = np.flatnonzero(changes == 1) ends = np.flatnonzero(changes == -1) # zero out short runs (loop over run count, not samples) short = (ends - starts) < min_speech_frames for s, e in zip(starts[short], ends[short]): v[s:e] = False # Fill short silence gaps between voiced (min_silence) min_silence_frames = int(np.ceil(min_silence_ms / hop_ms)) if min_silence_frames > 1: inv = ~v changes = np.diff(inv.astype(np.int8), prepend=0, append=0) s_starts = np.flatnonzero(changes == 1) s_ends = np.flatnonzero(changes == -1) short_gap = (s_ends - s_starts) < min_silence_frames for s, e in zip(s_starts[short_gap], s_ends[short_gap]): v[s:e] = True # ---- Convert voiced frame runs -> sample segments ---- changes = np.diff(v.astype(np.int8), prepend=0, append=0) f_starts = np.flatnonzero(changes == 1) f_ends = np.flatnonzero(changes == -1) if f_starts.size == 0: return np.zeros(0, dtype=np.float32), np.zeros(N, dtype=bool), np.zeros((0, 2), dtype=np.int64) # Map frame indices to sample indices; include full frame span segs = np.stack( [ f_starts * hop, np.minimum(N, f_ends * hop + win), # cap to original N ], axis=1, ).astype(np.int64) # ---- Build sample mask & concatenate voiced audio ---- mask = np.zeros(N, dtype=bool) for s, e in segs: mask[s:e] = True trimmed = x[mask] if normalize and trimmed.size: peak = np.max(np.abs(trimmed)) if peak > 0: trimmed = trimmed / peak return trimmed, mask, segs def _fast_trim_mono( x: np.ndarray, # shape: (samples,), float32/64 in [-1, 1] sr: int, # sample rate (Hz) thresh_db_rel: float = -35, # keep where RMS > max_RMS + thresh (dB) win_ms: float = 20.0, # moving RMS window size (ms) pad_ms: float = 20.0, # pad around kept regions (ms) min_keep_ms: float = 40.0, # drop kept bits shorter than this (ms) ) -> Tuple[np.ndarray, List[Tuple[int, int]]]: """ Ultra-fast silence trimmer for mono audio. No convolutions, all O(n). Returns (trimmed_audio, kept_spans) with kept_spans in original sample indices. """ assert x.ndim == 1, "Expected mono waveform of shape (samples,)" n = x.size if n == 0: return x[:0], [] # --- Moving RMS via cumulative sums (box filter), O(n) --- # Compute moving average of power over a window, then sqrt. win = max(1, int(round(sr * win_ms / 1000.0))) if win > n: win = n # power and cumulative sum (use float64 for numeric safety) sq = x.astype(np.float64) ** 2 csum = np.empty(n + 1, dtype=np.float64) csum[0] = 0.0 np.cumsum(sq, out=csum[1:]) # csum[k] = sum_{i thresh_db_rel # True = keep # --- Turn mask into spans, expand by pad, merge, drop short --- pad = max(0, int(round(sr * pad_ms / 1000.0))) min_keep = max(1, int(round(sr * min_keep_ms / 1000.0))) # Find rising/falling edges m = mask.astype(np.int8) edges = np.flatnonzero(np.diff(m, prepend=0, append=0)) # edges come in pairs [start0, end0, start1, end1, ...] starts = edges[::2] ends = edges[1::2] if starts.size == 0: return x[:0], [] # Expand by pad and clamp starts = np.maximum(0, starts - pad) ends = np.minimum(n, ends + pad) # Merge overlaps and drop short spans spans: List[Tuple[int, int]] = [] s_prev = int(starts[0]) e_prev = int(ends[0]) for s, e in zip(starts[1:], ends[1:]): s = int(s) e = int(e) if s <= e_prev: # overlap/adjacent -> merge e_prev = max(e_prev, e) else: if (e_prev - s_prev) >= min_keep: spans.append((s_prev, e_prev)) s_prev, e_prev = s, e # last span if (e_prev - s_prev) >= min_keep: spans.append((s_prev, e_prev)) if not spans: return x[:0], [] # --- Concatenate kept spans (one pass) --- parts = [x[a:b] for (a, b) in spans] y = np.concatenate(parts, axis=0).astype(x.dtype) return y, spans def _read_opus_window(path: str, start_s: float, total_len_s: float) -> "Audio": SAMPLE_RATE = 48000 buf_size = int(SAMPLE_RATE * total_len_s) float_arr = OpusFile(path=path).read( buf_size=buf_size, float_samples=True, from_position=int(max(0.0, start_s) * SAMPLE_RATE), ) return Audio.from_array_float(float_arr, sample_rate=SAMPLE_RATE, max_allowed_val=12) def _opus_duration_seconds(path: str) -> Optional[float]: """ Return the duration of an Opus file in seconds without decoding it. Uses libopusfile's op_pcm_total() which returns total PCM samples (at a fixed 48kHz sample rate for all Opus streams). Returns: float: duration in seconds, or None if duration can't be determined """ try: f = OpusFile(path=path) total_samples = f.pcm_total(-1) # -1 = across all links if total_samples < 0: return None return float(total_samples) / 48000 except Exception as e: print(f"Failed to read duration for {path}: {e}") return None finally: try: del f except NameError: pass def _extract_vox_segments( vox_path: str, vox_cond_duration_s: float = 30.0, min_segment_duration_s: float = 3.0, min_segments: int = 1, max_segments: int = 4, ) -> "Audio": """ Sample random vocal segments from an Opus file and return a single mono Audio made by concatenating all sampled segments in the order they were drawn. Assumptions/behavior: - segments may start anywhere (overlap allowed) and may be repeated - total concatenated duration <= vox_cond_duration_s - segment lengths are uniform in [min_segment_duration_s, per-segment-allowed-max] - no fades, no loudness changes, downmix to mono """ # assert the file is opus assert vox_path.endswith(".opus"), "File must be an Opus file" # --- 1) duration sanity --- vox_duration_s = _opus_duration_seconds(vox_path) if vox_duration_s is None or vox_duration_s <= 0: raise ValueError(f"Cannot determine duration for {vox_path}") # A segment can't be longer than the file itself seg_len_upper_bound = max(0.0, vox_duration_s) # If the file is shorter than requested min segment, cap the min at file length eff_min_seg = min(min_segment_duration_s, seg_len_upper_bound) if eff_min_seg <= 0: # zero-length or invalid file raise ValueError(f"Audio too short for segmenting: duration={vox_duration_s:.6f}s") # Ensure we can fit at least one segment under the total cap if vox_cond_duration_s < eff_min_seg: # Fall back to a single segment whose length is capped by both the file and the budget num_segments = 1 target_total = min(vox_cond_duration_s, seg_len_upper_bound) else: # --- 2) choose number of segments subject to total budget --- max_by_budget = int(math.floor(vox_cond_duration_s / eff_min_seg)) if max_by_budget < 1: max_by_budget = 1 num_segments = random.randint(min_segments, max_segments) num_segments = max(1, min(num_segments, max_by_budget)) target_total = vox_cond_duration_s # --- 3) draw segment lengths uniformly with a running budget --- # We keep a running "remaining" budget, but each individual segment length # is also clipped to the file's duration. lengths = [] remaining = target_total for i in range(num_segments): # Reserve minima for the remaining slots remaining_slots = num_segments - i - 1 min_needed_for_rest = remaining_slots * eff_min_seg max_len_this = min(seg_len_upper_bound, max(eff_min_seg, remaining - min_needed_for_rest)) min_len_this = eff_min_seg if max_len_this < min_len_this: # This can happen due to float rounding; clamp max_len_this = min_len_this L = random.uniform(min_len_this, max_len_this) L = max(eff_min_seg, min(L, seg_len_upper_bound)) lengths.append(L) remaining -= L # --- 4) open once; sample windows and downmix to mono; concatenate --- # Reuse a single OpusFile and seek for each segment for speed. of = OpusFile(path=vox_path) mono_chunks = [] for L in lengths: # start ∈ [0, vox_duration - L] max_start = max(0.0, vox_duration_s - L) start_s = 0.0 if max_start <= 0 else random.uniform(0.0, max_start) n_samps = int(round(L * 48000)) start_pcm = int(round(start_s * 48000)) # Read stereo float and downmix to mono arr = of.read( buf_size=n_samps, float_samples=True, from_position=start_pcm, ) # Expect (2, N) because the wrapper forces stereo in float path; handle mono just in case if arr.ndim == 2: # downmix L/R to mono mono = arr.mean(axis=0) else: mono = arr # already mono # If short due to stream edge, right-pad with zeros to requested length (keeps timing) if mono.shape[0] < n_samps: pad = np.zeros(n_samps - mono.shape[0], dtype=mono.dtype) mono = np.concatenate([mono, pad], axis=0) mono_chunks.append(mono.astype(np.float32, copy=False)) if not mono_chunks: # Should not happen, but guard anyway raise RuntimeError("No segments were produced") mono_all = np.concatenate(mono_chunks, axis=0) # Build final Audio (mono 1D array) audio = Audio.from_array_float(mono_all, sample_rate=48000, max_allowed_val=12) audio = audio.pad_to_length(vox_cond_duration_s) return audio def _extract_vox_segment( vox_path: str, vox_duration_s: float, vox_cond_duration_s: float = 30.0, min_segment_duration_s: float = 3.0, min_segments: int = 1, max_segments: int = 4, sample_rate: int = 48000, ) -> "Audio": """Extract vox segment from the file.""" # read the entire file if vox_path.endswith(".opus"): vox_duration_s = min(vox_duration_s, 8 * 60) vox_window = _read_opus_window(vox_path, 0.0, vox_duration_s) else: vox_window = Audio.from_file(vox_path, n_channels=2) mono_audio_array = (vox_window.array_float[0] + vox_window.array_float[1]) / 2 # strip silence (this is somewhat slow) # trimmed_audio, kept_spans = _fast_trim_mono(mono_audio_array, sample_rate) trimmed_audio, _, _ = _remove_silence_frames(mono_audio_array, sample_rate) # trimmed_audio = mono_audio_array if len(trimmed_audio) == 0: # If no audio after trimming, return empty audio return None # Convert kept_spans from original indices to trimmed audio indices # Since _fast_trim_mono returns trimmed audio, we need to work with the trimmed length trimmed_duration_s = len(trimmed_audio) / sample_rate # select the number of random segments num_segments = random.randint(min_segments, max_segments) # Calculate available duration for segments (min of vox_cond_duration_s and trimmed_duration_s) available_duration_s = min(vox_cond_duration_s, trimmed_duration_s) # Generate segment lengths that sum to available_duration_s # Each segment must be at least min_segment_duration_s segment_lengths = [] remaining_duration = available_duration_s for i in range(num_segments): if i == num_segments - 1: # Last segment gets all remaining duration segment_lengths.append(remaining_duration) else: # Calculate max possible length for this segment # Need to leave room for remaining segments (each needs min_segment_duration_s) max_length = remaining_duration - (num_segments - i - 1) * min_segment_duration_s min_length = min_segment_duration_s if max_length <= min_length: # Not enough duration left, use minimum segment_lengths.append(min_length) else: # Random length between min and max segment_lengths.append(random.uniform(min_length, max_length)) remaining_duration -= segment_lengths[-1] # Sample segments from the trimmed audio segments = [] current_pos = 0 for segment_length in segment_lengths: segment_samples = int(segment_length * sample_rate) # Ensure we don't go beyond the trimmed audio length if current_pos + segment_samples > len(trimmed_audio): segment_samples = len(trimmed_audio) - current_pos if segment_samples > 0: segment = trimmed_audio[current_pos : current_pos + segment_samples] segments.append(segment) current_pos += segment_samples else: break # shuffle the segments random.shuffle(segments) if not segments: # If no valid segments, return empty audio return None # concatenate the segments into one Audio object concatenated_audio = np.concatenate(segments) vox_window = Audio.from_array_float(concatenated_audio, sample_rate=sample_rate, max_allowed_val=12) # get a random segment and then pad to length, # first pick the size of the segment segment_length = random.uniform(min_segment_duration_s, vox_cond_duration_s) segment = vox_window.get_segment(from_s=0.0, to_s=segment_length) final_vox_window = segment.pad_to_length(vox_cond_duration_s) return final_vox_window def _build_artist_vox_mappings( metas: List[Dict], ) -> Tuple[Dict[str, List[int]], Dict[str, Optional[Dict[str, Any]]]]: """ Build mappings from artist IDs to metadata indices and vox stem paths with duration. Args: metas: List of metadata dictionaries Returns: Tuple of (artist_id_to_meta_idx, artist_id_to_vox_paths) - artist_id_to_meta_idx: Maps artist_id to list of meta indices - artist_id_to_vox_paths: Maps artist_id to dict with 'path' and 'duration_s' (or None) """ # Build mapping of artist_ids to meta indices artist_id_to_meta_idx = {} for idx, meta in enumerate(metas): if "artists" in meta and meta["artists"] is not None: if len(meta["artists"]) == 1: # only do if there is a single artist for artist_id in meta["artists"]: if artist_id not in artist_id_to_meta_idx: artist_id_to_meta_idx[artist_id] = [] artist_id_to_meta_idx[artist_id].append(idx) # Build mapping of artist_ids to vox stem paths with duration artist_id_to_vox_paths = {} for artist_id, meta_indices in artist_id_to_meta_idx.items(): for meta_idx in meta_indices: meta = metas[meta_idx] if meta.get("vox_stem") is not None: if artist_id not in artist_id_to_vox_paths: artist_id_to_vox_paths[artist_id] = [] artist_id_to_vox_paths[artist_id].append(meta["vox_stem"]) return artist_id_to_vox_paths def _extract_audio_segments( path: str, duration_s: float, audio_chunk_s: float, audio_ctx_s: float, use_ctx: bool, ) -> Tuple["Audio", Optional["Audio"], float, float]: """Extract target and context audio segments from the file. Extracts audio segments that are CROP_PADDING_S seconds longer on each side to allow for cropping after encoding. The target_start and target_end still represent the actual target window (without the padding) for text alignment. Returns: audio_target, audio_ctx, target_start, target_end """ audio_target_duration_s = audio_chunk_s + 2.0 * CROP_PADDING_S # Handle short files # For short files, we need symmetric padding: CROP_PADDING_S on left, original audio, then pad to full length if duration_s < audio_chunk_s: use_ctx = False window = Audio.from_file(path, n_channels=2) # Create left padding (CROP_PADDING_S seconds of silence) left_padding = Audio.from_silence(CROP_PADDING_S, window.sample_rate, window.n_channels) # Concatenate arrays: left_padding + original_audio left_samples = left_padding.array_float window_samples = window.array_float # Ensure both are 2D (channels, samples) and same channel count if left_samples.ndim == 1: left_samples = left_samples[np.newaxis, :] if window_samples.ndim == 1: window_samples = window_samples[np.newaxis, :] if left_samples.shape[0] != window_samples.shape[0]: # Convert mono to stereo if needed if left_samples.shape[0] == 1: left_samples = np.repeat(left_samples, window_samples.shape[0], axis=0) else: window_samples = np.repeat(window_samples, left_samples.shape[0], axis=0) concatenated_samples = np.concatenate([left_samples, window_samples], axis=1) audio_target = Audio.from_array_float( concatenated_samples, sample_rate=window.sample_rate, max_allowed_val=12, # Standard max_allowed_val for audio ) # Pad the rest on the right to reach the target duration audio_target = audio_target.pad_to_length(audio_target_duration_s) return audio_target, None, 0.0, audio_chunk_s # Choose window and extract segments # We still calculate target_start/target_end for the original chunk size (for text alignment) 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, ) if use_ctx: # When using context: we need to extract a larger window for target audio target_extract_start = max(0.0, target_start - CROP_PADDING_S) target_extract_end = min(duration_s, target_end + CROP_PADDING_S) target_extract_duration = target_extract_end - target_extract_start # Extract context with padding (same as target) ctx_extract_start = max(0.0, full_start - CROP_PADDING_S) ctx_extract_end = min(duration_s, full_start + audio_ctx_s + CROP_PADDING_S) ctx_extract_duration = ctx_extract_end - ctx_extract_start audio_ctx_duration_s = audio_ctx_s + 2.0 * CROP_PADDING_S if path.endswith(".opus"): ctx_window = _read_opus_window(path, ctx_extract_start, ctx_extract_duration) target_window = _read_opus_window(path, target_extract_start, target_extract_duration) else: ctx_window = Audio.from_file(path, n_channels=2).get_segment( from_s=ctx_extract_start, to_s=ctx_extract_end ) target_window = Audio.from_file(path, n_channels=2).get_segment( from_s=target_extract_start, to_s=target_extract_end ) audio_ctx = ctx_window.pad_to_length(audio_ctx_duration_s) audio_target = target_window.pad_to_length(audio_target_duration_s) else: # When not using context: extract target with padding on each side target_extract_start = max(0.0, target_start - CROP_PADDING_S) target_extract_end = min(duration_s, target_end + CROP_PADDING_S) target_extract_duration = target_extract_end - target_extract_start if path.endswith(".opus"): target_window = _read_opus_window(path, target_extract_start, target_extract_duration) else: target_window = Audio.from_file(path, n_channels=2).get_segment( from_s=target_extract_start, to_s=target_extract_end ) audio_ctx = None audio_target = target_window.pad_to_length(audio_target_duration_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, pad_prev_lines: int = 1, # how many lines before the first in-range line to include pad_next_lines: int = 1, # how many lines after the last in-range line to include joiner: str = "", # how to join parts (default keeps your original behavior) ) -> str: """ Extract lyrics fully contained within the given time window, with optional context padding. 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 pad_prev_lines: number of lines immediately preceding the first in-window line to include pad_next_lines: number of lines immediately following the last in-window line to include joiner: string used to join parts ("" preserves your original behavior) Returns: Concatenated text from lyrics fully contained within the window, optionally padded by neighboring lines. """ if not text_aligned: return "" in_range_indices: List[int] = [] for i, (s, e, _text) in enumerate(text_aligned): try: s_f = float(s) e_f = float(e) except (TypeError, ValueError): continue if (s_f >= float(window_start)) and (e_f <= float(window_end)): in_range_indices.append(i) if not in_range_indices: return "" # First and last in-window indices in the ORIGINAL list first_idx = in_range_indices[0] last_idx = in_range_indices[-1] # Compute padded bounds (clamped to list) pad_start = max(0, first_idx - pad_prev_lines) pad_end = min(len(text_aligned) - 1, last_idx + pad_next_lines) # Build final list of indices: [pad before] + [all in-range] + [pad after] # Note: this includes only immediate neighbors as requested; even if in-range # lines are non-contiguous, we *don’t* insert extra out-of-window lines in between. selected_indices: List[int] = [] selected_indices.extend(range(pad_start, first_idx)) selected_indices.extend(in_range_indices) selected_indices.extend(range(last_idx + 1, pad_end + 1)) parts: List[str] = [] for idx in selected_indices: _s, _e, t = text_aligned[idx] if t: parts.append(str(t)) return joiner.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, use_text_aligned_prob: float, text_drop_prob: float, stem_name: Optional[str] = None, ) -> Tuple[torch.Tensor, str]: """Process text data into tokenized tensor.""" text_aligned = meta.get("text_aligned") # Determine whether to use text_aligned, based on availability and probability has_text_aligned = ( text_aligned is not None and isinstance(text_aligned, list) and len(text_aligned) > 0 ) use_text_aligned = has_text_aligned and (random.random() <= use_text_aligned_prob) if use_text_aligned: text = _extract_lyrics_for_window(text_aligned, target_start, target_end) elif meta.get("text") is not None: text = meta["text"] else: text = "" # Extract text for the target window tags = meta.get("tags") or [] if stem_name: tags = tags + [stem_name, stem_name.lower(), "stem:" + stem_name] # append stem name to tags if stem_name.lower() != "vocals": text = "" # drop text if not vocals # check for stem captions stem_captions = meta.get("stems_captions", {}) if stem_captions: vocals_caption = stem_captions.get("Vocals", []) if vocals_caption: for caption_item in vocals_caption: if caption_item.get("prompt_type") == "voice_description_keywords": caption_str = caption_item.get("caption", "") if caption_str: tags.append(caption_str.strip()) break # 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), full_text def load_meta( meta: Dict[str, Any], tokenizer: Tokenizer, cond_text_len: int, audio_chunk_s: float, audio_ctx_s: float, audio_vox_s: float, is_training: bool, text_drop_prob: float = 0.1, use_text_aligned_prob: float = 0.8, target_loudness_db: Optional[float] = None, use_stem_prob: float = 0.0, use_vox_prob: float = 0.0, vox_paths: Optional[Dict[str, Any]] = [], ): """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 """ try: vox_path = None audio_vox = None duration_s = float(meta["duration_s"]) use_ctx = np.random.random() < 0.75 path = meta["local_filepath"] # normalize the audio to the target loudness audio_stats = meta.get("audio_stats", None) gain_db = 0.0 if audio_stats is not None: loudness_db = audio_stats.get("loudness", None) if loudness_db is not None and target_loudness_db is not None: gain_db = target_loudness_db - float(loudness_db) # clamp the gain_db between -12 and 12 gain_db = np.clip(gain_db, -12, 12) if duration_s > 8 * 60: raise ValueError(f"Duration is too long: {duration_s}s") # Validate required fields if not path: raise ValueError("local_filepath is empty or None") if not os.path.exists(path): raise FileNotFoundError(f"Audio file not found: {path}") except Exception as e: import traceback meta_id = meta.get("id", "") print(f"[ERROR] Failed to initialize load_meta for {meta_id}:") print(f" Exception: {type(e).__name__}: {e}") print(f" Traceback: {traceback.format_exc()}") raise # use vox conditioning, no ctx, no infill if len(vox_paths) > 0 and np.random.random() < use_vox_prob: vox_path = np.random.choice(vox_paths) use_ctx = False # use stem conditioning elif meta.get("stems") is not None and len(meta["stems"]) > 0 and np.random.random() < use_stem_prob: stems = meta["stems"] # this is a dict of stem_name -> stem_path # randomly sample a stem path stem_name = random.choice(list(stems.keys())) stem_path = stems[stem_name] if stem_path is not None and stem_path.strip(): path = stem_path # Extract audio segments try: 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, use_ctx=use_ctx, ) except Exception as e: import traceback meta_id = meta.get("id", "") print(f"[ERROR] Failed to extract audio segments for {meta_id}:") print(f" Path: {path}") print(f" Duration: {duration_s}s") print(f" Exception: {type(e).__name__}: {e}") print(f" Traceback: {traceback.format_exc()}") raise # Extract vox segment if vox_path is not None: audio_vox = _extract_vox_segments( vox_path=vox_path, vox_cond_duration_s=audio_vox_s, min_segments=1, max_segments=4, ) # Process text data try: text_codes, raw_text = _process_text( meta=meta, target_start=target_start, target_end=target_end, tokenizer=tokenizer, cond_text_len=cond_text_len, is_training=is_training, use_text_aligned_prob=use_text_aligned_prob, text_drop_prob=text_drop_prob, ) except Exception as e: import traceback meta_id = meta.get("id", "") print(f"[ERROR] Failed to process text for {meta_id}:") print(f" Exception: {type(e).__name__}: {e}") print(f" Traceback: {traceback.format_exc()}") raise # Create 24kHz mono version for semantic encoding # note: this is before headroom is applied to the audio target and ctx try: audio_target_24k = _resample_to_mert( audio_target.array_float, in_sample_rate=audio_target.sample_rate, out_sample_rate=24_000, ) except Exception as e: import traceback meta_id = meta.get("id", "") print(f"[ERROR] Failed to resample audio for {meta_id}:") print(f" Exception: {type(e).__name__}: {e}") print(f" Traceback: {traceback.format_exc()}") raise # apply headroom to the audio target and ctx but not the 24k version audio_target, _ = audio_target.apply_gain(gain_db) if audio_ctx is not None: audio_ctx, _ = audio_ctx.apply_gain(gain_db) # audio_target is padded with CROP_PADDING_S on each side for cropping after encoding expected_target_duration_s = audio_chunk_s + 2.0 * CROP_PADDING_S # Use tolerance for floating point comparison (1ms tolerance) duration_tolerance = 0.001 checks = [ ( abs(audio_target.duration_s - expected_target_duration_s) <= duration_tolerance, f"audio_target.duration_s: {audio_target.duration_s}, expected: {expected_target_duration_s} (audio_chunk_s: {audio_chunk_s} + {2.0 * CROP_PADDING_S} padding), diff: {abs(audio_target.duration_s - expected_target_duration_s)}", ), ( audio_ctx is None or abs(audio_ctx.duration_s - (audio_ctx_s + 2.0 * CROP_PADDING_S)) <= duration_tolerance, f"audio_ctx.duration_s: {getattr(audio_ctx, 'duration_s', None)}, expected: {audio_ctx_s + 2.0 * CROP_PADDING_S} (audio_ctx_s: {audio_ctx_s} + {2.0 * CROP_PADDING_S} padding), diff: {abs(getattr(audio_ctx, 'duration_s', 0) - (audio_ctx_s + 2.0 * CROP_PADDING_S)) if audio_ctx is not None else 0}", ), ( audio_vox is None or audio_vox.duration_s == audio_vox_s, f"audio_vox.duration_s: {getattr(audio_vox, 'duration_s', None)}, audio_vox_s: {audio_vox_s}", ), ( text_codes is None or len(text_codes) == cond_text_len, f"len(text_codes): {len(text_codes) if text_codes is not None else None}, cond_text_len: {cond_text_len}", ), ( audio_target_24k is None or audio_target_24k.duration_s == audio_target.duration_s, f"audio_target_24k.duration_s: {getattr(audio_target_24k, 'duration_s', None)}, audio_target.duration_s: {audio_target.duration_s}", ), ] for cond, msg in checks: assert cond, msg return audio_target, audio_ctx, audio_vox, text_codes, raw_text, audio_target_24k # ----------------------------- # Dataset & collate # ----------------------------- class DynamicDataset(IterableDataset): def __init__( self, metas, # Union[str, Iterable[Dict[str, Any]]] - file path or list of metas *, audio_chunk_s: float = 30.02, audio_ctx_s: float = 30.02, audio_vox_s: float = 30.02, cond_text_len: int = 1536, text_drop_prob: float = 0.1, use_stem_prob: float = 0.0, use_vox_prob: float = 0.0, use_text_aligned_prob: float = 0.8, is_training: bool = False, foreign_weight: float = 0.0, text_aligned_weight: float = 0.0, stem_weight: float = 0.0, tokenizer: Optional[Tokenizer] = None, tokenizer_path: str = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json", target_loudness_db: Optional[float] = None, ): # Load metas from file if string path is provided if isinstance(metas, str): self.metas = read_jsonl(metas) else: self.metas = list(metas) print_with_time_master(f"Loaded {len(self.metas)} metas") # Extract weights and apply language-based upweighting self.weights = [] for meta in self.metas: # Get base weight from meta, default to 0.0 if not present weight = meta.get("weight", 0.0) # note: we want 0 weight for covers # Ensure weight is a valid number try: weight = float(weight) if weight <= 0: weight = 1.0 # Default to 1.0 for invalid weights except (ValueError, TypeError): weight = 1.0 # Default to 1.0 for non-numeric weights # check if the meta has text_aligned has_text_aligned = ( meta.get("text_aligned") is not None and isinstance(meta["text_aligned"], list) and len(meta["text_aligned"]) > 0 ) # check if the meta has stems has_stems = ( meta.get("stems") is not None and isinstance(meta["stems"], dict) and len(meta["stems"]) > 0 ) # check if the meta has vox_stem has_vox_stem = ( meta.get("vox_stem") is not None and os.path.exists(meta["vox_stem"]) and meta["vox_stem"].strip() != "" ) # Check language and upweight non-English by foreign_weight lang = meta.get("lang", None) if lang is not None and lang != "en" and lang != "english": weight += foreign_weight if has_text_aligned: weight += text_aligned_weight if has_vox_stem: weight += stem_weight # cap weight at 10.0 weight = min(weight, 10.0) duration_s = float(meta.get("duration_s", 0.0)) if duration_s > 480 or duration_s < 10: weight = 0.0 # temporarily remove podcast audio if meta["id"].startswith("podcast_"): weight = 0.0 if meta["local_filepath"] is None: weight = 0.0 self.weights.append(weight) print_with_time_master(f"Extracted weights for {len(self.weights)} metas") # Build mappings for artist-based vox conditioning self.artist_id_to_vox_paths = _build_artist_vox_mappings(self.metas) print_with_time_master(f"Found {len(self.artist_id_to_vox_paths)} artists with vox stem paths") self.audio_chunk_s = float(audio_chunk_s) self.audio_ctx_s = float(audio_ctx_s) self.audio_vox_s = float(audio_vox_s) self.cond_text_len = int(cond_text_len) self.is_training = bool(is_training) self.text_drop_prob = float(text_drop_prob) self.use_stem_prob = float(use_stem_prob) self.use_vox_prob = float(use_vox_prob) self.foreign_weight = float(foreign_weight) self.text_aligned_weight = float(text_aligned_weight) self.stem_weight = float(stem_weight) self.target_loudness_db = target_loudness_db self.use_text_aligned_prob = float(use_text_aligned_prob) self.tokenizer = tokenizer or load_tokenizer(tokenizer_path) # Initialize random cache for weighted sampling self.random_cache = [] def _sample_meta(self): """ Randomly sample a meta from the metas list. Weighted sampling is slow for large weighted datasets so we sample 10k at a time. """ if len(self.random_cache) == 0: choices = random.choices(range(len(self.metas)), weights=self.weights, k=10000) self.random_cache.extend(choices) idx = self.random_cache.pop() return self.metas[idx] def __iter__(self): # Use weighted sampling instead of sequential iteration # This allows for proper weighted sampling across the entire dataset while True: # Infinite iteration for training try: # Sample a meta using weighted sampling meta = self._sample_meta() # Determine vox_paths based on the artist_id if available and unique vox_paths = [] artist_ids = meta.get("artists") if artist_ids is not None and len(artist_ids) == 1: artist_id = artist_ids[0] if artist_id in self.artist_id_to_vox_paths: vox_paths = self.artist_id_to_vox_paths[artist_id] # also check if this meta has a vox stem path if meta.get("vox_stem") is not None: vox_paths.append(meta["vox_stem"]) # load the meta audio_target, audio_ctx, audio_vox, text_codes, raw_text, 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, audio_vox_s=self.audio_vox_s, is_training=self.is_training, text_drop_prob=self.text_drop_prob, use_stem_prob=self.use_stem_prob, use_vox_prob=self.use_vox_prob, use_text_aligned_prob=self.use_text_aligned_prob, target_loudness_db=self.target_loudness_db, vox_paths=vox_paths, ) yield audio_target, audio_ctx, audio_vox, text_codes, raw_text, audio_target_24k except Exception as e: import traceback meta_id = meta.get("id", "") if "meta" in locals() else "" error_type = type(e).__name__ error_msg = str(e) if str(e) else "" tb_lines = traceback.format_exc().split("\n") # Get the last few lines of traceback for context relevant_tb = ( "\n".join(tb_lines[-6:-1]) if len(tb_lines) > 6 else "\n".join(tb_lines[:-1]) ) print(f"[WARN] Error loading meta {meta_id}:") print(f" Exception: {error_type}: {error_msg}") print(f" Traceback:\n{relevant_tb}") continue def collate_fn(batch: List[Tuple[Any, Optional[Any], torch.Tensor, Any, str, Any]]): audio_target_list = [item[0] for item in batch] audio_ctx_list = [item[1] for item in batch] audio_vox_list = [item[2] for item in batch] text_codes_list = [item[3] for item in batch] raw_text_list = [item[4] for item in batch] audio_target_24k_list = [item[5] for item in batch] # If your encoder expects tensors, you can add padding/stacking here instead. return ( audio_target_list, audio_ctx_list, audio_vox_list, text_codes_list, raw_text_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, semantic_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" assert semantic_encode_fn is not None, "Please pass semantic_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, loudness_normalize=False, ) 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, prefetch_factor=4, persistent_workers=True, ) seen = 0 t0 = time.time() for i, ( audio_target_list, audio_ctx_list, audio_vox_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_fn(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(): 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, ) 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, semantic_encode_fn=semantic_encode) if __name__ == "__main__": main()