| 128 | return quantized_states, codebook_indices |
| 129 | |
| 130 | def decode(self, encoding, is_codebook_indices=True): |
| 131 | if is_codebook_indices: |
| 132 | encoding = self.quantize(None, encoding) |
| 133 | T = None |
| 134 | if len(encoding.shape) == 5: |
| 135 | T = encoding.shape[1] |
| 136 | encoding = encoding.reshape(-1, *encoding.shape[2:]) |
| 137 | hidden_states = self.post_quant_conv(encoding) |
| 138 | reconstructed_pixel_values = self.decoder(hidden_states) |
| 139 | if T is not None: |
| 140 | reconstructed_pixel_values = reconstructed_pixel_values.reshape(-1, T, *reconstructed_pixel_values.shape[1:]) |
| 141 | return jnp.clip(reconstructed_pixel_values, -1, 1) |
| 142 | |
| 143 | def __call__(self, pixel_values): |
| 144 | encoding = self.encode(pixel_values)[1] |