import os import gc import re import math import copy import json import funcy import random import torch import tempfile import fasttext import numpy as np import collections import pandas as pd from tqdm import tqdm from joblib import Parallel, delayed from suno_utils.utils.display import capture_output from suno_utils.utils.lyrics import remove_speakers from suno_utils.utils.display import capture_output from suno_utils.utils.s3 import read_from_s3, check_s3_file_exists, open_from_s3 from suno_utils.harvest.youtube.constants.text_lang import BASE_TO_FASTTEXT_REMAP from suno_utils.utils.text import ( write_jsonl, read_jsonl, write_json, read_json, normalize_whitespace, ) # prep metas for each dataset with one row for each entry and (eg non-music will be filtered out) # genius_hq, youtube_music, freesound, jamendo, imslp, deezer, ytm_tagged # metas will contain text and tags: # text will be a simple string (eg for foreign etc) # text_segments will be a list of timestamped segments for aligned # tags will be (maybe empty) list of strings which could be any description (genre, title etc) # save this all to METAS_DIR LANG_ID_MODEL_FP = "s3://suno-data/georg/trained_models/chirp_v1/lid.176.bin" with capture_output(): text_lang_model = read_from_s3(LANG_ID_MODEL_FP, read_f=fasttext.load_model) def _get_text_lang(text): """get probability of input language for text""" text = text.replace("’", "'").lower() text = re.sub(r"\[.+?\]", " ", text) text = normalize_whitespace(text) out = text_lang_model.predict(text, k=1) lang_str = out[0][0] p_lang = out[1][0] lang = lang_str.split("__")[-1] lang = BASE_TO_FASTTEXT_REMAP.get(lang, lang) # if p_lang >= 0.8: # return lang return p_lang, lang def filter_by_audio_quality(audio_production_metas: list, percentile: float = 95): # compute the %1 and 99% percentile for the spectral centroid and loudness factor spectral_centroid = [ float(meta["features"]["spectral_centroid"]) for meta in audio_production_metas ] spectral_centroid = [value for value in spectral_centroid if np.isfinite(value)] loudness_factor = [ float(meta["features"]["loudness_factor"]) for meta in audio_production_metas ] loudness_factor = [value for value in loudness_factor if np.isfinite(value)] spectral_centroid_bot = np.percentile(spectral_centroid, 100 - percentile) spectral_centroid_top = np.percentile(spectral_centroid, percentile) loudness_factor_bot = np.percentile(loudness_factor, 100 - percentile) loudness_factor_top = np.percentile(loudness_factor, percentile) print( f"spectral centroid: {spectral_centroid_bot:.2f} - {spectral_centroid_top:.2f}" ) print(f"loudness factor: {loudness_factor_bot:.2f} - {loudness_factor_top:.2f}") # now filter the main mains and return a list of ids that pass the filter audio_quality_filtered_ids = [] for meta in tqdm(audio_production_metas): spectral_centroid_val = float(meta["features"]["spectral_centroid"]) loudness_factor_val = float(meta["features"]["loudness_factor"]) if ( spectral_centroid_val > spectral_centroid_bot and spectral_centroid_val < spectral_centroid_top and loudness_factor_val > loudness_factor_bot and loudness_factor_val < loudness_factor_top ): audio_quality_filtered_ids.append(meta["id"]) print( f"{len(audio_quality_filtered_ids):,}/{len(audio_production_metas):,} ids passed audio quality filter" ) return audio_quality_filtered_ids def _clean_tag(s): if s is None or not isinstance(s, str): return "" # remove symbold with special meaning s = re.sub(r"[\{\}\_\[\]]", " ", s) # squash whitespace return normalize_whitespace(s) def process_youtube_music(base_metas: list): # metas_plus = read_jsonl( # "/home/tony/Work/tony/FineTuning_chirp_v4/metadata/youtube_music_v9.jsonl" # ) # metas_plus_map = {meta["id"]: meta for meta in metas_plus} # load alignments alignments = read_jsonl( "/home/christian/code/christian/metadata/ytm_alignments_v11.jsonl" ) print(f"{len(alignments):,} alignments") aligned_lyrics_map = {} for meta_id, alignment in alignments: aligned_lyrics_map[meta_id] = alignment lang_counter = collections.Counter() print(f"{len(base_metas):,} clips") allowed_ids = set() for m in base_metas: if m["view_count"] < 1000: continue # if re.search(r"\blive\b", m["title"], flags=re.IGNORECASE): # continue allowed_ids.add(m["id"]) print(f"{len(allowed_ids):,} music clips") # 2,101,712 clips # 2,013,545 music clips metas = [] seen_ids = set() for m in base_metas: if m["id"] in seen_ids or m["id"] not in allowed_ids: continue if ( m["view_count"] < 1000 # default is 50 or m["duration_s"] < 1 * 60 # default is 100 or m["duration_s"] > 8 * 60 ): continue new_m = { "id": m["id"], "views": m["view_count"], "s3_filepath": m["s3_filepath"], "duration_s": m["duration_s"], } # get tags tags = [ m.get("genre_guess"), m.get("mood_guess"), ] tags = [t_clean for t in tags if len(t_clean := _clean_tag(t)) > 0] if len(tags) > 0: new_m["tags"] = list(set(tags)) # get private tags private_tags = [ m.get("title"), m.get("artists"), m.get("album"), ] private_tags = [ t_clean for t in private_tags if len(t_clean := _clean_tag(t)) > 0 ] private_tags += tags if len(private_tags) > 0: new_m["private_tags"] = list(set(private_tags)) new_m["dataset"] = "youtube_music" # get aligned lyrics text_segments = [] if m["id"] in aligned_lyrics_map: for mm in aligned_lyrics_map[m["id"]]: text_segments.append( { "text": remove_speakers(mm["text"]), "private_text": mm["text"], "start_s": mm["start_s"], "end_s": mm["end_s"], "vocal_start_s": mm["vocal_start_s"], "vocal_end_s": mm["vocal_end_s"], } ) new_m["text_segments"] = text_segments # get lyrics if ( "lang_guess" in m and "lyrics" in m and len(m["lyrics"]) <= 5120 # empirically verified ): lang_guess_1 = m["lang_guess"].split("-")[0].lower() lyrics = m["lyrics"] new_m["text"] = m["lyrics"] new_m["lang"] = lang_guess_1 metas.append(new_m) seen_ids.add(m["id"]) metas = sorted(metas, key=lambda x: "text" not in x) print(f"{len(metas):,} entries") print(f"{len([m for m in metas if 'text' in m]):,} entries with lyrics") print(f"{lang_counter['en']:,}/{len(metas):,} english") print(f"{lang_counter['foreign']:,}/{len(metas):,} foreign") # 2,013,545 entries # 674,466 entries with lyrics return metas def process_ytm_tagged(base_metas: list): print(f"{len(base_metas):,} clips") allowed_ids = set() for m in base_metas: if m["views"] < 100: continue allowed_ids.add(m["id"]) print(f"{len(allowed_ids):,} music clips") # 2,772,393 clips # 2,255,305 music clips metas = [] seen_ids = set() for m in base_metas: if m["id"] in seen_ids or m["id"] not in allowed_ids: continue metas.append(m) print(f"{len(metas):,} entries") return metas def process_pond5_music(base_metas: list): pd5_extra_metas = read_from_s3( "s3://suno-data/datasets/harvest/pond5_music/pond5_metas.jsonl", read_f=read_jsonl, ) print(len(pd5_extra_metas)) extra_info = {} seen_ids = set() genre_counter = collections.Counter() for m in tqdm(pd5_extra_metas): # if "s3_filepath" not in m: # continue if m["id"] in seen_ids: continue if m["duration"] < 60: continue if m["duration"] > (60 * 6): continue name = m.get("name", "") description = m.get("description", "") tags = m.get("tags", []) genre = m.get("genre", "") if not genre: continue genre_counter[genre.lower().strip()] += 1 if not description or not tags or not name: continue if len(description.strip().split()) < 10: continue if len(description.strip().split()) > 100: continue if not isinstance(tags, list) or len(tags) < 5 or len(tags) > 50: continue # if there are weird tags if max(len(tag) for tag in tags) > 30: continue extra_info[m["id"]] = { "name": name if isinstance(name, str) else "", "description": (description if isinstance(description, str) else ""), "tags": ( ([genre] if genre else []) + tags if isinstance(tags, list) else "" ), } seen_ids.add(m["id"]) print(len(seen_ids)) metas = [] total_duration = 0 n_c = 0 for m in tqdm(base_metas): # if n_c > 5: # break if m["id"] not in seen_ids: print("not in seen_ids") continue # if random.random() < 0.93: # continue n_c += 1 metas.append(m) total_duration += m["duration_s"] print(f"{len(metas):,} entries, total duration {total_duration/60/60:,.0f} hr") return metas def process_genius_hq(base_metas: list): metas = [] seen_ids = set() # filter based on audio quality # metas_plus = read_jsonl( # "/home/tony/Work/tony/FineTuning_chirp_v4/metadata/genius_hq_v9.jsonl" # ) # metas_plus_map = {meta["id"]: meta for meta in metas_plus} print("Loading audio production metas...") audio_production_metas = read_jsonl( "/home/christian/code/christian/metadata/genius_hq_metas_audio_production.jsonl" ) audio_quality_filtered_ids = filter_by_audio_quality(audio_production_metas) audio_quality_filtered_ids = set(audio_quality_filtered_ids) print("Loading alignments...") alignments = read_jsonl( "/home/christian/code/christian/metadata/genius_alignments_v11.jsonl" ) aligned_lyrics_map = {k: v for k, v in alignments} lang_counter = collections.Counter() for m in tqdm(base_metas): if ( m["genius_views"] < 50 or m["youtube_views"] < 100 # default is 50 or m["duration_s"] < 1 * 60 # default is 100 or m["duration_s"] > 8 * 60 or len(m["lyrics"]) > 6144 or len(m["lyrics"]) < 50 ): continue # filter based on audio quality if m["id"] not in audio_quality_filtered_ids: continue if m["id"] in seen_ids: continue new_m = { "id": m["id"], "original_id": m["genius_slug"], "views": m["youtube_views"], "lang": m["lang"].split("-")[0].lower(), } if new_m["lang"] == "en": lang_counter["en"] += 1 else: lang_counter["foreign"] += 1 # get aligned lyrics text_segments = [] if m["id"] in aligned_lyrics_map: for mm in aligned_lyrics_map[m["id"]]: text_segments.append( { "text": remove_speakers(mm["text"]), "private_text": mm["text"], "start_s": mm["start_s"], "end_s": mm["end_s"], "vocal_start_s": mm["vocal_start_s"], "vocal_end_s": mm["vocal_end_s"], } ) new_m["text_segments"] = text_segments new_m["duration_s"] = m["duration_s"] new_m["private_text"] = m["lyrics"] new_m["text"] = remove_speakers(m["lyrics"]) new_m["s3_filepath"] = m["audio_filepath"] new_m["dataset"] = "genius_hq" # add tags tags = [ normalize_whitespace(t_clean.replace("Genius", " ")) for t in m.get("tags_text", []) if len(t_clean := _clean_tag(t)) > 0 ] if len(tags) == 0: # we don't want to fine tune on things with no tags continue if len(tags) > 0: new_m["tags"] = list(set(tags)) # if any(t.lower() in blocked_tags for t in new_m["tags"]): # continue private_tag = _clean_tag(m["youtube_title"]) private_tags = tags if len(private_tag) > 0: private_tags += [private_tag] if len(private_tags) > 0: new_m["private_tags"] = list(set(private_tags)) metas.append(new_m) seen_ids.add(m["id"]) print(f"{len(metas):,} entries") print(f"{lang_counter['en']:,}/{len(metas):,} english") print(f"{lang_counter['foreign']:,}/{len(metas):,} foreign") return metas def process_jamendo(base_metas: str): # base_metas = read_from_s3( # "s3://suno-data/datasets/bundles/v1/jamendo/metas.jsonl", read_f=read_jsonl # ) metas = [] seen_ids = set() for m in base_metas: if m["id"] in seen_ids: continue metas.append(m) seen_ids.add(m["id"]) print(f"{len(metas):,} entries") # 55,609 entries return metas def process_imslp(base_metas_imslp: list): print(f"{len(base_metas_imslp):,} clips") blocked_tags = ["vocal", "chorus"] blocked_re = re.compile( r"(" + r")|(".join([s.lower() for s in blocked_tags]) + r")" ) allowed_ids = set() total_duration = 0 durations = [] for m in base_metas_imslp: tag_str = ( m.get("title", "") + " " + m.get("track_title", "") + " " + m.get("composer", "") + " " + m.get("language", "") + " " + m.get("recording_category", "") + " " + m.get("genre", "") + " " + m.get("instruments", "") ).lower() total_duration += m["duration_s"] durations.append(m["duration_s"]) if blocked_re.search(tag_str): continue if len(normalize_whitespace(tag_str)) == 0: continue allowed_ids.add(m["id"]) print(f"{len(allowed_ids):,} music clips") print(f"{total_duration / 3600 / 1000:0.2f}k hr") # 278,620 clips # 274,275 music clips genre_counter = collections.Counter() instruments_counter = collections.Counter() composer_counter = collections.Counter() for m in base_metas_imslp: genre_counter[m.get("genre", "").lower()] += 1 instruments_counter[m.get("instruments", "").lower()] += 1 composer_counter[m.get("composer", "").lower()] += 1 print(genre_counter) block_list = set( [ "simpson, daniel léo", "sousa, john philip", "various", "zhang, shuwen", "bartók, béla", "strauss, richard", "prokofiev, sergey", ] ) for composer in composer_counter: if composer_counter[composer] > 500: print(composer) composers_and_death = """ monteverdi, claudio - 1643 purcell, henry - 1695 corelli, arcangelo - 1713 bach, johann sebastian - 1750 telemann, georg philipp - 1767 scarlatti, domenico - 1757 handel, george frideric - 1759 mozart, wolfgang amadeus - 1791 beethoven, ludwig van - 1827 bach, carl philipp emanuel - 1788 gluck, christoph willibald - 1787 haydn, joseph - 1809 sor, fernando - 1839 donizetti, gaetano - 1848 schumann, robert - 1856 rossini, gioacchino - 1868 liszt, franz - 1886 wagner, richard - 1883 brahms, johannes - 1897 verdi, giuseppe - 1901 sullivan, arthur - 1900 debussy, claude - 1918 granados, enrique - 1916 mahler, gustav - 1911 puccini, giacomo - 1924 satie, erik - 1925 franck, césar - 1890 tchaikovsky, pyotr - 1893 gounod, charles - 1893 prokofiev, sergey - 1953 rachmaninoff, sergei - 1943 scriabin, aleksandr - 1915 paganini, niccolò - 1840 elgar, edward - 1934 villa-lobos, heitor - 1959 berlioz, hector - 1869 massenet, jules - 1912 rimsky-korsakov, nikolay - 1908 offenbach, jacques - 1880 franz lehár - 1948 saint-saëns, camille - 1921 sibelius, jean - 1957 franz schubert - 1828 manuel de falla - 1946 george gershwin - 1937 felix mendelssohn - 1847 michel rondeau - Unknown johann strauss jr. - 1899 john philip sousa - 1932 vincenzo bellini - 1835 georges bizet - 1875 georg philipp telemann - 1767 antonín dvořák - 1904 modest mussorgsky - 1881 antonio vivaldi - 1741 camille saint-saëns - 1921 jean sibelius - 1957 pierre-montan berton - 1780 bryan d. hoyt - Unknown claude le jeune - 1600 alexander nakarada - Unknown """ composers = [] allowed_composers = set() for l in composers_and_death.split("\n"): outs = l.split("-") if len(outs) < 2: continue composer = outs[0].strip() date = outs[1].strip() if not date.isnumeric(): continue if int(date) <= 1923: composers.append((int(date), composer)) allowed_composers.add(composer) print(sorted(allowed_composers)) # build metadata file metas = [] seen_ids = set() total_duration = 0 for m in base_metas_imslp: if m["id"] in seen_ids or m["id"] not in allowed_ids: continue if m.get("composer", "").lower() not in allowed_composers: continue new_m = {"id": m["id"]} # original tags spaced_tags = [] for t in [ m.get("genre", ""), m.get("title", ""), m.get("instruments", ""), ]: spaced_tags.extend(t.split(", ")) composer_tag = m.get("composer", "") if "," in composer_tag: composer_tag = " ".join(composer_tag.split(", ")[::-1]) spaced_tags.append(composer_tag) tags = [] for tag in spaced_tags: tag = tag.lower() no_pattern = r"no\. (\d+)" op_pattern = r"op\. (\d+)" bwv_pattern = r"bwv (\d+)" k_pattern = r"k\. (\d+)" l_pattern = r"l\. (\d+)" tag = re.sub(no_pattern, "", tag) tag = re.sub(op_pattern, "", tag) tag = re.sub(bwv_pattern, "", tag) tag = re.sub(k_pattern, "", tag) tag = re.sub(l_pattern, "", tag) if "voice" in tag.lower(): continue if "chorus" in tag.lower(): continue if len(tag.strip()) == 0: continue tag = tag.replace(" ", " ") tags.append(tag) if len(tags) > 0: new_m["tags"] = list(set(tags)) private_tags = [ _clean_tag(m.get("title")), _clean_tag(m.get("track_title", "")), _clean_tag(m.get("composer", "")), ] spaced_tags = [] for t in private_tags: spaced_tags.extend(t.split(", ")) private_tags = spaced_tags private_tags += tags private_tags = [t for t in private_tags if len(t) > 0] if len(private_tags) > 0: new_m["private_tags"] = list(set(private_tags)) new_m["s3_filepath"] = m.get("s3_filepath") metas.append(m) seen_ids.add(m["id"]) total_duration += m["duration_s"] print(f"{len(metas):,} entries, total duration {total_duration/60/60:,.0f} hr") return metas if __name__ == "__main__": datasets = [ ("imslp", 4_000), ("jamendo", 1_000), ("pond5_music", 2_000), ("ytm_tagged", 4_000), ("genius_hq", 80_000), ("youtube_music", 60_000), ] total_target_duration_hr = 0 for dataset, max_duration_hr in datasets: total_target_duration_hr += max_duration_hr print(f"Total target duration: {total_target_duration_hr:.2f} hours") new_dataset_metas = [] for dataset, max_duration_hr in datasets: print("Processing dataset: {}".format(dataset)) dataset_metas = read_jsonl( "/home/christian/code/christian/metadata/{}_metas.jsonl".format(dataset) ) print(f"{len(dataset_metas):,} entries") if dataset == "genius_hq": filtered_dataset_metas = process_genius_hq(dataset_metas) elif dataset == "youtube_music": filtered_dataset_metas = process_youtube_music(dataset_metas) elif dataset == "imslp": filtered_dataset_metas = process_imslp(dataset_metas) elif dataset == "jamendo": filtered_dataset_metas = process_jamendo(dataset_metas) elif dataset == "pond5_music": filtered_dataset_metas = process_pond5_music(dataset_metas) elif dataset == "ytm_tagged": filtered_dataset_metas = process_ytm_tagged(dataset_metas) else: raise ValueError(f"Unknown dataset: {dataset}") print(f"Filtered {len(dataset_metas)} to {len(filtered_dataset_metas)}") dataset_metas = filtered_dataset_metas # split into two lists metas_with_lyrics = [m for m in dataset_metas if "text" in m] metas_without_lyrics = [m for m in dataset_metas if "text" not in m] # filter down to max_duration_hr total_duration_s = 0 max_duration_s = ( max_duration_hr * 3600 if max_duration_hr is not None else float("inf") ) for meta_subset in [metas_with_lyrics, metas_without_lyrics]: meta_subset_indices = list(range(len(meta_subset))) # shuffle random.shuffle(meta_subset_indices) pbar = tqdm(meta_subset_indices) for idx in pbar: meta = meta_subset[idx] # duration_s = meta.get("duration_s", None) total_duration_s += meta["duration_s"] new_dataset_metas.append(meta) if idx % 100 == 0: pbar.set_description( f"Total duration: {total_duration_s / 3600:0.0f} hr" ) if max_duration_hr is not None and total_duration_s > max_duration_s: break # shuffle all the metas random.shuffle(new_dataset_metas) write_jsonl( new_dataset_metas, "/home/christian/code/christian/metadata/diffusion_mix_v2/metas.jsonl", )