| 397 | |
| 398 | |
| 399 | class ResnetBlock2D(nn.Module): |
| 400 | |
| 401 | def __init__( |
| 402 | self, |
| 403 | *, |
| 404 | in_channels, |
| 405 | out_channels=None, |
| 406 | conv_shortcut=False, |
| 407 | temb_channels=512, |
| 408 | groups=32, |
| 409 | groups_out=None, |
| 410 | eps=1e-6, |
| 411 | time_embedding_norm="default", |
| 412 | use_nin_shortcut=None, |
| 413 | ): |
| 414 | super().__init__() |
| 415 | self.in_channels = in_channels |
| 416 | self.out_channels = in_channels if out_channels is None else out_channels |
| 417 | self.use_conv_shortcut = conv_shortcut |
| 418 | self.time_embedding_norm = time_embedding_norm |
| 419 | |
| 420 | if groups_out is None: |
| 421 | groups_out = groups |
| 422 | |
| 423 | self.norm1 = torch.nn.GroupNorm(num_groups=groups, |
| 424 | num_channels=in_channels, |
| 425 | eps=eps, |
| 426 | affine=True) |
| 427 | |
| 428 | self.conv1 = torch.nn.Conv2d(in_channels, |
| 429 | out_channels, |
| 430 | kernel_size=3, |
| 431 | stride=1, |
| 432 | padding=1) |
| 433 | |
| 434 | if temb_channels is not None: |
| 435 | self.time_emb_proj = torch.nn.Conv2d(temb_channels, |
| 436 | out_channels, |
| 437 | kernel_size=1) |
| 438 | else: |
| 439 | self.time_emb_proj = None |
| 440 | |
| 441 | self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, |
| 442 | num_channels=out_channels, |
| 443 | eps=eps, |
| 444 | affine=True) |
| 445 | self.conv2 = torch.nn.Conv2d(out_channels, |
| 446 | out_channels, |
| 447 | kernel_size=3, |
| 448 | stride=1, |
| 449 | padding=1) |
| 450 | |
| 451 | self.nonlinearity = nn.SiLU() |
| 452 | |
| 453 | self.use_nin_shortcut = self.in_channels != self.out_channels if use_nin_shortcut is None else use_nin_shortcut |
| 454 | |
| 455 | self.conv_shortcut = None |
| 456 | if self.use_nin_shortcut: |