import base64
from io import BytesIO
from typing import Dict, Any

import torch
from PIL import Image
from diffusers import StableDiffusionPipeline


# helper decoder
def decode_base64_image(image_string):
    base64_image = base64.b64decode(image_string)
    buffer = BytesIO(base64_image)
    return Image.open(buffer)


class EndpointHandler:
    def __init__(self, path=""):
        self.pipe = StableDiffusionPipeline.from_pretrained("/repository/stable-diffusion-v1-5",
            torch_dtype=torch.float16, revision="fp16")
        self.pipe = self.pipe.to("cuda")

    def __call__(self, data: Any) -> Dict[str, str]:
        """
        Return predict value.
        :param data: A dictionary contains `inputs` and optional `image` field.
        :return: A dictionary with `image` field contains image in base64.
        """
        prompts = data.pop("inputs", None)
        encoded_image = data.pop("image", None)
        init_image = None
        if encoded_image:
            init_image = decode_base64_image(encoded_image)
            init_image.thumbnail((768, 768))

        image = self.pipe(prompts, init_image=init_image).images[0]
        buffered = BytesIO()
        image.save(buffered, format="png")
        img_str = base64.b64encode(buffered.getvalue())

        return {"image": img_str.decode()}