Upload model

image_processing_vqmodel.py

@@ -0,0 +1,63 @@
+import numpy as np
+import torch
+from PIL import Image
+from transformers.image_processing_utils import BaseImageProcessor
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class VQModelImageProcessor(BaseImageProcessor):  # type: ignore
+    def __init__(
+        self,
+        size: int = 256,
+        convert_rgb: bool = False,
+        resample: Image.Resampling = Image.Resampling.LANCZOS,
+    ) -> None:
+        self.size = size
+        self.convert_rgb = convert_rgb
+        self.resample = resample
+
+    def __call__(self, image: Image.Image) -> dict:
+        return self.preprocess(image)
+
+    def preprocess(self, image: Image.Image) -> dict:
+        width, height = image.size
+        size = (self.size, self.size)
+        image = image.resize(size, resample=self.resample)
+        image = image.convert("RGBA")
+
+        if self.convert_rgb:
+            # Paste RGBA image on white background
+            image_new = Image.new("RGB", image.size, (255, 255, 255))
+            image_new.paste(image, mask=image.split()[3])
+            image = image_new
+
+        return {
+            "image": self.to_tensor(image),
+            "width": width,
+            "height": height,
+        }
+
+    def to_tensor(self, image: Image.Image) -> torch.Tensor:
+        x = np.array(image) / 127.5 - 1.0
+        x = x.transpose(2, 0, 1).astype(np.float32)
+        return torch.as_tensor(x)
+
+    def postprocess(
+        self,
+        x: torch.Tensor,
+        width: int | None = None,
+        height: int | None = None,
+    ) -> Image.Image:
+        x_np = x.numpy().transpose(1, 2, 0)
+        x_np = (x_np + 1.0) * 127.5
+        x_np = np.clip(x_np, 0, 255).astype(np.uint8)
+        image = Image.fromarray(x_np)
+
+        # Resize image
+        width = width or self.size
+        height = height or self.size
+        image = image.resize((width, height), resample=self.resample)
+
+        return image

preprocessor_config.json

@@ -0,0 +1,9 @@
+{
+  "auto_map": {
+    "AutoImageProcessor": "image_processing_vqmodel.VQModelImageProcessor"
+  },
+  "convert_rgb": true,
+  "image_processor_type": "VQModelImageProcessor",
+  "resample": 1,
+  "size": 256
+}