| 142 | |
| 143 | |
| 144 | class AttnBlock(nn.Module): |
| 145 | def __init__(self, in_channels): |
| 146 | super().__init__() |
| 147 | self.in_channels = in_channels |
| 148 | |
| 149 | self.norm = Normalize(in_channels) |
| 150 | self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 151 | self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 152 | self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 153 | self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 154 | |
| 155 | def attention(self, h_: torch.Tensor) -> torch.Tensor: |
| 156 | h_ = self.norm(h_) |
| 157 | q = self.q(h_) |
| 158 | k = self.k(h_) |
| 159 | v = self.v(h_) |
| 160 | |
| 161 | b, c, h, w = q.shape |
| 162 | q, k, v = map(lambda x: rearrange(x, "b c h w -> b 1 (h w) c").contiguous(), (q, k, v)) |
| 163 | h_ = torch.nn.functional.scaled_dot_product_attention(q, k, v) # scale is dim ** -0.5 per default |
| 164 | # compute attention |
| 165 | |
| 166 | return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) |
| 167 | |
| 168 | def forward(self, x, **kwargs): |
| 169 | h_ = x |
| 170 | h_ = self.attention(h_) |
| 171 | h_ = self.proj_out(h_) |
| 172 | return x + h_ |
| 173 | |
| 174 | |
| 175 | class MemoryEfficientAttnBlock(nn.Module): |