| from typing import Dict, List, Any |
| import torch |
| from base64 import b64decode |
| from huggingface_hub import model_info |
| from diffusers import AutoencoderKL |
| from diffusers.image_processor import VaeImageProcessor |
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| kwargs = {"torch_dtype": torch.float16} |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" |
| self.dtype = kwargs["torch_dtype"] |
| self.vae = AutoencoderKL.from_pretrained(path, **kwargs).to(self.device).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) |
|
|
| |
| |
| 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] |
|
|
