(
self,
*,
in_channels,
out_channels=None,
conv_shortcut=False,
dropout,
temb_channels=512,
zq_ch=None,
add_conv=False,
gather_norm=False,
normalization=Normalize,
)
| 613 | |
| 614 | class ContextParallelResnetBlock3D(nn.Module): |
| 615 | def __init__( |
| 616 | self, |
| 617 | *, |
| 618 | in_channels, |
| 619 | out_channels=None, |
| 620 | conv_shortcut=False, |
| 621 | dropout, |
| 622 | temb_channels=512, |
| 623 | zq_ch=None, |
| 624 | add_conv=False, |
| 625 | gather_norm=False, |
| 626 | normalization=Normalize, |
| 627 | ): |
| 628 | super().__init__() |
| 629 | self.in_channels = in_channels |
| 630 | out_channels = in_channels if out_channels is None else out_channels |
| 631 | self.out_channels = out_channels |
| 632 | self.use_conv_shortcut = conv_shortcut |
| 633 | |
| 634 | self.norm1 = normalization( |
| 635 | in_channels, |
| 636 | zq_ch=zq_ch, |
| 637 | add_conv=add_conv, |
| 638 | gather=gather_norm, |
| 639 | ) |
| 640 | |
| 641 | self.conv1 = ContextParallelCausalConv3d( |
| 642 | chan_in=in_channels, |
| 643 | chan_out=out_channels, |
| 644 | kernel_size=3, |
| 645 | ) |
| 646 | if temb_channels > 0: |
| 647 | self.temb_proj = torch.nn.Linear(temb_channels, out_channels) |
| 648 | self.norm2 = normalization( |
| 649 | out_channels, |
| 650 | zq_ch=zq_ch, |
| 651 | add_conv=add_conv, |
| 652 | gather=gather_norm, |
| 653 | ) |
| 654 | self.dropout = torch.nn.Dropout(dropout) |
| 655 | self.conv2 = ContextParallelCausalConv3d( |
| 656 | chan_in=out_channels, |
| 657 | chan_out=out_channels, |
| 658 | kernel_size=3, |
| 659 | ) |
| 660 | if self.in_channels != self.out_channels: |
| 661 | if self.use_conv_shortcut: |
| 662 | self.conv_shortcut = ContextParallelCausalConv3d( |
| 663 | chan_in=in_channels, |
| 664 | chan_out=out_channels, |
| 665 | kernel_size=3, |
| 666 | ) |
| 667 | else: |
| 668 | self.nin_shortcut = Conv3d( |
| 669 | in_channels, |
| 670 | out_channels, |
| 671 | kernel_size=1, |
| 672 | stride=1, |
nothing calls this directly
no test coverage detected