(
self,
dim,
n_heads,
d_head,
dropout=0.0,
context_dim=None,
gated_ff=True,
checkpoint=True,
disable_self_attn=False,
attn_mode="softmax",
sdp_backend=None,
)
| 344 | } |
| 345 | |
| 346 | def __init__( |
| 347 | self, |
| 348 | dim, |
| 349 | n_heads, |
| 350 | d_head, |
| 351 | dropout=0.0, |
| 352 | context_dim=None, |
| 353 | gated_ff=True, |
| 354 | checkpoint=True, |
| 355 | disable_self_attn=False, |
| 356 | attn_mode="softmax", |
| 357 | sdp_backend=None, |
| 358 | ): |
| 359 | super().__init__() |
| 360 | assert attn_mode in self.ATTENTION_MODES |
| 361 | if attn_mode != "softmax" and not XFORMERS_IS_AVAILABLE: |
| 362 | print( |
| 363 | f"Attention mode '{attn_mode}' is not available. Falling back to native attention. " |
| 364 | f"This is not a problem in Pytorch >= 2.0. FYI, you are running with PyTorch version {torch.__version__}" |
| 365 | ) |
| 366 | attn_mode = "softmax" |
| 367 | elif attn_mode == "softmax" and not SDP_IS_AVAILABLE: |
| 368 | print("We do not support vanilla attention anymore, as it is too expensive. Sorry.") |
| 369 | if not XFORMERS_IS_AVAILABLE: |
| 370 | assert False, "Please install xformers via e.g. 'pip install xformers==0.0.16'" |
| 371 | else: |
| 372 | print("Falling back to xformers efficient attention.") |
| 373 | attn_mode = "softmax-xformers" |
| 374 | attn_cls = self.ATTENTION_MODES[attn_mode] |
| 375 | if version.parse(torch.__version__) >= version.parse("2.0.0"): |
| 376 | assert sdp_backend is None or isinstance(sdp_backend, SDPBackend) |
| 377 | else: |
| 378 | assert sdp_backend is None |
| 379 | self.disable_self_attn = disable_self_attn |
| 380 | self.attn1 = attn_cls( |
| 381 | query_dim=dim, |
| 382 | heads=n_heads, |
| 383 | dim_head=d_head, |
| 384 | dropout=dropout, |
| 385 | context_dim=context_dim if self.disable_self_attn else None, |
| 386 | backend=sdp_backend, |
| 387 | ) # is a self-attention if not self.disable_self_attn |
| 388 | self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) |
| 389 | self.attn2 = attn_cls( |
| 390 | query_dim=dim, |
| 391 | context_dim=context_dim, |
| 392 | heads=n_heads, |
| 393 | dim_head=d_head, |
| 394 | dropout=dropout, |
| 395 | backend=sdp_backend, |
| 396 | ) # is self-attn if context is none |
| 397 | self.norm1 = nn.LayerNorm(dim) |
| 398 | self.norm2 = nn.LayerNorm(dim) |
| 399 | self.norm3 = nn.LayerNorm(dim) |
| 400 | self.checkpoint = checkpoint |
| 401 | if self.checkpoint: |
| 402 | print(f"{self.__class__.__name__} is using checkpointing") |
| 403 |
nothing calls this directly
no test coverage detected