import gc import os import random import re import json import shutil import tempfile from tokenizers import Tokenizer from contextlib import contextmanager from typing import List, Optional, Iterable from collections import defaultdict import funcy import numpy as np import torch from helpers import ( dist_barrier, download_s3_file, get_filename, read_jsonl, write_jsonl, print_with_time_master, respell_random_words_in_text, ) @contextmanager def _download_from_s3_if_needed(maybe_s3_filepath): tmp_filepath = maybe_s3_filepath if maybe_s3_filepath.startswith("s3://"): temp_dir = tempfile.TemporaryDirectory() filename = get_filename(maybe_s3_filepath, keep_ext=True) tmp_filepath = os.path.join(temp_dir.name, filename) download_s3_file(maybe_s3_filepath, tmp_filepath) yield tmp_filepath def load_tokenizer( tokenizer_filepath="s3://suno-data/georg/models/tokenizers/tokenizer_60k.json", ): with _download_from_s3_if_needed(tokenizer_filepath) as tmp_fp: tokenizer = Tokenizer.from_file(tmp_fp) tokenizer.add_special_tokens(["\n"]) tokenizer.pad_idx = tokenizer.token_to_id("[PAD]") return tokenizer def _space_repl(m): s = m.group() n_newline = s.count("\n") if n_newline >= 2: return "\n\n" elif n_newline == 1: return "\n" return " " def _simplify_whitespace(text, retain_newlines=True): """simplify while respecting up to 2 newlines""" if retain_newlines: text = re.sub(r"\s+", _space_repl, text).strip() else: text = re.sub(r"\s+", " ", text).strip() return text CASE_AUGMENT_FUNCS = [ str.upper, str.lower, str.capitalize, str.title, ] def _augment_tag(s): # case augment if random.random() >= 0.8: s = random.choice(CASE_AUGMENT_FUNCS)(s) # other misc formatting if random.random() >= 0.5: s = s.replace("-", " ").strip() return s def _clean_tag(s, retain_newlines=False): s = s.replace("[", " ").replace("]", " ") return _simplify_whitespace(s, retain_newlines=retain_newlines) MAX_TAG_LEN = 256 MAX_TOT_TAGS_LEN = 512 def augment_text_training(tags: List[str], lyrics: str): # Clean tags and remove empty ones tags = [clean_tag for tag in tags if len(clean_tag := _clean_tag(tag, retain_newlines=False)) > 0] # Randomly shuffle and truncate tags 50% of the time random.shuffle(tags) if random.random() < 0.5: if len(tags) > 0: tags = tags[: random.randint(1, len(tags))] # Augment each tag tags = [_augment_tag(tag) for tag in tags] # Choose a random character to join tags merge_char = random.choice([", ", " ", "; ", ",", ";"]) # Join tags, truncate if necessary tags_str = f"{merge_char.join([tag[:MAX_TAG_LEN] for tag in tags])[:MAX_TOT_TAGS_LEN]}".strip() # Simplify whitespace in lyrics lyrics = _simplify_whitespace(lyrics, retain_newlines=True) # 5% chance to lowercase lyrics if random.random() < 0.05: lyrics = lyrics.lower() # 5% chance to remove all newlines from lyrics if random.random() < 0.05: lyrics = re.sub(r"\n+", " ", lyrics) lyrics = lyrics.strip() # Choose a random character to join tags and lyrics merge_char = random.choice(["\n\n", "\n", " "]) # Combine tags and lyrics text = "" if len(tags_str) > 0: text += f"[{tags_str}]" if len(lyrics) > 0: if len(text) > 0: merge_char = random.choice(["\n\n", "\n", " "]) text += merge_char text += lyrics return text def prepare_text_inference(tags: List[str], lyrics: str): tags = [clean_tag for tag in tags if len(clean_tag := _clean_tag(tag, retain_newlines=False)) > 0] tags_str = f"{', '.join([tag[:MAX_TAG_LEN] for tag in tags])[:MAX_TOT_TAGS_LEN]}".strip() lyrics = _simplify_whitespace(lyrics, retain_newlines=True) # Combine tags and lyrics text = "" if len(tags_str) > 0: text += f"[{tags_str}]" if len(lyrics) > 0: if len(text) > 0: text += "\n\n" text += lyrics return text def fold_tensor(x: torch.Tensor, patch_size: int) -> torch.Tensor: """Fold first dimension into last dimension with given patch size.""" batch, dim = x.shape assert ( batch % patch_size == 0 ), f"First dimension {batch} must be divisible by patch_size {patch_size}" # Ensure contiguous memory layout before reshaping x = x.contiguous() # Reshape to (new_batch, patch_size, dim) x = x.reshape(-1, patch_size, dim) # Flatten patch dimension into feature dimension return x.reshape(batch // patch_size, patch_size * dim) def unfold_tensor(x: torch.Tensor, original_dim: int) -> torch.Tensor: """Inverse operation of fold_tensor.""" batch, dim = x.shape patch_size = dim // original_dim # Ensure contiguous memory layout before reshaping x = x.contiguous() # Reshape back to (batch, patch_size, original_dim) x = x.reshape(batch, patch_size, original_dim) # Flatten first two dimensions return x.reshape(batch * patch_size, original_dim) class GeneralMemmapMapDataset(torch.utils.data.Dataset): def __init__( self, dataset_dir: str, mode: str = "pretraining", # "dpo" or "pretraining" vae_memmap_filename: str = "data_vae_val.bin", semantic_memmap_filename: str = "data_semantic_val.bin", metas_filename: str = "metas_val.jsonl", info_filename: str = None, vae_dim: int = 128, vae_n_tokens: int = 3000, vae_use_float16: bool = True, semantic_n_tokens: int = 750, semantic_pad_token: int = 4000, cond_text_len: int = 2560, semantic_noise_level: float = 0.0, ctx_len: Optional[int] = None, vae_scale_factor: float = 2.5, vae_rate_hz: int = 25, prev_vae_ctx: bool = True, infill_vae_ctx: bool = False, semantic_rate_hz: int = 25, aligned_text_prob: float = 0.5, semantic_dropout: bool = True, semantic_corruption_prob: float = 0.0, respell_augment_prob: float = 0.0, is_training: bool = False, patch_size: int = 1, foreign_weight: float = 1.0, scale_vae_ctx: bool = False, noise_ctx: float = 0.0, always_pad_semantic: bool = False, semantic_skip_factors: List[int] = [], always_skip_semantic: bool = False, use_fine_guidance: bool = False, beta_min: float = 0.05, beta_max: float = 1.0, max_masks: int = 3, shared_ctx: bool = False, ): """For use in training conditional diffusion model. When a metas file is provided, the metadata is loaded and returned with the data. This can be used for lyric and tags conditioning. """ super().__init__() self.dataset_dir = dataset_dir self.mode = mode self.vae_dim = vae_dim self.vae_memmap_filename = vae_memmap_filename self.semantic_memmap_filename = semantic_memmap_filename self.vae_n_tokens = vae_n_tokens self.semantic_n_tokens = semantic_n_tokens self.is_training = is_training self.metas_filename = metas_filename self.vae_use_float16 = vae_use_float16 self.vae_scale_factor = vae_scale_factor self.semantic_pad_token = semantic_pad_token self.cond_text_len = cond_text_len self.ctx_len = ctx_len self.prev_vae_ctx = prev_vae_ctx self.infill_vae_ctx = infill_vae_ctx self.aligned_text_prob = aligned_text_prob self.patch_size = patch_size self.foreign_weight = foreign_weight self.scale_vae_ctx = scale_vae_ctx self.always_pad_semantic = always_pad_semantic self.semantic_noise_level = semantic_noise_level self.vae_rate_hz = vae_rate_hz self.semantic_rate_hz = semantic_rate_hz self.noise_ctx = noise_ctx self.semantic_skip_factors = semantic_skip_factors self.always_skip_semantic = always_skip_semantic self.info_filename = info_filename self.use_fine_guidance = use_fine_guidance self.beta_min = beta_min self.beta_max = beta_max self.max_masks = max_masks self.respell_augment_prob = respell_augment_prob self.shared_ctx = shared_ctx self.semantic_dropout = semantic_dropout self.semantic_corruption_prob = semantic_corruption_prob # open vae memmap vae_data = np.memmap( os.path.join(dataset_dir, vae_memmap_filename), dtype=np.float16 if vae_use_float16 else np.float32, mode="r", ) vae_data = vae_data.reshape(-1, vae_n_tokens, vae_dim) self.vae_data = vae_data # open semantic memmap if semantic_memmap_filename is not None: semantic_data = np.memmap( os.path.join(dataset_dir, semantic_memmap_filename), dtype=np.uint16, mode="r", ) semantic_data = semantic_data.reshape(-1, semantic_n_tokens, 1) self.semantic_data = semantic_data[:, :, 0] assert vae_data.shape[0] == semantic_data.shape[0] # must have same number of rows else: self.semantic_data = None # print(f"vae_data.shape: {vae_data.shape}, semantic_data.shape: {semantic_data.shape}") # load metas self.metas = read_jsonl(os.path.join(dataset_dir, metas_filename)) assert len(self.metas) == self.vae_data.shape[0] # Load info.json if it exists to filter metas self.filtered_indices = None if self.info_filename is not None and os.path.exists( os.path.join(dataset_dir, self.info_filename) ): with open(os.path.join(dataset_dir, self.info_filename), "r") as f: info_data = json.load(f) # Create a set of all allowed IDs across all dataset keys allowed_ids = defaultdict(set) for dataset_name, id_list in info_data.items(): allowed_ids[dataset_name].update(id_list) # Create a mapping from original indices to filtered indices self.filtered_indices = [] for i, meta in enumerate(self.metas): dataset = meta.get("dataset") if dataset == "discogs_subset": dataset = "discogs" if dataset in allowed_ids: if meta.get("id") in allowed_ids[dataset]: self.filtered_indices.append(i) assert len(self.filtered_indices) > 0, "No samples left after filtering" print_with_time_master( f"Filtered dataset using {self.info_filename}: {len(self.filtered_indices)} / {len(self.metas)} samples kept" ) # split metas into english and non-english (foreign) if self.filtered_indices is not None: # Use filtered indices if available filtered_metas = [self.metas[i] for i in self.filtered_indices] self.metas_english = [meta for meta in filtered_metas if meta.get("text_lang", "en") == "en"] self.metas_foreign = [meta for meta in filtered_metas if meta.get("text_lang", "en") != "en"] # Create mappings from meta to original index for later use self.meta_to_idx = { id(meta): self.filtered_indices[i] for i, meta in enumerate(filtered_metas) } else: self.metas_english = [meta for meta in self.metas if meta.get("text_lang", "en") == "en"] self.metas_foreign = [meta for meta in self.metas if meta.get("text_lang", "en") != "en"] # Create mappings from meta to original index for later use self.meta_to_idx = {id(meta): i for i, meta in enumerate(self.metas)} self.n_english = len(self.metas_english) self.n_foreign = len(self.metas_foreign) self.effective_foreign = int(self.n_foreign * self.foreign_weight) self.total_effective_length = self.n_english + self.effective_foreign block_duration_s = 30.0 # hard coded for now but need io_hz to be set print_with_time_master(f"English: {self.n_english * block_duration_s / 3600.0:0.1f} hrs") print_with_time_master( f"Foreign: {self.n_foreign * block_duration_s / 3600.0:0.1f} hrs (foreign_weight: {self.foreign_weight}) -> {self.effective_foreign * block_duration_s / 3600.0:0.1f} hrs" ) self.tokenizer = load_tokenizer() def __len__(self): return self.vae_data.shape[0] def _prepare_conditioning(self, text_codes: torch.Tensor) -> dict: """Prepare beta values and unconditioned text for training. Returns: Dictionary containing beta_scale, beta, beta_uncond, and uncond_text_codes """ w_scales = [] w_conds = [] uncond_text_codess = [] for _ in range(self.max_masks): # Generate beta scale beta_scale = self.beta_min + (self.beta_max - self.beta_min) * torch.rand( 1, dtype=torch.bfloat16 ) # Sometimes negative beta_scale = torch.where(torch.rand_like(beta_scale) < 0.2, beta_scale * -1, beta_scale) w_scale = 1 / beta_scale - 1 # w ranges from 0 to 20, -2 to -21 # Prepare beta and unconditioned text w_cond = torch.zeros_like(text_codes, dtype=torch.bfloat16) uncond_text_codes = text_codes.clone() # Sample a span of text codes span_start = torch.randint(0, text_codes.shape[0] - 1, (1,)) span_end = torch.randint(span_start + 1, text_codes.shape[0], (1,)) # Apply beta scale to the span and create unconditioned text w_cond[span_start:span_end] = w_scale uncond_text_codes[span_start:span_end] = self.tokenizer.pad_idx w_scales.append(w_scale) w_conds.append(w_cond) uncond_text_codess.append(uncond_text_codes) w_scales = torch.stack(w_scales) w_conds = torch.stack(w_conds) uncond_text_codess = torch.stack(uncond_text_codess) num_masks = random.randint(1, self.max_masks) w_scales[num_masks:] = 0 w_conds[num_masks:] = 0 if random.random() < 0.1: # sft, maybe unncessary w_scales[:] = 0 w_conds[:] = 0 return { "uncond_text_codes": uncond_text_codess, "w_scale": w_scales, "w_cond": w_conds, } def __getitem__(self, idx=None): if self.mode != "dpo": # by default we ignore the index # if idx is None: # Determine if we should return a foreign sample based on weighted probability foreign_prob = self.effective_foreign / self.total_effective_length use_foreign = random.random() < foreign_prob if use_foreign: # Select from foreign samples real_idx = random.randrange(self.n_foreign) meta = self.metas_foreign[real_idx] else: # Select from English samples real_idx = random.randrange(self.n_english) meta = self.metas_english[real_idx] idx = self.metas.index(meta) # get original index # make sure we don't go out of bounds... if idx >= self.vae_data.shape[0]: print(f"Warning: data index will be out of bounds,idx: {idx}, len: {self.vae_data.shape[0]}") idx = idx % self.vae_data.shape[0] # get vae embeddings, ensure float32, apply scale factor vae_embeds = torch.from_numpy(self.vae_data[idx].copy()).float() * self.vae_scale_factor # patch the embeddings if self.patch_size > 1: vae_embeds = fold_tensor(vae_embeds, self.patch_size) vae_embeds = vae_embeds.permute(1, 0) # channels, seq_len info = {} info["idx"] = idx # save the sampled index for checks # construct a mask based on the number of tokens in the metadata # we only have to pad the vae tokens as we use the semantic pad token as input? n_vae_tokens = self.metas[idx]["n_vae_tokens"] # true where we compute loss, false where we don't padding_mask = torch.ones(self.vae_n_tokens // self.patch_size) # padding_mask[n_vae_tokens // self.patch_size :] = 0 info["padding_mask"] = padding_mask.bool() # semantic codes if self.semantic_data is not None: semantic_codes = torch.from_numpy(self.semantic_data[idx].copy()).long() if self.always_pad_semantic: semantic_codes[:] = self.semantic_pad_token # always pad semantic elif self.always_skip_semantic: n = random.choice(self.semantic_skip_factors) n_phase = random.randint(0, n - 1) for nn in range(n - 1): shifted_idx = (nn + n_phase) % n semantic_codes[shifted_idx::n] = self.semantic_pad_token elif self.is_training: do_augment = random.random() <= 0.1 if do_augment: # 10% chance to augment semantic dice_roll = random.random() if dice_roll <= 0.1 and self.semantic_dropout: # 10% chance semantic is all pad semantic_codes[:] = self.semantic_pad_token elif 0.1 <= dice_roll <= 0.2: # 10% chance semantic is incomplete (pad the end) dice_roll_2 = random.random() if dice_roll_2 <= 0.33: # 33% chance pad with random amount at the end n = random.randint(1, self.semantic_n_tokens - 1) elif dice_roll_2 <= 0.66: # 33% chance to pad exactly 15s at the end (15s chunk) n = self.semantic_rate_hz * 15 else: # 33% chance to pad exactly 25s at the end (5s chunk) n = self.semantic_rate_hz * 25 semantic_codes[-n:] = self.semantic_pad_token # to improve robustness, we will add noise to the target vae latents # for parts of the semantic that are padded # semantic_to_vae_factor = self.vae_rate_hz / self.semantic_rate_hz # vae_index = int(n * semantic_to_vae_factor) # vae_embeds[..., -vae_index:] = torch.randn_like(vae_embeds[..., -vae_index:]) elif 0.2 <= dice_roll <= 0.3 and len(self.semantic_skip_factors) > 0: # 10% chance semantic is skip n = random.choice(self.semantic_skip_factors) n_phase = random.randint(0, n - 1) for nn in range(n - 1): shifted_idx = (nn + n_phase) % n semantic_codes[shifted_idx::n] = self.semantic_pad_token elif 0.3 <= dice_roll <= 0.4: # 10% chance semantic is fully random mask p_mask = random.random() mask = torch.bernoulli(torch.full((self.semantic_n_tokens,), 1 - p_mask)) semantic_codes[mask == 0] = self.semantic_pad_token if random.random() <= self.semantic_corruption_prob: p = random.random() # select randomize token indicies to corrupt mask = torch.bernoulli(torch.full((self.semantic_n_tokens,), p)) # select random token to replace with random_tokens = torch.randint(0, self.semantic_pad_token, (self.semantic_n_tokens,)) semantic_codes[mask == 1] = random_tokens[mask == 1] else: semantic_codes = (torch.ones(vae_embeds.shape[1]) * 4000).long() info["semantic_codes"] = semantic_codes # Use aligned_text_prob parameter to determine whether to use aligned text if random.random() <= self.aligned_text_prob: lyrics = self.metas[idx].get("text_aligned", self.metas[idx].get("text", "")) else: lyrics = self.metas[idx].get("text", "") # ensure lyrics is a string lyrics = str(lyrics) # respell random words in lyrics if english if self.respell_augment_prob > 0.0 and self.metas[idx].get("text_lang", None) == "en": N = random.randint(1, 10) # respell 1-10 words lyrics = respell_random_words_in_text(lyrics, N) tags = self.metas[idx].get("tags", []) if self.is_training: if random.random() <= 0.1: text = "" else: text = augment_text_training(tags, lyrics) else: text = prepare_text_inference(tags, lyrics) do_infill = False if self.infill_vae_ctx: # this will be the same as previous # zero means no context, one means give context empty_ctx_vae_embeds = torch.zeros_like(vae_embeds) empty_ctx_mask = torch.zeros(empty_ctx_vae_embeds.shape[1]).bool() infill_ctx_vae_embeds = empty_ctx_vae_embeds.clone() infill_ctx_mask = empty_ctx_mask.clone() do_infill = random.random() <= 0.1 # 10% chance to use current sample as context if do_infill: infill_ctx_vae_embeds = torch.from_numpy(self.vae_data[idx].copy()).float() if self.scale_vae_ctx: infill_ctx_vae_embeds = infill_ctx_vae_embeds * self.vae_scale_factor if self.patch_size > 1: infill_ctx_vae_embeds = fold_tensor(infill_ctx_vae_embeds, self.patch_size) infill_ctx_vae_embeds = infill_ctx_vae_embeds.permute(1, 0) # channels, seq_len infill_ctx_mask = torch.ones(infill_ctx_vae_embeds.shape[1]).bool() # infill mask when True we are providing context for that portion # randomly choose one of three masking patterns: # 1. Left to middle - keep right side # 2. Right to middle - keep left side # 3. Middle section - mask middle pattern = random.randint(0, 3) seq_len = infill_ctx_vae_embeds.shape[1] min_tokens = 25 # for the mask, 1 where we want to replace, 0 where we want to keep if pattern == 0: # Left masked - keep right side n = random.randint(min_tokens, seq_len - 1) infill_ctx_mask[:n] = False # Set left portion to False (masked) infill_ctx_vae_embeds[..., :n] = empty_ctx_vae_embeds[..., :n] elif pattern == 1: # Right masked - keep left side n = random.randint(min_tokens, seq_len - 1) infill_ctx_mask[-n:] = False # Set right portion to False (masked) infill_ctx_vae_embeds[..., -n:] = empty_ctx_vae_embeds[..., -n:] elif pattern == 2: # Middle masked - keep edges # Ensure there's at least one token on each edge middle_start = random.randint(min_tokens, seq_len // 2) middle_end = random.randint(middle_start + 1, seq_len - 1) infill_ctx_mask[middle_start:middle_end] = False # Set middle to False (masked) infill_ctx_vae_embeds[..., middle_start:middle_end] = empty_ctx_vae_embeds[ ..., middle_start:middle_end ] else: # pattern == 3 # Edges masked - keep middle left_size = random.randint(min_tokens, seq_len // 3) # Left masked region right_start = random.randint( seq_len // 2, seq_len - 1 ) # Start of right masked region # Mask left edge infill_ctx_mask[:left_size] = False infill_ctx_vae_embeds[..., :left_size] = empty_ctx_vae_embeds[..., :left_size] # Mask right edge infill_ctx_mask[right_start:] = False infill_ctx_vae_embeds[..., right_start:] = empty_ctx_vae_embeds[..., right_start:] if ( self.is_training and self.noise_ctx > 0.0 and random.random() <= 0.8 ): # add noise to the infill context vae_noise = torch.randn_like(vae_embeds) * random.random() * self.noise_ctx infill_ctx_vae_embeds += vae_noise # also update padding mask to include infill context # don't compute loss where we have infill context padding_mask = infill_ctx_mask.bool() info["padding_mask"] = ~padding_mask # invert mask since we are providing context here info["infill_ctx_vae"] = infill_ctx_vae_embeds info["infill_ctx_mask"] = infill_ctx_mask do_prev_ctx = False if self.prev_vae_ctx: empty_ctx_vae_embeds = torch.zeros_like(vae_embeds) empty_ctx_mask = torch.zeros(empty_ctx_vae_embeds.shape[1]).bool() ctx_vae_embeds = empty_ctx_vae_embeds.clone() ctx_mask = empty_ctx_mask.clone() # has context (use 50% chance to avoid overusing context) if ( # "prev_context_id" in self.metas[idx] shouldn't need this check self.metas[idx].get("start_s", 0) > 0 and idx - 1 >= 0 and (random.random() <= 0.75 or not self.is_training) and not do_infill # don't use context if we are infilling ): do_prev_ctx = True # for DPO, we save this and this is relative row shift (should only apply to dpo memmaps) offset_rows = self.metas[idx].get("offset_rows", 1) assert offset_rows >= 0 if offset_rows > 1: # make sure that the info maches assert self.metas[idx]["id_x"] == self.metas[idx - offset_rows]["id_x"] ctx_vae_embeds = torch.from_numpy(self.vae_data[idx - offset_rows].copy()).float() if self.scale_vae_ctx: # scale the context vector before input to model ctx_vae_embeds = ctx_vae_embeds * self.vae_scale_factor if self.patch_size > 1: ctx_vae_embeds = fold_tensor(ctx_vae_embeds, self.patch_size) ctx_vae_embeds = ctx_vae_embeds.permute(1, 0) # channels, seq_len ctx_mask = torch.ones(ctx_vae_embeds.shape[1]).bool() if self.is_training and random.random() <= 0.2: # 50% use shorter context n = random.randint(1, self.ctx_len - 1) ctx_vae_embeds[..., :n] = empty_ctx_vae_embeds[..., :n] ctx_mask[:n] = empty_ctx_mask[:n] else: ctx_vae_embeds = torch.zeros_like(vae_embeds) ctx_mask = torch.zeros(ctx_vae_embeds.shape[1]).bool() # shape: (seq_len) if self.is_training and self.noise_ctx > 0.0 and random.random() <= 0.8: vae_noise = torch.randn_like(vae_embeds) * random.random() * self.noise_ctx if random.random() <= 0.2: # zero out the noise at the end of the sequence n = np.random.randint(1, vae_embeds.shape[1] // 2) vae_noise[..., -n:] = 0.0 ctx_vae_embeds += vae_noise info["ctx_vae"] = ctx_vae_embeds info["ctx_mask"] = ctx_mask # Build condition tensors text_codes = self.tokenizer.encode(text).ids[: self.cond_text_len] text_codes = text_codes + [self.tokenizer.pad_idx] * max(0, self.cond_text_len - len(text_codes)) text_codes = torch.tensor(text_codes).long() if self.use_fine_guidance: # randomly mask individual text codes if random.random() <= 0.2: dropout_rate = random.random() mask = torch.rand(len(text_codes)) < dropout_rate text_codes[mask] = self.tokenizer.pad_idx conditioning = self._prepare_conditioning(text_codes) info.update(conditioning) else: info["text_codes"] = text_codes # info["conditioning"] = None # deteremine the mode # there are 4 modes: # 0. no infill, no prev ctx (first chunk) # 1. no infill, prev ctx (non-first chunk) # 2. infill, no prev ctx (short context infill) # 3. infill, prev ctx (long context infill) ctx_mode = 0 if do_infill and do_prev_ctx: ctx_mode = 3 elif do_infill: ctx_mode = 2 elif do_prev_ctx: ctx_mode = 1 info["ctx_mode"] = ctx_mode # we need special handling for shared context if self.shared_ctx: ones_mask = torch.ones_like(info["ctx_mask"]) zero_mask = torch.zeros_like(info["ctx_mask"]) if do_infill: info["infill_ctx_mask"] = ones_mask else: info["infill_ctx_mask"] = zero_mask if do_prev_ctx: info["ctx_mask"] = ones_mask else: info["ctx_mask"] = zero_mask return (vae_embeds, info) from suno_utils.audio import Audio import random from copy import deepcopy from tqdm import tqdm from suno_utils.tasks.dac_vae_fixed_25hz import ( preload_models as preload_codec_models, encode as encode_audio, ) def random_lengthen_positions(audio: Audio, duration_s: float): max_repeat = 1 + (duration_s // audio.duration_s) n_repeats = random.randint(1, max_repeat) positions = [] for i in range(n_repeats): positions.append(random.randint(-5, duration_s)) return positions def random_lengthen(audio: Audio, duration_s: float, positions: list[int] = None): if positions is None: positions = random_lengthen_positions(audio, duration_s) # print(positions) arr = audio.array_float out_wav = np.zeros( ( arr.shape[0], int(duration_s * audio.sample_rate), ) ) for i, pos in enumerate(positions): start_idx = pos * audio.sample_rate if start_idx < 0: a = arr[:, start_idx:] start_idx = 0 else: a = arr dur = min(a.shape[-1], out_wav.shape[-1] - start_idx) out_wav[:, start_idx : start_idx + dur] += a[:, :dur] return Audio.from_array_float(out_wav, sample_rate=audio.sample_rate, max_allowed_val=1000) class StemGroup: def __init__(self, stems: List[dict], make_prefixes: bool = False): self.stems = stems # TODO: this might be slow assert all(isinstance(s["audio"], Audio) for s in stems) self.audio = Audio.sum([s["audio"] for s in stems]) self.vae_embed = None self.prefixes = [] if make_prefixes: for i in range(len(self.stems)): self.prefixes.append(StemGroup(self.stems[: i + 1])) def __len__(self): return len(self.stems) def __getitem__(self, idx): return self.stems[idx] def __add__(self, other): if isinstance(other, dict): return StemGroup(self.stems + [other]) elif isinstance(other, StemGroup): return StemGroup(self.stems + other.stems) else: raise ValueError(f"Cannot add {type(other)} to StemGroup") @property def is_karaoke(self): return all(s["dataset_source"] == "karaoke" for s in self.stems) class BundleDownloaderDataset(torch.utils.data.IterableDataset): def __init__( self, dataset_dir: str, metas_filename: str, duration_s: float = 30, max_multi_instruments: int = 20, ): self.duration_s = duration_s self.max_multi_instruments = max_multi_instruments # load metas # TODO: dont load everything, just the ids self.metas = read_jsonl(os.path.join(dataset_dir, metas_filename)) # filter karaoke stems self.karaoke_metas = [m for m in self.metas if m["dataset_source"] == "karaoke"] print(f"Loaded {len(self.metas)} metas") # make bundles self.bundles = defaultdict(list) for meta in self.karaoke_metas: self.bundles[meta["song_id"]].append(meta) print(f"Made {len(self.bundles)} bundles") def __len__(self): return len(self.metas) def __getitem__(self, meta): meta = deepcopy(meta) # avoid mutating the original meta, memory leak try: # audio = get_sample_oracle_file_segment(meta["s3_filepath"], max_duration_s=self.duration_s) local_filepath = f"/app2/suno/data/stems/{meta['id']}.opus" if os.path.exists(local_filepath): audio = Audio.read_opus(local_filepath) else: raise FileNotFoundError(f"File not found: {local_filepath}") # audio = Audio.from_silence(self.duration_s, 48000, 2, 2) if audio.duration_s < self.duration_s - 1: audio = random_lengthen(audio, self.duration_s) audio = audio.get_segment(0, self.duration_s) audio = audio.pad_to_length(self.duration_s) meta["audio"] = audio return meta except Exception as e: print(f"Error: {e}. Resampling.") return self.__next__() def make_bundle(self): if random.random() <= 1.0: # karaoke bundle song_id = random.choice(list(self.bundles.keys())) stems = self.bundles[song_id] stem_group = StemGroup([self[meta] for meta in stems], make_prefixes=True) assert stem_group.is_karaoke, stem_group else: # random multi-instrument bundle n = random.randint(1, self.max_multi_instruments) metas = random.sample(self.metas, n) stems = [self[meta] for meta in metas] stem_group = StemGroup(stems, make_prefixes=True) return stem_group def __iter__(self): while True: yield self.make_bundle() class StemMemmapMapDataset(torch.utils.data.IterableDataset): def __init__( self, dataset_dir: str, bundle_iter: Iterable[dict], vae_memmap_filename: str = "data_vae_val.bin", metas_filename: str = "metas_val.jsonl", stem_type_map_filename: str = "stem_type_map.json", vae_dim: int = 128, vae_n_tokens: int = 3000, semantic_n_tokens: int = 750, semantic_pad_token: int = 4000, vae_use_float16: bool = True, cond_text_len=1536, ctx_len: Optional[int] = None, vae_scale_factor: float = 2.5, is_training: bool = False, patch_size: int = 1, scale_vae_ctx: bool = False, max_stems: int = 256, max_multi_instruments: int = 20, ): """For use in training conditional diffusion model. When a metas file is provided, the metadata is loaded and returned with the data. This can be used for tags conditioning. """ super().__init__() self.bundle_iter = bundle_iter self.dataset_dir = dataset_dir self.vae_dim = vae_dim self.vae_memmap_filename = vae_memmap_filename self.vae_n_tokens = vae_n_tokens self.semantic_n_tokens = semantic_n_tokens self.semantic_pad_token = semantic_pad_token self.is_training = is_training self.metas_filename = metas_filename self.vae_use_float16 = vae_use_float16 self.vae_scale_factor = vae_scale_factor self.cond_text_len = cond_text_len self.ctx_len = ctx_len self.patch_size = patch_size self.scale_vae_ctx = scale_vae_ctx self.max_stems = max_stems self.max_multi_instruments = max_multi_instruments # load metas self.metas = read_jsonl(os.path.join(dataset_dir, metas_filename)) print(f"Loaded {len(self.metas)} metas") self.id_to_idx = {m["id"]: idx for idx, m in enumerate(self.metas)} self.metas_by_id = {m["id"]: m for m in self.metas} self.stem_type_map = json.load(open(stem_type_map_filename)) self.tokenizer = load_tokenizer() self.embed_cache = {} preload_codec_models( checkpoint_filepath="s3://suno-data/minz/models/dac_vae_tuned_25hz.pth", device="cuda", # TODO: fsdp device with rank compile=True, ) def __len__(self): return len(self.metas) def __iter__(self): while True: try: bundle = next(self.bundle_iter) bundle = self.get_items_for_bundle(bundle) for _ in range(4 * len(bundle.prefixes)): yield self.get_item(bundle) while len(self.embed_cache) > 100: self.embed_cache.popitem() except Exception as e: raise e print(f"Error: {e}. Resampling.") def get_items_for_bundle(self, bundle: StemGroup): ########## Figure out task and sample a stem ########## def embed_audio(audio: Audio): arr = torch.from_numpy(audio.array_float).cuda() return torch.from_numpy(encode_audio(arr, normalize_volume=False)) stems = bundle.stems for stem in stems: if stem["dataset_source"] == "splice": stem["tags"] = stem["tags"] + [stem["tag_category"]] elif stem["dataset_source"] == "karaoke": stem["tags"] = [stem["track_name"]] else: raise ValueError(f"Unknown dataset source: {stem['dataset_source']}") for stem in tqdm(stems, desc="Embedding stems", disable=True): if stem["id"] not in self.embed_cache: self.embed_cache[stem["id"]] = embed_audio(stem["audio"]) stem["vae_embed"] = self.embed_cache[stem["id"]] random.shuffle(stems) for prefix in bundle.prefixes: if prefix.vae_embed is None: prefix.vae_embed = embed_audio(prefix.audio) bundle.vae_embed = bundle.prefixes[-1].vae_embed return bundle # for i in range(len(prefix.stems)): # yield prefix, prefix.stems[i] def get_item(self, bundle: StemGroup): task_weights = { "extract": 2.0, "remove": 2.0, "add": 2.0, # "wet_to_dry": 1.0, # "dry_to_wet": 1.0, } if not bundle.is_karaoke: task_weights["add"] = 0.0 if len(bundle.prefixes) <= 1: task_weights["remove"] = 0.0 task_weights["add"] = 0.0 # sample task task = random.choices(list(task_weights.keys()), weights=list(task_weights.values()))[0] out_meta = None if task == "extract": prefix_idx = random.randint(0, len(bundle.prefixes) - 1) prefix = bundle.prefixes[prefix_idx] in_vae_embeds = prefix.vae_embed.clone() # clone to be safe out_meta = random.choice(prefix.stems) out_vae_embeds = out_meta["vae_embed"].clone() elif task == "remove": prefix_idx = random.randint(1, len(bundle.prefixes) - 1) prefix = bundle.prefixes[prefix_idx] in_vae_embeds = prefix.vae_embed.clone() out_meta = prefix.stems[-1] out_vae_embeds = bundle.prefixes[prefix_idx - 1].vae_embed.clone() elif task == "add": prefix_idx = random.randint(0, len(bundle.prefixes) - 2) prefix = bundle.prefixes[prefix_idx] in_vae_embeds = prefix.vae_embed.clone() out_meta = random.choice(bundle.stems[prefix_idx + 1 :]) out_vae_embeds = out_meta["vae_embed"].clone() # elif task == "wet_to_dry": # in_meta = random.choice(self.metas_wet) # out_meta = self.uuid_to_meta[in_meta["dry_uuid"]] # elif task == "dry_to_wet": # in_meta = random.choice(self.metas_dry) # out_meta = self.uuid_to_meta[in_meta["wet_uuid"]] # get vae embeddings, ensure float32, apply scale factor out_vae_embeds = out_vae_embeds * self.vae_scale_factor in_vae_embeds = in_vae_embeds * self.vae_scale_factor ########## Process embeddings ########## # crop ctx ctx_embeds = out_vae_embeds[: self.ctx_len] out_vae_embeds = out_vae_embeds[self.ctx_len :] in_vae_embeds = in_vae_embeds[self.ctx_len :] if random.random() <= 0.5: ctx_embeds = torch.zeros_like(ctx_embeds) ctx_mask = torch.zeros(ctx_embeds.shape[0]).bool() else: ctx_mask = torch.ones(ctx_embeds.shape[0]).bool() # patch the embeddings if self.patch_size > 1: out_vae_embeds = fold_tensor(out_vae_embeds, self.patch_size) ctx_embeds = fold_tensor(ctx_embeds, self.patch_size) in_vae_embeds = fold_tensor(in_vae_embeds, self.patch_size) out_vae_embeds = out_vae_embeds.permute(1, 0) # channels, seq_len ctx_embeds = ctx_embeds.permute(1, 0) # channels, seq_len in_vae_embeds = in_vae_embeds.permute(1, 0) # channels, seq_len # randomly noise context for less degradation if random.random() <= 0.5: in_vae_embeds = in_vae_embeds + torch.randn_like(in_vae_embeds) * 1.0 if random.random() <= 0.5: ctx_embeds = ctx_embeds + torch.randn_like(ctx_embeds) * 1.0 info = {} info["stem_ctx_vae"] = in_vae_embeds info["stem_ctx_mask"] = torch.ones(in_vae_embeds.shape[1]).bool() info["ctx_vae"] = ctx_embeds info["ctx_mask"] = ctx_mask # construct a mask based on the number of tokens in the metadata padding_mask = torch.ones(out_vae_embeds.shape[1]) # we dont use semantic codes for stems semantic_codes = torch.full((self.semantic_n_tokens,), self.semantic_pad_token) info["semantic_codes"] = semantic_codes info["padding_mask"] = padding_mask.bool() tags = out_meta["tags"] if self.is_training: text = augment_text_training(tags, "") text = f"{task} {prepare_text_inference(tags, '')}" if self.is_training and random.random() <= 0.1: text = "" # Build condition tensors text_codes = self.tokenizer.encode(text).ids[: self.cond_text_len] text_codes = text_codes + [self.tokenizer.pad_idx] * max(0, self.cond_text_len - len(text_codes)) info["text_codes"] = torch.tensor(text_codes).long() info["task"] = task return (out_vae_embeds, info) class DPOGeneralMemmapMapDataset(GeneralMemmapMapDataset): """DPO dataset assumes that data is prepared in pairs. The even indices are negative samples. The odd indices are positive samples. This is a simple wrapper around the GeneralMemmapMapDataset. When we load data, we load it in pairs. Depending on the sampling, but we only load the even incdices. The odd indices are then derived and loaded as well. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) def __getitem__(self, idx): idx = idx // 2 * 2 vae_embeds, info = super().__getitem__(idx) vae_embeds_2, info_2 = super().__getitem__(idx + 1) assert torch.equal(info["semantic_codes"], info_2["semantic_codes"]) assert torch.equal(info["text_codes"], info_2["text_codes"]) # assert torch.equal(info["ctx_mask"], info_2["ctx_mask"]) # masks should be the same # TODO: load the real infill ctx vae in the future info["infill_ctx_vae"] = torch.zeros_like(info["ctx_vae"]) info["infill_ctx_mask"] = torch.zeros_like(info["ctx_mask"]) info["ctx_vae"] = torch.cat([info["ctx_vae"], info_2["ctx_vae"]], dim=0) # we concat the two vae embeddings as we assume that the odd indices are the positive samples return torch.cat([vae_embeds, vae_embeds_2], dim=0), info def shard_data( input_dir: str, output_dir: str, memmap_filenames: List[str], vae_n_tokens_memmap: int = 3000, vae_dim: int = 128, vae_use_float16: bool = True, semantic_n_tokens_memmap: int = 750, codec_n_tokens_memmap: int = 750, codec_n_codebooks: int = 12, allow_shard_reuse: bool = False, ): """ Args: input_dir (str): directory where the memmap files are stored output_dir (str): directory where the sharded memmap files will be stored (on local disk) memmap_filenames (List[str]): list of filenames to shard vae_n_tokens_memmap (int): number of tokens in the VAE memmap vae_dim (int): dimension of the VAE memmap vae_use_float16 (bool): whether to use float16 for VAE memmap semantic_n_tokens_memmap (int): number of tokens in the semantic memmap codec_n_tokens_memmap (int): number of tokens in the codec memmap codec_n_codebooks (int): number of codebooks in the codec memmap allow_shard_reuse (bool): whether to allow reusing shards if they already exist on local disk """ ddp_rank = int(os.environ["RANK"]) ddp_local_rank = int(os.environ["LOCAL_RANK"]) world_size = torch.distributed.get_world_size() n_gpus_per_node = torch.cuda.device_count() master_process = ddp_rank == 0 dist_barrier() if ddp_local_rank == 0: # check for existing shards if reusing if allow_shard_reuse: # check each file exists for fn in memmap_filenames: if not os.path.exists(os.path.join(output_dir, fn)): raise ValueError( f"shard {os.path.join(output_dir, fn)} does not exist on rank {ddp_rank}" ) if not allow_shard_reuse: # remove existing shards if os.path.exists(output_dir): shutil.rmtree(output_dir, ignore_errors=False) os.makedirs(output_dir) os.chmod(output_dir, 0o774) # load from data_dir and shard based on fraction that node should receive from_frac = ddp_rank / world_size to_frac = (ddp_rank + n_gpus_per_node) / world_size assert 0 <= from_frac <= 1 assert 0 <= to_frac <= 1 check_sizes = [] vae_dtype = np.float16 if vae_use_float16 else np.float32 for fn in memmap_filenames: if "val" in fn: shutil.copyfile( os.path.join(input_dir, fn), os.path.join(output_dir, fn), ) continue if "info" in fn: shutil.copyfile( os.path.join(input_dir, fn), os.path.join(output_dir, fn), ) continue if "vae" in fn: data = np.memmap(os.path.join(input_dir, fn), dtype=vae_dtype, mode="r") out_data = np.memmap( os.path.join(output_dir, fn), dtype=vae_dtype, mode="w+", shape=(1, vae_n_tokens_memmap, vae_dim), ) data = data.reshape(-1, vae_n_tokens_memmap, vae_dim) elif "semantic" in fn: data = np.memmap(os.path.join(input_dir, fn), dtype=np.uint16, mode="r") data = data.reshape(-1, semantic_n_tokens_memmap) out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.uint16, mode="w+", shape=(1, semantic_n_tokens_memmap), ) elif "codec" in fn: data = np.memmap(os.path.join(input_dir, fn), dtype=np.uint16, mode="r") data = data.reshape(-1, codec_n_tokens_memmap, codec_n_codebooks) out_data = np.memmap( os.path.join(output_dir, codec_n_tokens_memmap), dtype=np.uint16, mode="w+", shape=(1, codec_n_tokens_memmap, codec_n_codebooks), ) elif "metas" in fn: # open jsonl file data = read_jsonl(os.path.join(input_dir, fn)) else: raise ValueError(f"{fn} unknown format") check_sizes.append(len(data)) assert len(set(check_sizes)) == 1 # check if all data have same total number of rows from_idx = int(round(from_frac * len(data))) to_idx = int(round(to_frac * len(data))) idx_chunks = list(funcy.chunks(100_000, list(range(from_idx, to_idx)))) n_offs = 0 # idx_chunk is a list of indices for n_chunk, idx_chunk in enumerate(idx_chunks): if master_process: print(f"writing shard {n_chunk + 1}/{len(idx_chunks)} for {fn}") if "vae" in fn: # Add debug print here total_elements = (n_offs + len(idx_chunk)) * vae_n_tokens_memmap * vae_dim # Get the correct itemsize from numpy dtype_size = np.dtype(vae_dtype).itemsize # print(f"Attempting to map {total_elements * dtype_size / (1024**3):.2f} GB") out_data = np.memmap( os.path.join(output_dir, fn), dtype=vae_dtype, mode="r+", shape=( n_offs + len(idx_chunk), vae_n_tokens_memmap, vae_dim, ), ) elif "semantic" in fn: out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.uint16, mode="r+", shape=(n_offs + len(idx_chunk), semantic_n_tokens_memmap), ) elif "codec" in fn: out_data = np.memmap( os.path.join(output_dir, fn), dtype=np.uint16, mode="r+", shape=( n_offs + len(idx_chunk), codec_n_tokens_memmap, codec_n_codebooks, ), ) elif "metas" in fn: shard_metas = [] else: raise ValueError(f"{fn} unknown format") # iterate over indices and store in out_data for n, idx in enumerate(idx_chunk): if "metas" in fn: shard_metas.append(data[idx]) else: out_data[n_offs + n] = data[idx] # write metas to disk if "metas" in fn: write_jsonl(shard_metas, os.path.join(output_dir, fn), do_append=True) del shard_metas else: out_data.flush() n_offs += len(idx_chunk) del out_data gc.collect() del data gc.collect() print(f"done sharding on rank {ddp_rank}") # loop until we're done (we could just do barrier instead as long as timeout long enough) dist_barrier()