(timesteps, dim, max_period=10000)
| 272 | return self.conv(x) |
| 273 | |
| 274 | def timestep_embedding(timesteps, dim, max_period=10000): |
| 275 | half = dim // 2 |
| 276 | freqs = (-math.log(max_period) * Tensor.arange(half) / half).exp() |
| 277 | args = timesteps * freqs |
| 278 | return Tensor.cat(args.cos(), args.sin()).reshape(1, -1) |
| 279 | |
| 280 | class UNetModel: |
| 281 | def __init__(self): |