| 277 | |
| 278 | |
| 279 | class VersatileAttention(Attention): |
| 280 | def __init__( |
| 281 | self, |
| 282 | attention_mode=None, |
| 283 | cross_frame_attention_mode=None, |
| 284 | temporal_position_encoding=False, |
| 285 | temporal_position_encoding_max_len=24, |
| 286 | *args, |
| 287 | **kwargs, |
| 288 | ): |
| 289 | super().__init__(*args, **kwargs) |
| 290 | assert attention_mode == "Temporal" |
| 291 | |
| 292 | self.attention_mode = attention_mode |
| 293 | self.is_cross_attention = kwargs["cross_attention_dim"] is not None |
| 294 | |
| 295 | self.pos_encoder = ( |
| 296 | PositionalEncoding( |
| 297 | kwargs["query_dim"], |
| 298 | dropout=0.0, |
| 299 | max_len=temporal_position_encoding_max_len, |
| 300 | ) |
| 301 | if (temporal_position_encoding and attention_mode == "Temporal") |
| 302 | else None |
| 303 | ) |
| 304 | |
| 305 | def extra_repr(self): |
| 306 | return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}" |
| 307 | |
| 308 | def set_use_memory_efficient_attention_xformers( |
| 309 | self, |
| 310 | use_memory_efficient_attention_xformers: bool, |
| 311 | attention_op: Optional[Callable] = None, |
| 312 | ): |
| 313 | if use_memory_efficient_attention_xformers: |
| 314 | if not is_xformers_available(): |
| 315 | raise ModuleNotFoundError( |
| 316 | ( |
| 317 | "Refer to https://github.com/facebookresearch/xformers for more information on how to install" |
| 318 | " xformers" |
| 319 | ), |
| 320 | name="xformers", |
| 321 | ) |
| 322 | elif not torch.cuda.is_available(): |
| 323 | raise ValueError( |
| 324 | "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is" |
| 325 | " only available for GPU " |
| 326 | ) |
| 327 | else: |
| 328 | try: |
| 329 | # Make sure we can run the memory efficient attention |
| 330 | _ = xformers.ops.memory_efficient_attention( |
| 331 | torch.randn((1, 2, 40), device="cuda"), |
| 332 | torch.randn((1, 2, 40), device="cuda"), |
| 333 | torch.randn((1, 2, 40), device="cuda"), |
| 334 | ) |
| 335 | except Exception as e: |
| 336 | raise e |