diffusers/FLUX.1-vae

1from typing import cast, Union2 3import PIL.Image4import torch5 6from diffusers import AutoencoderKL7from diffusers.image_processor import VaeImageProcessor8 9 10class EndpointHandler:11    def __init__(self, path=""):12        self.device = "cuda"13        self.dtype = torch.bfloat1614        self.vae = cast(AutoencoderKL, AutoencoderKL.from_pretrained(path, torch_dtype=self.dtype).to(self.device, self.dtype).eval())15 16        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)17        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)18 19    @staticmethod20    def _unpack_latents(latents, height, width, vae_scale_factor):21        batch_size, num_patches, channels = latents.shape22 23        # VAE applies 8x compression on images but we must also account for packing which requires24        # latent height and width to be divisible by 2.25        height = 2 * (int(height) // (vae_scale_factor * 2))26        width = 2 * (int(width) // (vae_scale_factor * 2))27 28        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)29        latents = latents.permute(0, 3, 1, 4, 2, 5)30 31        latents = latents.reshape(batch_size, channels // (2 * 2), height, width)32 33        return latents34 35    @torch.no_grad()36    def __call__(self, data) -> Union[torch.Tensor, PIL.Image.Image]:37        """38        Args:39            data (:obj:):40                includes the input data and the parameters for the inference.41        """42        tensor = cast(torch.Tensor, data["inputs"])43        parameters = cast(dict, data.get("parameters", {}))44        if tensor.ndim == 3 and ("height" not in parameters or "width" not in parameters):45            raise ValueError("Expected `height` and `width` in parameters.")46        height = cast(int, parameters.get("height", 0))47        width = cast(int, parameters.get("width", 0))48        do_scaling = cast(bool, parameters.get("do_scaling", True))49        output_type = cast(str, parameters.get("output_type", "pil"))50        partial_postprocess = cast(bool, parameters.get("partial_postprocess", False))51        if partial_postprocess and output_type != "pt":52            output_type = "pt"53 54        tensor = tensor.to(self.device, self.dtype)55        if tensor.ndim == 3:56            tensor = self._unpack_latents(tensor, height, width, self.vae_scale_factor)57 58        if do_scaling:59            tensor = (60                tensor / self.vae.config.scaling_factor61            ) + self.vae.config.shift_factor62 63        with torch.no_grad():64            image = cast(torch.Tensor, self.vae.decode(tensor, return_dict=False)[0])65 66        if partial_postprocess:67            image = (image * 0.5 + 0.5).clamp(0, 1)68            image = image.permute(0, 2, 3, 1).contiguous().float()69            image = (image * 255).round().to(torch.uint8)70        elif output_type == "pil":71            image = cast(PIL.Image.Image, self.image_processor.postprocess(image, output_type="pil")[0])72 73        return image