| 222 | |
| 223 | |
| 224 | class DownsamplingBlock(nn.Module): |
| 225 | config: VQGANConfig |
| 226 | block_idx: int |
| 227 | |
| 228 | @nn.compact |
| 229 | def __call__(self, hidden_states): |
| 230 | block_out = self.config.hidden_channels * self.config.channel_mult[self.block_idx] |
| 231 | for _ in range(self.config.num_res_blocks): |
| 232 | hidden_states = ResnetBlock( |
| 233 | block_out, dropout_prob=self.config.dropout |
| 234 | )(hidden_states) |
| 235 | if hidden_states.shape[1] in self.config.attn_resolutions: |
| 236 | hidden_states = AttnBlock()(hidden_states) |
| 237 | if self.block_idx != self.config.num_resolutions - 1: |
| 238 | hidden_states = Downsample(self.config.resample_with_conv)(hidden_states) |
| 239 | return hidden_states |
| 240 | |
| 241 | |
| 242 | class ResnetBlock(nn.Module): |