| 165 | |
| 166 | |
| 167 | class Decoder(nn.Module): |
| 168 | config: VQGANConfig |
| 169 | |
| 170 | @nn.compact |
| 171 | def __call__(self, hidden_states): |
| 172 | hidden_states = nn.Conv( |
| 173 | self.config.hidden_channels * self.config.channel_mult[self.config.num_resolutions - 1], |
| 174 | [3, 3] |
| 175 | )(hidden_states) |
| 176 | hidden_states = MidBlock( |
| 177 | self.config, self.config.no_attn_mid_block, self.config.dropout |
| 178 | )(hidden_states) |
| 179 | for i_level in reversed(range(self.config.num_resolutions)): |
| 180 | hidden_states = UpsamplingBlock(self.config, i_level)(hidden_states) |
| 181 | hidden_states = nn.GroupNorm()(hidden_states) |
| 182 | hidden_states = nn.silu(hidden_states) |
| 183 | hidden_states = nn.Conv(self.config.num_channels, [3, 3])(hidden_states) |
| 184 | return hidden_states |
| 185 | |
| 186 | |
| 187 | class VectorQuantizer(nn.Module): |