Update handler.py

handler.py

@@ -1,88 +1,72 @@
-from typing import Any, Dict, List, Tuple
+from typing import cast, Union
 
-import base64
+import PIL.Image
 import torch
 
 from diffusers import AutoencoderKL
-from safetensors.torch import _tobytes
+from diffusers.image_processor import VaeImageProcessor
 
-def prepare_tensor(tensor: torch.Tensor) -> Tuple[str, List[int], str]:
-    tensor_data = base64.b64encode(_tobytes(tensor, "inputs")).decode("utf-8")
-    shape = list(tensor.shape)
-    dtype = str(tensor.dtype).split(".")[-1]
-    return tensor_data, shape, dtype
-
-
-def unpack_tensor(tensor_data: bytes, shape: List[int], dtype: str) -> torch.Tensor:
-    tensor = tensor_data
-    DTYPE_MAP = {
-        "float16": torch.float16,
-        "float32": torch.float32,
-        "bfloat16": torch.bfloat16,
-    }
-    torch_dtype = DTYPE_MAP.get(dtype)
-    tensor = torch.frombuffer(bytearray(tensor), dtype=torch_dtype).reshape(shape)
-    return tensor
 
 class EndpointHandler:
     def __init__(self, path=""):
         self.device = "cuda"
         self.dtype = torch.float16
-        self.vae = AutoencoderKL.from_pretrained(path, torch_dtype=self.dtype).to(self.device, self.dtype).eval()
+        self.vae = cast(AutoencoderKL, AutoencoderKL.from_pretrained(path, torch_dtype=self.dtype).to(self.device, self.dtype).eval())
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
 
     @torch.no_grad()
-    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+    def __call__(self, data) -> Union[torch.Tensor, PIL.Image.Image]:
         """
         Args:
             data (:obj:):
                 includes the input data and the parameters for the inference.
         """
-        tensor_data = data["inputs"]
-        parameters = data.get("parameters", {})
-        if "shape" not in parameters:
-            raise ValueError("Expected `shape` in parameters.")
-        if "dtype" not in parameters:
-            raise ValueError("Expected `dtype` in parameters.")
-
-        shape = parameters.get("shape")
-        dtype = parameters.get("dtype")
-
-        tensor = unpack_tensor(tensor_data, shape, dtype)
+        tensor = cast(torch.Tensor, data["inputs"])
+        parameters = cast(dict, data.get("parameters", {}))
+        do_scaling = cast(bool, parameters.get("do_scaling", True))
+        output_type = cast(str, parameters.get("output_type", "pil"))
+        partial_postprocess = cast(bool, parameters.get("partial_postprocess", False))
+        if partial_postprocess and output_type != "pt":
+            output_type = "pt"
 
         tensor = tensor.to(self.device, self.dtype)
 
-        # unscale/denormalize the latents
-        # denormalize with the mean and std if available and not None
-        has_latents_mean = (
-            hasattr(self.vae.config, "latents_mean")
-            and self.vae.config.latents_mean is not None
-        )
-        has_latents_std = (
-            hasattr(self.vae.config, "latents_std")
-            and self.vae.config.latents_std is not None
-        )
-        if has_latents_mean and has_latents_std:
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, 4, 1, 1)
-                .to(tensor.device, tensor.dtype)
-            )
-            latents_std = (
-                torch.tensor(self.vae.config.latents_std)
-                .view(1, 4, 1, 1)
-                .to(tensor.device, tensor.dtype)
+        if do_scaling:
+            has_latents_mean = (
+                hasattr(self.vae.config, "latents_mean")
+                and self.vae.config.latents_mean is not None
             )
-            tensor = (
-                tensor * latents_std / self.vae.config.scaling_factor + latents_mean
+            has_latents_std = (
+                hasattr(self.vae.config, "latents_std")
+                and self.vae.config.latents_std is not None
             )
-        else:
-            tensor = tensor / self.vae.config.scaling_factor
+            if has_latents_mean and has_latents_std:
+                latents_mean = (
+                    torch.tensor(self.vae.config.latents_mean)
+                    .view(1, 4, 1, 1)
+                    .to(tensor.device, tensor.dtype)
+                )
+                latents_std = (
+                    torch.tensor(self.vae.config.latents_std)
+                    .view(1, 4, 1, 1)
+                    .to(tensor.device, tensor.dtype)
+                )
+                tensor = (
+                    tensor * latents_std / self.vae.config.scaling_factor + latents_mean
+                )
+            else:
+                tensor = tensor / self.vae.config.scaling_factor
 
         with torch.no_grad():
-            image: torch.Tensor = self.vae.decode(tensor, return_dict=False)[0]
+            image = cast(torch.Tensor, self.vae.decode(tensor, return_dict=False)[0])
 
-        image = (image * 0.5 + 0.5).clamp(0, 1)
-        image = image.permute(0, 2, 3, 1).contiguous().float()
-        image = (image * 255).round().to(torch.uint8)
+        if partial_postprocess:
+            image = (image * 0.5 + 0.5).clamp(0, 1)
+            image = image.permute(0, 2, 3, 1).contiguous().float()
+            image = (image * 255).round().to(torch.uint8)
+        elif output_type == "pil":
+            image = cast(PIL.Image.Image, self.image_processor.postprocess(image, output_type="pil")[0])
 
-        return _tobytes(image, "image")
+        return image