From 51d1908e941ebb4fd86380f3ab193247071e1999 Mon Sep 17 00:00:00 2001 From: James Betker Date: Sun, 5 Jun 2022 09:35:43 -0600 Subject: [PATCH] update --- codes/models/audio/music/transformer_diffusion8.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/codes/models/audio/music/transformer_diffusion8.py b/codes/models/audio/music/transformer_diffusion8.py index 9a9ae82b..acc0adfe 100644 --- a/codes/models/audio/music/transformer_diffusion8.py +++ b/codes/models/audio/music/transformer_diffusion8.py @@ -4,6 +4,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from models.arch_util import ResBlock from models.audio.music.music_quantizer2 import MusicQuantizer2 from models.diffusion.nn import timestep_embedding, normalization, zero_module, conv_nd, linear from models.diffusion.unet_diffusion import TimestepBlock @@ -45,13 +46,17 @@ class DietAttentionBlock(TimestepBlock): self.rms_scale_norm = RMSScaleShiftNorm(in_dim) self.proj = nn.Linear(in_dim, dim) self.attn = Attention(dim, heads=heads, causal=False, dropout=dropout) - self.ff = FeedForward(dim, in_dim, mult=1, dropout=dropout, zero_init_output=True) + self.res = ResBlock(channels=dim, dropout=dropout, out_channels=dim, dims=1) + self.ff = nn.Sequential(nn.LayerNorm(dim), + nn.GELU(), + FeedForward(dim, in_dim, mult=1, dropout=dropout, zero_init_output=True)) def forward(self, x, timestep_emb, rotary_emb): h = self.rms_scale_norm(x, norm_scale_shift_inp=timestep_emb) h = self.proj(h) k, _, _, _ = checkpoint(self.attn, h, None, None, None, None, None, rotary_emb) h = k + h + h = checkpoint(self.res, h.permute(0,2,1)).permute(0,2,1) h = checkpoint(self.ff, h) return h + x @@ -227,6 +232,7 @@ class TransformerDiffusionWithQuantizer(nn.Module): def get_grad_norm_parameter_groups(self): groups = { 'attention_layers': list(itertools.chain.from_iterable([lyr.attn.parameters() for lyr in self.diff.layers])), + 'res_layers': list(itertools.chain.from_iterable([lyr.res.parameters() for lyr in self.diff.layers])), 'ff_layers': list(itertools.chain.from_iterable([lyr.ff.parameters() for lyr in self.diff.layers])), 'quantizer_encoder': list(self.quantizer.encoder.parameters()), 'quant_codebook': [self.quantizer.quantizer.codevectors],