(self,
ch=32,
ch_mult=(4, 8, 16, 16),
num_res_blocks=2,
in_channels=3,
z_channels=16,
double_z=True,
down_sampling_layer=[1, 2],
resamp_with_conv=True,
version=1,
)
| 562 | |
| 563 | class VideoEncoder(nn.Module): |
| 564 | def __init__(self, |
| 565 | ch=32, |
| 566 | ch_mult=(4, 8, 16, 16), |
| 567 | num_res_blocks=2, |
| 568 | in_channels=3, |
| 569 | z_channels=16, |
| 570 | double_z=True, |
| 571 | down_sampling_layer=[1, 2], |
| 572 | resamp_with_conv=True, |
| 573 | version=1, |
| 574 | ): |
| 575 | super().__init__() |
| 576 | |
| 577 | temb_ch = 0 |
| 578 | |
| 579 | self.num_resolutions = len(ch_mult) |
| 580 | self.num_res_blocks = num_res_blocks |
| 581 | |
| 582 | # downsampling |
| 583 | self.conv_in = CausalConv(in_channels, ch, kernel_size=3) |
| 584 | self.down_sampling_layer = down_sampling_layer |
| 585 | |
| 586 | in_ch_mult = (1,) + tuple(ch_mult) |
| 587 | self.down = nn.ModuleList() |
| 588 | for i_level in range(self.num_resolutions): |
| 589 | block = nn.ModuleList() |
| 590 | attn = nn.ModuleList() |
| 591 | block_in = ch * in_ch_mult[i_level] |
| 592 | block_out = ch * ch_mult[i_level] |
| 593 | for i_block in range(self.num_res_blocks): |
| 594 | block.append( |
| 595 | Resnet3DBlock(in_channels=block_in, out_channels=block_out, temb_channels=temb_ch)) |
| 596 | block_in = block_out |
| 597 | down = nn.Module() |
| 598 | down.block = block |
| 599 | down.attn = attn |
| 600 | if i_level != self.num_resolutions - 1: |
| 601 | if i_level in self.down_sampling_layer: |
| 602 | down.downsample = Downsample3D(block_in, resamp_with_conv, stride=(2, 2, 2)) |
| 603 | else: |
| 604 | down.downsample = Downsample2D(block_in, resamp_with_conv, padding=0) #DIFF |
| 605 | self.down.append(down) |
| 606 | |
| 607 | # middle |
| 608 | self.mid = nn.Module() |
| 609 | self.mid.block_1 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) |
| 610 | self.mid.attn_1 = AttnBlock(block_in) |
| 611 | self.mid.block_2 = Resnet3DBlock(in_channels=block_in, out_channels=block_in, temb_channels=temb_ch) |
| 612 | |
| 613 | # end |
| 614 | self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in) |
| 615 | self.version = version |
| 616 | if version == 2: |
| 617 | channels = 4 * z_channels * 2 ** 3 |
| 618 | self.conv_patchify = ConvPixelUnshuffleDownSampleLayer3D(block_in, channels, kernel_size=3, factor=2) |
| 619 | self.shortcut_pathify = PixelUnshuffleChannelAveragingDownSampleLayer3D(block_in, channels, 2) |
| 620 | self.shortcut_out = PixelUnshuffleChannelAveragingDownSampleLayer3D(channels, 2 * z_channels if double_z else z_channels, 1) |
| 621 | self.conv_out = CausalConvChannelLast(channels, 2 * z_channels if double_z else z_channels, kernel_size=3) |
nothing calls this directly
no test coverage detected