import os import torch import torchaudio import numpy as np import pyloudnorm as pyln from ear.system import EarSystem from suno_utils.audio import Audio from suno_utils.utils.s3 import read_from_s3 import pandas as pd import json from tqdm import tqdm from tqdm.contrib.concurrent import process_map import time import argparse from multiprocessing import Pool # Set CUDA device device = "cuda" # neon/ear # pip install -e . # s3://suno-data/christian/ear/w5p4nhzn-epoch=27.ckpt checkpoint_filepath = "/app/suno/data/dpo/models/ear.ckpt" NUM_FRAMES = 131072 # def load_checkpoint(checkpoint_path): # return EarSystem.load_from_checkpoint(checkpoint_path, map_location="cpu") def load_ear_system(): ear_system = EarSystem.load_from_checkpoint(checkpoint_filepath, map_location="cpu") ear_system.to(device) ear_system.eval() return ear_system def compare_quality(ear_system, audio_a: torch.Tensor, audio_b: torch.Tensor): """Compare the quality of two audio files using the trained model. Parameters ---------- system : EarSystem the trained model audio_a : torch.Tensor audio tensor of shape (2, num_frames) audio_b : torch.Tensor audio tensor of shape (2, num_frames) Returns ------- pref : torch.Tensor preference prediction quant : torch.Tensor quantification prediction """ # move audio_a and audio_b to same device as system parameters audio_a = audio_a.to(ear_system.device) audio_b = audio_b.to(ear_system.device) # first, embed the audio that will be evaluated with torch.no_grad(): embeds_a = ear_system.embed(audio_a) embeds_b = ear_system.embed(audio_b) # aggregate embeddings over time with a moving mean of frame size embeds_a = torch.nn.functional.adaptive_avg_pool1d( embeds_a.permute(0, 2, 1), 137 ).permute(0, 2, 1) embeds_b = torch.nn.functional.adaptive_avg_pool1d( embeds_b.permute(0, 2, 1), 137 ).permute(0, 2, 1) # concat embeds into singular tensors embeds = torch.cat((embeds_a, embeds_b), dim=-1) # no run through the projection to make predictions with torch.no_grad(): pref_preds = ear_system.pref_classifier(embeds) quant_preds = ear_system.quant_classifier(embeds) # print(pref_preds.shape, quant_preds.shape) # get a final score by taking mean across seq of preds and chunks pref_preds = pref_preds.mean(dim=1).mean(dim=0) quant_preds = quant_preds.mean(dim=1).mean(dim=0) pref = torch.sigmoid(pref_preds) quant = torch.argmax(quant_preds, dim=0).float() return pref, quant def prepare_audio( audio: Audio, num_frames: int = NUM_FRAMES, start_s: float = None, end_s: float = None, ): sample_rate = audio.sample_rate audio = torch.from_numpy(audio.array_float) if audio.shape[0] != 2: audio = audio.repeat(2, 1) # crop audio based on metadata example if start_s is not None and end_s is not None: start_frame = int(start_s * sample_rate) end_frame = int(end_s * sample_rate) audio = audio[:, start_frame:end_frame] # if the file is long, only take part of it if audio.shape[-1] > (sample_rate * 120): audio = audio[:, : sample_rate * 120] # downmix and resample decoded audio to 24khz audio = torchaudio.functional.resample(audio, sample_rate, 24_000) # pad by repeating the signal if shorter than window if audio.shape[-1] < num_frames: pad_size = num_frames - audio.shape[-1] audio = torch.nn.functional.pad(audio, (1, pad_size), mode="replicate") # chunk into non-overlapping blocks of num_frames audio_chunks = [] num_chunks = audio.shape[-1] // num_frames for n in range(num_chunks): start_idx = n * num_frames end_idx = start_idx + num_frames audio_chunks.append(audio[:, start_idx:end_idx]) # loudness norm meter = pyln.Meter(24_000) for audio_chunk_idx in range(len(audio_chunks)): x_lufs_db = meter.integrated_loudness(audio.T.numpy()) if x_lufs_db == -float("inf"): gain_lin = 1.0 else: delta_lufs_db = -20.0 - x_lufs_db gain_lin = 10.0 ** (np.clip(delta_lufs_db, a_min=-120, a_max=48.0) / 20.0) audio_chunks[audio_chunk_idx] *= gain_lin # take the last chunk keeping the list # audio_chunks = audio_chunks[-1:] return torch.stack(audio_chunks) def process_request_batch(request_batch): try: ear_system = load_ear_system() results = {} for request_id, negative_id, positive_id in tqdm(request_batch): try: negative_audio = Audio.from_s3( f"s3://suno-data-uploads/studio/uploads/{negative_id}.mp3", n_channels=2, ) positive_audio = Audio.from_s3( f"s3://suno-data-uploads/studio/uploads/{positive_id}.mp3", n_channels=2, ) negative_audio = negative_audio.get_segment(0, 30) positive_audio = positive_audio.get_segment(0, 30) negative_prep = prepare_audio(negative_audio) positive_prep = prepare_audio(positive_audio) pref, quant = compare_quality(ear_system, negative_prep, positive_prep) results[request_id] = ( round(float(pref.cpu().numpy()[0]), 5), round(float(quant.cpu().numpy()), 5), ) except Exception as e: print(f"Error processing request {request_id}: {e}") results[request_id] = None return results except Exception as e: print(f"Error in process_request_batch: {e}") return {} def main(): total_job_n_gpus = 1 parser = argparse.ArgumentParser() parser.add_argument("--job_idx", type=int, default=0) args = parser.parse_args() print(f"CUDA_VISIBLE_DEVICES: {os.environ['CUDA_VISIBLE_DEVICES']}") start_time = time.time() # Load pretrained ear modelimport pandas as pd input_folder_path = "/home/tony/Data/Preference/up_v2" input_file_name = "interesting_clips_up_u_2_20241210_full.pkl" df = pd.read_pickle(f"{input_folder_path}/{input_file_name}") df = df.sort_values(by=["request_id", "preference"]) print("Preference data shape", df.shape) request_jobs = [] unique_request_ids = df["request_id"].unique() grouped = df.groupby("request_id") for request_id in tqdm(unique_request_ids): group = grouped.get_group(request_id) assert (group.iloc[0]["preference"], group.iloc[1]["preference"]) == (0, 1) negative_id, positive_id = group["s3_id"].tolist() request_jobs.append((str(request_id), negative_id, positive_id)) request_jobs = sorted(request_jobs) print(f"Total jobs: {len(request_jobs)}") # find the specific chunk request_jobs = request_jobs[ (len(request_jobs) // total_job_n_gpus) * args.job_idx : ( len(request_jobs) // total_job_n_gpus ) * (args.job_idx + 1) ] print(f"Chunked jobs: {len(request_jobs)}") # Example usage # test_batch = [ # ( # "0002c166-7355-4335-bc47-414ed0914f82", # "df9e8d5a-720e-4382-a10a-875b9b1e488a", # "30c532ed-9422-4db3-8f4c-2b56c3aed34a", # ) # ] # test_result = process_request_batch(test_batch) # print(test_result) with open(f"{input_folder_path}/pair_quality.json", "r") as f: known_results = json.load(f) print(f"Pre-filtered jobs: {len(request_jobs)}") request_jobs = [job for job in request_jobs if job[0] not in known_results.keys()] print(f"Total jobs: {len(request_jobs)}") # only when you debug # request_jobs = request_jobs[:200] # Using process_map from tqdm.contrib.concurrent for better parallelization # Split request_jobs into batches n_processes = 20 batch_size = len(request_jobs) // n_processes + 1 batches = [ request_jobs[i : i + batch_size] for i in range(0, len(request_jobs), batch_size) # for i in range(0, 10, batch_size) ] print(f"Total batches: {len(batches)}, batch size: {batch_size}") with Pool(processes=n_processes) as pool: results = pool.map(process_request_batch, batches) # Combine results from all batches combined_results = {} for batch_result in results: combined_results.update(batch_result) known_results.update(combined_results) # with open(f"{input_folder_path}/pair_quality_{args.job_idx}.json", "w") as f: with open(f"{input_folder_path}/pair_quality.json", "w") as f: json.dump(known_results, f, indent=4) print(f"DONE!!, total time: {round(time.time() - start_time, 2)}s") if __name__ == "__main__": # Make sure export is setup: export CUDA_VISIBLE_DEVICES=5 main()