(self,
in_channels,
out_channels=None,
temb_channels=512,
conv_shortcut=False,
)
| 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() |
nothing calls this directly
no test coverage detected