| 151 | # not to be confused with ResnetBlock |
| 152 | class ResBlock: |
| 153 | def __init__(self, channels, emb_channels, out_channels): |
| 154 | self.in_layers = [ |
| 155 | GroupNorm(32, channels), |
| 156 | Tensor.silu, |
| 157 | Conv2d(channels, out_channels, 3, padding=1) |
| 158 | ] |
| 159 | self.emb_layers = [ |
| 160 | Tensor.silu, |
| 161 | Linear(emb_channels, out_channels) |
| 162 | ] |
| 163 | self.out_layers = [ |
| 164 | GroupNorm(32, out_channels), |
| 165 | Tensor.silu, |
| 166 | lambda x: x, # needed for weights loading code to work |
| 167 | Conv2d(out_channels, out_channels, 3, padding=1) |
| 168 | ] |
| 169 | self.skip_connection = Conv2d(channels, out_channels, 1) if channels != out_channels else lambda x: x |
| 170 | |
| 171 | def __call__(self, x, emb): |
| 172 | h = x.sequential(self.in_layers) |