diff --git a/codes/models/audio/music/transformer_diffusion14.py b/codes/models/audio/music/transformer_diffusion14.py index f1b98180..4d28c7c5 100644 --- a/codes/models/audio/music/transformer_diffusion14.py +++ b/codes/models/audio/music/transformer_diffusion14.py @@ -4,65 +4,14 @@ import torch import torch.nn as nn import torch.nn.functional as F -from models.arch_util import AttentionBlock, TimestepEmbedSequential, build_local_attention_mask +from models.arch_util import TimestepEmbedSequential from models.audio.music.encoders import ResEncoder16x +from models.audio.music.transformer_diffusion13 import ConcatAttentionBlock from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear -from models.diffusion.unet_diffusion import TimestepBlock from trainer.networks import register_model from utils.util import checkpoint, print_network -class SubBlock(nn.Module): - def __init__(self, inp_dim, contraction_dim, blk_dim, heads, dropout, enable_attention_masking=False): - super().__init__() - self.enable_attention_masking = enable_attention_masking - self.dropout = nn.Dropout(p=dropout) - self.blk_emb_proj = nn.Conv1d(blk_dim, inp_dim, 1) - self.attn = AttentionBlock(inp_dim, out_channels=contraction_dim, num_heads=heads) - self.attnorm = nn.GroupNorm(8, contraction_dim) - self.ff = nn.Conv1d(inp_dim+contraction_dim, contraction_dim, kernel_size=3, padding=1) - self.ffnorm = nn.GroupNorm(8, contraction_dim) - if self.enable_attention_masking: - # All regions can attend to the first token, which will be the timestep embedding. Hence, fixed_region. - self.mask = build_local_attention_mask(n=4000, l=48, fixed_region=1) - self.mask_initialized = False - else: - self.mask = None - - def forward(self, x, blk_emb): - if self.mask is not None and not self.mask_initialized: - self.mask = self.mask.to(x.device) - self.mask_initialized = True - blk_enc = self.blk_emb_proj(blk_emb) - ah = self.dropout(self.attn(torch.cat([blk_enc, x], dim=-1), mask=self.mask)) - ah = ah[:,:,blk_emb.shape[-1]:] # Strip off the blk_emb and re-align with x. - ah = F.gelu(self.attnorm(ah)) - h = torch.cat([ah, x], dim=1) - hf = self.dropout(checkpoint(self.ff, h)) - hf = F.gelu(self.ffnorm(hf)) - h = torch.cat([h, hf], dim=1) - return h - - -class ConcatAttentionBlock(TimestepBlock): - def __init__(self, trunk_dim, contraction_dim, heads, dropout, enable_attention_masking=False): - super().__init__() - self.prenorm = nn.GroupNorm(8, trunk_dim) - self.block1 = SubBlock(trunk_dim, contraction_dim, trunk_dim, heads, dropout, - enable_attention_masking=enable_attention_masking) - self.block2 = SubBlock(trunk_dim+contraction_dim*2, contraction_dim, trunk_dim, heads, dropout, - enable_attention_masking=enable_attention_masking) - self.out = nn.Conv1d(contraction_dim*4, trunk_dim, kernel_size=1, bias=False) - self.out.weight.data.zero_() - - def forward(self, x, blk_emb): - h = self.prenorm(x) - h = self.block1(h, blk_emb) - h = self.block2(h, blk_emb) - h = self.out(h[:,x.shape[1]:]) - return h + x - - class TransformerDiffusion(nn.Module): """ A diffusion model composed entirely of stacks of transformer layers. Why would you do it any other way? @@ -102,13 +51,14 @@ class TransformerDiffusion(nn.Module): self.time_embed = nn.Sequential( linear(time_embed_dim, time_embed_dim), nn.SiLU(), - linear(time_embed_dim, model_channels), + linear(time_embed_dim, time_embed_dim//4), ) self.input_converter = nn.Conv1d(input_vec_dim, model_channels, 1) self.unconditioned_embedding = nn.Parameter(torch.randn(1,model_channels,1)) self.intg = nn.Conv1d(model_channels*2, model_channels, 1) - self.layers = TimestepEmbedSequential(*[ConcatAttentionBlock(model_channels, contraction_dim, num_heads, dropout, enable_attention_masking=True) for _ in range(num_layers)]) + self.layers = TimestepEmbedSequential(*[ConcatAttentionBlock(model_channels, contraction_dim, time_embed_dim//4, + num_heads, dropout) for _ in range(num_layers)]) self.out = nn.Sequential( normalization(model_channels), @@ -128,8 +78,10 @@ class TransformerDiffusion(nn.Module): def get_grad_norm_parameter_groups(self): attn1 = list(itertools.chain.from_iterable([lyr.block1.attn.parameters() for lyr in self.layers])) attn2 = list(itertools.chain.from_iterable([lyr.block2.attn.parameters() for lyr in self.layers])) - ff1 = list(itertools.chain.from_iterable([lyr.block1.ff.parameters() for lyr in self.layers])) - ff2 = list(itertools.chain.from_iterable([lyr.block2.ff.parameters() for lyr in self.layers])) + ff1 = list(itertools.chain.from_iterable([lyr.block1.ff1.parameters() for lyr in self.layers] + + [lyr.block1.ff2.parameters() for lyr in self.layers])) + ff2 = list(itertools.chain.from_iterable([lyr.block2.ff1.parameters() for lyr in self.layers] + + [lyr.block2.ff2.parameters() for lyr in self.layers])) blkout_layers = list(itertools.chain.from_iterable([lyr.out.parameters() for lyr in self.layers])) groups = { 'prenorms': list(itertools.chain.from_iterable([lyr.prenorm.parameters() for lyr in self.layers])), @@ -163,7 +115,7 @@ class TransformerDiffusion(nn.Module): code_emb = F.interpolate(code_emb, size=x.shape[-1], mode='nearest') with torch.autocast(x.device.type, enabled=self.enable_fp16): - blk_emb = self.time_embed(timestep_embedding(timesteps, self.time_embed_dim)).unsqueeze(-1) + blk_emb = self.time_embed(timestep_embedding(timesteps, self.time_embed_dim)) x = self.inp_block(x) x = self.intg(torch.cat([x, code_emb], dim=1))