import time import random import torch import lightning as L import torch.nn.functional as F import numpy as np from torch import nn from sklearn import metrics as skm from einops import rearrange def print_model_params(model): total_params = 0 trainable_params = 0 for _, parameter in model.named_parameters(): params = parameter.numel() total_params += params if parameter.requires_grad: trainable_params += params print(f"Total Params: {total_params:,}, Trainable_params: {trainable_params:,}") class DittoLitModule(L.LightningModule): def __init__( self, model, learning_rate=1e-4, ): super().__init__() self.lr = learning_rate self.model = model print_model_params(model) self.loss_function = nn.CrossEntropyLoss() self.save_hyperparameters(ignore=["model"]) ( self.inference_music_embeddings, self.inference_text_embeddings, self.inference_ids, self.start_ss, self.lyrics, ) = [], [], [], [], [] self.counter = 0 def get_metrics(self, music_z, text_z, logit_scale): metrics = {} # scaled dot product similarity logits_per_music = logit_scale * music_z @ text_z.t() logits_per_text = logits_per_music.t() labels = torch.arange(music_z.shape[0]).long().to(logit_scale.device) # get cross entropy loss loss = ( self.loss_function(logits_per_music, labels) + self.loss_function(logits_per_text, labels) ) / 2 metrics["loss"] = loss # get accuracy metrics["acc_music"] = skm.accuracy_score( labels.detach().cpu().numpy(), logits_per_music.argmax(dim=1).detach().cpu().numpy(), ) metrics["acc_text"] = skm.accuracy_score( labels.detach().cpu().numpy(), logits_per_text.argmax(dim=1).detach().cpu().numpy(), ) # get ranking metrics logits = { "music_to_text": logits_per_music.detach().cpu(), "text_to_music": logits_per_text.detach().cpu(), } ground_truth = torch.arange(len(text_z)).view(-1, 1) for name, logit in logits.items(): ranking = torch.argsort(logit, descending=True) preds = torch.where(ranking == ground_truth)[1] preds = preds.detach().cpu().numpy() metrics[f"{name}_mean_rank"] = preds.mean() + 1 metrics[f"{name}_mdeidan_rank"] = np.floor(np.median(preds)) + 1 for k in [1, 5, 10]: metrics[f"{name}_R@{k}"] = np.mean(preds < k) metrics[f"{name}_mAP@10"] = np.mean( np.where(preds < 10, 1 / (preds + 1), 0.0) ) return metrics def random_masking(self, x, mask_prob=0.125, mask_hop_s=0.5): """random masking of 500ms with given probability""" b, t = x.shape len_masking_raw = int(24000 * mask_hop_s) # get random mask indices start_indices = torch.rand(b, t // len_masking_raw) < mask_prob time_domain_masked_indices = torch.nonzero( start_indices.repeat_interleave(len_masking_raw, dim=1) ) # mask with random values masking_noise = ( torch.randn(time_domain_masked_indices.shape[0], dtype=x.dtype) * 0.1 ) # 0 mean 0.1 std x[tuple(time_domain_masked_indices.t())] = masking_noise.to(x.device) return x def sequence_masking(self, x, max_ratio=0.6): b, t = x.shape mask_len = random.randint(1, int(t * max_ratio)) masking_noise = torch.randn(b, mask_len, dtype=x.dtype) * 0.1 x[:, -mask_len:] = masking_noise.to(x.device) return x def step(self, batch, stage): # get batch wav = batch[0] text = batch[1] # random masking for data augmentation if stage == "train": rv = random.random() if random.random() >= 0.4: if rv >= 0.7: # 30% chance, random masking wav = self.random_masking(wav) else: # 30% chance, sequence masking wav = self.sequence_masking(wav) # forward outputs = self.model(wav, text) # gather multi-gpu outputs gathered_outputs = self.all_gather(outputs, sync_grads=True) music_emb = rearrange(gathered_outputs[0], "n b c -> (n b) c") text_emb = rearrange(gathered_outputs[1], "n b c -> (n b) c") logit_scale = outputs[2] # get metrics metrics = self.get_metrics(music_emb, text_emb, logit_scale) # log metrics self.log("loss_%s" % stage, metrics["loss"], prog_bar=True, sync_dist=True) self.log( "mAP10-t2m_%s" % stage, metrics["text_to_music_mAP@10"], prog_bar=True, sync_dist=True, ) self.log( "mAP10-m2t_%s" % stage, metrics["music_to_text_mAP@10"], prog_bar=True, sync_dist=True, ) return metrics def training_step(self, batch, batch_idx): metrics = self.step(batch, "train") return metrics["loss"] def validation_step(self, batch, batch_idx): metrics = self.step(batch, "validation") return metrics["loss"] def test_step(self, batch, batch_idx): # forward music_emb = self.model.music_to_latent(batch[0]) text_emb = self.model.text_to_latent(batch[1]) # gather multi-gpu results gathered_music_emb = self.all_gather(music_emb, sync_grads=True) gathered_text_emb = self.all_gather(text_emb, sync_grads=True) gathered_ids = self.all_gather(batch[2], sync_grads=True) gathered_start_s = self.all_gather(batch[3], sync_grads=True) self.inference_music_embeddings.append( gathered_music_emb.detach().cpu().numpy() ) self.inference_text_embeddings.append(gathered_text_emb.detach().cpu().numpy()) self.inference_ids.append([_id for _id in gathered_ids]) self.start_ss.append(gathered_start_s.detach().cpu().float().numpy()) self.lyrics.append([lyric for lyric in batch[1]]) self.counter += 1 # concatenate music_embs = np.concatenate(self.inference_music_embeddings) text_embs = np.concatenate(self.inference_text_embeddings) ids = np.concatenate(self.inference_ids) times = np.concatenate(self.start_ss) lyrics = np.concatenate(self.lyrics) # save embeddings music_emb_fn = ( "/app/suno/data/audio_mono_24khz/genius_hq/embeddings/music_emb.npy" ) text_emb_fn = ( "/app/suno/data/audio_mono_24khz/genius_hq/embeddings/text_emb.npy" ) ids_fn = "/app/suno/data/audio_mono_24khz/genius_hq/embeddings/ids.npy" times_fn = ( "/app/suno/data/audio_mono_24khz/genius_hq/embeddings/start_times.npy" ) lyrics_fn = "/app/suno/data/audio_mono_24khz/genius_hq/embeddings/lyrics.npy" np.save(open(music_emb_fn, "wb"), music_embs) np.save(open(text_emb_fn, "wb"), text_embs) np.save(open(ids_fn, "wb"), ids) np.save(open(times_fn, "wb"), times) np.save(open(lyrics_fn, "wb"), lyrics) def configure_optimizers(self): optimizer = torch.optim.AdamW( [ {"params": self.model.music_projection.parameters(), "lr": self.lr}, {"params": self.model.text_projection.parameters(), "lr": self.lr}, {"params": self.model.music_encoder.parameters(), "lr": self.lr / 10}, {"params": self.model.text_encoder.parameters(), "lr": self.lr / 10}, ], lr=self.lr, ) return [optimizer]