Method encode

evaluation/vqlm_demo/vqvae_muse.py:552–559

(self, pixel_values, return_loss=False)

        )

    def encode(self, pixel_values, return_loss=False):
        hidden_states = self.encoder(pixel_values)
        hidden_states = self.quant_conv(hidden_states)
        quantized_states, codebook_indices, codebook_loss = self.quantize(hidden_states, return_loss)
        output = (quantized_states, codebook_indices)
        if return_loss:
            output = output + (codebook_loss,)
        return output

    def decode(self, quantized_states):
        hidden_states = self.post_quant_conv(quantized_states)