Method __init__

model/DocDiff.py:324–330

(self, input_channels: int = 2, output_channels: int = 1, n_channels: int = 32,
                 ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4),
                 n_blocks: int = 1)

9
class DocDiff(nn.Module):
    def __init__(self, input_channels: int = 2, output_channels: int = 1, n_channels: int = 32,
                 ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4),
                 n_blocks: int = 1):
        super(DocDiff, self).__init__()
        self.denoiser = UNet(input_channels, output_channels, n_channels, ch_mults, n_blocks, is_noise=True)
        self.init_predictor = UNet(input_channels//2, output_channels, n_channels, ch_mults, n_blocks, is_noise=False)
        # self.init_predictor = UNet(input_channels, output_channels, 2 * n_channels, ch_mults, n_blocks)

    def forward(self, x, condition, t, diffusion):
        x_ = self.init_predictor(condition, t)