diff --git a/codes/models/audio/music/flat_diffusion.py b/codes/models/audio/music/flat_diffusion.py index 79c222bc..c143f6f1 100644 --- a/codes/models/audio/music/flat_diffusion.py +++ b/codes/models/audio/music/flat_diffusion.py @@ -1,13 +1,17 @@ +import os import random import torch import torch.nn as nn import torch.nn.functional as F +import torchvision from torch import autocast from models.arch_util import ResBlock from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear from models.diffusion.unet_diffusion import AttentionBlock, TimestepEmbedSequential, TimestepBlock +from scripts.audio.gen.use_mel2vec_codes import collapse_codegroups +from trainer.injectors.audio_injectors import normalize_mel from trainer.networks import register_model from utils.util import checkpoint @@ -107,6 +111,20 @@ class DiffusionLayer(TimestepBlock): return self.attn(y) +class NonTimestepResidualAttentionNorm(nn.Module): + def __init__(self, model_channels, dropout): + super().__init__() + self.resblk = ResBlock(dims=1, channels=model_channels, dropout=dropout) + self.attn = AttentionBlock(model_channels, num_heads=model_channels//64, relative_pos_embeddings=True) + self.norm = nn.GroupNorm(num_groups=8, num_channels=model_channels) + + def forward(self, x): + h = self.resblk(x) + h = self.norm(h) + h = self.attn(h) + return h + + class FlatDiffusion(nn.Module): def __init__( self, @@ -123,7 +141,6 @@ class FlatDiffusion(nn.Module): # Parameters for regularization. layer_drop=.1, unconditioned_percentage=.1, # This implements a mechanism similar to what is used in classifier-free training. - train_mel_head=False, ): super().__init__() @@ -137,6 +154,7 @@ class FlatDiffusion(nn.Module): self.layer_drop = layer_drop self.inp_block = conv_nd(1, in_channels, model_channels, 3, 1, 1) + # TODO: I'd really like to see if this could be ablated. It seems useless to me: why can't the embedding just learn this mapping directly? self.time_embed = nn.Sequential( linear(model_channels, model_channels), nn.SiLU(), @@ -150,32 +168,23 @@ class FlatDiffusion(nn.Module): self.embeddings = nn.ModuleList([nn.Embedding(in_vectors, model_channels//in_groups) for _ in range(in_groups)]) self.latent_conditioner = nn.Sequential( nn.Conv1d(in_latent_channels, model_channels, 3, padding=1), - AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), - AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), - AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), - AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), + NonTimestepResidualAttentionNorm(model_channels, dropout), + NonTimestepResidualAttentionNorm(model_channels, dropout), + nn.Conv1d(model_channels, model_channels, 3, padding=1), ) + self.latent_fade = nn.Parameter(torch.zeros(1,model_channels,1)) self.code_converter = nn.Sequential( - ResBlock(dims=1, channels=model_channels, dropout=dropout), - AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), - ResBlock(dims=1, channels=model_channels, dropout=dropout), - AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True), - ResBlock(dims=1, channels=model_channels, dropout=dropout), + NonTimestepResidualAttentionNorm(model_channels, dropout), + NonTimestepResidualAttentionNorm(model_channels, dropout), + NonTimestepResidualAttentionNorm(model_channels, dropout), + nn.Conv1d(model_channels, model_channels, 3, padding=1), ) - self.code_norm = normalization(model_channels) - self.contextual_embedder = nn.Sequential(nn.Conv1d(in_channels,model_channels,3,padding=1,stride=2), - nn.Conv1d(model_channels, model_channels*2,3,padding=1,stride=2), - AttentionBlock(model_channels*2, num_heads, relative_pos_embeddings=True, do_checkpoint=False), - AttentionBlock(model_channels*2, num_heads, relative_pos_embeddings=True, do_checkpoint=False), - AttentionBlock(model_channels*2, num_heads, relative_pos_embeddings=True, do_checkpoint=False), - AttentionBlock(model_channels*2, num_heads, relative_pos_embeddings=True, do_checkpoint=False), - AttentionBlock(model_channels*2, num_heads, relative_pos_embeddings=True, do_checkpoint=False)) + self.conditioning_embedder = nn.Sequential(nn.Conv1d(in_channels, model_channels // 2, 3, padding=1, stride=2), + nn.Conv1d(model_channels//2, model_channels,3,padding=1,stride=2), + NonTimestepResidualAttentionNorm(model_channels, dropout), + NonTimestepResidualAttentionNorm(model_channels, dropout), + NonTimestepResidualAttentionNorm(model_channels, dropout)) self.unconditioned_embedding = nn.Parameter(torch.randn(1,model_channels,1)) - self.conditioning_timestep_integrator = TimestepEmbedSequential( - DiffusionLayer(model_channels, dropout, num_heads), - DiffusionLayer(model_channels, dropout, num_heads), - DiffusionLayer(model_channels, dropout, num_heads), - ) self.integrating_conv = nn.Conv1d(model_channels*2, model_channels, kernel_size=1) self.mel_head = nn.Conv1d(model_channels, in_channels, kernel_size=3, padding=1) @@ -188,77 +197,78 @@ class FlatDiffusion(nn.Module): zero_module(conv_nd(1, model_channels, out_channels, 3, padding=1)), ) - if train_mel_head: - for m in [self.conditioning_timestep_integrator, self.integrating_conv, self.layers, - self.out]: - for p in m.parameters(): - p.requires_grad = False - p.DO_NOT_TRAIN = True + self.debug_codes = {} def get_grad_norm_parameter_groups(self): groups = { - 'minicoder': list(self.contextual_embedder.parameters()), - 'layers': list(self.layers.parameters()), + 'contextual_embedder': list(self.conditioning_embedder.parameters()), + 'layers': list(self.layers.parameters()) + list(self.integrating_conv.parameters()) + list(self.inp_block.parameters()), 'code_converters': list(self.embeddings.parameters()) + list(self.code_converter.parameters()) + list(self.latent_conditioner.parameters()), - 'timestep_integrator': list(self.conditioning_timestep_integrator.parameters()) + list(self.integrating_conv.parameters()), 'time_embed': list(self.time_embed.parameters()), } return groups - def timestep_independent(self, aligned_conditioning, conditioning_input, expected_seq_len, return_code_pred): - # Shuffle aligned_latent to BxCxS format - if is_latent(aligned_conditioning): - aligned_conditioning = aligned_conditioning.permute(0, 2, 1) + def timestep_independent(self, codes, conditioning_input, expected_seq_len, prenet_latent=None, return_code_pred=False): + cond_emb = self.conditioning_embedder(conditioning_input)[:, :, 0] - # Note: this block does not need to repeated on inference, since it is not timestep-dependent or x-dependent. - speech_conditioning_input = conditioning_input.unsqueeze(1) if len( - conditioning_input.shape) == 3 else conditioning_input - conds = [] - for j in range(speech_conditioning_input.shape[1]): - conds.append(self.contextual_embedder(speech_conditioning_input[:, j])) - conds = torch.cat(conds, dim=-1) - cond_emb = conds.mean(dim=-1) - cond_scale, cond_shift = torch.chunk(cond_emb, 2, dim=1) - if is_latent(aligned_conditioning): - code_emb = self.latent_conditioner(aligned_conditioning) - else: - code_emb = [embedding(aligned_conditioning[:, :, i]) for i, embedding in enumerate(self.embeddings)] - code_emb = torch.cat(code_emb, dim=-1).permute(0,2,1) + # Shuffle prenet_latent to BxCxS format + if prenet_latent is not None: + prenet_latent = prenet_latent.permute(0, 2, 1) + + code_emb = [embedding(codes[:, :, i]) for i, embedding in enumerate(self.embeddings)] + code_emb = torch.cat(code_emb, dim=-1).permute(0,2,1) + if prenet_latent is not None: + latent_conditioning = self.latent_conditioner(prenet_latent) + code_emb = code_emb + latent_conditioning * self.latent_fade unconditioned_batches = torch.zeros((code_emb.shape[0], 1, 1), device=code_emb.device) # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance. if self.training and self.unconditioned_percentage > 0: unconditioned_batches = torch.rand((code_emb.shape[0], 1, 1), device=code_emb.device) < self.unconditioned_percentage - code_emb = torch.where(unconditioned_batches, self.unconditioned_embedding.repeat(aligned_conditioning.shape[0], 1, 1), + code_emb = torch.where(unconditioned_batches, self.unconditioned_embedding.repeat(codes.shape[0], 1, 1), code_emb) + code_emb = self.code_converter(code_emb) + expanded_code_emb = F.interpolate(code_emb, size=expected_seq_len, mode='nearest') - expanded_code_emb = self.code_converter(expanded_code_emb) - expanded_code_emb = self.code_norm(expanded_code_emb) * (1 + cond_scale.unsqueeze(-1)) + cond_shift.unsqueeze(-1) - if not return_code_pred: - return expanded_code_emb + return expanded_code_emb, cond_emb else: - mel_pred = self.mel_head(expanded_code_emb) - # Multiply mel_pred by !unconditioned_branches, which drops the gradient on unconditioned branches. This is because we don't want that gradient being used to train parameters through the codes_embedder as it unbalances contributions to that network from the MSE loss. + # Perform the mel_head computation on the pre-exanded code embeddings, then interpolate it separately. + mel_pred = self.mel_head(code_emb) + mel_pred = F.interpolate(mel_pred, size=expected_seq_len, mode='nearest') + # Multiply mel_pred by !unconditioned_branches, which drops the gradient on unconditioned branches. + # This is because we don't want that gradient being used to train parameters through the codes_embedder as + # it unbalances contributions to that network from the MSE loss. mel_pred = mel_pred * unconditioned_batches.logical_not() - return expanded_code_emb, mel_pred + return expanded_code_emb, cond_emb, mel_pred - def forward(self, x, timesteps, aligned_conditioning=None, conditioning_input=None, precomputed_aligned_embeddings=None, conditioning_free=False, return_code_pred=False): + def forward(self, x, timesteps, + codes=None, conditioning_input=None, prenet_latent=None, + precomputed_code_embeddings=None, precomputed_cond_embeddings=None, + conditioning_free=False, return_code_pred=False): """ Apply the model to an input batch. + There are two ways to call this method: + 1) Specify codes, conditioning_input and optionally prenet_latent + 2) Specify precomputed_code_embeddings and precomputed_cond_embeddings, retrieved by calling timestep_independent yourself. + :param x: an [N x C x ...] Tensor of inputs. :param timesteps: a 1-D batch of timesteps. - :param aligned_conditioning: an aligned latent or sequence of tokens providing useful data about the sample to be produced. + :param codes: an aligned latent or sequence of tokens providing useful data about the sample to be produced. :param conditioning_input: a full-resolution audio clip that is used as a reference to the style you want decoded. - :param precomputed_aligned_embeddings: Embeddings returned from self.timestep_independent() + :param prenet_latent: optional latent vector aligned with codes derived from a prior network. + :param precomputed_code_embeddings: Code embeddings returned from self.timestep_independent() + :param precomputed_cond_embeddings: Conditional embeddings returned from self.timestep_independent() :param conditioning_free: When set, all conditioning inputs (including tokens and conditioning_input) will not be considered. :return: an [N x C x ...] Tensor of outputs. """ - assert precomputed_aligned_embeddings is not None or (aligned_conditioning is not None and conditioning_input is not None) - assert not (return_code_pred and precomputed_aligned_embeddings is not None) # These two are mutually exclusive. + if precomputed_code_embeddings is not None: + assert precomputed_cond_embeddings is not None, "Must specify both precomputed embeddings if one is specified" + assert codes is None and conditioning_input is None and prenet_latent is None, "Do not provide precomputed embeddings and the other parameters. It is unclear what you want me to do here." + assert not (return_code_pred and precomputed_code_embeddings is not None), "I cannot compute a code_pred output for you." unused_params = [] if conditioning_free: @@ -266,19 +276,17 @@ class FlatDiffusion(nn.Module): unused_params.extend(list(self.code_converter.parameters()) + list(self.code_embedding.parameters())) unused_params.extend(list(self.latent_conditioner.parameters())) else: - if precomputed_aligned_embeddings is not None: - code_emb = precomputed_aligned_embeddings + if precomputed_code_embeddings is not None: + code_emb = precomputed_code_embeddings + cond_emb = precomputed_cond_embeddings else: - code_emb, mel_pred = self.timestep_independent(aligned_conditioning, conditioning_input, x.shape[-1], True) - if is_latent(aligned_conditioning): - unused_params.extend(list(self.code_converter.parameters()) + list(self.code_embedding.parameters())) - else: + code_emb, cond_emb, mel_pred = self.timestep_independent(codes, conditioning_input, x.shape[-1], prenet_latent, True) + if prenet_latent is None: unused_params.extend(list(self.latent_conditioner.parameters())) unused_params.append(self.unconditioned_embedding) - time_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) - code_emb = self.conditioning_timestep_integrator(code_emb, time_emb) + blk_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels)) + cond_emb x = self.inp_block(x) x = torch.cat([x, code_emb], dim=1) x = self.integrating_conv(x) @@ -289,7 +297,7 @@ class FlatDiffusion(nn.Module): else: # First and last blocks will have autocast disabled for improved precision. with autocast(x.device.type, enabled=self.enable_fp16 and i != 0): - x = lyr(x, time_emb) + x = lyr(x, blk_emb) x = x.float() out = self.out(x) @@ -309,7 +317,7 @@ class FlatDiffusion(nn.Module): conditioning_input.shape) == 3 else conditioning_input conds = [] for j in range(speech_conditioning_input.shape[1]): - conds.append(self.contextual_embedder(speech_conditioning_input[:, j])) + conds.append(self.conditioning_embedder(speech_conditioning_input[:, j])) conds = torch.cat(conds, dim=-1) return conds.mean(dim=-1) @@ -320,13 +328,11 @@ def register_flat_diffusion(opt_net, opt): if __name__ == '__main__': clip = torch.randn(2, 256, 400) - aligned_latent = torch.randn(2,388,512) + aligned_latent = torch.randn(2,100,512) aligned_sequence = torch.randint(0,8,(2,100,8)) cond = torch.randn(2, 256, 400) ts = torch.LongTensor([600, 600]) - model = FlatDiffusion(512, layer_drop=.3, unconditioned_percentage=.5, train_mel_head=True) - # Test with latent aligned conditioning - #o = model(clip, ts, aligned_latent, cond) - # Test with sequence aligned conditioning + model = FlatDiffusion(512, layer_drop=.3, unconditioned_percentage=.5) o = model(clip, ts, aligned_sequence, cond, return_code_pred=True) + o = model(clip, ts, aligned_sequence, cond, aligned_latent) diff --git a/codes/train.py b/codes/train.py index bb2067ca..d9eaea4a 100644 --- a/codes/train.py +++ b/codes/train.py @@ -327,7 +327,7 @@ class Trainer: if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_music_gpt.yml') + parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_music_diffusion_flat.yml') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') args = parser.parse_args() opt = option.parse(args.opt, is_train=True) diff --git a/codes/trainer/eval/music_diffusion_fid.py b/codes/trainer/eval/music_diffusion_fid.py index aaacbfd6..1d0d0da5 100644 --- a/codes/trainer/eval/music_diffusion_fid.py +++ b/codes/trainer/eval/music_diffusion_fid.py @@ -167,10 +167,10 @@ class MusicDiffusionFid(evaluator.Evaluator): codegen = self.local_modules['codegen'].to(mel.device) codes = codegen.get_codes(mel) mel_norm = normalize_mel(mel) - precomputed = self.model.timestep_independent(aligned_conditioning=codes, conditioning_input=mel[:,:,:112], + precomputed_codes, precomputed_cond = self.model.timestep_independent(codes=codes, conditioning_input=mel_norm[:,:,:112], expected_seq_len=mel_norm.shape[-1], return_code_pred=False) gen_mel = self.diffuser.p_sample_loop(self.model, mel_norm.shape, noise=torch.zeros_like(mel_norm), - model_kwargs={'precomputed_aligned_embeddings': precomputed}) + model_kwargs={'precomputed_code_embeddings': precomputed_codes, 'precomputed_cond_embeddings': precomputed_cond}) gen_mel_denorm = denormalize_mel(gen_mel) output_shape = (1,16,audio.shape[-1]//16)