| 221 | return x.sequential(self.net) |
| 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 |
| 234 | x = self.attn2(self.norm2(x), context=context) + x |
| 235 | x = self.ff(self.norm3(x)) + x |
| 236 | return x |
| 237 | |
| 238 | class SpatialTransformer: |
| 239 | def __init__(self, channels, context_dim, n_heads, d_head): |