Allow non-512x512 with chameleon tokenizer

.gitignore

@@ -0,0 +1 @@
+__pycache__

app.py

@@ -130,11 +130,13 @@ class ChameleonVQImageRoundtripPipeline(ImageRoundtripPipeline):
     def roundtrip_image(self, image, output_type="pil"):
         # image = self.tokenizer._vqgan_input_from(image).to(device)
         image = self.preprocess(image).to(device)
+        _, _, im_height, im_width = image.shape
         _, _, [_, _, latents] = self.tokenizer._vq_model.encode(image)
-        # emb_dim = self._vq_model.quantize.embedding.weight.shape[-1]
-        output = self.tokenizer.pil_from_img_toks(latents)
+        scale = self.vae_scale_factor
+        shape = (1, im_height // scale, im_width // scale)
+        output = self.tokenizer.pil_from_img_toks(latents, shape=shape)
         # we actually do want this to be a grid, sorry!
-        latents = latents.reshape(1, 32, 32)
+        latents = latents.reshape(*shape)
 
         return (
             output,

chameleon/image_tokenizer.py

@@ -115,10 +115,11 @@ class ImageTokenizer:
 
         return pil_image
 
-    def pil_from_img_toks(self, img_tensor: torch.Tensor) -> PIL.Image:
+    # darknoon: added shape parameter
+    def pil_from_img_toks(self, img_tensor: torch.Tensor, shape = (1, 32, 32,)) -> PIL.Image:
         emb_dim = self._vq_model.quantize.embedding.weight.shape[-1]
         codebook_entry = self._vq_model.quantize.get_codebook_entry(
-            img_tensor, (1, 32, 32, emb_dim)
+            img_tensor, (*shape, emb_dim)
         )
         pixels = self._vq_model.decode(codebook_entry)
         return self._pil_from_chw_tensor(pixels[0])