| 266 | |
| 267 | |
| 268 | class DecoderBlock(nn.Module): |
| 269 | |
| 270 | def __init__( |
| 271 | self, |
| 272 | dim, |
| 273 | num_heads, |
| 274 | mlp_ratio=4.0, |
| 275 | qkv_bias=False, |
| 276 | drop=0.0, |
| 277 | attn_drop=0.0, |
| 278 | drop_path=0.0, |
| 279 | act_layer=nn.GELU, |
| 280 | norm_layer=nn.LayerNorm, |
| 281 | norm_mem=True, |
| 282 | rope=None, |
| 283 | ): |
| 284 | super().__init__() |
| 285 | self.norm1 = norm_layer(dim) |
| 286 | self.attn = Attention( |
| 287 | dim, |
| 288 | rope=rope, |
| 289 | num_heads=num_heads, |
| 290 | qkv_bias=qkv_bias, |
| 291 | attn_drop=attn_drop, |
| 292 | proj_drop=drop, |
| 293 | ) |
| 294 | self.cross_attn = CrossAttention( |
| 295 | dim, |
| 296 | rope=rope, |
| 297 | num_heads=num_heads, |
| 298 | qkv_bias=qkv_bias, |
| 299 | attn_drop=attn_drop, |
| 300 | proj_drop=drop, |
| 301 | ) |
| 302 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 303 | self.norm2 = norm_layer(dim) |
| 304 | self.norm3 = norm_layer(dim) |
| 305 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 306 | self.mlp = Mlp( |
| 307 | in_features=dim, |
| 308 | hidden_features=mlp_hidden_dim, |
| 309 | act_layer=act_layer, |
| 310 | drop=drop, |
| 311 | ) |
| 312 | self.norm_y = norm_layer(dim) if norm_mem else nn.Identity() |
| 313 | |
| 314 | def forward(self, x, y, xpos, ypos): |
| 315 | x = x + self.drop_path(self.attn(self.norm1(x), xpos)) |
| 316 | y_ = self.norm_y(y) |
| 317 | x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, y_, xpos, ypos)) |
| 318 | x = x + self.drop_path(self.mlp(self.norm3(x))) |
| 319 | return x, y |
| 320 | |
| 321 | |
| 322 | # patch embedding |