import os import json import funcy import boto3 import shutil import numpy as np from tqdm import tqdm from concurrent.futures import ThreadPoolExecutor, as_completed from suno_utils.utils.text import read_jsonl, write_jsonl from suno_utils.utils.s3 import read_from_s3, _verify_s3_filepath 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 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 USE_PAIRS = True # use pairs instead of single examples USE_AUDIO_QUALITY = False OUT_DATA_DIR = ( "/app/suno/data/diffusion_ft/upsample_100hz_v4_t_5_20241015_pairs_highpass" ) # OUT_DATA_DIR = ( # "/app/suno/data/diffusion_ft/upsample_100hz_v4_t_5_20241018_pairs_text_cfg" # ) if os.path.exists(OUT_DATA_DIR): print(f"out data dir {OUT_DATA_DIR} already exists, deleting") shutil.rmtree(OUT_DATA_DIR) os.makedirs(OUT_DATA_DIR, exist_ok=True) # christian/data/upsample_100hz_v4_t_5_20241018 bucket_name = "suno-data" base_dir = "christian/data/upsample_100z_v1" output_name = "v2" # # base_dir = "christian/data/upsample_100hz_v4_t_5_20241018" # output_name = "v2" base_metas_path = os.path.join("s3://", bucket_name, base_dir, "metas.jsonl") base_metas = read_from_s3(base_metas_path, read_f=read_jsonl) print(len(base_metas)) # s3 client s3 = boto3.client("s3") train_size = len(base_metas) - VAL_SIZE train_metas = base_metas[:train_size] val_metas = base_metas[train_size:] print("train/val split: ", len(train_metas), len(val_metas)) # find all files on s3 with the pattern filepaths = get_s3_files(bucket_name, f"{base_dir}/{output_name}") print("total files on s3: ", len(filepaths)) id_to_s3_paths = {} quality_scores = {} # create a dict with the id as key and the s3 paths a list of values for filepath in tqdm(filepaths): if "-text_cfg_" in filepath: meta_id = filepath.split("-text_cfg")[0].split("/")[-1] else: meta_id = filepath.split("/")[-1].split(".")[0].split("-")[:-1] meta_id = "-".join(meta_id) if meta_id not in id_to_s3_paths: id_to_s3_paths[meta_id] = set() s3_filepath_basename = filepath.split(".")[0] id_to_s3_paths[meta_id].add(s3_filepath_basename) print("total ids: ", len(id_to_s3_paths)) # load quality scores from local if USE_AUDIO_QUALITY: with open( # "/home/christian/code/christian/notebooks/diff_dpo/upsample_100z_v1_v2_quality_scores.json", "/home/christian/code/christian/notebooks/diff_dpo/upsample_100hz_v4_t_5_20241018_quality_scores.json", "r", ) as f: quality_scores = json.load(f) print("total quality scores: ", len(quality_scores)) # now iterate over the val, then train metas for dset_type in ["val", "train"]: metas = val_metas if dset_type == "val" else train_metas valid_metas = [] for meta in tqdm(metas): if meta["id"] in id_to_s3_paths: # check if the npz exists on s3 s3_filepath = ( f"s3://{bucket_name}/{list(id_to_s3_paths[meta['id']])[0]}.npz" ) try: _verify_s3_filepath(s3_filepath) valid_metas.append(meta) except Exception as e: print(f"error loading {s3_filepath}") continue print(f"total {dset_type} metas: ", len(valid_metas)) out_mm_vae_filepath = os.path.join(OUT_DATA_DIR, f"data_vae_{dset_type}.bin") out_metas_filepath = os.path.join(OUT_DATA_DIR, f"metas_{dset_type}.jsonl") out_mm_semantic_filepath = os.path.join( OUT_DATA_DIR, f"data_semantic_{dset_type}.bin" ) # initial write out_mm_semantic = np.memmap( out_mm_semantic_filepath, dtype=np.uint16, mode="w+", shape=(1), ) out_mm_vae = np.memmap( out_mm_vae_filepath, dtype=np.float16, mode="w+", shape=(1), ) n_offs_s = 0 n_offs_v = 0 CHUNK_SIZE = 100 # split valid metas into chunks of CHUNK_SIZE valid_metas_chunks = [ valid_metas[i : i + CHUNK_SIZE] for i in range(0, len(valid_metas), CHUNK_SIZE) ] print("total chunks: ", len(valid_metas_chunks)) for chunk_idx, meta_chunk in enumerate(tqdm(valid_metas_chunks)): arr_s_list = [] arr_v_list = [] new_metas = [] def process_meta(meta): if meta["id"] not in id_to_s3_paths: return None s3_paths = list(id_to_s3_paths[meta["id"]]) # when using pairs, select positive with highest score # and negative with lowest score # when using single example select based on highest score s3_path_list = [] # for s3_path in s3_paths: # if "_quality_scores" in s3_path: # continue # if USE_AUDIO_QUALITY: # score = float(quality_scores[meta["id"]][s3_path]) # s3_path_list.append((s3_path, score)) # elif "-text_cfg_" in s3_path: # text_cfg_value = s3_path.split("-text_cfg_")[-1] # value = int(text_cfg_value.split("-")[0]) # s3_path_list.append((s3_path, value)) # else: # s3_path_list.append((s3_path, 0)) found_highpass = False for s3_path in s3_paths: if "highpass" in s3_path: s3_path_list.append((s3_path, 0)) print(s3_path) found_highpass = True elif "quality_scores" in s3_path: continue else: s3_path_list.append((s3_path, 1)) if not found_highpass: return None # sort list based on score s3_path_list.sort(key=lambda x: x[1], reverse=True) if len(s3_path_list) < 2: return None if USE_PAIRS: positive_s3_path = s3_path_list[0][0] negative_s3_path = s3_path_list[1][0] selected_s3_paths = [positive_s3_path, negative_s3_path] else: positive_s3_path = s3_path_list[0][0] selected_s3_paths = [positive_s3_path] try: result = [] for s3_path_idx, s3_path in enumerate(selected_s3_paths): filepath = f"s3://{bucket_name}/{s3_path}.npz" data = read_from_s3(filepath, read_f=np.load) arr_s = data["semantic_codes"] arr_v = data["upsampled_latents"] if arr_s.size < SEMANTIC_MEMMAP_SIZE: return None if arr_v.size < VAE_MEMMAP_SIZE * VAE_DIM: return None result.append((arr_s, arr_v, meta)) return result except Exception as e: print(f"error loading {filepath}: {e}") return None with ThreadPoolExecutor(max_workers=16) as executor: futures = [executor.submit(process_meta, meta) for meta in meta_chunk] for future in as_completed(futures): result = future.result() # list of tuples (arr_s, arr_v, meta) if result: for arr_s, arr_v, meta in result: arr_s_list.append(arr_s) arr_v_list.append(arr_v) new_meta = meta.copy() new_meta["n_vae_tokens"] = VAE_MEMMAP_SIZE new_metas.append(new_meta) print(len(arr_s_list), len(arr_v_list), len(new_metas)) assert len(arr_s_list) == len(arr_v_list) == len(new_metas) # now write to the memmap # get a list of all the ids in the id_to_s3_paths to_write_len_s = SEMANTIC_MEMMAP_SIZE * len(arr_v_list) to_write_len_v = VAE_MEMMAP_SIZE * VAE_DIM * len(arr_v_list) out_mm_semantic = np.memmap( out_mm_semantic_filepath, dtype=np.uint16, mode="r+", shape=(n_offs_s + to_write_len_s,), ) out_mm_vae = np.memmap( out_mm_vae_filepath, dtype=np.float16, mode="r+", shape=(n_offs_v + to_write_len_v,), ) # write to memmap (has to happen sequentially) for new_meta, arr_s, arr_v in zip(new_metas, arr_s_list, arr_v_list): out_mm_semantic[n_offs_s : n_offs_s + arr_s.size] = arr_s.reshape( -1, ) out_mm_vae[n_offs_v : n_offs_v + arr_v.size] = arr_v.reshape( -1, ) n_offs_s += arr_s.size n_offs_v += arr_v.size # write it once out_mm_semantic.flush() out_mm_vae.flush() del out_mm_semantic, out_mm_vae write_jsonl(new_metas, out_metas_filepath, do_append=True)