(
self,
*,
in_channels,
out_channels=None,
conv_shortcut=False,
dropout,
temb_channels=512,
zq_ch=None,
add_conv=False,
gather_norm=False,
normalization=Normalize,
)
| 517 | |
| 518 | class ContextParallelResnetBlock3D(nn.Module): |
| 519 | def __init__( |
| 520 | self, |
| 521 | *, |
| 522 | in_channels, |
| 523 | out_channels=None, |
| 524 | conv_shortcut=False, |
| 525 | dropout, |
| 526 | temb_channels=512, |
| 527 | zq_ch=None, |
| 528 | add_conv=False, |
| 529 | gather_norm=False, |
| 530 | normalization=Normalize, |
| 531 | ): |
| 532 | super().__init__() |
| 533 | self.in_channels = in_channels |
| 534 | out_channels = in_channels if out_channels is None else out_channels |
| 535 | self.out_channels = out_channels |
| 536 | self.use_conv_shortcut = conv_shortcut |
| 537 | |
| 538 | self.norm1 = normalization( |
| 539 | in_channels, |
| 540 | zq_ch=zq_ch, |
| 541 | add_conv=add_conv, |
| 542 | gather=gather_norm, |
| 543 | ) |
| 544 | |
| 545 | self.conv1 = ContextParallelCausalConv3d( |
| 546 | chan_in=in_channels, |
| 547 | chan_out=out_channels, |
| 548 | kernel_size=3, |
| 549 | ) |
| 550 | if temb_channels > 0: |
| 551 | self.temb_proj = torch.nn.Linear(temb_channels, out_channels) |
| 552 | self.norm2 = normalization( |
| 553 | out_channels, |
| 554 | zq_ch=zq_ch, |
| 555 | add_conv=add_conv, |
| 556 | gather=gather_norm, |
| 557 | ) |
| 558 | self.dropout = torch.nn.Dropout(dropout) |
| 559 | self.conv2 = ContextParallelCausalConv3d( |
| 560 | chan_in=out_channels, |
| 561 | chan_out=out_channels, |
| 562 | kernel_size=3, |
| 563 | ) |
| 564 | if self.in_channels != self.out_channels: |
| 565 | if self.use_conv_shortcut: |
| 566 | self.conv_shortcut = ContextParallelCausalConv3d( |
| 567 | chan_in=in_channels, |
| 568 | chan_out=out_channels, |
| 569 | kernel_size=3, |
| 570 | ) |
| 571 | else: |
| 572 | self.nin_shortcut = Conv3d( |
| 573 | in_channels, |
| 574 | out_channels, |
| 575 | kernel_size=1, |
| 576 | stride=1, |
nothing calls this directly
no test coverage detected