| 201 | |
| 202 | class MOVQDecoder(nn.Module): |
| 203 | def __init__( |
| 204 | self, |
| 205 | *, |
| 206 | ch, |
| 207 | out_ch, |
| 208 | ch_mult=(1, 2, 4, 8), |
| 209 | num_res_blocks, |
| 210 | attn_resolutions, |
| 211 | dropout=0.0, |
| 212 | resamp_with_conv=True, |
| 213 | in_channels, |
| 214 | resolution, |
| 215 | z_channels, |
| 216 | give_pre_end=False, |
| 217 | zq_ch=None, |
| 218 | add_conv=False, |
| 219 | **ignorekwargs, |
| 220 | ): |
| 221 | super().__init__() |
| 222 | self.ch = ch |
| 223 | self.temb_ch = 0 |
| 224 | self.num_resolutions = len(ch_mult) |
| 225 | self.num_res_blocks = num_res_blocks |
| 226 | self.resolution = resolution |
| 227 | self.in_channels = in_channels |
| 228 | self.give_pre_end = give_pre_end |
| 229 | |
| 230 | # compute in_ch_mult, block_in and curr_res at lowest res |
| 231 | in_ch_mult = (1,) + tuple(ch_mult) |
| 232 | block_in = ch * ch_mult[self.num_resolutions - 1] |
| 233 | curr_res = resolution // 2 ** (self.num_resolutions - 1) |
| 234 | self.z_shape = (1, z_channels, curr_res, curr_res) |
| 235 | print("Working with z of shape {} = {} dimensions.".format(self.z_shape, np.prod(self.z_shape))) |
| 236 | |
| 237 | # z to block_in |
| 238 | self.conv_in = torch.nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) |
| 239 | |
| 240 | # middle |
| 241 | self.mid = nn.Module() |
| 242 | self.mid.block_1 = ResnetBlock( |
| 243 | in_channels=block_in, |
| 244 | out_channels=block_in, |
| 245 | temb_channels=self.temb_ch, |
| 246 | dropout=dropout, |
| 247 | zq_ch=zq_ch, |
| 248 | add_conv=add_conv, |
| 249 | ) |
| 250 | self.mid.attn_1 = AttnBlock(block_in, zq_ch, add_conv=add_conv) |
| 251 | self.mid.block_2 = ResnetBlock( |
| 252 | in_channels=block_in, |
| 253 | out_channels=block_in, |
| 254 | temb_channels=self.temb_ch, |
| 255 | dropout=dropout, |
| 256 | zq_ch=zq_ch, |
| 257 | add_conv=add_conv, |
| 258 | ) |
| 259 | |
| 260 | # upsampling |