from typing import Dict, List, Any
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from PIL import Image
from io import BytesIO
import base64
import json

class EndpointHandler():
    def __init__(self, path=""):
        model_id = "stabilityai/stable-diffusion-2-1"
        self.pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
        self.pipe = self.pipe.to("cuda")
    
    def __call__(self, data):

        inputs=data['inputs']
        text=inputs.pop('text',data)
        image = self.pipe(text).images[0] 
        img_data = image.tobytes()
        img_base64 = base64.b64encode(img_data)
        return json.dumps({'image':img_base64})