|
|
|
|
|
from PIL import Image |
|
from diffusers import ( |
|
StableDiffusionControlNetImg2ImgPipeline, |
|
ControlNetModel, |
|
DDIMScheduler, |
|
) |
|
from diffusers.utils import load_image |
|
import torch |
|
import openai |
|
from io import BytesIO |
|
import base64 |
|
import qrcode |
|
|
|
|
|
class EndpointHandler: |
|
def __init__( |
|
self, |
|
controlnet_path="DionTimmer/controlnet_qrcode-control_v11p_sd21", |
|
pipeline_path="stabilityai/stable-diffusion-2-1", |
|
): |
|
self.controlnet = ControlNetModel.from_pretrained( |
|
controlnet_path, torch_dtype=torch.float16 |
|
) |
|
|
|
self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( |
|
pipeline_path, |
|
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() |
|
|
|
@staticmethod |
|
def resize_for_condition_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 |
|
|
|
def __call__( |
|
self, |
|
prompt, |
|
negative_prompt, |
|
qrcode_data, |
|
guidance_scale, |
|
controlnet_conditioning_scale, |
|
strength, |
|
generator_seed, |
|
width, |
|
height, |
|
num_inference_steps, |
|
): |
|
openai.api_key = "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" |
|
|
|
qr = qrcode.QRCode( |
|
version=1, |
|
error_correction=qrcode.constants.ERROR_CORRECT_H, |
|
box_size=10, |
|
border=4, |
|
) |
|
qr.add_data(qrcode_data) |
|
qr.make(fit=True) |
|
img = qr.make_image(fill_color="black", back_color="white") |
|
|
|
|
|
basewidth = 768 |
|
wpercent = basewidth / float(img.size[0]) |
|
hsize = int((float(img.size[1]) * float(wpercent))) |
|
qrcode_image = img.resize((basewidth, hsize), Image.LANCZOS) |
|
|
|
response = openai.Image.create(prompt=prompt, n=1, size="1024x1024") |
|
image_url = response.data[0].url |
|
init_image = load_image(image_url) |
|
|
|
control_image = qrcode_image |
|
init_image = self.resize_for_condition_image(init_image, 768) |
|
|
|
generator = torch.manual_seed(generator_seed) |
|
|
|
image = self.pipe( |
|
prompt=prompt, |
|
negative_prompt=negative_prompt, |
|
image=init_image, |
|
control_image=control_image, |
|
width=width, |
|
height=height, |
|
guidance_scale=guidance_scale, |
|
controlnet_conditioning_scale=controlnet_conditioning_scale, |
|
generator=generator, |
|
strength=strength, |
|
num_inference_steps=num_inference_steps, |
|
) |
|
|
|
pil_image = image.images[0] |
|
buffered = BytesIO() |
|
pil_image.save(buffered, format="PNG") |
|
image_base64 = base64.b64encode(buffered.getvalue()).decode() |
|
|
|
return image_base64 |
|
|
