File size: 2,109 Bytes
314064f 9d02d34 314064f 9d02d34 314064f 1898dc9 314064f 9d02d34 1898dc9 9d02d34 314064f |
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 |
from typing import Dict, List, Any
import torch
from diffusers import StableDiffusionXLImg2ImgPipeline, DiffusionPipeline, AutoencoderKL
from PIL import Image
import base64
from io import BytesIO
from diffusers.image_processor import VaeImageProcessor
# 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=""):
self.smooth_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16
)
self.smooth_pipe.to("cuda")
def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
"""
:param data: A dictionary contains `inputs` and optional `image` field.
:return: A dictionary with `image` field contains image in base64.
"""
encoded_image = data.pop("image", None)
prompt = data.pop("prompt", "")
if encoded_image is not None:
image = self.decode_base64_image(encoded_image)
image_processor = VaeImageProcessor();
latents = image_processor.preprocess(image)
latents = latents.to(device="cuda")
vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
subfolder="vae", use_safetensors=True,
).to("cuda")
with torch.no_grad():
latents_dist = vae.encode(latents).latent_dist.sample() * vae.config.scaling_factor
self.smooth_pipe.enable_xformers_memory_efficient_attention()
out = self.smooth_pipe(prompt, image=latents_dist).images[0]
return out
# helper to decode input image
def decode_base64_image(self, image_string):
base64_image = base64.b64decode(image_string)
buffer = BytesIO(base64_image)
image = Image.open(buffer)
return image
|