import os import uuid import json import glob import torch import boto3 import random import resampy import torchaudio import numpy as np import soundfile as sf import concurrent.futures import pyloudnorm as pyln import multiprocessing as mp from tqdm import tqdm from time import perf_counter from typing import Optional, List from ear.utils import load_audio, apply_normalization from ear.system import EarSystem from suno_utils.audio import Audio from suno_utils.utils.text import ( write_jsonl, read_jsonl, write_json, read_json, normalize_whitespace, ) def remove_pad_tokens(arr, value: int = 2048): # Find the last index where the value is not equal to the specified value non_value_index = np.where(arr != value)[0] if non_value_index.size == 0: # If there are no values that are not equal to the specified value, return an empty array return np.array([], dtype=arr.dtype) # Get the last index where the value is not equal to the specified value last_non_value_index = non_value_index[-1] # Slice the array to remove trailing values return arr[: last_non_value_index + 1] def batch_evaluate_audio_quality_from_embeds( system: torch.nn.Module, eval_audio: torch.Tensor, ref_embeds: torch.Tensor ): """ 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 # 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 download_audio(s3_filepath: str, example_id: str, tmp_dir: str): filename = os.path.basename(s3_filepath) out_filepath = os.path.join(tmp_dir, f"{example_id}-{filename}") # only download the file if its not already downloaded if not os.path.isfile(out_filepath): os.system(f"aws s3 cp {s3_filepath} {out_filepath} > /dev/null 2>&1") return out_filepath def prepare_audio( filepath: str, num_frames: int, start_s: float = Optional[None], end_s: float = Optional[None], ): # start_time = perf_counter() # audio = Audio.from_s3(s3_filepath, n_channels=2) audio = Audio.from_file(filepath, n_channels=2) sample_rate = audio.sample_rate audio = torch.from_numpy(audio.array_float) if audio.shape[0] != 2: audio = audio.repeat(2, 1) # end_time = perf_counter() # elapsed_s = end_time - start_time # print(f"{elapsed_s:0.2f} - {filepath}") # 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] # check for valid duration # num_frames = audio.shape[-1] # if num_frames < 1: # print(filepath) # audio = torch.zeros(1, 524288) # meta_audio_dur_s = end_s - start_s # audio_dur_s = audio.shape[-1] / sample_rate # print(meta_audio_dur_s, audio_dur_s) # 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) # audio = resampy.resample(np.sum(audio, axis=0, keepdims=True), sample_rate, 24000) # audio = audio.mean(dim=0, keepdim=True) # 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") # take a central crop # if audio.shape[-1] > num_frames * 2: # start_idx = audio.shape[-1] // 2 # audio = audio[:, start_idx : start_idx + num_frames] # else: # take crop from the start # start_idx = 0 # audio = audio[:, start_idx : start_idx + num_frames] # 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 return torch.stack(audio_chunks) class AudioMetadataDataset(torch.utils.data.Dataset): def __init__(self, metas: List[dict], num_frames: int, tmp_dir: str): # we need to filter metas to only use ones with s3_filepath # but, we have to be careful to maintain the index of the original metadata filtered_metas = [] for meta_idx, meta in enumerate(metas): if "s3_filepath" in meta: meta["orig_idx"] = meta_idx filtered_metas.append(meta) num_filtered = len(filtered_metas) num_original = len(metas) percent_remaining = (num_filtered / num_original) * 100 print( f"{num_filtered}/{num_original} ({percent_remaining:0.2f}%) examples have s3_filepath." ) self.metas = filtered_metas self.num_frames = num_frames self.tmp_dir = tmp_dir def __len__(self): return len(self.metas) def __getitem__(self, idx): meta = self.metas[idx] filepath = download_audio(meta["s3_filepath"], meta["id"], self.tmp_dir) audio = prepare_audio( filepath, self.num_frames, meta.get("start_s"), meta.get("end_s") ) # 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 orig_idx, audio class AudioFileDataset(torch.utils.data.Dataset): def __init__(self, audio_files: List[str], num_frames: int): self.audio_files = audio_files self.num_frames = num_frames def __len__(self): return len(self.audio_files) def __getitem__(self, idx): audio_file = self.audio_files[idx] audio = prepare_audio(audio_file, self.num_frames, None, None) return audio_file, audio if __name__ == "__main__": num_compare = 5 num_frames = 131072 use_val = False # load pretrained ear model ckpt_path = "/home/christian/code/christian/checkpoints/w5p4nhzn-epoch=27.cpkt" if not os.path.isfile(ckpt_path): os.system( f"aws s3 cp s3://suno-data/christian/ear/w5p4nhzn-epoch=27.cpkt /home/christian/code/christian/checkpoints" ) system = EarSystem.load_from_checkpoint(ckpt_path) system.cuda() system.eval() metas_out_dir = "outputs/quality_metas/audio_2ch_48khz_lg" os.makedirs(metas_out_dir, exist_ok=True) if use_val: root_dir = "/app/suno/data/audio_2ch_48khz_lg/val" else: root_dir = "/app/suno/data/audio_2ch_48khz_lg/train" # load reference audio used for quality comparision # ref_dir = "/app/suno/christian/data/codec_audio/reference-audio-wav-mono-24khz/" # ref_filepaths = glob.glob(os.path.join(ref_dir, "*.input.wav")) # ref_filepaths = np.random.choice(ref_filepaths, num_compare) ref_filepaths = [ "/home/christian/audio/reference-audio-wav-mono-24khz/02 Dreams.wav", "/home/christian/audio/reference-audio-wav-mono-24khz/01 Mario Takes A Walk.wav", "/home/christian/audio/reference-audio-wav-mono-24khz/02 Freddie Freeloader.wav", "/home/christian/audio/reference-audio-wav-mono-24khz/09 Sounds Like Hallelujah.wav", "/home/christian/audio/reference-audio-wav-mono-24khz/03 Your New Aesthetic.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) tmp_dir = os.path.join(os.getcwd(), "tmp") os.makedirs(tmp_dir, exist_ok=True) # construct dataset # dataset = AudioMetadataDataset(metas, num_frames, tmp_dir) # note: this only works with batch_size = 1 # get all fileptahs in root_dir audio_subsets = glob.glob(os.path.join(root_dir, "**")) for audio_subset in audio_subsets: new_metas = [] new_metas_output_filepath = os.path.join( metas_out_dir, os.path.basename(root_dir), f"{os.path.basename(audio_subset)}_quality_metas.jsonl", ) os.makedirs(os.path.dirname(new_metas_output_filepath), exist_ok=True) audio_subset_filepaths = glob.glob(os.path.join(audio_subset, "*.wav")) print(audio_subset, len(audio_subset_filepaths)) dataset = AudioFileDataset(audio_subset_filepaths, num_frames) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=32) for batch_idx, batch in enumerate(tqdm(dataloader)): audio_file, audios = batch audios = audios.squeeze(0) audios = audios.cuda() prefs, quants, scores = batch_evaluate_audio_quality_from_embeds( system, audios, ref_embeds ) # add this meta to existing new meta new_meta = { "audio_file": audio_file, "audio_quality": { "score": f"{scores.item():0.2f}", "preference": f"{prefs.item():0.2f}", "quantification": f"{quants.item():0.2f}", }, } new_metas.append(new_meta) write_jsonl(new_metas, new_metas_output_filepath)