MCPcopy Create free account
hub / github.com/Royalvice/DocDiff / MiddleBlock

Class MiddleBlock

model/DocDiff.py:151–174  ·  view source on GitHub ↗

### Middle block It combines a `ResidualBlock`, `AttentionBlock`, followed by another `ResidualBlock`. This block is applied at the lowest resolution of the U-Net.

Source from the content-addressed store, hash-verified

149
150
151class MiddleBlock(nn.Module):
152 """
153 ### Middle block
154 It combines a `ResidualBlock`, `AttentionBlock`, followed by another `ResidualBlock`.
155 This block is applied at the lowest resolution of the U-Net.
156 """
157
158 def __init__(self, n_channels: int, time_channels: int, is_noise: bool = True):
159 super().__init__()
160 self.res1 = ResidualBlock(n_channels, n_channels, time_channels, is_noise=is_noise)
161 self.dia1 = nn.Conv2d(n_channels, n_channels, 3, 1, dilation=2, padding=get_pad(16, 3, 1, 2))
162 self.dia2 = nn.Conv2d(n_channels, n_channels, 3, 1, dilation=4, padding=get_pad(16, 3, 1, 4))
163 self.dia3 = nn.Conv2d(n_channels, n_channels, 3, 1, dilation=8, padding=get_pad(16, 3, 1, 8))
164 self.dia4 = nn.Conv2d(n_channels, n_channels, 3, 1, dilation=16, padding=get_pad(16, 3, 1, 16))
165 self.res2 = ResidualBlock(n_channels, n_channels, time_channels, is_noise=is_noise)
166
167 def forward(self, x: torch.Tensor, t: torch.Tensor):
168 x = self.res1(x, t)
169 x = self.dia1(x)
170 x = self.dia2(x)
171 x = self.dia3(x)
172 x = self.dia4(x)
173 x = self.res2(x, t)
174 return x
175
176
177class Upsample(nn.Module):

Callers 1

__init__Method · 0.85

Calls

no outgoing calls

Tested by

no test coverage detected