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