(self, channels, context_dim, n_heads, d_head)
| 237 | |
| 238 | class SpatialTransformer: |
| 239 | def __init__(self, channels, context_dim, n_heads, d_head): |
| 240 | self.norm = GroupNorm(32, channels) |
| 241 | assert channels == n_heads * d_head |
| 242 | self.proj_in = Conv2d(channels, n_heads * d_head, 1) |
| 243 | self.transformer_blocks = [BasicTransformerBlock(channels, context_dim, n_heads, d_head)] |
| 244 | self.proj_out = Conv2d(n_heads * d_head, channels, 1) |
| 245 | |
| 246 | def __call__(self, x, context=None): |
| 247 | b, c, h, w = x.shape |
nothing calls this directly
no test coverage detected