| 321 | |
| 322 | 9 |
| 323 | class DocDiff(nn.Module): |
| 324 | def __init__(self, input_channels: int = 2, output_channels: int = 1, n_channels: int = 32, |
| 325 | ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4), |
| 326 | n_blocks: int = 1): |
| 327 | super(DocDiff, self).__init__() |
| 328 | self.denoiser = UNet(input_channels, output_channels, n_channels, ch_mults, n_blocks, is_noise=True) |
| 329 | self.init_predictor = UNet(input_channels//2, output_channels, n_channels, ch_mults, n_blocks, is_noise=False) |
| 330 | # self.init_predictor = UNet(input_channels, output_channels, 2 * n_channels, ch_mults, n_blocks) |
| 331 | |
| 332 | def forward(self, x, condition, t, diffusion): |
| 333 | x_ = self.init_predictor(condition, t) |
| 334 | residual = x - x_ |
| 335 | noisy_image, noise_ref = diffusion.noisy_image(t, residual) |
| 336 | x__ = self.denoiser(torch.cat((noisy_image, x_.clone().detach()), dim=1), t) |
| 337 | return x_, x__, noisy_image, noise_ref |
| 338 | |
| 339 | |
| 340 | class EMA(): |
no outgoing calls
no test coverage detected