| 242 | |
| 243 | |
| 244 | class BasicTransformerBlock(nn.Module): |
| 245 | ATTENTION_MODES = { |
| 246 | "softmax": CrossAttention, # vanilla attention |
| 247 | "softmax-xformers": MemoryEfficientCrossAttention |
| 248 | } |
| 249 | def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True, |
| 250 | disable_self_attn=False): |
| 251 | super().__init__() |
| 252 | attn_mode = "softmax-xformers" if XFORMERS_IS_AVAILBLE else "softmax" |
| 253 | assert attn_mode in self.ATTENTION_MODES |
| 254 | attn_cls = self.ATTENTION_MODES[attn_mode] |
| 255 | self.disable_self_attn = disable_self_attn |
| 256 | self.attn1 = attn_cls(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout, |
| 257 | context_dim=context_dim if self.disable_self_attn else None) # is a self-attention if not self.disable_self_attn |
| 258 | self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) |
| 259 | self.attn2 = attn_cls(query_dim=dim, context_dim=context_dim, |
| 260 | heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none |
| 261 | self.norm1 = nn.LayerNorm(dim) |
| 262 | self.norm2 = nn.LayerNorm(dim) |
| 263 | self.norm3 = nn.LayerNorm(dim) |
| 264 | self.checkpoint = checkpoint |
| 265 | |
| 266 | def forward(self, x, context=None): |
| 267 | return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint) |
| 268 | |
| 269 | def _forward(self, x, context=None): |
| 270 | x = self.attn1(self.norm1(x), context=context if self.disable_self_attn else None) + x |
| 271 | x = self.attn2(self.norm2(x), context=context) + x |
| 272 | x = self.ff(self.norm3(x)) + x |
| 273 | return x |
| 274 | |
| 275 | |
| 276 | class SpatialTransformer(nn.Module): |