| import torch | |
| from PIL import Image | |
| import qrcode | |
| from pathlib import Path | |
| from multiprocessing import cpu_count | |
| import requests | |
| import io | |
| import os | |
| from PIL import Image | |
| from diffusers import ( | |
| StableDiffusionPipeline, | |
| StableDiffusionControlNetImg2ImgPipeline, | |
| ControlNetModel, | |
| DDIMScheduler, | |
| DPMSolverMultistepScheduler, | |
| DEISMultistepScheduler, | |
| HeunDiscreteScheduler, | |
| EulerDiscreteScheduler, | |
| ) | |
| SAMPLER_MAP = { | |
| "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"), | |
| "DPM++ Karras": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True), | |
| "Heun": lambda config: HeunDiscreteScheduler.from_config(config), | |
| "Euler": lambda config: EulerDiscreteScheduler.from_config(config), | |
| "DDIM": lambda config: DDIMScheduler.from_config(config), | |
| "DEIS": lambda config: DEISMultistepScheduler.from_config(config), | |
| } | |
| def resize_for_condition_image(input_image: 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 | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| qrcode_generator = qrcode.QRCode( | |
| version=1, | |
| error_correction=qrcode.ERROR_CORRECT_H, | |
| box_size=10, | |
| border=4, | |
| ) | |
| controlnet = ControlNetModel.from_pretrained( | |
| "DionTimmer/controlnet_qrcode-control_v1p_sd15", torch_dtype=torch.float16 | |
| ) | |
| pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( | |
| "runwayml/stable-diffusion-v1-5", | |
| controlnet=controlnet, | |
| safety_checker=None, | |
| torch_dtype=torch.float16, | |
| ).to("cuda") | |
| pipe.enable_xformers_memory_efficient_attention() | |
| def __call__inference(self, | |
| qr_code_content: str, | |
| prompt: str, | |
| negative_prompt: str, | |
| guidance_scale: float = 10.0, | |
| controlnet_conditioning_scale: float = 2.0, | |
| strength: float = 0.8, | |
| seed: int = -1, | |
| init_image: Image.Image | None = None, | |
| qrcode_image: Image.Image | None = None, | |
| use_qr_code_as_init_image = True, | |
| sampler = "DPM++ Karras SDE", | |
| ): | |
| if prompt is None or prompt == "": | |
| raise gr.Error("Prompt is required") | |
| if qrcode_image is None and qr_code_content == "": | |
| raise gr.Error("QR Code Image or QR Code Content is required") | |
| pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config) | |
| generator = torch.manual_seed(seed) if seed != -1 else torch.Generator() | |
| if qr_code_content != "" or qrcode_image.size == (1, 1): | |
| qr = qrcode.QRCode( | |
| version=1, | |
| error_correction=qrcode.constants.ERROR_CORRECT_H, | |
| box_size=10, | |
| border=4, | |
| ) | |
| qr.add_data(qr_code_content) | |
| qr.make(fit=True) | |
| qrcode_image = qr.make_image(fill_color="black", back_color="white") | |
| if init_image is None: | |
| if use_qr_code_as_init_image: | |
| init_image = qrcode_image.convert("RGB") | |
| resolution = controlnet.config.resolution | |
| qrcode_image = resize_for_condition_image(qrcode_image, resolution) | |
| if init_image is not None: | |
| init_image = init_image.convert("RGB") | |
| init_image = resize_for_condition_image(init_image, resolution) | |
| init_image = torch.nn.functional.interpolate( | |
| torch.nn.functional.to_tensor(init_image).unsqueeze(0), | |
| size=(resolution, resolution), | |
| mode="bilinear", | |
| align_corners=False, | |
| )[0].unsqueeze(0) | |
| else: | |
| init_image = torch.zeros( | |
| (1, 3, resolution, resolution), device=pipe.device | |
| ).to(dtype=torch.float32) | |
| with torch.no_grad(): | |
| result_image = pipe( | |
| qr_code_condition=qrcode_image, | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| init_image=init_image, | |
| strength=strength, | |
| guidance_scale=guidance_scale, | |
| controlnet_conditioning_scale=controlnet_conditioning_scale, | |
| disable_progress_bar=True, | |
| seed=generator, | |
| ) | |
| result_image = ( | |
| result_image.clamp(-1, 1).squeeze().permute(1, 2, 0).numpy() * 255 | |
| ) | |
| result_image = Image.fromarray(result_image.astype("uint8")) | |
| return result_image | |