| 504 | |
| 505 | |
| 506 | class SpatialTransformer(nn.Module): |
| 507 | |
| 508 | def __init__( |
| 509 | self, |
| 510 | in_channels, |
| 511 | n_heads, |
| 512 | d_head, |
| 513 | depth=1, |
| 514 | context_dim=None, |
| 515 | ): |
| 516 | super().__init__() |
| 517 | self.n_heads = n_heads |
| 518 | self.d_head = d_head |
| 519 | self.in_channels = in_channels |
| 520 | inner_dim = n_heads * d_head |
| 521 | self.norm = torch.nn.GroupNorm(num_groups=32, |
| 522 | num_channels=in_channels, |
| 523 | eps=1e-6, |
| 524 | affine=True) |
| 525 | |
| 526 | self.proj_in = nn.Conv2d(in_channels, |
| 527 | inner_dim, |
| 528 | kernel_size=1, |
| 529 | stride=1, |
| 530 | padding=0) |
| 531 | |
| 532 | self.transformer_blocks = nn.ModuleList([ |
| 533 | BasicTransformerBlock(inner_dim, |
| 534 | n_heads, |
| 535 | d_head, |
| 536 | context_dim=context_dim) |
| 537 | for d in range(depth) |
| 538 | ]) |
| 539 | |
| 540 | self.proj_out = nn.Conv2d(inner_dim, |
| 541 | in_channels, |
| 542 | kernel_size=1, |
| 543 | stride=1, |
| 544 | padding=0) |
| 545 | |
| 546 | def forward(self, hidden_states, context=None): |
| 547 | batch, channel, height, weight = hidden_states.shape |
| 548 | residual = hidden_states |
| 549 | hidden_states = self.norm(hidden_states) |
| 550 | hidden_states = self.proj_in(hidden_states) |
| 551 | hidden_states = hidden_states.view(batch, channel, 1, height * weight) |
| 552 | for block in self.transformer_blocks: |
| 553 | hidden_states = block(hidden_states, context=context) |
| 554 | hidden_states = hidden_states.view(batch, channel, height, weight) |
| 555 | hidden_states = self.proj_out(hidden_states) |
| 556 | return hidden_states + residual |
| 557 | |
| 558 | |
| 559 | class BasicTransformerBlock(nn.Module): |