Class DownBlock

model/DocDiff.py:119–131

### Down block This combines `ResidualBlock` and `AttentionBlock`. These are used in the first half of U-Net at each resolution.

class DownBlock(nn.Module):
    """
    ### Down block
    This combines `ResidualBlock` and `AttentionBlock`. These are used in the first half of U-Net at each resolution.
    """

    def __init__(self, in_channels: int, out_channels: int, time_channels: int, is_noise: bool = True):
        super().__init__()
        self.res = ResidualBlock(in_channels, out_channels, time_channels, is_noise=is_noise)

    def forward(self, x: torch.Tensor, t: torch.Tensor):
        x = self.res(x, t)
        return x


class UpBlock(nn.Module):