Transformer block for image-like data. First, project the input (aka embedding) and reshape to b, t, d. Then apply standard transformer action. Finally, reshape to image NEW: use_linear for more efficiency instead of the 1x1 convs
| 274 | |
| 275 | |
| 276 | class SpatialTransformer(nn.Module): |
| 277 | """ |
| 278 | Transformer block for image-like data. |
| 279 | First, project the input (aka embedding) |
| 280 | and reshape to b, t, d. |
| 281 | Then apply standard transformer action. |
| 282 | Finally, reshape to image |
| 283 | NEW: use_linear for more efficiency instead of the 1x1 convs |
| 284 | """ |
| 285 | def __init__(self, in_channels, n_heads, d_head, |
| 286 | depth=1, dropout=0., context_dim=None, |
| 287 | disable_self_attn=False, use_linear=False, |
| 288 | use_checkpoint=True): |
| 289 | super().__init__() |
| 290 | if exists(context_dim) and not isinstance(context_dim, list): |
| 291 | context_dim = [context_dim] |
| 292 | self.in_channels = in_channels |
| 293 | inner_dim = n_heads * d_head |
| 294 | self.norm = Normalize(in_channels) |
| 295 | if not use_linear: |
| 296 | self.proj_in = nn.Conv2d(in_channels, |
| 297 | inner_dim, |
| 298 | kernel_size=1, |
| 299 | stride=1, |
| 300 | padding=0) |
| 301 | else: |
| 302 | self.proj_in = nn.Linear(in_channels, inner_dim) |
| 303 | |
| 304 | self.transformer_blocks = nn.ModuleList( |
| 305 | [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d], |
| 306 | disable_self_attn=disable_self_attn, checkpoint=use_checkpoint) |
| 307 | for d in range(depth)] |
| 308 | ) |
| 309 | if not use_linear: |
| 310 | self.proj_out = zero_module(nn.Conv2d(inner_dim, |
| 311 | in_channels, |
| 312 | kernel_size=1, |
| 313 | stride=1, |
| 314 | padding=0)) |
| 315 | else: |
| 316 | self.proj_out = zero_module(nn.Linear(in_channels, inner_dim)) |
| 317 | self.use_linear = use_linear |
| 318 | |
| 319 | def forward(self, x, context=None): |
| 320 | # note: if no context is given, cross-attention defaults to self-attention |
| 321 | if not isinstance(context, list): |
| 322 | context = [context] |
| 323 | b, c, h, w = x.shape |
| 324 | x_in = x |
| 325 | x = self.norm(x) |
| 326 | if not self.use_linear: |
| 327 | x = self.proj_in(x) |
| 328 | x = rearrange(x, 'b c h w -> b (h w) c').contiguous() |
| 329 | if self.use_linear: |
| 330 | x = self.proj_in(x) |
| 331 | for i, block in enumerate(self.transformer_blocks): |
| 332 | x = block(x, context=context[i]) |
| 333 | if self.use_linear: |