| 783 | |
| 784 | |
| 785 | class MergedRescaleEncoder(nn.Module): |
| 786 | def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks, |
| 787 | attn_resolutions, dropout=0.0, resamp_with_conv=True, |
| 788 | ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1): |
| 789 | super().__init__() |
| 790 | intermediate_chn = ch * ch_mult[-1] |
| 791 | self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult, |
| 792 | z_channels=intermediate_chn, double_z=False, resolution=resolution, |
| 793 | attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv, |
| 794 | out_ch=None) |
| 795 | self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn, |
| 796 | mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth) |
| 797 | |
| 798 | def forward(self, x): |
| 799 | x = self.encoder(x) |
| 800 | x = self.rescaler(x) |
| 801 | return x |
| 802 | |
| 803 | |
| 804 | class MergedRescaleDecoder(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected