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

class EndpointHandler():
    def __init__(self, path=""):
        model_id = "stabilityai/stable-diffusion-2-1"
        self.pipe = StableDiffusionPipeline.from_pretrained(model_id)
        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)
        img = self.pipe(text).images[0] 
        img.save("./1.png")
        with open('./1.png','rb') as img_file:
            encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
        return {'image':encoded_string}