(self, x, context=None)
| 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): |
nothing calls this directly
no outgoing calls
no test coverage detected