import torch import torch.nn as nn from torch.utils.checkpoint import checkpoint from .mask_dit import DiTBlock, FinalBlock, UDiT from .modules import ( film_modulate, PatchEmbed, PE_wrapper, TimestepEmbedder, RMSNorm, ) class AudioDiTBlock(DiTBlock): """ A modified DiT block with time aligned context add to latent. """ def __init__( self, dim, ta_context_dim, ta_context_norm=False, context_dim=None, num_heads=8, mlp_ratio=4., qkv_bias=False, qk_scale=None, qk_norm=None, act_layer='gelu', norm_layer=nn.LayerNorm, ta_context_fusion='add', time_fusion='none', ada_sola_rank=None, ada_sola_alpha=None, skip=False, skip_norm=False, rope_mode='none', context_norm=False, use_checkpoint=False ): super().__init__( dim=dim, context_dim=context_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm, act_layer=act_layer, norm_layer=norm_layer, time_fusion=time_fusion, ada_sola_rank=ada_sola_rank, ada_sola_alpha=ada_sola_alpha, skip=skip, skip_norm=skip_norm, rope_mode=rope_mode, context_norm=context_norm, use_checkpoint=use_checkpoint ) self.ta_context_fusion = ta_context_fusion self.ta_context_norm = ta_context_norm if self.ta_context_fusion == "add": self.ta_context_projection = nn.Linear(ta_context_dim, dim) self.ta_context_norm = norm_layer( ta_context_dim ) if self.ta_context_norm else nn.Identity() elif self.ta_context_fusion == "concat": self.ta_context_projection = nn.Linear(ta_context_dim + dim, dim) self.ta_context_norm = norm_layer( ta_context_dim + dim ) if self.ta_context_norm else nn.Identity() def forward( self, x, time_aligned_context, time_token=None, time_ada=None, skip=None, context=None, x_mask=None, context_mask=None, extras=None ): if self.use_checkpoint: return checkpoint( self._forward, x, time_aligned_context, time_token, time_ada, skip, context, x_mask, context_mask, extras, use_reentrant=False ) else: return self._forward( x, time_aligned_context, time_token, time_ada, skip, context, x_mask, context_mask, extras, ) def _forward( self, x, time_aligned_context, time_token=None, time_ada=None, skip=None, context=None, x_mask=None, context_mask=None, extras=None ): B, T, C = x.shape # # time aligned context # if self.ta_context_fusion == "add": # time_aligned_context = self.ta_context_projection( # self.ta_context_norm(time_aligned_context) # ) # x = x + time_aligned_context # elif self.ta_context_fusion == "concat": # cat = torch.cat([x, time_aligned_context], dim=-1) # cat = self.ta_context_norm(cat) # x = self.ta_context_projection(cat) # skip connection if self.skip_linear is not None: assert skip is not None cat = torch.cat([x, skip], dim=-1) cat = self.skip_norm(cat) x = self.skip_linear(cat) #print('skip') #print(x) if self.use_adanorm: time_ada = self.adaln(time_token, time_ada) (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = time_ada.chunk(6, dim=1) # self attention if self.use_adanorm: x_norm = film_modulate( self.norm1(x), shift=shift_msa, scale=scale_msa ) x = x + (1-gate_msa) * self.attn( x_norm, context=None, context_mask=x_mask, extras=extras ) else: # TODO diffusion timestep input is not fused here x = x + self.attn( self.norm1(x), context=None, context_mask=x_mask, extras=extras ) # time aligned context fusion if self.ta_context_fusion == "add": time_aligned_context = self.ta_context_projection( self.ta_context_norm(time_aligned_context) ) x = x + time_aligned_context elif self.ta_context_fusion == "concat": cat = torch.cat([x, time_aligned_context], dim=-1) cat = self.ta_context_norm(cat) x = self.ta_context_projection(cat) # cross attention if self.use_context: assert context is not None x = x + self.cross_attn( x=self.norm2(x), context=self.norm_context(context), context_mask=context_mask, extras=extras ) # mlp if self.use_adanorm: x_norm = film_modulate( self.norm3(x), shift=shift_mlp, scale=scale_mlp ) x = x + (1-gate_mlp) * self.mlp(x_norm) else: x = x + self.mlp(self.norm3(x)) return x class AudioUDiT(UDiT): def __init__( self, img_size=224, patch_size=16, in_chans=3, input_type='2d', out_chans=None, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4, qkv_bias=False, qk_scale=None, qk_norm=None, act_layer='gelu', norm_layer='layernorm', context_norm=False, use_checkpoint=False, time_fusion='token', ada_sola_rank=None, ada_sola_alpha=None, cls_dim=None, ta_context_dim=768, ta_context_fusion='concat', ta_context_norm=True, context_dim=768, context_fusion='concat', context_max_length=128, context_pe_method='sinu', pe_method='abs', rope_mode='none', use_conv=True, skip=True, skip_norm=True ): nn.Module.__init__(self) self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models # input self.in_chans = in_chans self.input_type = input_type if self.input_type == '2d': num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size) elif self.input_type == '1d': num_patches = img_size // patch_size self.patch_embed = PatchEmbed( patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, input_type=input_type ) out_chans = in_chans if out_chans is None else out_chans self.out_chans = out_chans # position embedding self.rope = rope_mode self.x_pe = PE_wrapper( dim=embed_dim, method=pe_method, length=num_patches ) # time embed self.time_embed = TimestepEmbedder(embed_dim) self.time_fusion = time_fusion self.use_adanorm = False # cls embed if cls_dim is not None: self.cls_embed = nn.Sequential( nn.Linear(cls_dim, embed_dim, bias=True), nn.SiLU(), nn.Linear(embed_dim, embed_dim, bias=True), ) else: self.cls_embed = None # time fusion if time_fusion == 'token': # put token at the beginning of sequence self.extras = 2 if self.cls_embed else 1 self.time_pe = PE_wrapper( dim=embed_dim, method='abs', length=self.extras ) elif time_fusion in ['ada', 'ada_single', 'ada_sola', 'ada_sola_bias']: self.use_adanorm = True # aviod repetitive silu for each adaln block self.time_act = nn.SiLU() self.extras = 0 self.time_ada_final = nn.Linear( embed_dim, 2 * embed_dim, bias=True ) if time_fusion in ['ada_single', 'ada_sola', 'ada_sola_bias']: # shared adaln self.time_ada = nn.Linear(embed_dim, 6 * embed_dim, bias=True) else: self.time_ada = None else: raise NotImplementedError # context # use a simple projection self.use_context = False self.context_cross = False self.context_max_length = context_max_length self.context_fusion = 'none' if context_dim is not None: self.use_context = True self.context_embed = nn.Sequential( nn.Linear(context_dim, embed_dim, bias=True), nn.SiLU(), nn.Linear(embed_dim, embed_dim, bias=True), ) self.context_fusion = context_fusion if context_fusion == 'concat' or context_fusion == 'joint': self.extras += context_max_length self.context_pe = PE_wrapper( dim=embed_dim, method=context_pe_method, length=context_max_length ) # no cross attention layers context_dim = None elif context_fusion == 'cross': self.context_pe = PE_wrapper( dim=embed_dim, method=context_pe_method, length=context_max_length ) self.context_cross = True context_dim = embed_dim else: raise NotImplementedError self.use_skip = skip # norm layers if norm_layer == 'layernorm': norm_layer = nn.LayerNorm elif norm_layer == 'rmsnorm': norm_layer = RMSNorm else: raise NotImplementedError self.in_blocks = nn.ModuleList([ AudioDiTBlock( dim=embed_dim, ta_context_dim=ta_context_dim, ta_context_fusion=ta_context_fusion, ta_context_norm=ta_context_norm, context_dim=context_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm, act_layer=act_layer, norm_layer=norm_layer, time_fusion=time_fusion, ada_sola_rank=ada_sola_rank, ada_sola_alpha=ada_sola_alpha, skip=False, skip_norm=False, rope_mode=self.rope, context_norm=context_norm, use_checkpoint=use_checkpoint ) for i in range(depth // 2) ]) self.mid_block = AudioDiTBlock( dim=embed_dim, ta_context_dim=ta_context_dim, context_dim=context_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm, act_layer=act_layer, norm_layer=norm_layer, time_fusion=time_fusion, ada_sola_rank=ada_sola_rank, ada_sola_alpha=ada_sola_alpha, ta_context_fusion=ta_context_fusion, ta_context_norm=ta_context_norm, skip=False, skip_norm=False, rope_mode=self.rope, context_norm=context_norm, use_checkpoint=use_checkpoint ) self.out_blocks = nn.ModuleList([ AudioDiTBlock( dim=embed_dim, ta_context_dim=ta_context_dim, context_dim=context_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, qk_norm=qk_norm, act_layer=act_layer, norm_layer=norm_layer, time_fusion=time_fusion, ada_sola_rank=ada_sola_rank, ada_sola_alpha=ada_sola_alpha, ta_context_fusion=ta_context_fusion, ta_context_norm=ta_context_norm, skip=skip, skip_norm=skip_norm, rope_mode=self.rope, context_norm=context_norm, use_checkpoint=use_checkpoint ) for i in range(depth // 2) ]) # FinalLayer block self.use_conv = use_conv self.final_block = FinalBlock( embed_dim=embed_dim, patch_size=patch_size, img_size=img_size, in_chans=out_chans, input_type=input_type, norm_layer=norm_layer, use_conv=use_conv, use_adanorm=self.use_adanorm ) self.initialize_weights() def forward( self, x, timesteps, time_aligned_context, context, x_mask=None, context_mask=None, cls_token=None, controlnet_skips=None, ): # make it compatible with int time step during inference if timesteps.dim() == 0: timesteps = timesteps.expand(x.shape[0] ).to(x.device, dtype=torch.long) x = self.patch_embed(x) x = self.x_pe(x) B, L, D = x.shape if self.use_context: context_token = self.context_embed(context) context_token = self.context_pe(context_token) if self.context_fusion == 'concat' or self.context_fusion == 'joint': x, x_mask = self._concat_x_context( x=x, context=context_token, x_mask=x_mask, context_mask=context_mask ) context_token, context_mask = None, None else: context_token, context_mask = None, None time_token = self.time_embed(timesteps) if self.cls_embed: cls_token = self.cls_embed(cls_token) time_ada = None time_ada_final = None if self.use_adanorm: if self.cls_embed: time_token = time_token + cls_token time_token = self.time_act(time_token) time_ada_final = self.time_ada_final(time_token) if self.time_ada is not None: time_ada = self.time_ada(time_token) else: time_token = time_token.unsqueeze(dim=1) if self.cls_embed: cls_token = cls_token.unsqueeze(dim=1) time_token = torch.cat([time_token, cls_token], dim=1) time_token = self.time_pe(time_token) x = torch.cat((time_token, x), dim=1) if x_mask is not None: x_mask = torch.cat([ torch.ones(B, time_token.shape[1], device=x_mask.device).bool(), x_mask ], dim=1) time_token = None skips = [] for blk in self.in_blocks: x = blk( x=x, time_aligned_context=time_aligned_context, time_token=time_token, time_ada=time_ada, skip=None, context=context_token, x_mask=x_mask, context_mask=context_mask, extras=self.extras ) if self.use_skip: skips.append(x) x = self.mid_block( x=x, time_aligned_context=time_aligned_context, time_token=time_token, time_ada=time_ada, skip=None, context=context_token, x_mask=x_mask, context_mask=context_mask, extras=self.extras ) for blk in self.out_blocks: if self.use_skip: skip = skips.pop() if controlnet_skips: # add to skip like u-net controlnet skip = skip + controlnet_skips.pop() else: skip = None if controlnet_skips: # directly add to x x = x + controlnet_skips.pop() x = blk( x=x, time_aligned_context=time_aligned_context, time_token=time_token, time_ada=time_ada, skip=skip, context=context_token, x_mask=x_mask, context_mask=context_mask, extras=self.extras ) x = self.final_block(x, time_ada=time_ada_final, extras=self.extras) return x