QRDiffuser.py

import torch
import gradio as gr
from PIL import Image
import qrcode
import os

from diffusers import (
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    UniPCMultistepScheduler,
    DEISMultistepScheduler,
    HeunDiscreteScheduler,
    EulerDiscreteScheduler,
    EulerAncestralDiscreteScheduler,
)

controlnet = ControlNetModel.from_pretrained(
    "monster-labs/control_v1p_sd15_qrcode_monster",
    torch_dtype=torch.float16,
)

pipe = StableDiffusionControlNetPipeline.from_pretrained(
    #"runwayml/stable-diffusion-v1-5",
    "SG161222/Realistic_Vision_V3.0_VAE",
    controlnet=controlnet,
    safety_checker=None,
    torch_dtype=torch.float16,
).to("cuda")
#pipe.enable_xformers_memory_efficient_attention()
pipe.enable_attention_slicing(1)
pipe.enable_model_cpu_offload()
#pipe.enable_vae_tiling()
pipe.enable_vae_slicing()
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)

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 a": lambda config: EulerAncestralDiscreteScheduler.from_config(config),
    "Euler": lambda config: EulerDiscreteScheduler.from_config(config),
    "DDIM": lambda config: DDIMScheduler.from_config(config),
    "DEIS": lambda config: DEISMultistepScheduler.from_config(config),
}

boxsize=16
def create_code(content: str, errorCorrection: str):
    match errorCorrection:
        case "L 7%":
            errCorr = qrcode.constants.ERROR_CORRECT_L
        case "M 15%":
            errCorr = qrcode.constants.ERROR_CORRECT_M
        case "Q 25%":
            errCorr = qrcode.constants.ERROR_CORRECT_Q
        case "H 30%":
            errCorr = qrcode.constants.ERROR_CORRECT_H
    
    qr = qrcode.QRCode(
        version=1,
        error_correction=errCorr,
        box_size=boxsize,
        border=0,
    )
    qr.add_data(content)
    qr.make(fit=True)
    img = qr.make_image(fill_color="black", back_color="white")

    # find smallest image size multiple of 256 that can fit qr
    offset_min = 8 * boxsize
    w, h = img.size
    w = (w + 255 + offset_min) // 256 * 256
    h = (h + 255 + offset_min) // 256 * 256
    if w > 1024:
        raise gr.Error("QR code is too large, please use a shorter content")
    bg = Image.new('L', (w, h), 128)

    # align on 16px grid
    coords = ((w - img.size[0]) // 2 // boxsize * boxsize,
              (h - img.size[1]) // 2 // boxsize * boxsize)
    bg.paste(img, coords)
    return bg


def inference(
    qr_code_content: str,
    errorCorrection: str,
    prompt: str,
    negative_prompt: str,
    inferenceSteps: float,
    guidance_scale: float = 10.0,
    controlnet_conditioning_scale: float = 2.0,
    seed: int = -1,
    sampler="Euler a",
):
    if prompt is None or prompt == "":
        raise gr.Error("Prompt is required")

    if qr_code_content is None or qr_code_content == "":
        raise gr.Error("QR Code Content is required")

    pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)

    generator = torch.manual_seed(seed) if seed != -1 else torch.Generator()

    print("Generating QR Code from content")
    qrcode_image = create_code(qr_code_content, errorCorrection)

    # hack due to gradio examples
    init_image = qrcode_image
    init_image.save("c:\\temp\\qr.jpg")

    out = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=qrcode_image,
        width=qrcode_image.width,
        height=qrcode_image.height,
        guidance_scale=float(guidance_scale),
        controlnet_conditioning_scale=float(controlnet_conditioning_scale),
        generator=generator,
        num_inference_steps=inferenceSteps,
    )
    return out.images[0]


css = """
#result_image {
    display: flex;
    place-content: center;
    align-items: center;
}
#result_image > img {
    height: auto;
    max-width: 100%;
    width: revert;
}
"""

with gr.Blocks(css=css) as blocks:

    with gr.Row():
        with gr.Column():
            qr_code_content = gr.Textbox(
                label="QR Code Content or URL",
                info="The text you want to encode into the QR code",
                value="",
            )
            errorCorrection = gr.Dropdown(
                label="QR Code Error Correction Level",
                choices=["L 7%", "M 15%", "Q 25%", "H 30%"],
                value="H 30%"
            )

            prompt = gr.Textbox(
                label="Prompt",
                info="Prompt that guides the generation towards",
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                value="ugly, disfigured, low quality, blurry, nsfw",
                info="Prompt that guides the generation away from",
            )
            inferenceSteps = gr.Slider(
                minimum=10.0,
                maximum=60.0,
                step=1,
                value=20,
                label="Inference Steps",
                info="More steps give better image but longer runtime",
            )

            with gr.Accordion(
                label="Params: The generated QR Code functionality is largely influenced by the parameters detailed below",
                open=True,
            ):
                controlnet_conditioning_scale = gr.Slider(
                    minimum=0.5,
                    maximum=2.5,
                    step=0.01,
                    value=1.5,
                    label="Controlnet Conditioning Scale",
                    info="""Controls the readability/creativity of the QR code.
                    High values: The generated QR code will be more readable.
                    Low values: The generated QR code will be more creative.
                    """
                )
                guidance_scale = gr.Slider(
                    minimum=0.0,
                    maximum=25.0,
                    step=0.25,
                    value=7,
                    label="Guidance Scale",
                    info="Controls the amount of guidance the text prompt guides the image generation"
                )
                sampler = gr.Dropdown(choices=list(
                    SAMPLER_MAP.keys()), value="Euler a", label="Sampler")
                seed = gr.Number(
                    minimum=-1,
                    maximum=9999999999,
                    value=-1,
                    label="Seed",
                    info="Seed for the random number generator. Set to -1 for a random seed"
                )
            with gr.Row():
                run_btn = gr.Button("Run")
        with gr.Column():
            result_image = gr.Image(label="Result Image", elem_id="result_image")
    run_btn.click(
        inference,
        inputs=[
            qr_code_content,
            errorCorrection,
            prompt,
            negative_prompt,
            inferenceSteps,
            guidance_scale,
            controlnet_conditioning_scale,
            seed,
            sampler,
        ],
        outputs=[result_image],
    )

blocks.queue(concurrency_count=1, max_size=20, api_open=False)
blocks.launch(share=bool(os.environ.get("SHARE", True)), show_api=False)