# filter audio dataset to remove low quality examples 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 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) # 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.sum(dim=0, keepdim=True), 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, filepaths: List[str], num_frames: int): # 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 print(f"Found {len(filepaths)} filepaths") self.filepaths = filepaths self.num_frames = num_frames def __len__(self): return len(self.filepaths) def __getitem__(self, idx): filepath = self.filepaths[idx] audio = prepare_audio(filepath, self.num_frames) # 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 return audio if __name__ == "__main__": manifest_filepath = "" # create a new json file with labels for aaudio files in manifest # we can later use this to filter out audio # load manifest dataset = AudioMetadataDataset(filepaths, 131072) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=32) # load pretrained ear model ckpt_path = "/home/christian/code/christian/checkpoints/0bwk1wd7-7.ckpt" if not os.path.isfile(ckpt_path): os.system( f"aws s3 cp s3://suno-data/christian/ear/0bwk1wd7-7.ckpt /home/christian/code/christian/checkpoints" ) system = EarSystem.load_from_checkpoint(ckpt_path) system.cuda() system.eval() for batch in dataloader: audio = batch score, pref, quant = batch_evaluate_audio_quality_from_embeds( system, eval_audio, ref_embeds )