| 320 | |
| 321 | |
| 322 | class UpsamplingBlock(nn.Module): |
| 323 | config: VQGANConfig |
| 324 | block_idx: int |
| 325 | |
| 326 | @nn.compact |
| 327 | def __call__(self, hidden_states): |
| 328 | block_out = self.config.hidden_channels * self.config.channel_mult[self.block_idx] |
| 329 | for _ in range(self.config.num_res_blocks + 1): |
| 330 | hidden_states = ResnetBlock( |
| 331 | block_out, dropout_prob=self.config.dropout |
| 332 | )(hidden_states) |
| 333 | if hidden_states.shape[1] in self.config.attn_resolutions: |
| 334 | hidden_states = AttnBlock()(hidden_states) |
| 335 | if self.block_idx != 0: |
| 336 | hidden_states = Upsample(self.config.resample_with_conv)(hidden_states) |
| 337 | return hidden_states |
| 338 | |
| 339 | |
| 340 | class MidBlock(nn.Module): |