File size: 1,398 Bytes
a7eb3c4 |
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 43 44 45 46 47 |
import base64
from io import BytesIO
from typing import Dict, List, Any
import torch
from PIL import Image
from diffusers import StableDiffusionPipeline
REPO_ID = "runwayml/stable-diffusion-v1-5"
# 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("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16,
revision="fp16")
self.pipe = self.pipe.to("cuda")
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
"""
Args:
data (:obj:):
includes the input data and the parameters for the inference.
Return:
A :obj:`dict`:. base64 encoded image
"""
prompts = data.pop("prompts", None)
encoded_image = data.pop("image", None)
init_image = decode_base64_image(encoded_image)
init_image.thumbnail((768, 768))
image = self.pipe(prompts, init_image=init_image).images[0]
# encode image as base 64
buffered = BytesIO()
image.save(buffered, format="png")
img_str = base64.b64encode(buffered.getvalue())
# post process the prediction
return {"image": img_str.decode()}
|