| 257 | |
| 258 | |
| 259 | class Block(nn.Module): |
| 260 | def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, |
| 261 | drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, |
| 262 | window_size=None, window=False, rel_pos_spatial=False, prompt=None): |
| 263 | super().__init__() |
| 264 | self.norm1 = norm_layer(dim) |
| 265 | if not window: |
| 266 | self.attn = Attention( |
| 267 | dim, num_heads=num_heads, qkv_bias=qkv_bias, |
| 268 | window_size=window_size, rel_pos_spatial=rel_pos_spatial) |
| 269 | else: |
| 270 | self.attn = WindowAttention( |
| 271 | dim, num_heads=num_heads, qkv_bias=qkv_bias, |
| 272 | window_size=window_size, rel_pos_spatial=rel_pos_spatial |
| 273 | ) |
| 274 | # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here |
| 275 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() |
| 276 | self.norm2 = norm_layer(dim) |
| 277 | mlp_hidden_dim = int(dim * mlp_ratio) |
| 278 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer) |
| 279 | |
| 280 | def forward(self, x, H, W, mask=None): |
| 281 | x = x + self.drop_path(self.attn(self.norm1(x), H, W)) |
| 282 | x = x + self.drop_path(self.mlp(self.norm2(x))) |
| 283 | return x |
| 284 | |
| 285 | |
| 286 | class PatchEmbed(nn.Module): |