import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
import base64
from io import BytesIO
from PIL import Image
import json

class EndpointHandler():
    def __init__(self, path=""):
        model_id = "timbrooks/instruct-pix2pix"
        self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None)
        self.pipe.to("cuda")
        self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
    
    def __call__(self, data):

        info=data['inputs']
        image=info.pop("image",data)
        prompt=info.pop("prompt",data)
        image=base64.b64decode(image)
        raw_images = Image.open(BytesIO(image)).convert('RGB')
        
        images = self.pipe(prompt, image=raw_images, num_inference_steps=25, image_guidance_scale=1).images
        return json.dumps({'image':images[0]})
    


if __name__=="__main__":
    my_handler=EndpointHandler(path='.')