import numpy as np import os import tqdm from collections import defaultdict, Counter import pandas as pd from suno_utils.utils.text import write_jsonl, write_json SEMANTIC_N_CODEBOOKS = 4 COARSE_RATE_HZ = 25 SEMANTIC_CODEBOOK_SIZE = 4000 SEMANTIC_PAD_TOKEN = SEMANTIC_CODEBOOK_SIZE SEMANTIC_RATE_HZ = 25 COARSE_CODEBOOK_SIZE = 2048 COARSE_N_CODEBOOKS = 12 COARSE_PAD_TOKEN = COARSE_CODEBOOK_SIZE N_TOKENS_AUDIO = 25 * 8 * 60 # max 8 mins of audio def load_semantic_codes( s3_id: str, semantic_code_dir: str = "/app2/suno/data/semantic_code/crow/" ) -> np.ndarray: """Load semantic codes from separate directory for 4RVQ.""" semantic_path = os.path.join(semantic_code_dir, f"{s3_id}.npz") if not os.path.exists(semantic_path): raise FileNotFoundError(f"Semantic codes not found: {semantic_path}") codes = np.load(semantic_path)["codes"] semantic_codes = codes[:, :4] # Load first 4 codebooks (indices 0-3) return semantic_codes.astype(np.uint16) def load_conditioning_codes( s3_id: str, semantic_code_dir: str, task_name: str, start_s: float = None, end_s: float = None, max_duration_s: int = 240, ) -> np.ndarray: """Load conditioning semantic codes for a specific task. Args: s3_id: Clip ID to load semantic_code_dir: Directory containing semantic codes task_name: Task name for error messages start_s: Start time in seconds (optional) end_s: End time in seconds (optional) max_duration_s: Maximum duration for this conditioning type Returns: Trimmed semantic codes array """ try: codes = load_semantic_codes(s3_id, semantic_code_dir) # Apply time-based trimming if specified rate = 25 # semantic rate is 25 Hz for v5.0 if start_s is not None: start_idx = int(start_s * rate) codes = codes[start_idx:] if end_s is not None: end_idx = ( int(end_s * rate) if start_s is None else int((end_s - (start_s or 0)) * rate) ) codes = codes[:end_idx] # Apply max duration trimming (keep the start) max_tokens = int(max_duration_s * rate) if codes.shape[0] > max_tokens: codes = codes[:max_tokens] return codes except Exception as e: raise ValueError( f"Failed to load conditioning codes for {task_name} from {s3_id}: {e}" ) def make_dataset( input_df, output_data_dir: str, is_val=False, t_data_memmap=N_TOKENS_AUDIO, semantic_code_dir="/app2/suno/data/semantic_code/crow/", # directory for 4RVQ semantic codes ): print(f"t_data_memmap is set to: {t_data_memmap}") dset_type = "val" if is_val else "tr" out_mmap_path = os.path.join(output_data_dir, f"data_{dset_type}.bin") out_metas_path = os.path.join(output_data_dir, f"meta_{dset_type}.jsonl") out_info_filepath = os.path.join(output_data_dir, f"info_{dset_type}.json") def reshift(arr): sem_start_idx = 0 sem_end_idx = len(arr) - 1 semantic_arr = arr[:, :SEMANTIC_N_CODEBOOKS] # get array segments arr_s = semantic_arr[sem_start_idx:sem_end_idx, :].copy() assert arr_s.max() <= SEMANTIC_PAD_TOKEN # concat and stack arr_s = np.pad( arr_s, ((0, t_data_memmap - len(arr_s)), (0, 0)), constant_values=SEMANTIC_PAD_TOKEN, mode="constant", ) arr = arr_s arr = arr.astype(np.uint16) assert arr.shape == (t_data_memmap, SEMANTIC_N_CODEBOOKS) return arr # gather the data _ = np.memmap(out_mmap_path, dtype=np.uint16, mode="w+", shape=(1,)) n_offs = 0 curr_idx = 0 tot_duration_dict = defaultdict(int) datasets_info = defaultdict(dict) total_different_prompts = 0 total_different_tags = 0 total_different_neg_tags = 0 negative_prompt = "" negative_tag = "" negative_neg_tag = "" total_task_counter = Counter() total_error_task_counter = Counter() gender_counter = Counter() control_slider_counter = Counter() neg_tags_counter = Counter() prev_skip = None for i, row in tqdm.tqdm(input_df.iterrows(), total=len(input_df)): # we need to alternate between preference: neg, pos if i - 1 == prev_skip: print( f"WTF --> {i}, skip, preference: {row['preference']}, {row['s3_id']}, task: {row.get('task', '')}." ) continue # print(i, row) assert row["preference"] == (i % 2 == 1) is_positive = row["preference"] # need to use prompt_text -- prompt could be edited; # but this only works with recent data? since Aug 2024? current_prompt = row["prompt_text"] if not pd.isna(row["prompt_text"]) else "" current_tag = row["tags"] if not pd.isna(row["tags"]) else "" current_neg_tag = ( row["negative_tags"] if not pd.isna(row["negative_tags"]) else "" ) current_vocal_gender = ( row["vocal_gender"] if not pd.isna(row["vocal_gender"]) else "" ) # let's take this with a grain of salt # cause param_experiment could mask_control_slider current_control_sliders = ( row["control_sliders"] if not pd.isna(row["control_sliders"]) else {} ) control_slider_dict = {} if isinstance(current_control_sliders, dict): tag_weight = current_control_sliders.get("style_weight", None) audio_weight = current_control_sliders.get("audio_weight", None) mixed_weirdness = current_control_sliders.get("weirdness_constraint", None) if tag_weight is not None: control_slider_dict["tag_strength"] = tag_weight if audio_weight is not None: control_slider_dict["audio_strength"] = audio_weight if mixed_weirdness is not None: control_slider_dict["weird_strength"] = mixed_weirdness # this is a necessary augmentation if current_vocal_gender == "m": current_vocal_gender = "male" elif current_vocal_gender == "f": current_vocal_gender = "female" gender_counter[ current_vocal_gender if current_vocal_gender else "unspecified" ] += 1 control_slider_counter[ "has_control_slider" if control_slider_dict else "no_control_slider" ] += 1 neg_tags_counter["has_neg_tags" if current_neg_tag else "no_neg_tags"] += 1 if is_positive: if current_prompt != negative_prompt: # print("different prompts", negative_prompt, current_prompt) # ppl probably won't change negative... current_prompt = negative_prompt total_different_prompts += 1 if current_tag != negative_tag: # print("different tags", negative_tag, current_tag) # ppl probably won't change negative... current_tag = negative_tag total_different_tags += 1 if current_neg_tag != negative_neg_tag: current_neg_tag = negative_neg_tag total_different_neg_tags += 1 else: negative_prompt = current_prompt negative_tag = current_tag negative_neg_tag = current_neg_tag # Try to load semantic codes from separate directory try: arr = load_semantic_codes(row["s3_id"], semantic_code_dir) except Exception as e: print(f"Failed to load semantic codes for {row['s3_id']}: {e}") # Skip both samples in the pair when semantic loading fails if i % 2 == 1 and prev_skip != i - 1: print(f"WTF --> {i}, semantic loading failed for {row['s3_id']}.") prev_skip = i total_error_task_counter["semantic_loading"] += 1 continue # All NPZ files (including conditioning arrays) are loaded from semantic_code_dir local_path = os.path.join(semantic_code_dir, f"{row['s3_id']}.npz") if arr.shape[0] > t_data_memmap: print(f"Overflow at {i}, {row['s3_id']}: {arr.shape[0]} > {t_data_memmap}") arr = arr[:t_data_memmap, :] assert arr.shape[0] <= t_data_memmap # can't be longer than 8 mins :) if arr.shape[1] != 4: print( f"weird, semantic codes for {row['s3_id']}, with {arr.shape[1]} codebooks, expected 4" ) # Try to handle different shapes gracefully if arr.shape[1] < 4: # Pad with zeros if fewer codebooks arr = np.pad(arr, ((0, 0), (0, 4 - arr.shape[1])), constant_values=0) else: # Truncate if more codebooks arr = arr[:, :4] current_task = row.get("task", "") add_meta = { "dataset": f"perference_{int(row['preference'])}", "id": row["s3_id"], # this is the row s3_id "start_s": row["total_start_s"] if row["total_start_s"] >= 0 else 0, "vocal_start_s": None, # these are unfortunately missing for now "vocal_end_s": None, # these are unfortunately missing for now "tags": [current_tag], # tags is a list, do you know :) "neg_tags": current_neg_tag, # but we save negative tags as a string "control_tags": "" if pd.isna(row.get("control_tags", "")) else row.get("control_tags", ""), # save the control tags "gender": current_vocal_gender, # save the gender "control": control_slider_dict, # save the OG inputs, augment in training "text": current_prompt, # replace this with the fixed prompt "generated_start_index": 0, "user_id": row["user_id"], } # default gens, arr is what we look for. However... # all the conditioning shit # order is now: future + cover + artist known_dpo_tasks = [ "cover", "artist_consistency", "infill", "extend", "upload_extend", "cover_extend", "artist_extend", "artist_cover", "playlist_condition", "overpainting", "underpainting", "stem_condition", "", ] try: # Load metadata and conditioning IDs from NPZ if they exist temp_npz = None if os.path.exists(local_path): temp_npz = np.load(local_path) # Extract conditioning IDs from metadata metadata = row.get("metadata", {}) if isinstance(metadata, str): import json metadata = json.loads(metadata) # Set history_max_duration for crow/auk models (v5.0) - 240 seconds (4 minutes) history_max_duration = 240 # For extend tasks, reduce to half if current_task in [ "extend", "upload_extend", "artist_extend", "cover_extend", ]: history_max_duration = history_max_duration // 2 # 120 seconds max_condition_tokens = history_max_duration * 25 if current_task == "cover": # Load cover audio semantic codes with smart cropping cover_id = row.get("edited_clip_id") or metadata.get("edited_clip_id") if not cover_id: raise ValueError(f"No cover_id found for cover task") # Load full cover codes to check length cover_prompt_arr = load_semantic_codes(cover_id, semantic_code_dir) original_cover_length = int(cover_prompt_arr.shape[0] / 25) # Apply smart trimming logic from modal_runner if original_cover_length > history_max_duration: cover_duration_cutoff = max( 60, 2 * history_max_duration - original_cover_length ) else: cover_duration_cutoff = history_max_duration cover_prompt_arr = cover_prompt_arr[: int(cover_duration_cutoff * 25)] assert cover_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([cover_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = cover_prompt_arr.shape[0] add_meta["task"] = "cover" elif current_task == "artist_consistency": # Load artist audio semantic codes with smart cropping artist_id = row.get("edited_clip_id") or metadata.get("edited_clip_id") if not artist_id: raise ValueError(f"No artist_id found for artist_consistency task") # Load full artist codes to check length artist_prompt_arr = load_semantic_codes(artist_id, semantic_code_dir) original_artist_length = int(artist_prompt_arr.shape[0] / 25) # Apply smart trimming logic from modal_runner if original_artist_length > history_max_duration: artist_duration_cutoff = max( 60, 2 * history_max_duration - original_artist_length ) else: artist_duration_cutoff = history_max_duration artist_prompt_arr = artist_prompt_arr[ : int(artist_duration_cutoff * 25) ] assert artist_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([artist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "artist_consistency" elif current_task == "overpainting": # Load overpainting audio semantic codes overpainting_id = metadata.get("overpainting_clip_id") or row.get( "overpainting_clip_id" ) if not overpainting_id: raise ValueError(f"No overpainting_id found for overpainting task") # Crop to history_max_duration artist_prompt_arr = load_conditioning_codes( overpainting_id, semantic_code_dir, "overpainting", max_duration_s=history_max_duration, ) assert artist_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([artist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "overpainting" elif current_task == "underpainting": # Load underpainting audio semantic codes underpainting_id = metadata.get("underpainting_clip_id") or row.get( "underpainting_clip_id" ) if not underpainting_id: raise ValueError( f"No underpainting_id found for underpainting task" ) # Crop to history_max_duration artist_prompt_arr = load_conditioning_codes( underpainting_id, semantic_code_dir, "underpainting", max_duration_s=history_max_duration, ) assert artist_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([artist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "underpainting" elif current_task == "stem_condition": # Load stem audio semantic codes stem_id = metadata.get("stem_clip_id") or row.get("stem_clip_id") if not stem_id: raise ValueError(f"No stem_id found for stem_condition task") # Crop to history_max_duration stem_prompt_arr = load_conditioning_codes( stem_id, semantic_code_dir, "stem", max_duration_s=history_max_duration, ) assert stem_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([stem_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = stem_prompt_arr.shape[0] add_meta["task"] = "stem_condition" elif current_task == "playlist_condition": # For playlist, need to load multiple clip IDs and concatenate # This requires playlist_arr_ids from metadata or temp_npz if temp_npz and "playlist_arr_ids" in temp_npz: playlist_ids = [str(pid) for pid in temp_npz["playlist_arr_ids"]] elif metadata.get("playlist_clip_ids"): playlist_ids = metadata["playlist_clip_ids"] else: raise ValueError( f"No playlist_ids found for playlist_condition task" ) # Load and concatenate all playlist clips (line 798: keep 60s each) playlist_arrays = [] actual_lengths = [] for playlist_id in playlist_ids: playlist_clip_arr = load_conditioning_codes( playlist_id, semantic_code_dir, "playlist", max_duration_s=60 ) playlist_arrays.append(playlist_clip_arr) actual_lengths.append(playlist_clip_arr.shape[0]) playlist_prompt_arr = np.concatenate(playlist_arrays, axis=0) assert playlist_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([playlist_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = playlist_prompt_arr.shape[0] add_meta["task"] = "playlist_condition" add_meta["playlist_arr_len"] = actual_lengths elif current_task == "artist_cover": # Load both artist and cover semantic codes with balancing logic artist_id = metadata.get("artist_clip_id") or row.get("artist_clip_id") cover_id = ( metadata.get("cover_clip_id") or row.get("cover_clip_id") or row.get("edited_clip_id") ) if not artist_id or not cover_id: raise ValueError( f"Missing artist_id or cover_id for artist_cover task" ) # Load both without trimming first artist_prompt_arr = load_semantic_codes(artist_id, semantic_code_dir) cover_prompt_arr = load_semantic_codes(cover_id, semantic_code_dir) # Apply token balancing logic from modal_runner (lines 821-840) n_artist_tokens = artist_prompt_arr.shape[0] n_cover_tokens = cover_prompt_arr.shape[0] if n_artist_tokens + n_cover_tokens > max_condition_tokens: # Keep minimum 60s (1500 tokens) for artist n_artist_tokens = min(25 * 60, n_artist_tokens) # Use remaining for cover n_cover_tokens = max_condition_tokens - n_artist_tokens artist_prompt_arr = artist_prompt_arr[:n_artist_tokens] cover_prompt_arr = cover_prompt_arr[:n_cover_tokens] assert artist_prompt_arr.shape[0] <= t_data_memmap assert cover_prompt_arr.shape[0] <= t_data_memmap # the order is now: artist, cover, generated arr = np.concatenate([artist_prompt_arr, cover_prompt_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = ( artist_prompt_arr.shape[0] + cover_prompt_arr.shape[0] ) add_meta["cover_start_index"] = artist_prompt_arr.shape[0] add_meta["task"] = "artist_cover" elif ( current_task == "infill" or current_task == "infill_outro" or current_task == "infill_intro" ): # Load infill clip using edited_clip_id infill_clip_id = metadata.get("edited_clip_id") or row.get( "edited_clip_id" ) if not infill_clip_id: raise ValueError(f"No edited_clip_id found for infill task") # Load the full clip arr = load_semantic_codes(infill_clip_id, semantic_code_dir) assert arr.shape[1] == 4 # expecting 4RVQ here too # Get include_history_s and include_future_s to determine cropping include_history_s = metadata.get("include_history_s", 0) include_future_s = metadata.get("include_future_s", 0) # Convert to tokens (25 Hz) history_tokens = int(include_history_s * 25) future_tokens = int(include_future_s * 25) # Total length of the array total_length = arr.shape[0] # Crop off the history and future portions # The structure is: [history context] [generated part] [future context] if history_tokens > 0 or future_tokens > 0: # Remove history from start, future from end start_idx = history_tokens end_idx = ( total_length - future_tokens if future_tokens > 0 else total_length ) arr = arr[start_idx:end_idx, :] # Set indices for the infill structure # After cropping, the generated part starts at 0 # and future starts at (total - future_tokens) add_meta["generated_start_index"] = 0 if future_tokens > 0 and arr.shape[0] > future_tokens: # Future starts where generation ends add_meta["future_start_index"] = arr.shape[0] - future_tokens else: add_meta["future_start_index"] = arr.shape[0] # crop to max duration to avoid overflow context window if arr.shape[0] > t_data_memmap: print(f"Overflow at {i}: {arr.shape[0]} > {t_data_memmap}") arr = arr[:t_data_memmap, :] # Adjust future_start_index if we cropped if add_meta["future_start_index"] > t_data_memmap: add_meta["future_start_index"] = t_data_memmap add_meta["task"] = "infill" add_meta["infill_lyrics"] = metadata.get("infill_lyrics", "") or "" elif ( current_task == "extend" or current_task == "upload_extend" or current_task == "artist_extend" or current_task == "cover_extend" ): add_meta["task"] = "extend" # Load history clip semantic codes history_id = ( row.get("continued_parent") or metadata.get("history", [{}])[-1].get("id") if metadata.get("history") else None ) if not history_id: raise ValueError(f"No history_id found for extend task") # Load full history, then keep the LAST history_max_duration seconds history_arr = load_semantic_codes(history_id, semantic_code_dir) # Keep the END of the history (line 1102 in modal_runner) history_arr = history_arr[-int(history_max_duration * 25) :] assert history_arr.shape[1] == 4 arr = np.concatenate([history_arr, arr], axis=0) # crop to max duration arr = arr[:t_data_memmap, :] add_meta["generated_start_index"] = history_arr.shape[0] # Get history text if available history_text = "" if temp_npz and "history_text" in temp_npz: history_text = str(temp_npz["history_text"]) add_meta["text"] = ( history_text + "\n" + add_meta["text"] if history_text else add_meta["text"] ) if current_task == "artist_extend": add_meta["task"] = "artist_extend" artist_id = metadata.get("artist_clip_id") or row.get( "artist_clip_id" ) if not artist_id: raise ValueError(f"No artist_id for artist_extend task") # Load artist with smart cropping artist_prompt_arr = load_semantic_codes( artist_id, semantic_code_dir ) original_artist_length = int(artist_prompt_arr.shape[0] / 25) if original_artist_length > history_max_duration: artist_duration_cutoff = max( 60, 2 * history_max_duration - original_artist_length ) else: artist_duration_cutoff = history_max_duration artist_prompt_arr = artist_prompt_arr[ : int(artist_duration_cutoff * 25) ] assert artist_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([artist_prompt_arr, arr], axis=0) add_meta["history_start_index"] = artist_prompt_arr.shape[0] elif current_task == "cover_extend": add_meta["task"] = "cover_extend" cover_id = ( metadata.get("cover_clip_id") or row.get("cover_clip_id") or row.get("edited_clip_id") ) if not cover_id: raise ValueError(f"No cover_id for cover_extend task") # Load cover with smart cropping cover_prompt_arr = load_semantic_codes(cover_id, semantic_code_dir) original_cover_length = int(cover_prompt_arr.shape[0] / 25) if original_cover_length > history_max_duration: cover_duration_cutoff = max( 60, 2 * history_max_duration - original_cover_length ) else: cover_duration_cutoff = history_max_duration cover_prompt_arr = cover_prompt_arr[ : int(cover_duration_cutoff * 25) ] assert cover_prompt_arr.shape[0] <= t_data_memmap arr = np.concatenate([cover_prompt_arr, arr], axis=0) add_meta["history_start_index"] = cover_prompt_arr.shape[0] arr = arr[:t_data_memmap, :] except Exception as E: # there are a bunch of sth wrong with some data... print(f"{i}, {E}, {current_task}, {local_path}.") # don't keep the junks # they tend to be pairs anyways... # if not we wll find them and kick them out... if i % 2 == 1 and prev_skip != i - 1: print(f"WTF --> {i}, {E}, {current_task}, {local_path}.") prev_skip = i total_error_task_counter[current_task] += 1 continue # raise E assert add_meta.get("task", "") in known_dpo_tasks arr_duration = arr.shape[0] / 25 arr = reshift(arr) # print("after shift and pad", arr.shape) arr = arr.reshape( -1, ) # print(arr.shape) out_mm = np.memmap( out_mmap_path, dtype=np.uint16, mode="r+", shape=(n_offs + arr.size,), ) out_mm[n_offs : n_offs + arr.size] = arr # print(f"offset is: {n_offs}") # break # write it once out_mm.flush() del out_mm add_metas = [] add_metas.append(add_meta) tot_duration_dict[row["preference"]] += arr_duration total_task_counter[add_meta.get("task", "gen")] += arr_duration # print(add_metas) write_jsonl( add_metas, os.path.join(out_metas_path), do_append=bool(n_offs != 0), ) if "idx_list" not in datasets_info[add_meta["dataset"]]: datasets_info[add_meta["dataset"]]["idx_list"] = [curr_idx] else: datasets_info[add_meta["dataset"]]["idx_list"].append(curr_idx) curr_idx += 1 n_offs += arr.size write_json(datasets_info, out_info_filepath) print( f"Total {curr_idx} clips, {total_different_prompts} different prompts, {total_different_tags} different tags, {total_different_neg_tags} different negative tags" ) for k, v in tot_duration_dict.items(): print(f"{round(v / 60 / 60):,} hours of {k}") for k, v in total_task_counter.items(): print(f"{k}: {round(v / 60 / 60, 1)} hours") for k, v in total_error_task_counter.items(): print(f"🚨 Error {k}: {v}") print("\n--- Gender Distribution ---") total_gender_count = sum(gender_counter.values()) for gender, count in sorted(gender_counter.items()): percentage = (count / total_gender_count * 100) if total_gender_count > 0 else 0 print(f" {gender}: {count:,} ({percentage:.1f}%)") print("\n--- Negative Tags Usage ---") total_neg_tags = sum(neg_tags_counter.values()) for tag_status, count in sorted(neg_tags_counter.items()): percentage = (count / total_neg_tags * 100) if total_neg_tags > 0 else 0 print(f" {tag_status}: {count:,} ({percentage:.1f}%)") print("\n--- Control Slider Usage ---") for slider, count in sorted(control_slider_counter.items()): percentage = (count / curr_idx * 100) if curr_idx > 0 else 0 print(f" {slider}: {count:,} ({percentage:.1f}% of clips)") print("Done")