| 115 | self.post_quant_conv = nn.Conv(self.config.z_channels, [1, 1]) |
| 116 | |
| 117 | def encode(self, pixel_values): |
| 118 | T = None |
| 119 | if len(pixel_values.shape) == 5: # video |
| 120 | T = pixel_values.shape[1] |
| 121 | pixel_values = pixel_values.reshape(-1, *pixel_values.shape[2:]) |
| 122 | hidden_states = self.encoder(pixel_values) |
| 123 | hidden_states = self.quant_conv(hidden_states) |
| 124 | quantized_states, codebook_indices = self.quantize(hidden_states) |
| 125 | if T is not None: |
| 126 | quantized_states = quantized_states.reshape(-1, T, *quantized_states.shape[1:]) |
| 127 | codebook_indices = codebook_indices.reshape(-1, T, *codebook_indices.shape[1:]) |
| 128 | return quantized_states, codebook_indices |
| 129 | |
| 130 | def decode(self, encoding, is_codebook_indices=True): |
| 131 | if is_codebook_indices: |