import os os.environ["CUDA_VISIBLE_DEVICES"] = "2" import math import torch import funcy import einsum import torchaudio import numpy as np from torch import nn, einsum import torch.optim as optim from tqdm import tqdm from torch.nn.functional import mse_loss from dac.model.dac4 import DAC from suno_utils.utils.s3 import read_from_s3 from suno_utils.utils.text import read_jsonl def load_vae( device: str = "cuda", checkpoint_filepath: str = "s3://suno-data/christian/100hz_vae_peaq_kl_0.005.pth", ): # checkpoint_filepath = "/app/suno/christian/checkpoints/dac/100hz_vae_peaq_kl_0.005/best/dac/weights.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 } model_100hz = DAC(**sd["metadata"]["kwargs"]) model_100hz.load_state_dict(sd["state_dict"]) model_100hz.eval() model_100hz.to(device) return model_100hz # load pretrained VAE model vae_model = load_vae() # Define the noise schedule and sampling loop def get_alphas_sigmas(t): """Returns the scaling factors for the clean image (alpha) and for the noise (sigma), given a timestep.""" return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2) def alpha_sigma_to_t(alpha, sigma): """Returns a timestep, given the scaling factors for the clean image and for the noise.""" return torch.atan2(sigma, alpha) / math.pi * 2 def t_to_alpha_sigma(t): """Returns the scaling factors for the clean image and for the noise, given a timestep.""" return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2) class ScaledSinusoidalEmbedding(nn.Module): def __init__(self, dim, theta=10000): super().__init__() assert (dim % 2) == 0, "dimension must be divisible by 2" self.scale = nn.Parameter(torch.ones(1) * dim**-0.5) half_dim = dim // 2 freq_seq = torch.arange(half_dim).float() / half_dim inv_freq = theta**-freq_seq self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, x, pos=None, seq_start_pos=None): seq_len, device = x.shape[1], x.device if pos is None: pos = torch.arange(seq_len, device=device) if seq_start_pos is not None: pos = pos - seq_start_pos[..., None] emb = einsum("i, j -> i j", pos, self.inv_freq) emb = torch.cat((emb.sin(), emb.cos()), dim=-1) return emb * self.scale class SinusoidalPositionalEmbedding(nn.Module): def __init__(self, d_model, max_len=5000): super().__init__() self.d_model = d_model # Create a long enough position vector position = torch.arange(max_len).unsqueeze(1) div_term = torch.exp( torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model) ) pe = torch.zeros(max_len, d_model) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer("pe", pe) def forward(self, x): """ Args: x: Tensor, shape [batch_size, seq_len, embed_dim] Returns: Tensor, shape [batch_size, seq_len, embed_dim] """ seq_len = x.size(1) return x + self.pe[:seq_len].unsqueeze(0) # Define a causal Transformer-based autoregressive model class SimpleARModel(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim, n_heads, n_layers): super(SimpleARModel, self).__init__() self.embedding = nn.Linear(input_dim, hidden_dim) encoder_layer = nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=n_heads) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers) self.linear = nn.Linear(hidden_dim, output_dim) # self.pos_embedding = ScaledSinusoidalEmbedding(hidden_dim) self.pos_embedding = SinusoidalPositionalEmbedding(hidden_dim) self.infer_embed = nn.Parameter(torch.randn(1, 1, input_dim)) def forward(self, x: torch.Tensor): # x = self.embedding(x) + self.pos_embedding(x) x = self.pos_embedding(self.embedding(x)) seq_len = x.size(1) mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(x.device) x = self.transformer(x, mask=mask, is_causal=True) z = self.linear(x) return z class AdaLN(nn.Module): def __init__(self, num_features, conditioning_dim): super(AdaLN, self).__init__() self.norm = nn.LayerNorm(num_features) self.fc_scale = nn.Linear(conditioning_dim, num_features) self.fc_shift = nn.Linear(conditioning_dim, num_features) def forward(self, x, z): normalized_x = self.norm(x) scale = self.fc_scale(z) shift = self.fc_shift(z) return scale * normalized_x + shift class ResidualBlock(nn.Module): def __init__(self, input_dim, hidden_dim, conditioning_dim): super(ResidualBlock, self).__init__() self.adaln1 = AdaLN(input_dim, conditioning_dim) self.linear1 = nn.Linear(input_dim, hidden_dim) self.silu = nn.SiLU() self.adaln2 = AdaLN(hidden_dim, conditioning_dim) self.linear2 = nn.Linear(hidden_dim, input_dim) def forward(self, x, z): residual = x x = self.adaln1(x, z) x = self.linear1(x) x = self.silu(x) x = self.adaln2(x, z) x = self.linear2(x) return x + residual # Define the diffusion model (small MLP) class DenoisingMLP(nn.Module): def __init__( self, input_dim, hidden_dim: int, conditioning_dim: int, num_blocks: int = 3, ): super(DenoisingMLP, self).__init__() self.blocks = nn.ModuleList( [ ResidualBlock(input_dim, hidden_dim, conditioning_dim) for _ in range(num_blocks) ] ) self.time_embedding = nn.Linear(1, conditioning_dim) def forward(self, x, z, t): """ x (torch.Tensor): input tensor of shape (batch_size, seq_len, input_dim) z (torch.Tensor): conditioning tensor of shape (batch_size, seq_len, conditioning_dim) t (torch.Tensor): time tensor of shape (batch_size, seq_len) """ t_emb = self.time_embedding(t.unsqueeze(-1)) z = z + t_emb for block in self.blocks: x = block(x, z) return x # Define the diffusion loss function def diffusion_loss(z: torch.Tensor, x: torch.Tensor, diffusion_model: nn.Module): # here x should not include the infer embed # but z should be generated using sequence with infer embed at the start # Draw uniformly distributed continuous timesteps # t = self.rng.draw(reals.shape[0])[:, 0].to(self.device) # torch.manual_seed(42) t = torch.rand((x.size(0), x.size(1)), device=x.device) # Replace 1% of t with ones to ensure training on terminal SNR t = torch.where(torch.rand_like(t) < 0.01, torch.ones_like(t), t) # Calculate the noise schedule parameters for those timesteps alphas, sigmas = get_alphas_sigmas(t) alphas = alphas.unsqueeze(-1) sigmas = sigmas.unsqueeze(-1) # combine noise with inputs noise = torch.randn_like(x) noised_inputs = x * alphas + noise * sigmas targets = noise * alphas - x * sigmas # Calculate the velocity v = diffusion_model(noised_inputs, z, t) # don't compute loss for the first timestep v = v[:, 1:, :] targets = targets[:, 1:, :] return mse_loss(v, targets) # Example dataset class SimpleDataset(torch.utils.data.Dataset): def __init__(self, size, seq_length, feature_dim): self.data = torch.randn(size, seq_length, feature_dim) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx] class VAEMemmapDataset(torch.utils.data.Dataset): def __init__( self, vae_memmap_path: str, vae_metas_path: str = None, vae_dim: int = 128, n_tokens_memmap: int = 1000, ): """For use in training unconditional diffusion model. When a metas file is provided, the metadata is loaded and returned with the data. This can be used for lyric conditioning, etc. """ super().__init__() self.vae_metas_path = vae_metas_path self.vae_dim = vae_dim self.n_tokens_memmap = n_tokens_memmap vae_data = np.memmap(vae_memmap_path, dtype=np.float32, mode="r") vae_data = vae_data.reshape(-1, vae_dim, n_tokens_memmap) self.vae_data = vae_data print(f"Found {vae_data.shape[0]} examples.") if vae_metas_path is not None: self.metas = read_jsonl(vae_metas_path) assert len(self.metas) == self.vae_data.shape[0] print("Loaded metadata for", len(self.metas), "examples.") else: self.metas = None def __len__(self): return self.vae_data.shape[0] def __getitem__(self, idx: int): info = {} idx = idx % self.vae_data.shape[0] idx = np.random.randint(0, 2) vae_embeds = torch.from_numpy(self.vae_data[idx, ...].copy()).float() info["idx"] = idx info["seconds_start"] = 0 info["seconds_total"] = 10.0 if self.metas is not None: info["lyrics"] = self.metas[idx]["lyrics"] return (vae_embeds.permute(1, 0), info) def trange(*args, **kwargs): """Shortcut for tqdm(range(*args), **kwargs).""" return tqdm(range(*args), **kwargs) @torch.no_grad() def sample( model: nn.Module, diffusion_model: nn.Module, max_seq_len: int, num_steps: int, device="cuda", ): """Draws samples from a model given starting noise. v-diffusion""" bs = 1 current_context = model.infer_embed.view(1, 1, -1) # initial context (bs, seq_len, embed_dim) # sequential generation loop for i in trange(max_seq_len): # given the set of current embeddings, predict next conditioning embedding with AR model z = model(current_context) z = z[:, -1:, :] # use the last embedding # start with noise for the next timestep embedding x0 = torch.randn(bs, 1, current_context.shape[-1], device=device) # Reverse the diffusion process # Pre-calculate all timesteps and noise levels timesteps = torch.linspace(1, 0, num_steps + 1, device=device)[:-1] alphas, sigmas = get_alphas_sigmas(timesteps) for j in range(num_steps): t = timesteps[j].expand(bs) alpha = alphas[j].view(1, 1, 1) sigma = sigmas[j].view(1, 1, 1) with torch.no_grad(): v = diffusion_model(x0, z, t) pred_x0 = (x0 - sigma * v) / alpha pred_x0 = torch.clamp(pred_x0, -4, 4) x0 = pred_x0 # Update the context with the generated sample current_context = torch.cat([current_context, x0], dim=1) return current_context def generate( model: nn.Module, diffusion_model: nn.Module, vae_model: nn.Module, max_seq_len: int = 10, num_steps: int = 100, device: str = "cuda", ): pred_latents = sample( model, diffusion_model, max_seq_len=max_seq_len, num_steps=num_steps, device=device, ) # decode latents back to audio with torch.no_grad(): pred_audio = vae_model.decode(pred_latents.permute(0, 2, 1)) return pred_audio # Training loop def train( model, diffusion_model, dataloader, optimizer, diffusion_optimizer, epochs: str, ): model.train() diffusion_model.train() overall_losses = [] for epoch in range(epochs): epoch_losses = [] pbar = tqdm(dataloader) for batch in pbar: vae_embeds, metadata = batch optimizer.zero_grad() diffusion_optimizer.zero_grad() # only use the first 10 tokens vae_embeds = vae_embeds[:, :10, :] vae_embeds = vae_embeds.cuda() # concat the infer embed batch_infer_embed = model.infer_embed.view(1, 1, -1).repeat( vae_embeds.shape[0], 1, 1 ) vae_embeds_with_infer = torch.cat( [ batch_infer_embed, vae_embeds, ], dim=1, ) z = model(vae_embeds_with_infer) # Autoregressive model's output z = z[:, :-1, :] # remove the last embedding loss = diffusion_loss(z, vae_embeds, diffusion_model) loss.backward() optimizer.step() diffusion_optimizer.step() epoch_losses.append(loss.item()) pbar.set_description( f"Epoch {epoch + 1}/{epochs}, Loss: {np.mean(epoch_losses):.4f}" ) # if epoch % 10 == 0: print(f"Epoch {epoch + 1}/{epochs}, Loss: {np.mean(epoch_losses):.4f}") # at the end of epoch sample from the model pred_audio = generate( model, diffusion_model, vae_model=vae_model, max_seq_len=vae_embeds.shape[1] - 1, num_steps=100, device="cuda", ) # save the audio save_path = f"outputs/ar-diffusion/output_epoch_{epoch}.wav" torchaudio.save(save_path, pred_audio.squeeze(0).cpu(), 48000) # decode a training example with torch.no_grad(): pred_audio = vae_model.decode(vae_embeds[0:1, ...].permute(0, 2, 1)) save_path = f"outputs/ar-diffusion/input_epoch_{epoch}.wav" torchaudio.save(save_path, pred_audio.squeeze(0).cpu(), 48000) if __name__ == "__main__": os.makedirs("outputs/ar-diffusion", exist_ok=True) # Define the model and the diffusion model vae_dim = 128 hidden_dim = 1024 cond_dim = 1024 n_heads = 8 n_layers = 32 batch_size = 16 lr = 8e-4 epochs = 100 n_blocks = 9 # create the AR model and the diffusion model model = SimpleARModel(vae_dim, hidden_dim, cond_dim, n_heads, n_layers).cuda() diffusion_model = DenoisingMLP( vae_dim, hidden_dim, cond_dim, num_blocks=n_blocks ).cuda() # print number of parameters in millions for ar model and diffusion model ar_params = sum(p.numel() for p in model.parameters()) / 1e6 diffusion_params = sum(p.numel() for p in diffusion_model.parameters()) / 1e6 print(f"AR Model: {ar_params:.2f}M, Diffusion Model: {diffusion_params:.2f}M") # Define the optimizer and the diffusion optimizer optimizer = optim.AdamW( model.parameters(), weight_decay=0.02, betas=(0.9, 0.95), lr=lr, ) diffusion_optimizer = optim.AdamW( diffusion_model.parameters(), weight_decay=0.02, betas=(0.9, 0.95), lr=lr, ) # Define the dataset and the dataloader dataset = VAEMemmapDataset( vae_memmap_path="/app/suno/christian/data/suno_diffusion_tiktok_covers_lyrics/vae_train.bin", # vae_metas_path="/app/suno/christian/data/suno_diffusion_genius_hq_lyrics/val_metas.jsonl", ) dataloader = torch.utils.data.DataLoader( dataset, batch_size=batch_size, shuffle=True, num_workers=4 ) # Train the model train( model, diffusion_model, dataloader, optimizer, diffusion_optimizer, epochs, )