import os import gc import time import json import glob import torch import funcy import IPython import random import tempfile import torchaudio import numpy as np import pyloudnorm as pyln import multiprocessing from tqdm import tqdm from itertools import repeat from dac.model.dac4 import DAC from suno_utils.audio import Audio from suno_utils.utils.s3 import read_from_s3 from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor from suno_utils.utils.text import write_jsonl, read_jsonl # load the semantic model from suno_utils.tasks.mert_25 import ( preload_models as preload_semantic_models, encode as encode_semantic, EMBEDDING_RATE as SEMANTIC_HZ, ) print("loading semantic model...") 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" _ = preload_semantic_models( checkpoint_filepath=semantic_model_filepath, centroids_filepath=semantic_clusters_filepath, device="cuda", ) def corrupt_audio(audio, sr): audio_corrupt = audio.clone() # 10% chance to add highpass if np.random.uniform() < 0.25: hpf_cutoff = np.random.uniform(20, 1000) audio_corrupt = torchaudio.functional.highpass_biquad( audio_corrupt, sr, hpf_cutoff ) # 10% chance to add lowpass if np.random.uniform() < 0.25: lpf_cutoff = np.random.uniform(1000, 20000) audio_corrupt = torchaudio.functional.lowpass_biquad( audio_corrupt, sr, lpf_cutoff ) # 10% chance to make mono if np.random.uniform() < 0.7: audio_corrupt = audio_corrupt.mean(dim=0, keepdim=True) audio_corrupt = audio_corrupt.repeat(2, 1) # 10% chance to add noise if np.random.uniform() < 0.2: noise_level_db = -np.random.uniform(42, 96) noise_signal = torch.randn_like(audio_corrupt) * 10 ** (noise_level_db / 20) # apply filter to noise signal lpf_cutoff = np.random.uniform(1000, 20000) noise_signal = torchaudio.functional.lowpass_biquad( noise_signal, sr, lpf_cutoff ) audio_corrupt = audio_corrupt + noise_signal # 5% chance to add contrast if np.random.uniform() < 0.05: contrast = np.random.uniform(0, 100.0) audio_corrupt = torchaudio.functional.contrast( audio_corrupt, enhancement_amount=contrast ) # 10% add codec artifacts if np.random.uniform() < 1.0: compression = np.random.choice([32, 64, 128]) audio_corrupt = torchaudio.functional.apply_codec( audio_corrupt, sr, "mp3", compression ) # 10% chance to add distortion if np.random.uniform() < 0.3: drive_db = np.random.uniform(6, 24) else: drive_db = 0 drive_lin = 10 ** (drive_db / 20) audio_corrupt *= drive_lin if np.random.uniform() < 0.5: audio_corrupt = torch.tanh(audio_corrupt) else: audio_corrupt = audio_corrupt.clamp(-1, 1) return audio_corrupt def load_audio(meta): """Just load the audio file - I/O bound operation""" try: meta_id = meta["id"] # skip if npz file exists # npz_filepath = os.path.join(out_dir, f"{meta_id}.npz") # if os.path.exists(npz_filepath): # return None start_time = time.time() audio, sr = read_from_s3(meta["audio_filepath"], read_f=torchaudio.load) print(f"loaded {meta['id']} in {time.time() - start_time:.2f} seconds") # crop to max 60s audio = audio[:, : int(48000 * 60.01)] # ensure stereo if audio.shape[0] == 1: audio = audio.repeat(2, 1) elif audio.shape[0] > 2: audio = audio[:2, :] # resample to 48k audio = torchaudio.functional.resample(audio, sr, 48000) # crop to first 30 seconds audio_48k = audio[:, : int(48000 * 30.01)] # Apply highpass filter audio_48k_corrupted = corrupt_audio(audio_48k, 48000) # Normalize for VAE audio_48k_normalized = audio_48k # / audio_48k.abs().max().clamp(min=1e-5) audio_48k_corrupted_normalized = ( audio_48k_corrupted # / audio_48k_corrupted.abs().max().clamp(min=1e-5) ) # Create 24k versions for semantic audio_24k = torchaudio.functional.resample(audio_48k, 48000, 24000) audio_24k_corrupted = torchaudio.functional.resample( audio_48k_corrupted, 48000, 24000 ) audio_24k_mono = audio_24k.mean(dim=0).unsqueeze(0) audio_24k_corrupted_mono = audio_24k_corrupted.mean(dim=0).unsqueeze(0) return { "id": meta["id"], "audio_48k": audio_48k_normalized, "audio_48k_corrupted": audio_48k_corrupted_normalized, "audio_24k_mono": audio_24k_mono, "audio_24k_corrupted_mono": audio_24k_corrupted_mono, } except Exception as e: print(f"Error loading {meta['id']}: {str(e)}") return None def process_batch(batch_data, codec_model, out_dir, num_vae_tokens: int): """Process a batch through models and save results""" if not batch_data: return # Stack for VAE processing original_batch = torch.stack([data["audio_48k"] for data in batch_data]).cuda() corrupted_batch = torch.stack( [data["audio_48k_corrupted"] for data in batch_data] ).cuda() # Process through VAE with torch.no_grad(): vae_original = codec_model.encode(original_batch)["z"].cpu() vae_corrupted = codec_model.encode(corrupted_batch)["z"].cpu() # Clear GPU memory del original_batch, corrupted_batch torch.cuda.empty_cache() # Process semantic codes semantic_original = encode_semantic([d["audio_24k_mono"] for d in batch_data])[0] semantic_corrupted = encode_semantic( [d["audio_24k_corrupted_mono"] for d in batch_data] )[0] # Save results for idx, data in enumerate(batch_data): vae_orig = vae_original[idx, :, :num_vae_tokens] vae_corrupt = vae_corrupted[idx, :, :num_vae_tokens] np.savez( os.path.join(out_dir, f"{data['id']}.npz"), vae_latents_original=vae_orig, vae_latents_corrupted=vae_corrupt, semantic_codes=semantic_original.astype(np.uint16)[:, 0], semantic_codes_corrupted=semantic_corrupted.astype(np.uint16)[:, 0], ) # print(f"saved {len(batch_data)} files") # Clear memory del vae_original, vae_corrupted, semantic_original, semantic_corrupted torch.cuda.empty_cache() def clear_memory(): """Aggressive memory cleanup""" gc.collect() torch.cuda.empty_cache() if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() def main( out_dir: str, num_threads: int = 16, batch_size: int = 8, num_samples: int = 10000, ): VAE_RATE_HZ = 25 # load the vae model device = "cuda:0" # device = "cpu" if VAE_RATE_HZ == 100: checkpoint_filepath = "s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth" # checkpoint_filepath = "/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.pth" num_vae_tokens = 3000 elif VAE_RATE_HZ == 25: checkpoint_filepath = "s3://suno-data/christian/25hz_vae_peaq_kl_0.005.pth" num_vae_tokens = 750 else: raise ValueError(f"VAE_RATE_HZ must be 100 or 25, got {VAE_RATE_HZ}") load_f = funcy.partial(torch.load, map_location="cpu") if checkpoint_filepath.startswith("s3://"): sd = read_from_s3(checkpoint_filepath, read_f=load_f) else: sd = load_f(checkpoint_filepath) sd["metadata"]["kwargs"] = { k: v for k, v in sd["metadata"]["kwargs"].items() if k in DAC.__init__.__code__.co_varnames } codec_model = DAC(**sd["metadata"]["kwargs"]) codec_model.load_state_dict(sd["state_dict"]) codec_model.eval() codec_model.to(device) base_metas = ( "/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl" ) metas = read_jsonl(base_metas) print(f"loaded {len(metas):,} metas") # also load the lyric alignments genius_alignments_filepath = ( "/home/tony/Work/tony/hoot/tmp/genius_hq_alignments_t30_v1.jsonl" ) if os.path.exists(genius_alignments_filepath): print("loading genius alignments") genius_alignments = read_jsonl(genius_alignments_filepath, progress=False) genius_alignments_map = {a[0]: a[1] for a in genius_alignments} else: genius_alignments_map = {} # merge alignments into single map alignments_map = {**genius_alignments_map} print(f"loaded {len(alignments_map):,} alignments") # create metas map metas_map = {meta["id"]: meta for meta in metas} print(f"loaded {len(metas_map):,} metas") # Setup os.makedirs(out_dir, exist_ok=True) random.seed(10) random.shuffle(metas) subset_metas = metas[:num_samples] subset_metas_with_alignments = [] # add alignments to metas for meta in subset_metas: if meta["id"] not in alignments_map: continue meta["alignments"] = alignments_map[meta["id"]] subset_metas_with_alignments.append(meta) # Move model to GPU codec_model.to("cuda:0") # do this in chunks chunk_size = num_threads for i in tqdm(range(0, len(subset_metas_with_alignments), chunk_size)): chunk_metas = subset_metas_with_alignments[i : i + chunk_size] # load the audio files in parallel with ThreadPoolExecutor(max_workers=num_threads) as executor: results = list( executor.map(load_audio, chunk_metas), ) print(f"loaded {len(results):,} audio files") # process the batch results = [r for r in results if r is not None] print(f"processing {len(results):,} audio files") # use batch_size to process in chunks for i in range(0, len(results), batch_size): batch_results = results[i : i + batch_size] process_batch(batch_results, codec_model, out_dir, num_vae_tokens) if __name__ == "__main__": out_dir = "/home/christian/data/dpo/genius_hq_corrupt_dpo_25hz_30_v2_npz" num_samples = 10000 num_threads = 64 batch_size = 8 main(out_dir, num_threads, batch_size, num_samples)