(self, dim, context_dim, n_heads, d_head)
| 222 | |
| 223 | class BasicTransformerBlock: |
| 224 | def __init__(self, dim, context_dim, n_heads, d_head): |
| 225 | self.attn1 = CrossAttention(dim, dim, n_heads, d_head) |
| 226 | self.ff = FeedForward(dim) |
| 227 | self.attn2 = CrossAttention(dim, context_dim, n_heads, d_head) |
| 228 | self.norm1 = LayerNorm(dim) |
| 229 | self.norm2 = LayerNorm(dim) |
| 230 | self.norm3 = LayerNorm(dim) |
| 231 | |
| 232 | def __call__(self, x, context=None): |
| 233 | x = self.attn1(self.norm1(x)) + x |
nothing calls this directly
no test coverage detected