| 278 | return Tensor.cat(args.cos(), args.sin()).reshape(1, -1) |
| 279 | |
| 280 | class UNetModel: |
| 281 | def __init__(self): |
| 282 | self.time_embed = [ |
| 283 | Linear(320, 1280), |
| 284 | Tensor.silu, |
| 285 | Linear(1280, 1280), |
| 286 | ] |
| 287 | self.input_blocks = [ |
| 288 | [Conv2d(4, 320, kernel_size=3, padding=1)], |
| 289 | [ResBlock(320, 1280, 320), SpatialTransformer(320, 768, 8, 40)], |
| 290 | [ResBlock(320, 1280, 320), SpatialTransformer(320, 768, 8, 40)], |
| 291 | [Downsample(320)], |
| 292 | [ResBlock(320, 1280, 640), SpatialTransformer(640, 768, 8, 80)], |
| 293 | [ResBlock(640, 1280, 640), SpatialTransformer(640, 768, 8, 80)], |
| 294 | [Downsample(640)], |
| 295 | [ResBlock(640, 1280, 1280), SpatialTransformer(1280, 768, 8, 160)], |
| 296 | [ResBlock(1280, 1280, 1280), SpatialTransformer(1280, 768, 8, 160)], |
| 297 | [Downsample(1280)], |
| 298 | [ResBlock(1280, 1280, 1280)], |
| 299 | [ResBlock(1280, 1280, 1280)] |
| 300 | ] |
| 301 | self.middle_block = [ |
| 302 | ResBlock(1280, 1280, 1280), |
| 303 | SpatialTransformer(1280, 768, 8, 160), |
| 304 | ResBlock(1280, 1280, 1280) |
| 305 | ] |
| 306 | self.output_blocks = [ |
| 307 | [ResBlock(2560, 1280, 1280)], |
| 308 | [ResBlock(2560, 1280, 1280)], |
| 309 | [ResBlock(2560, 1280, 1280), Upsample(1280)], |
| 310 | [ResBlock(2560, 1280, 1280), SpatialTransformer(1280, 768, 8, 160)], |
| 311 | [ResBlock(2560, 1280, 1280), SpatialTransformer(1280, 768, 8, 160)], |
| 312 | [ResBlock(1920, 1280, 1280), SpatialTransformer(1280, 768, 8, 160), Upsample(1280)], |
| 313 | [ResBlock(1920, 1280, 640), SpatialTransformer(640, 768, 8, 80)], # 6 |
| 314 | [ResBlock(1280, 1280, 640), SpatialTransformer(640, 768, 8, 80)], |
| 315 | [ResBlock(960, 1280, 640), SpatialTransformer(640, 768, 8, 80), Upsample(640)], |
| 316 | [ResBlock(960, 1280, 320), SpatialTransformer(320, 768, 8, 40)], |
| 317 | [ResBlock(640, 1280, 320), SpatialTransformer(320, 768, 8, 40)], |
| 318 | [ResBlock(640, 1280, 320), SpatialTransformer(320, 768, 8, 40)], |
| 319 | ] |
| 320 | self.out = [ |
| 321 | GroupNorm(32, 320), |
| 322 | Tensor.silu, |
| 323 | Conv2d(320, 4, kernel_size=3, padding=1) |
| 324 | ] |
| 325 | |
| 326 | def __call__(self, x, timesteps=None, context=None): |
| 327 | # TODO: real time embedding |
| 328 | t_emb = timestep_embedding(timesteps, 320) |
| 329 | emb = t_emb.sequential(self.time_embed) |
| 330 | |
| 331 | |
| 332 | |
| 333 | def run(x, bb): |
| 334 | if isinstance(bb, ResBlock): x = bb(x, emb) |
| 335 | elif isinstance(bb, SpatialTransformer): x = bb(x, context) |
| 336 | else: x = bb(x) |
| 337 | return x |