Uses xformers efficient implementation, see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 Note: this is a single-head self-attention operation
| 176 | |
| 177 | |
| 178 | class MemoryEfficientAttnBlock(nn.Module): |
| 179 | """ |
| 180 | Uses xformers efficient implementation, |
| 181 | see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 |
| 182 | Note: this is a single-head self-attention operation |
| 183 | """ |
| 184 | def __init__(self, in_channels, natten_kernel_size=-1, use_null_attention=False): |
| 185 | super().__init__() |
| 186 | self.in_channels = in_channels |
| 187 | |
| 188 | self.norm = Normalize(in_channels) |
| 189 | conv_kwargs = dict(kernel_size=1, stride=1, padding=0) |
| 190 | self.q = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 191 | self.k = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 192 | self.v = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 193 | self.proj_out = nn.Conv2d(in_channels, in_channels, **conv_kwargs) |
| 194 | |
| 195 | if natten_kernel_size > -1: |
| 196 | assert NATTEN_IS_AVAILBLE, "natten_kernel_size > -1 but natten is not available" |
| 197 | assert (natten_kernel_size % 2) == 1, 'natten_kernel_size must be odd' |
| 198 | self.natten_kernel_size = natten_kernel_size |
| 199 | self.use_null_attention = use_null_attention |
| 200 | |
| 201 | def forward(self, x): |
| 202 | h_ = x |
| 203 | h_ = self.norm(h_) |
| 204 | q = self.q(h_) |
| 205 | k = self.k(h_) |
| 206 | v = self.v(h_) |
| 207 | |
| 208 | # compute null attention by discarding q,k |
| 209 | if self.use_null_attention: |
| 210 | out = self.proj_out(v) |
| 211 | return x+out |
| 212 | |
| 213 | # compute attention |
| 214 | if NATTEN_IS_AVAILBLE and self.natten_kernel_size > -1: |
| 215 | q, k, v = map(lambda x: rearrange(x, 'b c h w -> b 1 h w c'), (q, k, v)) |
| 216 | qk = natten.functional.natten2dqk(q, k, self.natten_kernel_size, 1) |
| 217 | a = torch.softmax(qk, dim=-1) |
| 218 | out = natten.functional.natten2dav(a, v, self.natten_kernel_size, 1) |
| 219 | out = rearrange(out, 'b 1 h w c -> b c h w') |
| 220 | |
| 221 | else: |
| 222 | b,c,h,w = q.shape |
| 223 | q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v)) |
| 224 | out = nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0) |
| 225 | out = rearrange(out, 'b (h w) c -> b c h w', h=h, w=w) |
| 226 | |
| 227 | out = self.proj_out(out) |
| 228 | return x+out |
| 229 | |
| 230 | |
| 231 | def make_attn(in_channels, attn_type="vanilla", natten_kernel_size=-1, use_null_attention=False): |