| 212 | self.attention_op: Optional[Any] = None |
| 213 | |
| 214 | def forward(self, x, context=None, mask=None): |
| 215 | q = self.to_q(x) |
| 216 | context = default(context, x) |
| 217 | k = self.to_k(context) |
| 218 | v = self.to_v(context) |
| 219 | |
| 220 | b, _, _ = q.shape |
| 221 | q, k, v = map( |
| 222 | lambda t: t.unsqueeze(3) |
| 223 | .reshape(b, t.shape[1], self.heads, self.dim_head) |
| 224 | .permute(0, 2, 1, 3) |
| 225 | .reshape(b * self.heads, t.shape[1], self.dim_head) |
| 226 | .contiguous(), |
| 227 | (q, k, v), |
| 228 | ) |
| 229 | |
| 230 | # actually compute the attention, what we cannot get enough of |
| 231 | out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) |
| 232 | |
| 233 | if exists(mask): |
| 234 | raise NotImplementedError |
| 235 | out = ( |
| 236 | out.unsqueeze(0) |
| 237 | .reshape(b, self.heads, out.shape[1], self.dim_head) |
| 238 | .permute(0, 2, 1, 3) |
| 239 | .reshape(b, out.shape[1], self.heads * self.dim_head) |
| 240 | ) |
| 241 | return self.to_out(out) |
| 242 | |
| 243 | |
| 244 | class BasicTransformerBlock(nn.Module): |