| 150 | |
| 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) |
| 173 | emb_out = emb.sequential(self.emb_layers) |
| 174 | h = h + emb_out.reshape(*emb_out.shape, 1, 1) |
| 175 | h = h.sequential(self.out_layers) |
| 176 | ret = self.skip_connection(x) + h |
| 177 | return ret |
| 178 | |
| 179 | class CrossAttention: |
| 180 | def __init__(self, query_dim, context_dim, n_heads, d_head): |