| 709 | |
| 710 | |
| 711 | class MergedRescaleDecoder(nn.Module): |
| 712 | def __init__(self, z_channels, out_ch, resolution, num_res_blocks, attn_resolutions, ch, ch_mult=(1,2,4,8), |
| 713 | dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1): |
| 714 | super().__init__() |
| 715 | tmp_chn = z_channels*ch_mult[-1] |
| 716 | self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_resolutions=attn_resolutions, dropout=dropout, |
| 717 | resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks, |
| 718 | ch_mult=ch_mult, resolution=resolution, ch=ch) |
| 719 | self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn, |
| 720 | out_channels=tmp_chn, depth=rescale_module_depth) |
| 721 | |
| 722 | def forward(self, x): |
| 723 | x = self.rescaler(x) |
| 724 | x = self.decoder(x) |
| 725 | return x |
| 726 | |
| 727 | |
| 728 | class Upsampler(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected