mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
James Betker 2022-06-05 09:35:43 -06:00
parent f9ebcf11d8
commit 51d1908e94
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
codes/models/audio/music/transformer_diffusion8.py
View File

@ -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],