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
| 474 | |
| 475 | |
| 476 | class SpatialTransformer(nn.Module): |
| 477 | """ |
| 478 | Transformer block for image-like data. |
| 479 | First, project the input (aka embedding) |
| 480 | and reshape to b, t, d. |
| 481 | Then apply standard transformer action. |
| 482 | Finally, reshape to image |
| 483 | NEW: use_linear for more efficiency instead of the 1x1 convs |
| 484 | """ |
| 485 | |
| 486 | def __init__( |
| 487 | self, |
| 488 | in_channels, |
| 489 | n_heads, |
| 490 | d_head, |
| 491 | depth=1, |
| 492 | dropout=0.0, |
| 493 | context_dim=None, |
| 494 | disable_self_attn=False, |
| 495 | use_linear=False, |
| 496 | attn_type="softmax", |
| 497 | use_checkpoint=True, |
| 498 | # sdp_backend=SDPBackend.FLASH_ATTENTION |
| 499 | sdp_backend=None, |
| 500 | ): |
| 501 | super().__init__() |
| 502 | print(f"constructing {self.__class__.__name__} of depth {depth} w/ {in_channels} channels and {n_heads} heads") |
| 503 | from omegaconf import ListConfig |
| 504 | |
| 505 | if exists(context_dim) and not isinstance(context_dim, (list, ListConfig)): |
| 506 | context_dim = [context_dim] |
| 507 | if exists(context_dim) and isinstance(context_dim, list): |
| 508 | if depth != len(context_dim): |
| 509 | print( |
| 510 | f"WARNING: {self.__class__.__name__}: Found context dims {context_dim} of depth {len(context_dim)}, " |
| 511 | f"which does not match the specified 'depth' of {depth}. Setting context_dim to {depth * [context_dim[0]]} now." |
| 512 | ) |
| 513 | # depth does not match context dims. |
| 514 | assert all( |
| 515 | map(lambda x: x == context_dim[0], context_dim) |
| 516 | ), "need homogenous context_dim to match depth automatically" |
| 517 | context_dim = depth * [context_dim[0]] |
| 518 | elif context_dim is None: |
| 519 | context_dim = [None] * depth |
| 520 | self.in_channels = in_channels |
| 521 | inner_dim = n_heads * d_head |
| 522 | self.norm = Normalize(in_channels) |
| 523 | if not use_linear: |
| 524 | self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0) |
| 525 | else: |
| 526 | self.proj_in = nn.Linear(in_channels, inner_dim) |
| 527 | |
| 528 | self.transformer_blocks = nn.ModuleList( |
| 529 | [ |
| 530 | BasicTransformerBlock( |
| 531 | inner_dim, |
| 532 | n_heads, |
| 533 | d_head, |