import os import glob import json import funcy import torch import torchaudio import numpy as np import soundfile as sf import torch.nn as nn import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm from dac.model.dac4 import DAC from suno_utils.utils.s3 import read_from_s3, download_s3_files from suno_utils.utils.text import read_jsonl, write_jsonl from transformers import Wav2Vec2FeatureExtractor from transformers import WhisperProcessor, WhisperForConditionalGeneration from suno_utils.models.mert.modeling_MERT import MERTModel, MERTConfig from suno_utils.tasks.mert_25 import ClusterModel def load_models(device="cpu"): # ---------------- load vae model for embedding ---------------- print("Loading VAE model...") checkpoint_filepath = "s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth" load_f = funcy.partial(torch.load, map_location="cpu") if checkpoint_filepath.startswith("s3://"): sd = read_from_s3(checkpoint_filepath, read_f=load_f) else: sd = load_f(checkpoint_filepath) sd["metadata"]["kwargs"] = { k: v for k, v in sd["metadata"]["kwargs"].items() if k in DAC.__init__.__code__.co_varnames } vae_model = DAC(**sd["metadata"]["kwargs"]) vae_model.load_state_dict(sd["state_dict"]) vae_model.eval() vae_model.to(device) # ---------------- load semantic model ---------------- print("Loading semantic model...") centroids_filepath = "s3://suno-data/georg/models/semantic/mert_25_2x4k.npy" checkpoint_filepath = "s3://suno-data/georg/models/semantic/mert_25.pt" # MERT preprocess_config = "/home/christian/code/glockenspiel/suno_utils/suno_utils/models/mert/preprocessor_config.json.py" with open(preprocess_config) as f: mert_processor = Wav2Vec2FeatureExtractor(**json.load(f)) mert_config = "/home/christian/code/glockenspiel/suno_utils/suno_utils/models/mert/config_small_25hz.json.py" with open(mert_config) as f: cfg = MERTConfig(**json.load(f)) mert_model = MERTModel(cfg) # assert processor.sampling_rate == SAMPLE_RATE _torch_load_p = funcy.partial(torch.load, map_location="cpu") if checkpoint_filepath.startswith("s3://"): sd = read_from_s3(checkpoint_filepath, read_f=_torch_load_p) else: sd = _torch_load_p(checkpoint_filepath) mert_model.load_state_dict(sd) del sd mert_model.eval() mert_model.to(device) if centroids_filepath.startswith("s3://"): cluster_centers = read_from_s3(centroids_filepath, read_f=np.load) else: cluster_centers = np.load(centroids_filepath) cluster_model = ClusterModel(cluster_centers) cluster_model.eval() cluster_model.to(device) # ---------------- load acoustic codec model ---------------- print("Loading acoustic codec model...") # filter args checkpoint_filepath = "s3://suno-data/georg/models/codec/dac_2c_25x12.pt" load_f = funcy.partial(torch.load, map_location="cpu") if checkpoint_filepath.startswith("s3://"): sd = read_from_s3(checkpoint_filepath, read_f=load_f) else: sd = load_f(checkpoint_filepath) sd["metadata"]["kwargs"] = { k: v for k, v in sd["metadata"]["kwargs"].items() if k in DAC.__init__.__code__.co_varnames } codec_model = DAC(**sd["metadata"]["kwargs"]) codec_model.load_state_dict(sd["state_dict"]) codec_model.eval() codec_model.to(device) # ---------------- load whisper model ---------------- print("Loading whisper model...") whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base") whisper_model = WhisperForConditionalGeneration.from_pretrained( "openai/whisper-base" ) whisper_model.to(device) return ( vae_model, mert_model, mert_processor, cluster_model, codec_model, whisper_processor, whisper_model, ) class S3AudioLoader(torch.utils.data.Dataset): def __init__(self, metas: list, block_size: int = 8388608, max_length_s: int = 600): self.metas = metas self.block_size = block_size self.max_length_s = max_length_s self.tmp_dir = "/mnt/localdisk/cjs" os.makedirs(self.tmp_dir, exist_ok=True) def __len__(self): return len(self.metas) def __getitem__(self, idx): meta = self.metas[idx] meta_id = meta["id"] text_tags = meta["tags_text"] s3_filepath = meta["audio_filepath"] audio_filepath = os.path.join(self.tmp_dir, os.path.basename(s3_filepath)) # check if file is already downloaded if not os.path.exists(audio_filepath): try: download_s3_files([s3_filepath], [audio_filepath]) except Exception as e: print(f"Error downloading {s3_filepath}: {e}") return ( torch.tensor(False), torch.tensor(False), torch.tensor(False), torch.tensor(False), torch.tensor(False), ) try: # load the file data, sr = sf.read(audio_filepath) audio = torch.from_numpy(data).float() os.remove(audio_filepath) except Exception as e: print(f"Error loading {audio_filepath}: {e}") return ( torch.tensor(False), torch.tensor(False), torch.tensor(False), torch.tensor(False), torch.tensor(False), ) max_length_samples = self.max_length_s * sr if audio.shape[-1] > max_length_samples: audio = audio[:, :max_length_samples] # resample to 48khz audio_48khz = torchaudio.functional.resample(audio, sr, 48000) # ensure 48khz audio is stereo if audio_48khz.shape[0] == 1: audio_48khz = audio_48khz.repeat(2, 1) elif audio_48khz.shape[0] > 2: audio_48khz = audio_48khz[:2] # resample to 24khz audio_24khz = torchaudio.functional.resample(audio.mean(dim=0), sr, 24000) # resample to 16khz # audio_16khz = torchaudio.functional.resample(audio.mean(dim=0), sr, 16000) # remove the audio file os.remove(audio_filepath) if audio_48khz.shape[-1] > self.block_size: audio_48khz = torch.split(audio_48khz, self.block_size, dim=-1) else: audio_48khz = [audio_48khz] if audio_24khz.shape[-1] > self.block_size: # chunk the audio into blocks of block_size audio_24khz = torch.split(audio_24khz, self.block_size, dim=-1) else: audio_24khz = [audio_24khz] # if audio_16khz.shape[-1] > self.block_size: # # chunk the audio into blocks of block_size # audio_16khz = torch.split(audio_16khz, self.block_size, dim=-1) return audio_48khz, audio_24khz, meta_id, s3_filepath, text_tags # --------- distributed setup ------------ def setup(): dist.init_process_group("nccl") def cleanup(): dist.destroy_process_group() def run_inference(out_dir: str): local_rank = int(os.environ["LOCAL_RANK"]) global_rank = int(os.environ["RANK"]) world_size = int(os.environ["WORLD_SIZE"]) setup() torch.cuda.set_device(local_rank) # load models ( vae_model, mert_model, mert_processor, cluster_model, codec_model, whisper_processor, whisper_model, ) = load_models(device=local_rank) # wrap models in DDP # vae_model = DDP(vae_model, device_ids=[local_rank]) # codec_model = DDP(codec_model, device_ids=[local_rank]) # mert_model = DDP(mert_model, device_ids=[local_rank]) # cluster_model = DDP(cluster_model, device_ids=[local_rank]) # load dataset dataset = S3AudioLoader(metas) sampler = DistributedSampler(dataset, num_replicas=world_size, rank=global_rank) dataloader = DataLoader( dataset, batch_size=1, sampler=sampler, shuffle=False, num_workers=4 ) with torch.no_grad(): for batch in dataloader: audio_48khz, audio_24khz, meta_id, s3_filepath, text_tags = batch if audio_48khz.item() == False: continue meta_id = meta_id[0] text_tags = list(text_tags[0]) s3_filepath = s3_filepath[0] # check if we have the alignment segments = alignemnts_map.get(meta_id, None) if segments is None: print(f"Skipping {meta_id} as no alignment found") continue # move each tensor to device audio_48khz = [block.to(local_rank) for block in audio_48khz] audio_24khz = [block.cuda(local_rank) for block in audio_24khz] # -------------- first embed with VAE -------------- out_filename = f"{meta_id}.npz" out_filepath = os.path.join(out_dir, "vae", out_filename) # if file does not exist, embed and save if not os.path.exists(out_filepath): # embed entire audio stream vae_latents_list = [] with torch.no_grad(): for audio_block in audio_48khz: vae_latents = vae_model.encode(audio_block)["z"] vae_latents = vae_latents.cpu().numpy() print(vae_latents.shape) vae_latents_list.append(vae_latents) vae_latents = np.concatenate(vae_latents_list, axis=-1) print(vae_latents.shape) # save npz of embeddings to disk using id as filename np.savez(out_filepath, vae_latents=vae_latents) # -------------- then embed with semantic model -------------- out_filepath = os.path.join(out_dir, "semantic", out_filename) N_LAYER_EMBED = 7 n_codebooks = 2 semantic_codes_list = [] for audio_block in audio_24khz: attention_mask = torch.ones( audio_block.shape, dtype=torch.int32, device=local_rank ) outputs = mert_model( input_values=audio_block, attention_mask=attention_mask, output_hidden_states=True, ) hidden_states = outputs.hidden_states[N_LAYER_EMBED] # b, t, d semantic_codes = cluster_model.encode( hidden_states, n_codebooks=n_codebooks ) print(semantic_codes.shape) semantic_codes_list.append(semantic_codes.cpu().numpy()) semantic_codes = np.concatenate(semantic_codes_list, axis=1) print(semantic_codes.shape) # if file does not exist, embed and save # if not os.path.exists(out_filepath): # semantic_codes = semantic_encode(audio_24khz) # # has shape (seq_len, 2) # semantic_codes = np.concatenate(semantic_codes, axis=0) # # save npz of embeddings to disk using id as filename np.savez(out_filepath, semantic_codes=semantic_codes) # -------------- then embed with acoustic codec model -------------- out_filepath = os.path.join(out_dir, "codec", out_filename) # if file does not exist, embed and save if not os.path.exists(out_filepath): # embed entire audio stream with torch.no_grad(): codec_codes_list = [] for audio_block in audio_48khz: codec_codes = codec_model.encode(audio_block)["codes"] codec_codes = codec_codes.cpu().numpy() print(codec_codes.shape) codec_codes_list.append(codec_codes) codec_codes = np.concatenate(codec_codes_list, axis=-1) print(codec_codes.shape) # save npz of embeddings to disk using id as filename np.savez(out_filepath, codec_codes=codec_codes) out_filename = f"{meta_id}.json" out_filepath = os.path.join(out_dir, "meta", out_filename) # -------------- then embed with whisper model -------------- if False: # audio_16khz = [block.numpy() for block in audio_16khz] # list of numpy tensors # if file does not exist, embed and save input_features = whisper_processor( audio_16khz, sampling_rate=16000, return_tensors="pt", ).input_features input_features = input_features.to(local_rank) # Generate token ids predicted_ids = whisper_model.generate( input_features, # condition_on_previous_text=False, ) # Decode token ids to text transcriptions = whisper_processor.batch_decode( predicted_ids, skip_special_tokens=True, ) # combine transcriptions into a single string ranscriptions = " ".join(transcriptions) meta = { "id": meta_id, "s3_filepath": s3_filepath, # "transcription": transcriptions, "segments": segments, "tags": text_tags, } print(meta) # save transcriptions to disk using id as filename with open(out_filepath, "w") as f: json.dump(meta, f, indent=2) cleanup() if __name__ == "__main__": dataset_name = "genius_hq" out_dir = f"/app/suno/christian/data/{dataset_name}/" os.makedirs(out_dir, exist_ok=True) os.makedirs(os.path.join(out_dir, "vae"), exist_ok=True) os.makedirs(os.path.join(out_dir, "codec"), exist_ok=True) os.makedirs(os.path.join(out_dir, "semantic"), exist_ok=True) os.makedirs(os.path.join(out_dir, "meta"), exist_ok=True) metas_filepath = f"metadata/{dataset_name}_metas.jsonl" alignments_filepath = f"metadata/{dataset_name}_alignments_v6.jsonl" # load the main metas and alignments metas = read_jsonl(metas_filepath, progress=True) alignments = read_jsonl(alignments_filepath, progress=True) alignemnts_map = {} for alignment in alignments: alignment_id = alignment[0] segments = alignment[1] alignemnts_map[alignment_id] = segments # launch run_inference(out_dir)