File size: 1,460 Bytes
a7eb3c4 eddf5bd a7eb3c4 b1bfbf4 a7eb3c4 b1bfbf4 a7eb3c4 322fe05 |
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 48 |
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", use_auth_token="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")
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("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]
# encode image as base 64
buffered = BytesIO()
image.save(buffered, format="png")
# post process the prediction
return {"image": buffered.getvalue()}
|