| 746 | |
| 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) |
| 775 | for block in self.res_block1: |
| 776 | x = block(x, None) |
| 777 | x = torch.nn.functional.interpolate(x, size=(int(round(x.shape[2]*self.factor)), int(round(x.shape[3]*self.factor)))) |
| 778 | x = self.attn(x) |
| 779 | for block in self.res_block2: |
| 780 | x = block(x, None) |
| 781 | x = self.conv_out(x) |
| 782 | return x |
| 783 | |
| 784 | |
| 785 | class MergedRescaleEncoder(nn.Module): |