File size: 1,325 Bytes
059fad7 b1e2b1f 059fad7 b1e2b1f e5e883d 059fad7 b1e2b1f 059fad7 b1e2b1f 059fad7 b1e2b1f 059fad7 b1e2b1f 059fad7 b1e2b1f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import base64
from io import BytesIO
from typing import Dict, Any
import torch
from PIL import Image
from diffusers import StableDiffusionPipeline
# helper decoder
def decode_base64_image(image_string):
base64_image = base64.b64decode(image_string)
buffer = BytesIO(base64_image)
return Image.open(buffer)
class EndpointHandler:
def __init__(self, path=""):
self.pipe = StableDiffusionPipeline.from_pretrained("/repository/stable-diffusion-v1-5",
torch_dtype=torch.float16, revision="fp16")
self.pipe = self.pipe.to("cuda")
def __call__(self, data: Any) -> Dict[str, str]:
"""
Return predict value.
:param data: A dictionary contains `inputs` and optional `image` field.
:return: A dictionary with `image` field contains image in base64.
"""
prompts = data.pop("inputs", None)
encoded_image = data.pop("image", None)
init_image = None
if encoded_image:
init_image = decode_base64_image(encoded_image)
init_image.thumbnail((768, 768))
image = self.pipe(prompts, init_image=init_image).images[0]
buffered = BytesIO()
image.save(buffered, format="png")
img_str = base64.b64encode(buffered.getvalue())
return {"image": img_str.decode()}
|