| 335 | |
| 336 | |
| 337 | class DecoderBlock(nn.Module): |
| 338 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=True, proj_bias=True, mlp_bias=True, drop=0., attn_drop=0., |
| 339 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, gated_mlp=False, qk_norm=False, allow_zero_attn=False): |
| 340 | super().__init__() |
| 341 | self.norm1 = norm_layer(dim) |
| 342 | |
| 343 | if not qk_norm: |
| 344 | self.self_attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, proj_bias=proj_bias, attn_drop=attn_drop, proj_drop=drop, allow_zero_attn=allow_zero_attn) |
| 345 | self.cross_attn = CrossAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, proj_bias=proj_bias, attn_drop=attn_drop, proj_drop=drop, allow_zero_attn=allow_zero_attn) |
| 346 | else: |
| 347 | self.self_attn = NormAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, proj_bias=proj_bias, norm_layer=norm_layer, attn_drop=attn_drop, proj_drop=drop, allow_zero_attn=allow_zero_attn) |
| 348 | self.cross_attn = NormCrossAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias, proj_bias=proj_bias, norm_layer=norm_layer, attn_drop=attn_drop, proj_drop=drop, allow_zero_attn=allow_zero_attn) |
| 349 | |
| 350 | |
| 351 | self.query_norm = norm_layer(dim) |
| 352 | self.context_norm = norm_layer(dim) |
| 353 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 354 | self.norm2 = norm_layer(dim) |
| 355 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 356 | |
| 357 | if not gated_mlp: |
| 358 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, bias=mlp_bias, drop=drop) |
| 359 | else: |
| 360 | self.mlp = GatedMlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, bias=mlp_bias) |
| 361 | |
| 362 | def forward(self, x, context, sa_mask=None, xa_mask=None): |
| 363 | x = x + self.drop_path(self.self_attn(self.norm1(x), sa_mask)) |
| 364 | x = x + self.drop_path(self.cross_attn(self.query_norm(x), self.context_norm(context), xa_mask)) |
| 365 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 366 | return x |
| 367 | |
| 368 | |
| 369 | class CrossAttentionBlock(nn.Module): |