| 305 | |
| 306 | # Also ConvResblock |
| 307 | class Upsample(nn.Module): |
| 308 | def __init__(self, in_channels=1024) -> None: |
| 309 | super().__init__() |
| 310 | self.f_t = nn.Linear(1280, in_channels * 2) |
| 311 | |
| 312 | self.gn_1 = nn.GroupNorm(32, in_channels) |
| 313 | self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) |
| 314 | self.gn_2 = nn.GroupNorm(32, in_channels) |
| 315 | |
| 316 | self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) |
| 317 | |
| 318 | def forward(self, x, t) -> torch.Tensor: |
| 319 | x_skip = x |
| 320 | |
| 321 | t = self.f_t(F.silu(t)) |
| 322 | t_1, t_2 = t.chunk(2, dim=1) |
| 323 | t_1 = t_1.unsqueeze(2).unsqueeze(3) + 1 |
| 324 | t_2 = t_2.unsqueeze(2).unsqueeze(3) |
| 325 | |
| 326 | gn_1 = F.silu(self.gn_1(x)) |
| 327 | upsample = F.upsample_nearest(gn_1, scale_factor=2) |
| 328 | f_1 = self.f_1(upsample) |
| 329 | gn_2 = self.gn_2(f_1) |
| 330 | |
| 331 | f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2))) |
| 332 | |
| 333 | return f_2 + F.upsample_nearest(x_skip, scale_factor=2) |
| 334 | |
| 335 | |
| 336 | class ConvUNetVAE(nn.Module): |