File size: 850 Bytes
bfe2873 13639c4 bfe2873 afa593f bfe2873 dd3b112 baaad9e f17ad12 dd3b112 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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} |