bpmatt/QR-code-AI-art-generator

1import torch2import gradio as gr3from PIL import Image4import qrcode5from pathlib import Path6from multiprocessing import cpu_count7import requests8import io9import os10from PIL import Image11import spaces12 13from diffusers import (14    StableDiffusionControlNetImg2ImgPipeline,15    ControlNetModel,16    DDIMScheduler,17    DPMSolverMultistepScheduler,18    DEISMultistepScheduler,19    HeunDiscreteScheduler,20    EulerDiscreteScheduler,21)22 23qrcode_generator = qrcode.QRCode(24    version=1,25    error_correction=qrcode.ERROR_CORRECT_H,26    box_size=10,27    border=4,28)29 30# Load models to CPU at startup (ZeroGPU moves to GPU inside @spaces.GPU())31controlnet = ControlNetModel.from_pretrained(32    "DionTimmer/controlnet_qrcode-control_v1p_sd15", torch_dtype=torch.float1633)34 35pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(36    "stable-diffusion-v1-5/stable-diffusion-v1-5",37    controlnet=controlnet,38    safety_checker=None,39    torch_dtype=torch.float16,40)41 42 43def resize_for_condition_image(input_image: Image.Image, resolution: int):44    input_image = input_image.convert("RGB")45    W, H = input_image.size46    k = float(resolution) / min(H, W)47    H *= k48    W *= k49    H = int(round(H / 64.0)) * 6450    W = int(round(W / 64.0)) * 6451    img = input_image.resize((W, H), resample=Image.LANCZOS)52    return img53 54 55SAMPLER_MAP = {56    "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),57    "DPM++ Karras": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True),58    "Heun": lambda config: HeunDiscreteScheduler.from_config(config),59    "Euler": lambda config: EulerDiscreteScheduler.from_config(config),60    "DDIM": lambda config: DDIMScheduler.from_config(config),61    "DEIS": lambda config: DEISMultistepScheduler.from_config(config),62}63 64@spaces.GPU()65def inference(66    qr_code_content: str,67    prompt: str,68    negative_prompt: str,69    guidance_scale: float = 10.0,70    controlnet_conditioning_scale: float = 2.0,71    strength: float = 0.8,72    seed: int = -1,73    init_image: Image.Image | None = None,74    qrcode_image: Image.Image | None = None,75    use_qr_code_as_init_image = True,76    sampler = "DPM++ Karras SDE",77):78    if prompt is None or prompt == "":79        raise gr.Error("Prompt is required")80 81    if qrcode_image is None and qr_code_content == "":82        raise gr.Error("QR Code Image or QR Code Content is required")83 84    pipe.to("cuda")85    pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)86 87    generator = torch.manual_seed(seed) if seed != -1 else torch.Generator()88 89    if qr_code_content != "" or qrcode_image.size == (1, 1):90        print("Generating QR Code from content")91        qr = qrcode.QRCode(92            version=1,93            error_correction=qrcode.constants.ERROR_CORRECT_H,94            box_size=10,95            border=4,96        )97        qr.add_data(qr_code_content)98        qr.make(fit=True)99 100        qrcode_image = qr.make_image(fill_color="black", back_color="white")101        qrcode_image = resize_for_condition_image(qrcode_image, 768)102    else:103        print("Using QR Code Image")104        qrcode_image = resize_for_condition_image(qrcode_image, 768)105 106    # hack due to gradio examples107    init_image = qrcode_image108 109    out = pipe(110        prompt=prompt,111        negative_prompt=negative_prompt,112        image=qrcode_image,113        control_image=qrcode_image,  # type: ignore114        width=768,  # type: ignore115        height=768,  # type: ignore116        guidance_scale=float(guidance_scale),117        controlnet_conditioning_scale=float(controlnet_conditioning_scale),  # type: ignore118        generator=generator,119        strength=float(strength),120        num_inference_steps=40,121    )122    return out.images[0]  # type: ignore123 124 125with gr.Blocks() as blocks:126    gr.Markdown(127        """128# QR Code AI Art Generator129 130## 💡 How to generate beautiful QR codes131 132We use the QR code image as the initial image **and** the control image, which allows you to generate 133QR Codes that blend in **very naturally** with your provided prompt.134The strength parameter defines how much noise is added to your QR code and the noisy QR code is then guided towards both your prompt and the QR code image via Controlnet.135Use a high strength value between 0.8 and 0.95 and choose a conditioning scale between 0.6 and 2.0.136This mode arguably achieves the asthetically most appealing QR code images, but also requires more tuning of the controlnet conditioning scale and the strength value. If the generated image 137looks way to much like the original QR code, make sure to gently increase the *strength* value and reduce the *conditioning* scale. Also check out the examples below.138 139model: https://huggingface.co/DionTimmer/controlnet_qrcode-control_v1p_sd15140 141<a href="https://huggingface.co/spaces/huggingface-projects/QR-code-AI-art-generator?duplicate=true" style="display: inline-block;margin-top: .5em;margin-right: .25em;" target="_blank">142<img style="margin-bottom: 0em;display: inline;margin-top: -.25em;" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a> for no queue on your own hardware.</p>143                """144    )145 146    with gr.Row():147        with gr.Column():148            qr_code_content = gr.Textbox(149                label="QR Code Content",150                info="QR Code Content or URL",151                value="",152            )153            with gr.Accordion(label="QR Code Image (Optional)", open=False):154                qr_code_image = gr.Image(155                    label="QR Code Image (Optional). Leave blank to automatically generate QR code",156                    type="pil",157                )158 159            prompt = gr.Textbox(160                label="Prompt",161                info="Prompt that guides the generation towards",162            )163            negative_prompt = gr.Textbox(164                label="Negative Prompt",165                value="ugly, disfigured, low quality, blurry, nsfw",166            )167            use_qr_code_as_init_image = gr.Checkbox(label="Use QR code as init image", value=True, interactive=False, info="Whether init image should be QR code. Unclick to pass init image or generate init image with Stable Diffusion 2.1")168 169            with gr.Accordion(label="Init Images (Optional)", open=False, visible=False) as init_image_acc:170                init_image = gr.Image(label="Init Image (Optional). Leave blank to generate image with SD 2.1", type="pil")171 172 173            with gr.Accordion(174                label="Params: The generated QR Code functionality is largely influenced by the parameters detailed below",175                open=True,176            ):177                controlnet_conditioning_scale = gr.Slider(178                    minimum=0.0,179                    maximum=5.0,180                    step=0.01,181                    value=1.1,182                    label="Controlnet Conditioning Scale",183                )184                strength = gr.Slider(185                    minimum=0.0, maximum=1.0, step=0.01, value=0.9, label="Strength"186                )187                guidance_scale = gr.Slider(188                    minimum=0.0,189                    maximum=50.0,190                    step=0.25,191                    value=7.5,192                    label="Guidance Scale",193                )194                sampler = gr.Dropdown(choices=list(SAMPLER_MAP.keys()), value="DPM++ Karras SDE", label="Sampler")195                seed = gr.Slider(196                    minimum=-1,197                    maximum=9999999999,198                    step=1,199                    value=2313123,200                    label="Seed",201                    randomize=True,202                )203            with gr.Row():204                run_btn = gr.Button("Run")205        with gr.Column():206            result_image = gr.Image(label="Result Image")207    run_btn.click(208        inference,209        inputs=[210            qr_code_content,211            prompt,212            negative_prompt,213            guidance_scale,214            controlnet_conditioning_scale,215            strength,216            seed,217            init_image,218            qr_code_image,219            use_qr_code_as_init_image,220            sampler,221        ],222        outputs=[result_image],223        concurrency_limit=1224    )225 226    gr.Examples(227        examples=[228            [229                "https://huggingface.co/",230                "A sky view of a colorful lakes and rivers flowing through the desert",231                "ugly, disfigured, low quality, blurry, nsfw",232                7.5,233                1.3,234                0.9,235                5392011833,236                "DPM++ Karras SDE",237            ],238            [239                "https://huggingface.co/",240                "Bright sunshine coming through the cracks of a wet, cave wall of big rocks",241                "ugly, disfigured, low quality, blurry, nsfw",242                7.5,243                1.11,244                0.9,245                2523992465,246                "DPM++ Karras SDE",247            ],248            [249                "https://huggingface.co/",250                "Sky view of highly aesthetic, ancient greek thermal baths  in beautiful nature",251                "ugly, disfigured, low quality, blurry, nsfw",252                7.5,253                1.5,254                0.9,255                2523992465,256                "DPM++ Karras SDE",257            ],258        ],259        inputs=[260            qr_code_content,261            prompt,262            negative_prompt,263            guidance_scale,264            controlnet_conditioning_scale,265            strength,266            seed,267            sampler,268        ],269    )270 271blocks.queue(max_size=20)272blocks.launch(show_api=False)273