(self, x)
| 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) |
| 270 | .permute(0, 2, 1, 3) |
| 271 | .reshape(B, out.shape[1], C) |
| 272 | ) |
| 273 | out = rearrange(out, 'b (h w) c -> b c h w', b=B, h=H, w=W, c=C) |
| 274 | out = self.proj_out(out) |
| 275 | return x+out |
| 276 | |
| 277 | |
| 278 | class MemoryEfficientCrossAttentionWrapper(MemoryEfficientCrossAttention): |
nothing calls this directly
no outgoing calls
no test coverage detected