| 263 | return self.op(x) |
| 264 | |
| 265 | class Upsample: |
| 266 | def __init__(self, channels): |
| 267 | self.conv = Conv2d(channels, channels, 3, padding=1) |
| 268 | |
| 269 | def __call__(self, x): |
| 270 | bs,c,py,px = x.shape |
| 271 | x = x.reshape(bs, c, py, 1, px, 1).expand(bs, c, py, 2, px, 2).reshape(bs, c, py*2, px*2) |
| 272 | return self.conv(x) |
| 273 | |
| 274 | def timestep_embedding(timesteps, dim, max_period=10000): |
| 275 | half = dim // 2 |