| 184 | |
| 185 | |
| 186 | class UpBlock(nn.Module): |
| 187 | def __init__( |
| 188 | self, |
| 189 | in_channels: int, |
| 190 | prev_out_channels: int, |
| 191 | out_channels: int, |
| 192 | num_layers: int = 1, |
| 193 | upsample: bool = True, |
| 194 | attention: bool = True, |
| 195 | attention_heads: int = 16, |
| 196 | skip_scale: float = 1, |
| 197 | ): |
| 198 | super().__init__() |
| 199 | |
| 200 | nets = [] |
| 201 | attns = [] |
| 202 | for i in range(num_layers): |
| 203 | cin = in_channels if i == 0 else out_channels |
| 204 | cskip = prev_out_channels if (i == num_layers - 1) else out_channels |
| 205 | |
| 206 | nets.append(ResnetBlock(cin + cskip, out_channels, skip_scale=skip_scale)) |
| 207 | if attention: |
| 208 | attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale)) |
| 209 | else: |
| 210 | attns.append(None) |
| 211 | self.nets = nn.ModuleList(nets) |
| 212 | self.attns = nn.ModuleList(attns) |
| 213 | |
| 214 | self.upsample = None |
| 215 | if upsample: |
| 216 | self.upsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) |
| 217 | |
| 218 | def forward(self, x, xs): |
| 219 | |
| 220 | for attn, net in zip(self.attns, self.nets): |
| 221 | res_x = xs[-1] |
| 222 | xs = xs[:-1] |
| 223 | x = torch.cat([x, res_x], dim=1) |
| 224 | x = net(x) |
| 225 | if attn: |
| 226 | x = attn(x) |
| 227 | |
| 228 | if self.upsample: |
| 229 | x = F.interpolate(x, scale_factor=2.0, mode='nearest') |
| 230 | x = self.upsample(x) |
| 231 | |
| 232 | return x |
| 233 | |
| 234 | |
| 235 | # it could be asymmetric! |