from diffusers import AutoPipelineForImage2Image
import torch
from typing import Dict, Any
from PIL import Image
from io import BytesIO
import base64


class EndpointHandler():

    def __init__(self, path="."):
        if torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"
        self._pipe = AutoPipelineForImage2Image.from_pretrained(path, torch_dtype=torch.float16).to(device)

    def __call__(self, data: Dict[str, Any]) -> list[Dict[str, Any]]:
        inputs = data.pop("inputs", data)

        params = {"prompt": inputs.get("prompt", ""),
                  "image": Image.open(BytesIO(base64.b64decode(inputs['image']))),
                  "strength": float(inputs.get("strength", 0.3)),
                  "guidance_scale": float(inputs.get("guidance_scale", 10)),
                  "height": 768,
                  "width": 768}

        img: Image = self._pipe(**params).images[0]

        stream = BytesIO()
        img.save(stream, format="jpeg")
        res = {"status": 200,
               "image": base64.b64encode(stream.getvalue()).decode("utf8")
               }
        return res

if __name__ == "__main__":
    h = EndpointHandler()
    v = h({})
    print(v)