| 244 | |
| 245 | |
| 246 | class ConvResblock(nn.Module): |
| 247 | def __init__(self, in_features=320, out_features=320) -> None: |
| 248 | super().__init__() |
| 249 | self.f_t = nn.Linear(1280, out_features * 2) |
| 250 | |
| 251 | self.gn_1 = nn.GroupNorm(32, in_features) |
| 252 | self.f_1 = nn.Conv2d(in_features, out_features, kernel_size=3, padding=1) |
| 253 | |
| 254 | self.gn_2 = nn.GroupNorm(32, out_features) |
| 255 | self.f_2 = nn.Conv2d(out_features, out_features, kernel_size=3, padding=1) |
| 256 | |
| 257 | skip_conv = in_features != out_features |
| 258 | self.f_s = nn.Conv2d(in_features, out_features, kernel_size=1, padding=0) if skip_conv else nn.Identity() |
| 259 | |
| 260 | def forward(self, x, t): |
| 261 | x_skip = x |
| 262 | t = self.f_t(F.silu(t)) |
| 263 | t = t.chunk(2, dim=1) |
| 264 | t_1 = t[0].unsqueeze(dim=2).unsqueeze(dim=3) + 1 |
| 265 | t_2 = t[1].unsqueeze(dim=2).unsqueeze(dim=3) |
| 266 | |
| 267 | gn_1 = F.silu(self.gn_1(x)) |
| 268 | f_1 = self.f_1(gn_1) |
| 269 | |
| 270 | gn_2 = self.gn_2(f_1) |
| 271 | |
| 272 | return self.f_s(x_skip) + self.f_2(F.silu(gn_2 * t_1 + t_2)) |
| 273 | |
| 274 | |
| 275 | # Also ConvResblock |