|
from typing import Any |
|
import torch, base64 |
|
from PIL import Image |
|
from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel, DDIMScheduler |
|
from diffusers.utils import load_image |
|
from io import BytesIO |
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
self.controlnet = ControlNetModel.from_pretrained("DionTimmer/controlnet_qrcode-control_v11p_sd21", torch_dtype=torch.float16) |
|
self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16) |
|
|
|
self.pipe.enable_xformers_memory_efficient_attention() |
|
self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config) |
|
self.pipe.enable_model_cpu_offload() |
|
|
|
def __call__(self, data): |
|
""" |
|
data args: |
|
inputs (:obj: `str`) |
|
date (:obj: `str`) |
|
Return: |
|
A :obj:`list` | `dict`: will be serialized and returned |
|
""" |
|
|
|
inputs = data.pop("inputs", data) |
|
params = data.pop("parameters", data) |
|
prompt = params.get("prompt") |
|
negative_prompt = params.get("negative_prompt") |
|
|
|
def resize_image(input_image: Image, resolution: int): |
|
input_image = input_image.convert("RGB") |
|
W, H = input_image.size |
|
k = float(resolution) / min(H, W) |
|
H *= k |
|
W *= k |
|
H = int(round(H / 64.0)) * 64 |
|
W = int(round(W / 64.0)) * 64 |
|
img = input_image.resize((W, H), resample=Image.LANCZOS) |
|
return img |
|
|
|
orriginal_qr_code_image = load_image(inputs) |
|
img_path = 'https://images.squarespace-cdn.com/content/v1/59413d96e6f2e1c6837c7ecd/1536503659130-R84NUPOY4QPQTEGCTSAI/15fe1e62172035.5a87280d713e4.png' |
|
|
|
|
|
init_image = load_image(img_path) |
|
condition_image = resize_image(orriginal_qr_code_image, 768) |
|
init_image = resize_image(init_image, 768) |
|
generator = torch.manual_seed(123121231) |
|
image = self.pipe(prompt=prompt or "a bilboard in NYC with a qrcode", |
|
negative_prompt=negative_prompt or "ugly, disfigured, low quality, blurry, nsfw, worst quality, illustration, drawing", |
|
image=init_image, |
|
control_image=condition_image, |
|
width=768, |
|
height=768, |
|
guidance_scale=20, |
|
controlnet_conditioning_scale=2.5, |
|
generator=generator, |
|
strength=0.9, |
|
num_inference_steps=150, |
|
) |
|
|
|
image = image.images[0] |
|
buffered = BytesIO() |
|
image.save(buffered, format="JPEG") |
|
img_str = base64.b64encode(buffered.getvalue()) |
|
|
|
return {"image": img_str.decode()} |
