(
self,
dim,
n_heads,
d_head,
dropout=0.0,
context_dim=None,
gated_ff=True,
checkpoint=True,
attn_mode="softmax",
)
| 439 | } |
| 440 | |
| 441 | def __init__( |
| 442 | self, |
| 443 | dim, |
| 444 | n_heads, |
| 445 | d_head, |
| 446 | dropout=0.0, |
| 447 | context_dim=None, |
| 448 | gated_ff=True, |
| 449 | checkpoint=True, |
| 450 | attn_mode="softmax", |
| 451 | ): |
| 452 | super().__init__() |
| 453 | assert attn_mode in self.ATTENTION_MODES |
| 454 | attn_cls = self.ATTENTION_MODES[attn_mode] |
| 455 | self.attn1 = attn_cls( |
| 456 | query_dim=dim, |
| 457 | heads=n_heads, |
| 458 | dim_head=d_head, |
| 459 | dropout=dropout, |
| 460 | context_dim=context_dim, |
| 461 | ) |
| 462 | self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff) |
| 463 | self.norm1 = nn.LayerNorm(dim) |
| 464 | self.norm2 = nn.LayerNorm(dim) |
| 465 | self.checkpoint = checkpoint |
| 466 | |
| 467 | def forward(self, x, context=None): |
| 468 | return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint) |
nothing calls this directly
no test coverage detected