| 245 | |
| 246 | |
| 247 | class VisualAttentionBlock(nn.Module): |
| 248 | def __init__( |
| 249 | self, |
| 250 | d_model: int, |
| 251 | n_head: int, |
| 252 | mlp_ratio: float = 4.0, |
| 253 | act_layer: Callable = nn.GELU, |
| 254 | norm_layer: Callable = nn.LayerNorm, |
| 255 | is_cross_attention: bool = False, |
| 256 | ): |
| 257 | super().__init__() |
| 258 | |
| 259 | self.ln_1 = norm_layer(d_model) |
| 260 | if is_cross_attention: |
| 261 | self.ln_1_kv = norm_layer(d_model) |
| 262 | |
| 263 | self.ln_2 = norm_layer(d_model) |
| 264 | mlp_width = int(d_model * mlp_ratio) |
| 265 | self.attn = VisualAttention(d_model, n_head) |
| 266 | self.mlp = nn.Sequential(OrderedDict([ |
| 267 | ("c_fc", nn.Linear(d_model, mlp_width)), |
| 268 | ("gelu", act_layer()), |
| 269 | ("c_proj", nn.Linear(mlp_width, d_model)) |
| 270 | ])) |
| 271 | |
| 272 | def attention( |
| 273 | self, |
| 274 | q_x: torch.Tensor, |
| 275 | k_x: Optional[torch.Tensor] = None, |
| 276 | v_x: Optional[torch.Tensor] = None, |
| 277 | attn_mask: Optional[torch.Tensor] = None, |
| 278 | ): |
| 279 | k_x = k_x if k_x is not None else q_x |
| 280 | v_x = v_x if v_x is not None else q_x |
| 281 | |
| 282 | attn_mask = attn_mask.to(q_x.dtype) if attn_mask is not None else None |
| 283 | return self.attn(q_x, k_x, v_x, attn_mask=attn_mask) |
| 284 | |
| 285 | def forward( |
| 286 | self, |
| 287 | q_x: torch.Tensor, |
| 288 | k_x: Optional[torch.Tensor] = None, |
| 289 | v_x: Optional[torch.Tensor] = None, |
| 290 | attn_mask: Optional[torch.Tensor] = None, |
| 291 | ): |
| 292 | k_x = self.ln_1_kv(k_x) if hasattr(self, "ln_1_kv") and k_x is not None else None |
| 293 | v_x = self.ln_1_kv(v_x) if hasattr(self, "ln_1_kv") and v_x is not None else None |
| 294 | |
| 295 | x = q_x + self.attention(q_x=self.ln_1(q_x), k_x=k_x, v_x=v_x, attn_mask=attn_mask) |
| 296 | x = x + self.mlp(self.ln_2(x)) |
| 297 | return x |
| 298 | |
| 299 | |
| 300 | class TransformerBlock(nn.Module): |