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, ) N_TOKENS_MEMMAP = 6016 SEMANTIC_N_CODEBOOKS = 1 COARSE_N_CODEBOOKS = 12 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) # 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 if __name__ == "__main__": num_compare = 5 num_frames = 131072 use_val = True num_examples = 5000 # 50000 # when set to None, use all examples shuffle = True # random subset # load pretrained ear model # 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() # load metadata data_dir = "/app/suno/data/chirp_v4_ft/multi/" out_dir = "/app/suno/data/chirp_v4_ft_sm/multi/" if use_val: metas_filename = "metas_val.jsonl" memmap_filepath = os.path.join("/app/suno/data/chirp_v4/multi/", "data_val.bin") out_mm_filepath = os.path.join(out_dir, "data_val_quality.bin") new_metas_output_filepath = os.path.join(out_dir, "metas_val_quality.jsonl") else: metas_filename = "metas_tr.jsonl" memmap_filepath = os.path.join("/app/suno/data/chirp_v4/multi/", "data_tr.bin") out_mm_filepath = os.path.join(out_dir, "data_tr_quality.bin") new_metas_output_filepath = os.path.join(out_dir, "metas_tr_quality.jsonl") # load memmap mm = np.memmap(os.path.join(memmap_filepath), dtype=np.uint16, mode="r") mm = mm.reshape(-1, N_TOKENS_MEMMAP, SEMANTIC_N_CODEBOOKS + COARSE_N_CODEBOOKS) # load reference audio used for quality comparision # ref_dir = "/app/suno/christian/data/codec_audio/reference-audio-wav_24khz/" # ref_filepaths = glob.glob(os.path.join(ref_dir, "*.input.wav")) # ref_filepaths = np.random.choice(ref_filepaths, num_compare) 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) tmp_dir = os.path.join(os.getcwd(), "tmp") os.makedirs(tmp_dir, exist_ok=True) # construct dataset metas_filepath = os.path.join(data_dir, metas_filename) print(f"Reading jsonl metadata from {metas_filepath}...") metas = read_jsonl(metas_filepath) dataset = AudioMetadataDataset(metas, num_frames, tmp_dir) # note: this only works with batch_size = 1 dataloader = torch.utils.data.DataLoader( dataset, batch_size=1, num_workers=32, shuffle=shuffle ) assert len(metas) == mm.shape[0] # metas and memmap must have same shape if num_examples is None: num_examples = len(metas) # new_metas = metas.copy() # create a copy of original list new_metas = [] new_arrays = [] # create output memmap example_size = N_TOKENS_MEMMAP * (COARSE_N_CODEBOOKS + SEMANTIC_N_CODEBOOKS) out_mm = np.memmap( out_mm_filepath, dtype=np.uint16, mode="w+", shape=(num_examples * example_size), ) mm_write_idx = 0 count = 0 for batch_idx, batch in enumerate(tqdm(dataloader)): if count >= num_examples: break orig_idx, audios = batch audios = audios.squeeze(0) audios = audios.cuda() prefs, quants, scores = batch_evaluate_audio_quality_from_embeds( system, audios, ref_embeds ) # now add metas with quality labels to new metas for output_idx, orig_index in enumerate(orig_idx): new_meta = metas[orig_index] # add this meta to existing new meta new_meta["audio_quality"] = { "score": f"{scores.item():0.2f}", "preference": f"{prefs.item():0.2f}", "quantification": f"{quants.item():0.2f}", } print(new_meta["audio_quality"]) # append new_metas.append(new_meta) # add to memmap array = np.reshape(mm[orig_index, ...], (-1)) print(count, mm.shape, array.shape) out_mm[mm_write_idx : mm_write_idx + example_size] = array mm_write_idx += example_size count += 1 # save metadata write_jsonl(new_metas, new_metas_output_filepath) out_mm.flush()