| 333 | return x, y, None, None |
| 334 | |
| 335 | class CustomDecoderBlock(nn.Module): |
| 336 | |
| 337 | def __init__( |
| 338 | self, |
| 339 | dim, |
| 340 | num_heads, |
| 341 | mlp_ratio=4.0, |
| 342 | qkv_bias=False, |
| 343 | drop=0.0, |
| 344 | attn_drop=0.0, |
| 345 | drop_path=0.0, |
| 346 | act_layer=nn.GELU, |
| 347 | norm_layer=nn.LayerNorm, |
| 348 | norm_mem=True, |
| 349 | rope=None, |
| 350 | ): |
| 351 | super().__init__() |
| 352 | self.norm1 = norm_layer(dim) |
| 353 | self.attn = Attention( |
| 354 | dim, |
| 355 | rope=rope, |
| 356 | num_heads=num_heads, |
| 357 | qkv_bias=qkv_bias, |
| 358 | attn_drop=attn_drop, |
| 359 | proj_drop=drop, |
| 360 | ) |
| 361 | self.cross_attn = CrossAttention( |
| 362 | dim, |
| 363 | rope=rope, |
| 364 | num_heads=num_heads, |
| 365 | qkv_bias=qkv_bias, |
| 366 | attn_drop=attn_drop, |
| 367 | proj_drop=drop, |
| 368 | ) |
| 369 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 370 | self.norm2 = norm_layer(dim) |
| 371 | self.norm3 = norm_layer(dim) |
| 372 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 373 | self.mlp = Mlp( |
| 374 | in_features=dim, |
| 375 | hidden_features=mlp_hidden_dim, |
| 376 | act_layer=act_layer, |
| 377 | drop=drop, |
| 378 | ) |
| 379 | self.norm_y = norm_layer(dim) if norm_mem else nn.Identity() |
| 380 | self.norm_z = norm_layer(dim) if norm_mem else nn.Identity() |
| 381 | |
| 382 | def forward(self, x, y, z, xpos, ypos): |
| 383 | x = x + self.drop_path(self.attn(self.norm1(x), xpos)) |
| 384 | y_ = self.norm_y(y) |
| 385 | z_ = self.norm_z(z) |
| 386 | x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, z_, xpos, ypos)) |
| 387 | x = x + self.drop_path(self.mlp(self.norm3(x))) |
| 388 | return x, y |
| 389 | |
| 390 | |
| 391 | class ModLN(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected