| 298 | |
| 299 | |
| 300 | class CustomDecoderBlock(nn.Module): |
| 301 | |
| 302 | def __init__( |
| 303 | self, |
| 304 | dim, |
| 305 | num_heads, |
| 306 | mlp_ratio=4.0, |
| 307 | qkv_bias=False, |
| 308 | drop=0.0, |
| 309 | attn_drop=0.0, |
| 310 | drop_path=0.0, |
| 311 | act_layer=nn.GELU, |
| 312 | norm_layer=nn.LayerNorm, |
| 313 | norm_mem=True, |
| 314 | rope=None, |
| 315 | ): |
| 316 | super().__init__() |
| 317 | self.norm1 = norm_layer(dim) |
| 318 | self.attn = Attention( |
| 319 | dim, |
| 320 | rope=rope, |
| 321 | num_heads=num_heads, |
| 322 | qkv_bias=qkv_bias, |
| 323 | attn_drop=attn_drop, |
| 324 | proj_drop=drop, |
| 325 | ) |
| 326 | self.cross_attn = CrossAttention( |
| 327 | dim, |
| 328 | rope=rope, |
| 329 | num_heads=num_heads, |
| 330 | qkv_bias=qkv_bias, |
| 331 | attn_drop=attn_drop, |
| 332 | proj_drop=drop, |
| 333 | ) |
| 334 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() |
| 335 | self.norm2 = norm_layer(dim) |
| 336 | self.norm3 = norm_layer(dim) |
| 337 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 338 | self.mlp = Mlp( |
| 339 | in_features=dim, |
| 340 | hidden_features=mlp_hidden_dim, |
| 341 | act_layer=act_layer, |
| 342 | drop=drop, |
| 343 | ) |
| 344 | self.norm_y = norm_layer(dim) if norm_mem else nn.Identity() |
| 345 | self.norm_z = norm_layer(dim) if norm_mem else nn.Identity() |
| 346 | |
| 347 | def forward(self, x, y, z, xpos, ypos): |
| 348 | x = x + self.drop_path(self.attn(self.norm1(x), xpos)) |
| 349 | y_ = self.norm_y(y) |
| 350 | z_ = self.norm_z(z) |
| 351 | x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, z_, xpos, ypos)) |
| 352 | x = x + self.drop_path(self.mlp(self.norm3(x))) |
| 353 | return x, y |
| 354 | |
| 355 | |
| 356 | class ModLN(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected