# THIS ONE ACTUALLY RUNS HOOT ON AUDIOS import os import json import argparse import tqdm from tqdm.contrib.concurrent import thread_map import numpy as np from suno_utils.tasks.hoot import ( encode_filepaths, ctc_align, load_model, clean_text, decode_logits, preload_models, ) from suno_utils.utils.metrics import get_cer def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--dataset", type=str) parser.add_argument("--start_index", type=int, default=0) args = parser.parse_args() return args if __name__ == "__main__": avialbe_device = f"cuda:{os.environ['CUDA_VISIBLE_DEVICES']}" print(avialbe_device) input_args = parse_args() if input_args.dataset == "genius": INPUT_PATH = "/home/tony/Data/Hoot/genius_metas.json" DOWNLOAD_PATH = "/app/suno/data/hoot/audios/genius" OUTPUT_PATH = "/app/suno/data/hoot/alignments/genius_v5_fix" elif input_args.dataset == "discogs": INPUT_PATH = "/home/tony/Data/Hoot/discogs_full_metas.json" DOWNLOAD_PATH = "/app/suno/data/audios/discogs" OUTPUT_PATH = "/app/suno/data/hoot/alignments/discogs_full_v5_fix" elif input_args.dataset == "ytm": INPUT_PATH = "/home/tony/Data/Hoot/ytm_metas.json" DOWNLOAD_PATH = "/app/suno/data/audios/ytm" OUTPUT_PATH = "/app/suno/data/hoot/alignments/ytm_v5_fix" elif input_args.dataset == "deezer": INPUT_PATH = "/home/tony/Data/Hoot/deezer_metas.json" DOWNLOAD_PATH = "/app/suno/data/audios/deezer" OUTPUT_PATH = "/app/suno/data/hoot/alignments/deezer_v5_fix" else: raise ValueError("Unknown dataset") print("start loading metadata file...") with open(INPUT_PATH, "r") as fp: metas = json.load(fp) print(f"check {input_args.dataset} metas size", len(metas)) # keep the same order as the download... if input_args.dataset == "discogs": metas.sort(key=lambda x: x["id"]) # this is stupid but if lyrics doesn't exist it is text :( for meta in tqdm.tqdm(metas): if "lyrics" not in meta: meta["lyrics"] = meta["text"] # decode logits uses model...this is really stupid but they are loaded into different keys _ = load_model( checkpoint_filepath="/home/tony/Data/checkpoints/hoot/2025-02-13_05-07-20/15k_ckpt.pt", tokenizer_filepath="/home/tony/Work/tony/hoot/tokenizers/v5/tokenizer_spe_bpe_v20481/tokenizer.model", ) # encode uses mode list preload_models( checkpoint_filepath="/home/tony/Data/checkpoints/hoot/2025-02-13_05-07-20/15k_ckpt.pt", tokenizer_filepath="/home/tony/Work/tony/hoot/tokenizers/v5/tokenizer_spe_bpe_v20481/tokenizer.model", ) align_metas = [] # will do this in batches of 100 n_threads = 16 # A100 8, A10 16 batch_size = 100 print(f"Total jobs {len(metas) / batch_size}.") for i in tqdm.tqdm(range(input_args.start_index, len(metas), batch_size)): start_index = i end_index = min(start_index + batch_size, len(metas)) if start_index > len(metas): continue output_json_path = os.path.join(OUTPUT_PATH, f"batch_{start_index}.json") # if os.path.exists(output_json_path): # continue file_size = os.stat(output_json_path).st_size if file_size < 1000: print("busted job", start_index) os.remove(output_json_path) else: continue # create an empyt file for place holder, in case there are racing conditions with open(output_json_path, "w") as fp: fp.write("") sub_metas = metas[start_index:end_index] # map to the local file paths sub_audio_paths = [] for meta in sub_metas: # if do local ... # audio_local_path = os.path.join( # DOWNLOAD_PATH, f"{meta.get('original_id', meta.get('id'))}.mp3" # ) # if os.path.exists(audio_local_path): # sub_audio_paths.append(audio_local_path) # else: # print(f"file {audio_local_path} does not exist, skipping...") # if do s3 ... sub_audio_paths.append(meta["audio_filepath"]) # print(f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}, start encoding set {i}") # batch size 192 will use ~ 20 GB, be safe outputs = encode_filepaths( sub_audio_paths, batch_size=512, # A100 256, A10 128 return_logits=True, dataloader_num_workers=n_threads, force_threads=True, ) # assert len(outputs) == len(sub_metas) # print(f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}, finish encoding set {i}") def get_meta_output_based_on_id(test_id): # print('working on test_id', test_id) try: meta = metas[test_id] meta_id = meta["id"] out = outputs[test_id - start_index] if out is None: # somehow the logit space is empty, could be missing file print("index ", test_id, " is empty, skipping...") return None # we need to lower case cause the tokenzier doesn't know upper case... basic_cleaned_lyrics = meta["lyrics"].lower() # pass in the lyrics for a boost of performance decoded_preds = decode_logits(out, basic_cleaned_lyrics) true_text_norm = clean_text(meta["lyrics"]) # print("pre compute cer for test_id", test_id) cer_val = round(get_cer(true_text_norm, decoded_preds), 3) # print("pre align for test_id", test_id) # print("decoded_preds, ", decoded_preds ) # print("basic_cleaned_lyrics", basic_cleaned_lyrics) word_timings = ctc_align(out.astype(np.float16), basic_cleaned_lyrics) aligned_meta = { "id": meta_id, "duration_s": meta["duration_s"], "alignment": word_timings, "cer_val": cer_val, "lyrics": meta["lyrics"], # keep the original, cased, lyrics "pred_lyrics": decoded_preds, } return aligned_meta except Exception as e: print("oops", e) return None aligned_outputs = thread_map( get_meta_output_based_on_id, list(range(start_index, end_index)), max_workers=n_threads, chunksize=1, disable=True, # print or not ) # print(f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}, finish align set {i}") # aligned_outputs = [] # for x in tqdm.tqdm(list(range(start_index, end_index))): # aligned_outputs.append(get_meta_output_based_on_id(x)) align_metas = [ aligned_output for aligned_output in aligned_outputs if aligned_output is not None ] # # DEBUG # for i, test_id in tqdm.tqdm(enumerate(range(start_index, end_index))): # print(test_id) # meta = metas[test_id] # meta_id = meta["id"] # out = outputs[test_id - start_index] # basic_cleaned_lyrics = meta["lyrics"].lower() # # pass in the lyrics for a boost of performance # decoded_preds = decode_logits(out, basic_cleaned_lyrics) # true_text_norm = clean_text(meta["lyrics"]) # # print("pre compute cer for test_id", test_id) # cer_val = round(get_cer(true_text_norm, decoded_preds), 3) # # print("pre align for test_id", test_id) # # print("decoded_preds, ", decoded_preds ) # # print("basic_cleaned_lyrics", basic_cleaned_lyrics) # word_timings = ctc_align(out.astype(np.float16), basic_cleaned_lyrics) # aligned_meta = { # "id": meta_id, # "duration_s": meta["duration_s"], # "alignment": word_timings, # "cer_val": cer_val, # "lyrics": basic_cleaned_lyrics, # "pred_lyrics": decoded_preds, # } # align_metas.append(aligned_meta) with open(output_json_path, "w") as fp: json.dump(align_metas, fp) print("DONE!!!")