import os import sys import torch import argparse import torchaudio import numpy as np import pandas as pd import pyloudnorm as pyln from tqdm import tqdm from ear.utils import load_audio from suno_utils.audio import Audio from ear.system import EarSystem from suno_boost.utils import apply_normalization def closer_to_target(value_a: float, value_b: float, target: float = 0.0): diff1 = abs(value_a - target) diff2 = abs(value_b - target) if diff1 < diff2: return "a" elif diff2 < diff1: return "b" else: return "Both values are equally close to the target" def get_audio_stats(x: np.ndarray, sample_rate: int, measure_true_peak: bool = False): """Measure singal processing audio quality statisics on audio array. Args: x (np.ndarray): Audio array of shape (2, samples) sample_rate (int): Audio sample rate. Returns: dict containing audio features """ meter = pyln.Meter(sample_rate) peak_lin = np.max(np.abs(x)) peak_db = 20 * np.log10(peak_lin + 1e-8) lufs_db = meter.integrated_loudness(x) dc_offset = np.mean(x) rms = np.sqrt(np.mean(x**2)) crest_factor = peak_lin / (rms + 1e-8) # for a fair comparision (loudness normalize to -16 dBLUFS and recompute) gain_lin = 10 ** ((-16.0 - lufs_db) / 20) x_norm = x * gain_lin peak_lin_norm = np.max(np.abs(x)) rms_norm = np.sqrt(np.mean(x_norm**2)) crest_factor_norm = peak_lin_norm / (rms_norm + 1e-8) # compute the compression factor by peak normalizing and then x_peak_norm = x / np.clip(peak_lin, a_min=1e-8, a_max=None) compression_factor = meter.integrated_loudness(x_peak_norm) # resample 4x to get true-peak (this is slow) if measure_true_peak: x_up = resampy.resample(x, sample_rate, int(4 * sample_rate), axis=0) true_peak_lin = np.max(np.abs(x_up)) true_peak_db = 20 * np.log10(true_peak_lin + 1e-8) result = { "peak_lin": peak_lin, "peak_db": peak_db, # "true_peak_lin" : true_peak_lin, # "true_peak_db" : true_peak_db, "lufs_db": lufs_db, "dc_offset": dc_offset, "rms": rms, "crest_factor": crest_factor, "rms_norm": rms_norm, "crest_factor_norm": crest_factor_norm, "compression_factor": compression_factor, } return result def compare_audio_stats(stat_dict_a: dict, stat_dict_b: dict): """Given audio statisitics for two audios, decide which is better.""" # compare the DC Offset, which is closer to 0? dc_offset_a = stat_dict_a["dc_offset"] dc_offset_b = stat_dict_b["dc_offset"] dc_offset_winner = closer_to_target(dc_offset_a, dc_offset_b, target=0.0) # compare the loudness, which is closer to -16 dB LUFS loudness_a = stat_dict_a["lufs_db"] loudness_b = stat_dict_b["lufs_db"] loudness_winner = closer_to_target(loudness_a, loudness_b, target=-16.0) # check if the peaks go beyond 1.0 peak_db_a = stat_dict_a["peak_db"] peak_db_b = stat_dict_b["peak_db"] # if one generation has peaks beyond 0 dB and the other doesn't if peak_db_a > 0.0 and peak_db_b < 0.0: peak_db_winner = "b" elif peak_db_a < 0.0 and peak_db_b > 0.0: peak_db_winner = "a" else: peak_db_winner = None test_results = [dc_offset_winner, loudness_winner, peak_db_winner] # count number of wins for A and B a_wins = 0 b_wins = 0 for test_result in test_results: if test_result == "a": a_wins += 1 elif test_result == "b": b_wins += 1 else: pass # no winner if a_wins > b_wins: return "a" else: return "b" if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--csv_path", default="/home/tony/Data/Preference/7b_v0/interesting_clips_v3_processed.csv", type=str, ) parser.add_argument( "--root_dir", type=str, default="/app/suno/christian/data/interesting_clips_v3_processed", ) parser.add_argument("--num_examples", type=int, default=100000) args = parser.parse_args() # load pretrained model # music fm models ckpt_path = "/app/suno/christian/ear-logs/ear/1l1qrmqw/checkpoints/epoch=63-step=359552.ckpt" # ckpt_path = "./checkpoints/epoch=3-step=22472.ckpt" # ckpt_path = "./checkpoints/epoch=10-step=61798.ckpt" # panns models # ckpt_path = ( # "/app/suno/christian/ear-logs/ear/rkkkl0ck/checkpoints/epoch=13-step=78652.ckpt" # system = EarSystem.load_from_checkpoint(ckpt_path) system.eval() ear_version = os.path.basename(ckpt_path).replace(".ckpt", "") out_csv_path = f"/app/suno/christian/preference/interesting_clips_v3_processed+ear={ear_version}.csv" # load csv df = pd.read_csv(args.csv_path) print(df.shape) # get pairs df = df.sort_values(by=["request_id", "preference"]) print(df.head()) if system.hparams.sample_rate == 48000: num_frames = 262144 elif system.hparams.sample_rate == 24000: num_frames = 131072 new_df = df.copy() # add column with preference new_df["ear_preference"] = None if (args.num_examples * 2) > len(df): indices = np.arange(0, len(df), step=2) else: indices = np.arange(0, args.num_examples * 2, step=2) same = [] count = 0 pbar = tqdm(indices) for idx in pbar: # idx = np.random.randint(0, len(df)) # idx += idx % 2 neg_row = df.iloc[idx] pos_row = df.iloc[idx + 1] # print(pos_row["request_id"], pos_row["preference"], pos_row["s3_id"]) # print(neg_row["request_id"], neg_row["preference"], neg_row["s3_id"]) # get local paths pos_filepath = os.path.join(args.root_dir, pos_row["s3_id"] + ".mp3") neg_filepath = os.path.join(args.root_dir, neg_row["s3_id"] + ".mp3") try: pos_audio = load_audio( pos_filepath, num_frames=num_frames, target_sample_rate=system.hparams.sample_rate, ) neg_audio = load_audio( neg_filepath, num_frames=num_frames, target_sample_rate=system.hparams.sample_rate, ) except KeyboardInterrupt: print("Execution stopped by user.") sys.exit() except: # print("Skipping since audio load failed...") continue if pos_audio.shape[-1] != neg_audio.shape[-1]: continue # run inference with torch.no_grad(): # run inference # run forward pref_preds, quant_preds = system.forward( neg_audio.unsqueeze(0), pos_audio.unsqueeze(0), ) f_pref_preds = pref_preds.mean(dim=1).squeeze(1) f_quant_preds = quant_preds.mean(dim=1).squeeze(1) f_pref = torch.sigmoid(f_pref_preds) f_quant = torch.argmax(f_quant_preds) # pref = f_pref # quant = f_quant # run reverse pref_preds, quant_preds = system.forward( pos_audio.unsqueeze(0), neg_audio.unsqueeze(0), ) r_pref_preds = pref_preds.mean(dim=1).squeeze(1) r_quant_preds = quant_preds.mean(dim=1).squeeze(1) r_pref = torch.sigmoid(r_pref_preds) r_quant = torch.argmax(r_quant_preds) print(f_pref, r_pref) pref = (f_pref + (1 - r_pref)) / 2 quant = (f_quant + r_quant) / 2 # print(pref, quant) if pref > 0.5: same.append(True) pos_ear_preference = True neg_ear_preference = False else: same.append(False) pos_ear_preference = False neg_ear_preference = True # add preference row new_df.at[idx, "ear_preference"] = pos_ear_preference new_df.at[idx + 1, "ear_preference"] = neg_ear_preference count += 1 pbar.set_description( f"agree {np.mean(same)*100.0:0.2f} % ({count}) pref: {pref.item():0.2f} qaunt: {quant.item():0.2f}/32" ) print(new_df.head()) new_df.to_csv(out_csv_path)