File size: 3,425 Bytes
b402ddd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
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}
model_data = model_info(path)
has_model_index = any(
file.rfilename == "model_index.json" for file in model_data.siblings
)
if has_model_index:
kwargs["subfolder"] = "vae"
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)
# 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]
|