import os import torch import torchaudio import numpy as np import pandas as pd import pyloudnorm as pyln from tqdm import tqdm from ear.utils import load_audio, apply_normalization from ear.system import EarSystem from suno_utils.utils.text import read_jsonl from suno_utils.audio import Audio from suno_utils.tasks.dac_2c_12cb import ( encode as codec_encode, decode as codec_decode, EMBEDDING_RATE as CODEC_EMBEDDING_RATE, ) from suno_utils.tasks.dac_2c_12cb import preload_models as preload_codec_models def batch_evaluate_audio_quality_from_embeds( system: torch.nn.Module, eval_audio: torch.Tensor, ref_embeds: torch.Tensor, num_chunks: None, ): """ Args: system (torch.nn.Module): eval_audio (torch.Tensor): Audio to evaluate with shape (bs, 1, seq_len) ref_embeds (torch.Tensor): Embedded reference audio (num_refs, seq_len, embed_dim) """ # print("eval_audio", eval_audio.shape) bs, chs, eval_seq_len = eval_audio.shape num_refs, ref_seq_len, embed_dim = ref_embeds.shape # randomly sample N chunks max_chunks = min(8, eval_audio.shape[0]) chunk_idx = torch.randint(0, eval_audio.shape[0], [max_chunks]) eval_audio = eval_audio[chunk_idx, ...] # first, embed the audio that will be evaluated with torch.no_grad(): eval_embeds = system.embed(eval_audio) # eval_embeds has shape (bs, embed_dim) # ref_embeds has shape (num_refs, embed_dim) # now copy the eval and reference embeds to evaluate against all ref_embeds = ref_embeds.repeat(bs, 1, 1) eval_embeds = eval_embeds.repeat(num_refs, 1, 1) # concat embeds into singular tensors embeds = torch.cat((eval_embeds, ref_embeds), dim=-1) # print("embeds", embeds.shape) # no run through the projection to make predictions with torch.no_grad(): pref_preds = system.pref_classifier(embeds) quant_preds = system.quant_classifier(embeds) # print(pref_preds.shape, quant_preds.shape) # get a final score by taking mean across seq of preds pref_preds = pref_preds.mean(dim=1).squeeze(1) quant_preds = quant_preds.mean(dim=1).squeeze(1) pref = torch.sigmoid(pref_preds) quant = torch.argmax(quant_preds, dim=1).float() # aggregate predictions across the reference recordings prefs = pref.view(bs, -1).mean() quants = quant.view(bs, -1).mean() scores = -((prefs * 2) - 1) * (quants + 1) return prefs, quants, scores def prepare_audio(audio_filepath: str, num_frames: int = 131072): audio, sample_rate = torchaudio.load(audio_filepath) if sample_rate != 24_000: audio = torchaudio.functional.resample(audio, sample_rate, 24_000) # 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]) if len(audio_chunks) < 1: return # 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 *= gain_lin audio_chunks[audio_chunk_idx] *= gain_lin return torch.stack(audio_chunks) def decode_npz_from_csv(csv_filepath: str, npz_dir: str, audio_dir: str): df = pd.read_csv(csv_filepath) num_examples = len(df.index) for index in tqdm(range(num_examples)): example = df.iloc[index] file_id = example["id"] npz_filepath = os.path.join(npz_dir, f"{file_id}.npz") audio_filepath = os.path.join(audio_dir, f"{file_id}.wav") if os.path.isfile(audio_filepath): print(audio_filepath, "already decoded.") continue if not os.path.isfile(npz_filepath): print(npz_filepath, "not found!") os.system( f"aws s3 cp s3://suno-data-uploads/studio/uploads/{file_id}.npz tmp/{file_id}.npz" ) npz_filepath = os.path.join("tmp", f"{file_id}.npz") data = np.load(npz_filepath) if "v3.5_raw" in data: data = data["v3.5_raw"] elif "v3.0_raw" in data: data = data["v3.0_raw"] else: print("no valid data in", file_id) continue audio = codec_decode(data[:, 1:]) sample_rate = audio.sample_rate audio = torch.from_numpy(audio.array_float) torchaudio.save(audio_filepath, audio, sample_rate) class PreferenceDataset(torch.utils.data.Dataset): def __init__( self, csv_filepath: str, audio_dir: str = "/home/christian/data/13b_audio/" ): # load csv file self.audio_dir = audio_dir self.df = pd.read_csv(csv_filepath) self.num_examples = len(self.df.index) self.indices = [i for i in range(0, self.num_examples, 2)] print(self.num_examples) def __len__(self): return len(self.indices) def __getitem__(self, idx): df_index = self.indices[idx] a_example = self.df.iloc[df_index] b_example = self.df.iloc[df_index + 1] a_id = a_example["id"] a_audio_filepath = os.path.join(self.audio_dir, f"{a_id}.wav") b_id = b_example["id"] b_audio_filepath = os.path.join(self.audio_dir, f"{b_id}.wav") # check that both files exist if not os.path.isfile(a_audio_filepath) or not os.path.isfile(b_audio_filepath): a_audio = None b_audio = None valid = False else: a_audio = prepare_audio(a_audio_filepath) b_audio = prepare_audio(b_audio_filepath) valid = True if a_audio is None or b_audio is None: a_audio = torch.zeros(1) b_audio = torch.zeros(1) valid = False # randomly sample N chunks # max_chunks = min(8, audio.shape[0]) # chunk_idx = torch.randint(0, audio.shape[0], [max_chunks]) # audio = audio[chunk_idx, ...] # os.remove(filepath) # delete audio file from tmp directory # orig_idx = torch.tensor(meta["orig_idx"]) return df_index, a_audio, b_audio, valid if __name__ == "__main__": # load pretrained ear model num_frames = 131072 ckpt_path = "/app/suno/christian/checkpoints/ear/0bwk1wd7-6.ckpt" system = EarSystem.load_from_checkpoint(ckpt_path) system.cuda() system.eval() # setup references ref_filepaths = [ "/app/suno/christian/reference-audio-wav-24khz/02 Dreams.wav", "/app/suno/christian/reference-audio-wav-24khz/01 Mario Takes A Walk.wav", "/app/suno/christian/reference-audio-wav-24khz/02 Freddie Freeloader.wav", "/app/suno/christian/reference-audio-wav-24khz/09 Sounds Like Hallelujah.wav", "/app/suno/christian/reference-audio-wav-24khz/03 Your New Aesthetic.wav", "/app/suno/christian/reference-audio-wav-24khz/01 J.S. Bach Suite No.1, S.1007, G major - I. Prelude.wav", "/app/suno/christian/reference-audio-wav-24khz/08 Get Lucky.wav", "/app/suno/christian/reference-audio-wav-24khz/03 Ocean Waves (O Mar).wav", "/app/suno/christian/reference-audio-wav-24khz/02 Take Five.wav", "/app/suno/christian/reference-audio-wav-24khz/03 Guess I_m Doing Fine.wav", ] ref_audios = [ load_audio( filepath, num_frames=num_frames, target_sample_rate=system.hparams.sample_rate, ) for filepath in ref_filepaths ] ref_audios = torch.stack(ref_audios) print("ref_audios", ref_audios.shape) ref_audio = ref_audios.cuda() # first precompute the reference embeddings ref_embeds = system.embed(ref_audios) print("ref_embeds", ref_embeds.shape) # load DAC codec _ = preload_codec_models( checkpoint_filepath="s3://suno-data/georg/models/codec/dac_2c_25x12.pt", ) # source csv npz_dir = "/app/suno/data/dpo/13b_npz/" audio_dir = "/home/christian/data/13b_audio" preference_csv = "/home/tony/Data/Preference/13b_v0/interesting_clips_20240528.csv" # first, decode all npz to audio decode_npz_from_csv(preference_csv, npz_dir, audio_dir) # then create a dataset that loads dataset = PreferenceDataset(preference_csv, audio_dir) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=64) # load our own copy of the preference data df = pd.read_csv(preference_csv) out_preference_csv = "/home/christian/data/interesting_clips_20240528_ear.csv" for idx, batch in enumerate(tqdm(dataloader)): df_index, a_audio, b_audio, valid = batch print(valid.item()) # print(df_index, a_audio.shape, b_audio.shape) if not valid.item(): continue # directory compare audios # first, embed the audio that will be evaluated with torch.no_grad(): a_embeds = system.embed(a_audio.squeeze(0).cuda()) b_embeds = system.embed(b_audio.squeeze(0).cuda()) # ensure equal length min_length = min(a_embeds.shape[1], b_embeds.shape[1]) min_bs = min(a_embeds.shape[0], b_embeds.shape[0]) a_embeds = a_embeds[:min_bs, :min_length, :] b_embeds = b_embeds[:min_bs, :min_length, :] # concat embeds into singular tensors # embeds = torch.cat((cycled_embeds, orig_embeds), dim=-1) embeds = torch.cat((a_embeds, b_embeds), dim=-1) # print("embeds", embeds.shape) # no run through the projection to make predictions with torch.no_grad(): pref_preds = system.pref_classifier(embeds) # quant_preds = system.quant_classifier(embeds) # get a final score by taking mean across seq of preds pref_preds = torch.sigmoid(pref_preds.mean()) # if pref is greater than 0,5 then b is better than a if pref_preds.item() > 0.5: a_ear_preference = False b_ear_preference = True else: a_ear_preference = True b_ear_preference = False df.loc[df_index, "ear_preference"] = a_ear_preference df.loc[df_index + 1, "ear_preference"] = b_ear_preference df.to_csv(out_preference_csv)