| 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} |