File size: 2,522 Bytes
6642c60 513b0b2 6642c60 53385b6 6642c60 6ad32f2 6642c60 513b0b2 35c6947 6642c60 9fc68a1 6642c60 9fc68a1 6642c60 a632781 513b0b2 c7e0853 513b0b2 6642c60 9fc68a1 6642c60 513b0b2 a9108ee c7e0853 9fc68a1 513b0b2 492876f c7e0853 9fc68a1 6642c60 9fc68a1 6642c60 9fc68a1 |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
from typing import Dict, List, Any
import torch
import os
import PIL
from PIL import Image
from torch import autocast
from diffusers import StableDiffusionPipeline,EulerDiscreteScheduler
import base64
from io import BytesIO
# set device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if device.type != 'cuda':
raise ValueError("need to run on GPU")
class EndpointHandler():
def __init__(self, path=""):
# load the optimized model
self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16,low_cpu_mem_usage=False)
self.pipe.scheduler = EulerDiscreteScheduler.from_config(self.pipe.scheduler.config)
self.pipe = self.pipe.to(device)
def __call__(self, data: Any) -> Dict[str, str]:
"""
Args:
data (Any): Includes the input data and the parameters for the inference.
Returns:
Dict[str, str]: Dictionary with the base64 encoded image.
"""
inputs = data.pop("inputs", data)
# positive_prompt = data.pop("positive_prompt", None)
negative_prompt = data.pop("negative_prompt", None)
height = data.pop("height", 512)
width = data.pop("width", 512)
inference_steps = data.pop("inference_steps", 25)
guidance_scale = data.pop("guidance_scale", 7.5)
# Run inference pipeline
with autocast(device.type):
if negative_prompt is None:
print(str(inputs), str(height), str(width), str(guidance_scale))
image = self.pipe(prompt=inputs, height=height, width=width, guidance_scale=float(guidance_scale),num_inference_steps=inference_steps)
image = image.images[0]
else:
print(str(inputs), str(height), str(negative_prompt), str(width), str(guidance_scale))
image = self.pipe(prompt=inputs, negative_prompt=negative_prompt, height=height, width=width, guidance_scale=float(guidance_scale),num_inference_steps=inference_steps)
image = image.images[0]
# Encode image as base64
buffered = BytesIO()
image.save(buffered, format="JPEG")
img_str = base64.b64encode(buffered.getvalue())
# Postprocess the prediction
return {"image": img_str.decode()}
def decode_base64_image(self, image_string):
base64_image = base64.b64decode(image_string)
buffer = BytesIO(base64_image)
image = Image.open(buffer)
return image
|