import os import json import boto3 import torch import shutil import tempfile import torchaudio import numpy as np from tqdm import tqdm from suno_utils.utils.text import read_jsonl, write_jsonl from suno_utils.utils.s3 import read_from_s3, _verify_s3_filepath from ear.utils import load_audio, apply_normalization from ear.system import EarSystem def get_s3_files(bucket_name, prefix, max_keys: int = 100000): all_files = [] continuation_token = None while True: # Prepare the arguments for the request list_kwargs = { "Bucket": bucket_name, "Prefix": prefix, # List objects under this prefix, or leave blank for all objects } if continuation_token: list_kwargs["ContinuationToken"] = continuation_token # Make the request to list objects response = s3.list_objects_v2(**list_kwargs) # Collect the file keys all_files += [obj["Key"] for obj in response.get("Contents", [])] # Check if more results are available if response.get("IsTruncated"): # True if there are more results to fetch continuation_token = response["NextContinuationToken"] else: break # No more results to fetch return all_files def list_s3_directories(bucket_name: str, prefix: str = ""): """ List all directories (prefixes) in an S3 bucket. Args: bucket_name (str): Name of the S3 bucket prefix (str): Optional prefix to filter results (like a directory path) Returns: List[str]: List of directory paths (prefixes) """ s3_client = boto3.client("s3") directories = set() # Use paginator to handle buckets with many objects paginator = s3_client.get_paginator("list_objects_v2") page_iterator = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter="/") # Collect all prefixes (directories) for page in page_iterator: # Get common prefixes (directories) if "CommonPrefixes" in page: for prefix_obj in page["CommonPrefixes"]: directories.add(prefix_obj["Prefix"]) # Also check Contents for any directory-like objects if "Contents" in page: for obj in page["Contents"]: key = obj["Key"] # If the key contains a slash, add the directory part if "/" in key: directory = key.rsplit("/", 1)[0] + "/" directories.add(directory) return sorted(list(directories)) if __name__ == "__main__": # load the base metas VAE_DIM = 128 VAE_RATE_HZ = 100 VAE_MEMMAP_SIZE = 3000 SEMANTIC_VOCAB_SIZE = 4000 SEMANTIC_MEMMAP_SIZE = 750 VAL_SIZE = 100 # 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() NUM_FRAMES = 131072 # 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) def evaluate_ear(audio: torch.Tensor, ref_embeds: torch.Tensor): bs = audio.shape[0] num_refs = ref_embeds.shape[0] # first, embed the audio that will be evaluated with torch.no_grad(): eval_embeds = system.embed(audio) # aggregate the eval_embed eval_embeds = eval_embeds.mean(dim=1, keepdim=True) ref_embeds = ref_embeds.mean(dim=1, keepdim=True) # 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 # -------------------------------------------------------------------------- bucket_name = "suno-data" # base_dir = "christian/data/upsample_100z_v1" # output_name = "v2" base_dir = "christian/data/upsample_v4_t_5_20241018" output_name = "25hz_20241031_v1/" # s3 client s3 = boto3.client("s3") # find all files on s3 with the pattern dir_paths = list_s3_directories(bucket_name, f"{base_dir}/{output_name}") print("total dirs: ", len(dir_paths)) for dir_path in tqdm(dir_paths): # check if the directory has a quality_scores.json file # s3_filepath = f"{dir_path}quality_scores.json" # if s3.head_object(bucket_name, s3_filepath): # print(f"skipping {dir_path} because it already has quality scores") # continue # get all the mp3 filepaths in this directory all_filepaths = get_s3_files(bucket_name, f"{dir_path}") mp3_filepaths = [fp for fp in all_filepaths if fp.endswith(".mp3")] # create a json file with the scores for this example and save to s3 scores = {} for mp3_filepath in mp3_filepaths: base_name = mp3_filepath.strip("/").split("/")[-1] try: full_mp3_filepath = f"s3://{bucket_name}/{mp3_filepath}" data = read_from_s3(full_mp3_filepath, read_f=torchaudio.load) except Exception as e: print(f"error loading {full_mp3_filepath}: {e}") continue audio, sr = data audio = audio.cuda() # evaluate the audio score = evaluate_ear(audio.unsqueeze(0), ref_embeds) scores[base_name] = score.item() print(scores) # save to temporary file and then upload to s3 with tempfile.TemporaryDirectory() as td: temp_filepath = os.path.join(td, "scores.json") with open(temp_filepath, "w") as f: json.dump(scores, f) s3_filepath = f"{dir_path}quality_scores.json" s3.upload_file(temp_filepath, bucket_name, s3_filepath)