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
| 173 | |
| 174 | |
| 175 | class MemoryEfficientAttnBlock(nn.Module): |
| 176 | """ |
| 177 | Uses xformers efficient implementation, |
| 178 | see https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 |
| 179 | Note: this is a single-head self-attention operation |
| 180 | """ |
| 181 | |
| 182 | # |
| 183 | def __init__(self, in_channels): |
| 184 | super().__init__() |
| 185 | self.in_channels = in_channels |
| 186 | |
| 187 | self.norm = Normalize(in_channels) |
| 188 | self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 189 | self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 190 | self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 191 | self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) |
| 192 | self.attention_op: Optional[Any] = None |
| 193 | |
| 194 | def attention(self, h_: torch.Tensor) -> torch.Tensor: |
| 195 | h_ = self.norm(h_) |
| 196 | q = self.q(h_) |
| 197 | k = self.k(h_) |
| 198 | v = self.v(h_) |
| 199 | |
| 200 | # compute attention |
| 201 | B, C, H, W = q.shape |
| 202 | q, k, v = map(lambda x: rearrange(x, "b c h w -> b (h w) c"), (q, k, v)) |
| 203 | |
| 204 | q, k, v = map( |
| 205 | lambda t: t.unsqueeze(3) |
| 206 | .reshape(B, t.shape[1], 1, C) |
| 207 | .permute(0, 2, 1, 3) |
| 208 | .reshape(B * 1, t.shape[1], C) |
| 209 | .contiguous(), |
| 210 | (q, k, v), |
| 211 | ) |
| 212 | out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) |
| 213 | |
| 214 | out = out.unsqueeze(0).reshape(B, 1, out.shape[1], C).permute(0, 2, 1, 3).reshape(B, out.shape[1], C) |
| 215 | return rearrange(out, "b (h w) c -> b c h w", b=B, h=H, w=W, c=C) |
| 216 | |
| 217 | def forward(self, x, **kwargs): |
| 218 | h_ = x |
| 219 | h_ = self.attention(h_) |
| 220 | h_ = self.proj_out(h_) |
| 221 | return x + h_ |
| 222 | |
| 223 | |
| 224 | class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention): |