Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+images/fix-fp16.png filter=lfs diff=lfs merge=lfs -text
+images/fix-fp32.png filter=lfs diff=lfs merge=lfs -text
+images/orig-fp32.png filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,32 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.18.0.dev0",
+  "_name_or_path": ".",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 512,
+  "scaling_factor": 0.13025,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "force_upcast": false
+}

diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b909373b28f2137098b0fd9dbc6f97f8410854f31f84ddc9fa04b077b0ace2c
+size 334643238

handler.py

@@ -0,0 +1,85 @@
+from typing import Dict, List, Any
+import torch
+from base64 import b64decode
+from diffusers import AutoencoderKL
+from diffusers.image_processor import VaeImageProcessor
+
+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_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]]]:
+        """
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        """
+        tensor = data["inputs"]
+        tensor = b64decode(tensor.encode("utf-8"))
+        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.")
+
+        DTYPE_MAP = {
+            "float16": torch.float16,
+            "float32": torch.float32,
+            "bfloat16": torch.bfloat16,
+        }
+
+        shape = parameters.get("shape")
+        dtype = DTYPE_MAP.get(parameters.get("dtype"))
+        tensor = torch.frombuffer(bytearray(tensor), dtype=dtype).reshape(shape)
+
+        needs_upcasting = (
+            self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+        )
+        if needs_upcasting:
+            self.vae = self.vae.to(torch.float32)
+            tensor = tensor.to(self.device, torch.float32)
+        else:
+            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)
+            )
+            tensor = (
+                tensor * latents_std / self.vae.config.scaling_factor + latents_mean
+            )
+        else:
+            tensor = tensor / self.vae.config.scaling_factor
+
+        with torch.no_grad():
+            image = self.vae.decode(tensor, return_dict=False)[0]
+
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float16)
+
+        image = self.image_processor.postprocess(image, output_type="pil")
+
+        return image[0]

requirements.txt

@@ -0,0 +1,2 @@
+huggingface_hub
+diffusers