| 103 | |
| 104 | |
| 105 | class VQGANModel(nn.Module): |
| 106 | config: VQGANConfig |
| 107 | |
| 108 | def setup(self): |
| 109 | self.encoder = Encoder(self.config) |
| 110 | self.decoder = Decoder(self.config) |
| 111 | self.quantize = VectorQuantizer( |
| 112 | self.config.num_embeddings, self.config.quantized_embed_dim |
| 113 | ) |
| 114 | self.quant_conv = nn.Conv(self.config.quantized_embed_dim, [1, 1]) |
| 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: |
| 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] |
| 145 | recon = self.decode(encoding) |
| 146 | return recon |
| 147 | |
| 148 | |
| 149 | class Encoder(nn.Module): |