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]