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,86 @@
+from typing import Any, Dict, List, Tuple
+
+import base64
+import torch
+
+from diffusers import AutoencoderKL
+from safetensors.torch import _tobytes
+
+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: str, shape: List[int], dtype: str) -> torch.Tensor:
+    tensor = base64.b64decode(tensor_data.encode("utf-8"))
+    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()
+
+    @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 = 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 = 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]
+
+        tensor_data, shape, dtype = prepare_tensor(image)
+
+        return {"tensor": tensor_data, "shape": shape, "dtype": dtype}

requirements.txt

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