|
|
import torch |
|
|
from diffusers import ( |
|
|
StableDiffusionControlNetImg2ImgPipeline, |
|
|
ControlNetModel, |
|
|
DDIMScheduler, |
|
|
) |
|
|
|
|
|
from PIL import Image |
|
|
|
|
|
|
|
|
class QRControlNet: |
|
|
|
|
|
def __init__(self, device: str = "cuda"): |
|
|
|
|
|
torch_dtype = torch.float16 if device == "cuda" else torch.float32 |
|
|
controlnet = ControlNetModel.from_pretrained( |
|
|
"DionTimmer/controlnet_qrcode-control_v1p_sd15", torch_dtype=torch_dtype |
|
|
) |
|
|
|
|
|
pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained( |
|
|
"runwayml/stable-diffusion-v1-5", |
|
|
controlnet=controlnet, |
|
|
safety_checker=None, |
|
|
torch_dtype=torch_dtype, |
|
|
).to(device) |
|
|
|
|
|
if device == "cuda": |
|
|
pipe.enable_xformers_memory_efficient_attention() |
|
|
pipe.enable_model_cpu_offload() |
|
|
|
|
|
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) |
|
|
self.pipe = pipe |
|
|
|
|
|
def generate_image( |
|
|
self, |
|
|
source_image: Image, |
|
|
control_image: Image, |
|
|
prompt: str, |
|
|
negative_prompt: str, |
|
|
img_size=512, |
|
|
num_inference_steps: int = 50, |
|
|
guidance_scale: int = 20, |
|
|
controlnet_conditioning_scale: float = 3.0, |
|
|
strength=0.9, |
|
|
seed=42, |
|
|
**kwargs |
|
|
): |
|
|
|
|
|
width = height = img_size |
|
|
generator = torch.manual_seed(seed) |
|
|
|
|
|
image = self.pipe( |
|
|
prompt=prompt, |
|
|
negative_prompt=negative_prompt, |
|
|
image=source_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, |
|
|
**kwargs |
|
|
) |
|
|
|
|
|
return image.images[0] |
|
|
|
