| 366 | return x |
| 367 | |
| 368 | class DecoderBlock(nn.Module): |
| 369 | |
| 370 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0., |
| 371 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, norm_mem=True, rope=None, arch_mode='VanillaDust3r', rope_mode='full_3d'): |
| 372 | super().__init__() |
| 373 | self.norm1 = norm_layer(dim) |
| 374 | self.attn = Attention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) |
| 375 | self.cross_attn = CrossAttention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) |
| 376 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 377 | self.norm2 = norm_layer(dim) |
| 378 | self.norm3 = norm_layer(dim) |
| 379 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 380 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) |
| 381 | self.norm_y = norm_layer(dim) if norm_mem else nn.Identity() |
| 382 | self.arch_mode = arch_mode |
| 383 | self.rope_mode = rope_mode |
| 384 | |
| 385 | |
| 386 | def forward(self, x, y, xpos, ypos): |
| 387 | infer_mode = 'train' if self.training else 'eval' |
| 388 | x = x + self.drop_path(self.attn(self.norm1(x), xpos)) |
| 389 | y_ = self.norm_y(y) |
| 390 | x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, y_, xpos, ypos)) |
| 391 | x = x + self.drop_path(self.mlp(self.norm3(x))) |
| 392 | return x, y |
| 393 | |
| 394 | class LoRA_DecoderBlock(nn.Module): |
| 395 | def __init__( |