| 240 | return x+h_ |
| 241 | |
| 242 | class MemoryEfficientAttnBlock(nn.Module): |
| 243 | def __init__(self, in_channels): |
| 244 | super().__init__() |
| 245 | self.in_channels = in_channels |
| 246 | |
| 247 | self.norm = Normalize(in_channels) |
| 248 | self.q = torch.nn.Conv2d(in_channels, |
| 249 | in_channels, |
| 250 | kernel_size=1, |
| 251 | stride=1, |
| 252 | padding=0) |
| 253 | self.k = torch.nn.Conv2d(in_channels, |
| 254 | in_channels, |
| 255 | kernel_size=1, |
| 256 | stride=1, |
| 257 | padding=0) |
| 258 | self.v = torch.nn.Conv2d(in_channels, |
| 259 | in_channels, |
| 260 | kernel_size=1, |
| 261 | stride=1, |
| 262 | padding=0) |
| 263 | self.proj_out = torch.nn.Conv2d(in_channels, |
| 264 | in_channels, |
| 265 | kernel_size=1, |
| 266 | stride=1, |
| 267 | padding=0) |
| 268 | self.attention_op: Optional[Any] = None |
| 269 | |
| 270 | |
| 271 | def forward(self, x): |
| 272 | h_ = x |
| 273 | h_ = self.norm(h_) |
| 274 | q = self.q(h_) |
| 275 | k = self.k(h_) |
| 276 | v = self.v(h_) |
| 277 | |
| 278 | # compute attention |
| 279 | b,c,h,w = q.shape |
| 280 | q, k, v = map( |
| 281 | lambda t:t.reshape(b, t.shape[1], t.shape[2]*t.shape[3], 1) |
| 282 | .squeeze(3) |
| 283 | .permute(0,2,1) |
| 284 | .contiguous(), |
| 285 | (q, k, v), |
| 286 | ) |
| 287 | |
| 288 | # actually compute the attention, what we cannot get enough of |
| 289 | out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, scale=(int(c)**(-0.5)), op=self.attention_op) |
| 290 | |
| 291 | h_ = ( |
| 292 | out.permute(0,2,1) |
| 293 | .unsqueeze(3) |
| 294 | .reshape(b, c, h, w) |
| 295 | ) |
| 296 | |
| 297 | h_ = self.proj_out(h_) |
| 298 | |
| 299 | return x+h_ |