| 690 | |
| 691 | |
| 692 | class MergedRescaleEncoder(nn.Module): |
| 693 | def __init__(self, in_channels, ch, resolution, out_ch, num_res_blocks, |
| 694 | attn_resolutions, dropout=0.0, resamp_with_conv=True, |
| 695 | ch_mult=(1,2,4,8), rescale_factor=1.0, rescale_module_depth=1): |
| 696 | super().__init__() |
| 697 | intermediate_chn = ch * ch_mult[-1] |
| 698 | self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult, |
| 699 | z_channels=intermediate_chn, double_z=False, resolution=resolution, |
| 700 | attn_resolutions=attn_resolutions, dropout=dropout, resamp_with_conv=resamp_with_conv, |
| 701 | out_ch=None) |
| 702 | self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn, |
| 703 | mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth) |
| 704 | |
| 705 | def forward(self, x): |
| 706 | x = self.encoder(x) |
| 707 | x = self.rescaler(x) |
| 708 | return x |
| 709 | |
| 710 | |
| 711 | class MergedRescaleDecoder(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected