import random from typing import Iterator import numpy as np from pathlib import Path import json import torch import torch.nn.functional as F from torch.utils.data import IterableDataset, DataLoader from modules.fine import FineConfig from tqdm import tqdm class CodesDir: meta: dict[str, dict] def __init__(self, dir: Path): """ Run glockenspiel/suno_utils/suno_utils/scripts/gpt/concat_codes.py to generate the codes.npy and metas.jsonl files. """ self.dir = dir assert self.dir.exists(), self.dir self.codes_path = self.dir / "codes.npy" self.metas_path = self.dir / "metas.jsonl" assert self.codes_path.exists(), self.codes_path assert self.metas_path.exists(), self.metas_path self.meta = {} with open(self.metas_path, "r") as f: for line in tqdm(f, desc="Loading metas"): song_meta = json.loads(line) self.meta[song_meta["id"]] = song_meta self.is_mfcc = "mfcc" in self.codes_path.name self.dtype = np.float16 if self.is_mfcc else np.int16 self.n_codebooks = self.meta[list(self.meta.keys())[0]]["arr_shape"][1] self.codes = np.memmap(self.codes_path, dtype=self.dtype, mode="r").reshape(-1, self.n_codebooks) self.lengths = {k: v["arr_shape"][0] for k, v in self.meta.items()} # for sampling a random song self.lottery = [] for song_id, length in self.lengths.items(): self.lottery.extend([song_id] * (length // (25 * 10))) print(f"len(self.lottery): {len(self.lottery)}") def __len__(self): return len(self.meta) def get_song_codes(self, song_id: str, offset: int = 0, length: int = None): """(T, n_codebooks)""" if song_id not in self.meta: raise ValueError(f"song_id {song_id} not found") song_meta = self.meta[song_id] offset = song_meta["arr_offset"] + offset shape = song_meta["arr_shape"] length = length or shape[0] length = min(length, shape[0]) codes = self.codes[offset : offset + length].copy() return codes def song_samples(self, song_id: str): return self.lengths[song_id] def sample_song(self, split: str = "train"): # return self.train_lengths[0][0] split_idx = int(len(self.lottery) * 0.9) if split == "train": song_idx = random.randint(0, split_idx) elif split == "val": song_idx = random.randint(split_idx, len(self.lottery) - 1) return self.lottery[song_idx] class FineDataset(IterableDataset): def __init__( self, cfg: FineConfig, data_coarse: CodesDir, data_fine: CodesDir, split: str = "train", ): super().__init__() self.cfg = cfg self.split = split self.data_coarse = data_coarse self.data_fine = data_fine assert len(self.data_coarse) == len(self.data_fine), ( len(self.data_coarse), len(self.data_fine), ) def __iter__(self) -> Iterator: while True: coarse_codes, fine_codes = self.sample_codes() ( coarse_offset, x_coarse_arr, x_fine_arr, y_fine_arr, ) = self.build_training_sample(coarse_codes, fine_codes) x_coarse_arr = torch.from_numpy(x_coarse_arr) x_fine_arr = torch.from_numpy(x_fine_arr) y_fine_arr = torch.from_numpy(y_fine_arr) x, y = self.pad_and_stack(x_coarse_arr, x_fine_arr, y_fine_arr) yield torch.tensor([coarse_offset]), x, y # @profile def get_batch(self, batch_size: int, device: torch.device = None, noise_level=0.0): coarse_offset = None x_coarse_list = [] x_fine_list = [] y_fine_list = [] for _ in range(batch_size): coarse_codes, fine_codes = self.sample_codes() # self.build_training_sample(coarse_codes, fine_codes) ( coarse_offset, x_coarse_arr, x_fine_arr, y_fine_arr, ) = self.build_training_sample(coarse_codes, fine_codes, noise_level=noise_level) x_coarse_list.append(torch.from_numpy(x_coarse_arr)) x_fine_list.append(torch.from_numpy(x_fine_arr)) y_fine_list.append(torch.from_numpy(y_fine_arr)) x_coarse_arr = torch.stack(x_coarse_list) x_fine_arr = torch.stack(x_fine_list) y_fine_arr = torch.stack(y_fine_list) xs, ys = self.pad_and_stack(x_coarse_arr, x_fine_arr, y_fine_arr) if device == "cuda": # pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True) xs, ys = ( xs.pin_memory().to(device, non_blocking=True), ys.pin_memory().to(device, non_blocking=True), ) else: xs, ys = xs.to(device), ys.to(device) return coarse_offset, xs, ys def sample_codes(self): song_id = self.data_coarse.sample_song(split=self.split) while song_id not in self.data_fine.meta: # song is not in fine dataset, sample again song_id = self.data_coarse.sample_song(split=self.split) coarse_len = self.data_coarse.song_samples(song_id) # sample a chunk from coarse codes coarse_start = random.randint(0, coarse_len) fine_start = int(coarse_start * self.cfg.fine_rate_hz / self.cfg.coarse_rate_hz) coarse_codes = self.data_coarse.get_song_codes(song_id, coarse_start, self.cfg.coarse_samples) fine_codes = self.data_fine.get_song_codes(song_id, fine_start, self.cfg.fine_samples) return coarse_codes, fine_codes def build_training_sample(self, coarse_codes, fine_codes, noise_level=0.0): cfg = self.cfg coarse_codes = coarse_codes.T fine_codes = fine_codes.T # randomly add noise to coarse codes if noise_level > 0.0: # sample the actual noise level noise_level = np.random.uniform(0.0, noise_level) rand = np.random.rand(*coarse_codes.shape) coarse_codes[rand < noise_level] = np.random.randint( 0, cfg.coarse_codebook_size, size=(rand < noise_level).sum() ) # mask if cfg.coarse_masked_samples > 0: period = cfg.coarse_mask_period mask_len = cfg.coarse_masked_samples for i in range(0, coarse_codes.shape[-1], period): unmasked = period - mask_len coarse_codes[:, i + unmasked : i + period] = cfg.coarse_pad_token # pad if necessary if coarse_codes.shape[-1] < cfg.coarse_samples: print(f"coarse_codes.shape: {coarse_codes.shape}") print(f"cfg.coarse_samples: {cfg.coarse_samples}") coarse_codes = np.pad( coarse_codes, ((0, 0), (0, cfg.coarse_samples - coarse_codes.shape[-1])), "constant", constant_values=cfg.coarse_pad_token, ) if fine_codes.shape[-1] < cfg.fine_samples: fine_codes = np.pad( fine_codes, ((0, 0), (0, self.cfg.fine_samples - fine_codes.shape[-1])), "constant", constant_values=self.cfg.fine_pad_token, ) # build coarse x_coarse_arr = coarse_codes # build fine x_fine_arr = np.full( (cfg.fine_n_codebooks, cfg.t_fine - 1), cfg.fine_pad_token, dtype=np.int64, ) for n in range(cfg.fine_n_codebooks): offs = n * cfg.fine_shift_factor x_fine_arr[n, offs : offs + cfg.fine_samples] = fine_codes[n] y_fine_arr = x_fine_arr.copy() # combine audio and add x with infer token x_fine_arr = np.concatenate( [ np.array([[cfg.fine_infer_token]] * cfg.fine_n_codebooks), x_fine_arr[:, :-1], ], axis=-1, ) coarse_offset = x_coarse_arr.shape[-1] return coarse_offset, x_coarse_arr, x_fine_arr, y_fine_arr def pad_and_stack(self, x_coarse_arr, x_fine_arr, y_fine_arr): cfg = self.cfg # pad x to full codebooks x_coarse = F.pad( x_coarse_arr, (0, 0, 0, cfg.fine_n_codebooks), "constant", cfg.fine_pad_token, ) x_fine = F.pad( x_fine_arr, (0, 0, cfg.coarse_n_codebooks, 0), "constant", cfg.coarse_pad_token, ) # print(f"x_coarse.shape: {x_coarse.shape}") # print(f"x_fine.shape: {x_fine.shape}") x = torch.cat([x_coarse, x_fine], dim=-1) assert x.shape[-2:] == ( cfg.coarse_n_codebooks + cfg.fine_n_codebooks, cfg.block_size, ), (x.shape, cfg.block_size) return x, y_fine_arr if __name__ == "__main__": # test load_from_npz dac_12_dir = Path("/app/suno/victor/music_sample/concat/dac_2c_25_12") dac_100_dir = Path("/app/suno/victor/music_sample/concat/dac_2c_100x16") assert dac_12_dir.exists(), dac_12_dir coarse_dir = CodesDir(dac_12_dir) # coarse_dir = CodesDir(dac_12_dir) fine_dir = CodesDir(dac_12_dir) dataset = FineDataset( FineConfig(coarse_n_codebooks=12, fine_n_codebooks=12), coarse_dir, fine_dir, split="val", ) coarse_offset, xs, ys = dataset.get_batch(1, "cuda") print(coarse_offset, xs.shape, ys.shape) # test noise _ = dataset.get_batch(1, "cuda", noise_level=0.1) # test iter for s in dataset: print(s[0].shape, s[1].shape, s[2].shape) break # test dataloader dataloader = DataLoader(dataset, batch_size=64, num_workers=4, pin_memory=True) total = 100 for i, s in tqdm(enumerate(dataloader), total=total): if i == 0: print(s[0].shape, s[1].shape, s[2].shape) if i > total: break # # benchmark time to load 1000 batches # import time # start = time.time() # for _ in tqdm(range(100)): # dataset.get_batch(64, "cpu") # end = time.time() # test strided strided_config = FineConfig( coarse_n_codebooks=12, fine_n_codebooks=12, coarse_mask_period=25, coarse_masked_samples=12, ) strided_dataset = FineDataset( strided_config, coarse_dir, fine_dir, split="val", ) coarse_offset, xs, ys = strided_dataset.get_batch(1, "cuda") print(coarse_offset, xs.shape, ys.shape) assert xs[:, 0, 20] == strided_config.coarse_pad_token assert xs[:, 0, 10] != strided_config.coarse_pad_token