| 496 | return x |
| 497 | |
| 498 | class Resnet3DBlock(nn.Module): |
| 499 | def __init__(self, |
| 500 | in_channels, |
| 501 | out_channels=None, |
| 502 | temb_channels=512, |
| 503 | conv_shortcut=False, |
| 504 | ): |
| 505 | super().__init__() |
| 506 | |
| 507 | self.in_channels = in_channels |
| 508 | out_channels = in_channels if out_channels is None else out_channels |
| 509 | self.out_channels = out_channels |
| 510 | |
| 511 | self.norm1 = BaseGroupNorm(num_groups=32, num_channels=in_channels) |
| 512 | self.conv1 = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) |
| 513 | if temb_channels > 0: |
| 514 | self.temb_proj = nn.Linear(temb_channels, out_channels) |
| 515 | |
| 516 | self.norm2 = BaseGroupNorm(num_groups=32, num_channels=out_channels) |
| 517 | self.conv2 = CausalConvAfterNorm(out_channels, out_channels, kernel_size=3) |
| 518 | |
| 519 | assert conv_shortcut is False |
| 520 | self.use_conv_shortcut = conv_shortcut |
| 521 | if self.in_channels != self.out_channels: |
| 522 | if self.use_conv_shortcut: |
| 523 | self.conv_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=3) |
| 524 | else: |
| 525 | self.nin_shortcut = CausalConvAfterNorm(in_channels, out_channels, kernel_size=1) |
| 526 | |
| 527 | def forward(self, x, temb=None, is_init=True): |
| 528 | x = x.permute(0,2,3,4,1).contiguous() |
| 529 | |
| 530 | h = self.norm1(x, zero_pad=True, act_silu=True, pad_size=2) |
| 531 | h = self.conv1(h) |
| 532 | if temb is not None: |
| 533 | h = h + self.temb_proj(nn.functional.silu(temb))[:, :, None, None] |
| 534 | |
| 535 | x = self.nin_shortcut(x) if self.in_channels != self.out_channels else x |
| 536 | |
| 537 | h = self.norm2(h, zero_pad=True, act_silu=True, pad_size=2) |
| 538 | x = self.conv2(h, residual=x) |
| 539 | |
| 540 | x = x.permute(0,4,1,2,3) |
| 541 | return x |
| 542 | |
| 543 | |
| 544 | class Downsample3D(nn.Module): |