(
self,
in_channels,
n_heads,
d_head,
depth=1,
dropout=0.0,
context_dim=None,
disable_self_attn=False,
use_linear=False,
attn_type="softmax",
use_checkpoint=True,
# sdp_backend=SDPBackend.FLASH_ATTENTION
sdp_backend=None,
)
| 484 | """ |
| 485 | |
| 486 | def __init__( |
| 487 | self, |
| 488 | in_channels, |
| 489 | n_heads, |
| 490 | d_head, |
| 491 | depth=1, |
| 492 | dropout=0.0, |
| 493 | context_dim=None, |
| 494 | disable_self_attn=False, |
| 495 | use_linear=False, |
| 496 | attn_type="softmax", |
| 497 | use_checkpoint=True, |
| 498 | # sdp_backend=SDPBackend.FLASH_ATTENTION |
| 499 | sdp_backend=None, |
| 500 | ): |
| 501 | super().__init__() |
| 502 | print(f"constructing {self.__class__.__name__} of depth {depth} w/ {in_channels} channels and {n_heads} heads") |
| 503 | from omegaconf import ListConfig |
| 504 | |
| 505 | if exists(context_dim) and not isinstance(context_dim, (list, ListConfig)): |
| 506 | context_dim = [context_dim] |
| 507 | if exists(context_dim) and isinstance(context_dim, list): |
| 508 | if depth != len(context_dim): |
| 509 | print( |
| 510 | f"WARNING: {self.__class__.__name__}: Found context dims {context_dim} of depth {len(context_dim)}, " |
| 511 | f"which does not match the specified 'depth' of {depth}. Setting context_dim to {depth * [context_dim[0]]} now." |
| 512 | ) |
| 513 | # depth does not match context dims. |
| 514 | assert all( |
| 515 | map(lambda x: x == context_dim[0], context_dim) |
| 516 | ), "need homogenous context_dim to match depth automatically" |
| 517 | context_dim = depth * [context_dim[0]] |
| 518 | elif context_dim is None: |
| 519 | context_dim = [None] * depth |
| 520 | self.in_channels = in_channels |
| 521 | inner_dim = n_heads * d_head |
| 522 | self.norm = Normalize(in_channels) |
| 523 | if not use_linear: |
| 524 | self.proj_in = nn.Conv2d(in_channels, inner_dim, kernel_size=1, stride=1, padding=0) |
| 525 | else: |
| 526 | self.proj_in = nn.Linear(in_channels, inner_dim) |
| 527 | |
| 528 | self.transformer_blocks = nn.ModuleList( |
| 529 | [ |
| 530 | BasicTransformerBlock( |
| 531 | inner_dim, |
| 532 | n_heads, |
| 533 | d_head, |
| 534 | dropout=dropout, |
| 535 | context_dim=context_dim[d], |
| 536 | disable_self_attn=disable_self_attn, |
| 537 | attn_mode=attn_type, |
| 538 | checkpoint=use_checkpoint, |
| 539 | sdp_backend=sdp_backend, |
| 540 | ) |
| 541 | for d in range(depth) |
| 542 | ] |
| 543 | ) |
nothing calls this directly
no test coverage detected