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
| 210 | return x+h_ |
| 211 | |
| 212 | class MemoryEfficientAttnBlock(nn.Module): |
| 213 | """ |
| 214 | Uses xformers efficient implementation, |
| 215 | see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 |
| 216 | Note: this is a single-head self-attention operation |
| 217 | """ |
| 218 | # |
| 219 | def __init__(self, in_channels): |
| 220 | super().__init__() |
| 221 | self.in_channels = in_channels |
| 222 | |
| 223 | self.norm = Normalize(in_channels) |
| 224 | self.q = torch.nn.Conv2d(in_channels, |
| 225 | in_channels, |
| 226 | kernel_size=1, |
| 227 | stride=1, |
| 228 | padding=0) |
| 229 | self.k = torch.nn.Conv2d(in_channels, |
| 230 | in_channels, |
| 231 | kernel_size=1, |
| 232 | stride=1, |
| 233 | padding=0) |
| 234 | self.v = torch.nn.Conv2d(in_channels, |
| 235 | in_channels, |
| 236 | kernel_size=1, |
| 237 | stride=1, |
| 238 | padding=0) |
| 239 | self.proj_out = torch.nn.Conv2d(in_channels, |
| 240 | in_channels, |
| 241 | kernel_size=1, |
| 242 | stride=1, |
| 243 | padding=0) |
| 244 | self.attention_op: Optional[Any] = None |
| 245 | |
| 246 | def forward(self, x): |
| 247 | h_ = x |
| 248 | h_ = self.norm(h_) |
| 249 | q = self.q(h_) |
| 250 | k = self.k(h_) |
| 251 | v = self.v(h_) |
| 252 | |
| 253 | # compute attention |
| 254 | B, C, H, W = q.shape |
| 255 | q, k, v = map(lambda x: rearrange(x, 'b c h w -> b (h w) c'), (q, k, v)) |
| 256 | |
| 257 | q, k, v = map( |
| 258 | lambda t: t.unsqueeze(3) |
| 259 | .reshape(B, t.shape[1], 1, C) |
| 260 | .permute(0, 2, 1, 3) |
| 261 | .reshape(B * 1, t.shape[1], C) |
| 262 | .contiguous(), |
| 263 | (q, k, v), |
| 264 | ) |
| 265 | out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) |
| 266 | |
| 267 | out = ( |
| 268 | out.unsqueeze(0) |
| 269 | .reshape(B, 1, out.shape[1], C) |