| 802 | |
| 803 | |
| 804 | class MergedRescaleDecoder(nn.Module): |
| 805 | def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8), |
| 806 | dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1): |
| 807 | super().__init__() |
| 808 | tmp_chn = z_channels*ch_mult[-1] |
| 809 | self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout, |
| 810 | resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks, |
| 811 | ch_mult=ch_mult, resolution=resolution, ch=ch) |
| 812 | self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn, |
| 813 | out_channels=tmp_chn, depth=rescale_module_depth) |
| 814 | |
| 815 | def forward(self, x): |
| 816 | x = self.rescaler(x) |
| 817 | x = self.decoder(x) |
| 818 | return x |
| 819 | |
| 820 | |
| 821 | class Upsampler(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected