| | |
| |
|
| | 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 |
| |
|
