Method get_code

models/modeling_magvitv2.py:423–427

(self, pixel_values)

        return output

    def get_code(self, pixel_values):
        hidden_states = self.encoder(pixel_values)
        codebook_indices = self.quantize.get_indices(self.quantize(hidden_states)['z']).reshape(pixel_values.shape[0], -1)

        return codebook_indices

    def decode_code(self, codebook_indices, shape=None):
        z_q = self.quantize.get_codebook_entry(codebook_indices, shape=shape)