| 236 | return x |
| 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 |
| 248 | x_in = x |
| 249 | x = self.norm(x) |
| 250 | x = self.proj_in(x) |
| 251 | x = x.reshape(b, c, h*w).permute(0,2,1) |
| 252 | for block in self.transformer_blocks: |
| 253 | x = block(x, context=context) |
| 254 | x = x.permute(0,2,1).reshape(b, c, h, w) |
| 255 | ret = self.proj_out(x) + x_in |
| 256 | return ret |
| 257 | |
| 258 | class Downsample: |
| 259 | def __init__(self, channels): |