(self, factor, in_channels, mid_channels, out_channels, depth=2)
| 747 | |
| 748 | class LatentRescaler(nn.Module): |
| 749 | def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2): |
| 750 | super().__init__() |
| 751 | # residual block, interpolate, residual block |
| 752 | self.factor = factor |
| 753 | self.conv_in = nn.Conv2d(in_channels, |
| 754 | mid_channels, |
| 755 | kernel_size=3, |
| 756 | stride=1, |
| 757 | padding=1) |
| 758 | self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, |
| 759 | out_channels=mid_channels, |
| 760 | temb_channels=0, |
| 761 | dropout=0.0) for _ in range(depth)]) |
| 762 | self.attn = AttnBlock(mid_channels) |
| 763 | self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels, |
| 764 | out_channels=mid_channels, |
| 765 | temb_channels=0, |
| 766 | dropout=0.0) for _ in range(depth)]) |
| 767 | |
| 768 | self.conv_out = nn.Conv2d(mid_channels, |
| 769 | out_channels, |
| 770 | kernel_size=1, |
| 771 | ) |
| 772 | |
| 773 | def forward(self, x): |
| 774 | x = self.conv_in(x) |
nothing calls this directly
no test coverage detected