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