| 277 | self.attention_op: Optional[Any] = None |
| 278 | |
| 279 | def forward( |
| 280 | self, |
| 281 | x, |
| 282 | context=None, |
| 283 | mask=None, |
| 284 | additional_tokens=None, |
| 285 | n_times_crossframe_attn_in_self=0, |
| 286 | ): |
| 287 | if additional_tokens is not None: |
| 288 | # get the number of masked tokens at the beginning of the output sequence |
| 289 | n_tokens_to_mask = additional_tokens.shape[1] |
| 290 | # add additional token |
| 291 | x = torch.cat([additional_tokens, x], dim=1) |
| 292 | q = self.to_q(x) |
| 293 | context = default(context, x) |
| 294 | k = self.to_k(context) |
| 295 | v = self.to_v(context) |
| 296 | |
| 297 | if n_times_crossframe_attn_in_self: |
| 298 | # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439 |
| 299 | assert x.shape[0] % n_times_crossframe_attn_in_self == 0 |
| 300 | # n_cp = x.shape[0]//n_times_crossframe_attn_in_self |
| 301 | k = repeat( |
| 302 | k[::n_times_crossframe_attn_in_self], |
| 303 | "b ... -> (b n) ...", |
| 304 | n=n_times_crossframe_attn_in_self, |
| 305 | ) |
| 306 | v = repeat( |
| 307 | v[::n_times_crossframe_attn_in_self], |
| 308 | "b ... -> (b n) ...", |
| 309 | n=n_times_crossframe_attn_in_self, |
| 310 | ) |
| 311 | |
| 312 | b, _, _ = q.shape |
| 313 | q, k, v = map( |
| 314 | lambda t: t.unsqueeze(3) |
| 315 | .reshape(b, t.shape[1], self.heads, self.dim_head) |
| 316 | .permute(0, 2, 1, 3) |
| 317 | .reshape(b * self.heads, t.shape[1], self.dim_head) |
| 318 | .contiguous(), |
| 319 | (q, k, v), |
| 320 | ) |
| 321 | |
| 322 | # actually compute the attention, what we cannot get enough of |
| 323 | out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) |
| 324 | |
| 325 | # TODO: Use this directly in the attention operation, as a bias |
| 326 | if exists(mask): |
| 327 | raise NotImplementedError |
| 328 | out = ( |
| 329 | out.unsqueeze(0) |
| 330 | .reshape(b, self.heads, out.shape[1], self.dim_head) |
| 331 | .permute(0, 2, 1, 3) |
| 332 | .reshape(b, out.shape[1], self.heads * self.dim_head) |
| 333 | ) |
| 334 | if additional_tokens is not None: |
| 335 | # remove additional token |
| 336 | out = out[:, n_tokens_to_mask:] |