rg321/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 Image11 12device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')13 14from diffusers import (15    StableDiffusionPipeline,16    StableDiffusionControlNetImg2ImgPipeline,17    ControlNetModel,18    DDIMScheduler,19    DPMSolverMultistepScheduler,20    DEISMultistepScheduler,21    HeunDiscreteScheduler,22    EulerDiscreteScheduler,23)24 25qrcode_generator = qrcode.QRCode(26    version=1,27    error_correction=qrcode.ERROR_CORRECT_H,28    box_size=10,29    border=4,30)31 32controlnet = ControlNetModel.from_pretrained(33    "DionTimmer/controlnet_qrcode-control_v1p_sd15", torch_dtype=torch.float1634)35 36pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(37    "runwayml/stable-diffusion-v1-5",38    controlnet=controlnet,39    safety_checker=None,40    torch_dtype=torch.float16,41).to(device)42pipe.enable_xformers_memory_efficient_attention()43 44 45def resize_for_condition_image(input_image: Image.Image, resolution: int):46    input_image = input_image.convert("RGB")47    W, H = input_image.size48    k = float(resolution) / min(H, W)49    H *= k50    W *= k51    H = int(round(H / 64.0)) * 6452    W = int(round(W / 64.0)) * 6453    img = input_image.resize((W, H), resample=Image.LANCZOS)54    return img55 56 57SAMPLER_MAP = {58    "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),59    "DPM++ Karras": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True),60    "Heun": lambda config: HeunDiscreteScheduler.from_config(config),61    "Euler": lambda config: EulerDiscreteScheduler.from_config(config),62    "DDIM": lambda config: DDIMScheduler.from_config(config),63    "DEIS": lambda config: DEISMultistepScheduler.from_config(config),64}65 66 67def inference(68    qr_code_content: str,69    prompt: str,70    negative_prompt: str,71    guidance_scale: float = 10.0,72    controlnet_conditioning_scale: float = 2.0,73    strength: float = 0.8,74    seed: int = -1,75    init_image: Image.Image | None = None,76    qrcode_image: Image.Image | None = None,77    use_qr_code_as_init_image = True,78    sampler = "DPM++ Karras SDE",79):80    if prompt is None or prompt == "":81        raise gr.Error("Prompt is required")82 83    if qrcode_image is None and qr_code_content == "":84        raise gr.Error("QR Code Image or QR Code Content is required")85 86    pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)87 88    generator = torch.manual_seed(seed) if seed != -1 else torch.Generator()89 90    if qr_code_content != "" or qrcode_image.size == (1, 1):91        print("Generating QR Code from content")92        qr = qrcode.QRCode(93            version=1,94            error_correction=qrcode.constants.ERROR_CORRECT_H,95            box_size=10,96            border=4,97        )98        qr.add_data(qr_code_content)99        qr.make(fit=True)100 101        qrcode_image = qr.make_image(fill_color="black", back_color="white")102        qrcode_image = resize_for_condition_image(qrcode_image, 768)103    else:104        print("Using QR Code Image")105        qrcode_image = resize_for_condition_image(qrcode_image, 768)106 107    # hack due to gradio examples108    init_image = qrcode_image109 110    out = pipe(111        prompt=prompt,112        negative_prompt=negative_prompt,113        image=qrcode_image,114        control_image=qrcode_image,  # type: ignore115        width=768,  # type: ignore116        height=768,  # type: ignore117        guidance_scale=float(guidance_scale),118        controlnet_conditioning_scale=float(controlnet_conditioning_scale),  # type: ignore119        generator=generator,120        strength=float(strength),121        num_inference_steps=40,122    )123    return out.images[0]  # type: ignore124 125 126with gr.Blocks() as blocks:127    gr.Markdown(128        """129# QR Code AI Art Generator130 131## 💡 How to generate beautiful QR codes132 133We use the QR code image as the initial image **and** the control image, which allows you to generate 134QR Codes that blend in **very naturally** with your provided prompt.135The 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.136Use a high strength value between 0.8 and 0.95 and choose a conditioning scale between 0.6 and 2.0.137This 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 138looks 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.139 140model: https://huggingface.co/DionTimmer/controlnet_qrcode-control_v1p_sd15141 142<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">143<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>144                """145    )146 147    with gr.Row():148        with gr.Column():149            qr_code_content = gr.Textbox(150                label="QR Code Content",151                info="QR Code Content or URL",152                value="",153            )154            with gr.Accordion(label="QR Code Image (Optional)", open=False):155                qr_code_image = gr.Image(156                    label="QR Code Image (Optional). Leave blank to automatically generate QR code",157                    type="pil",158                )159 160            prompt = gr.Textbox(161                label="Prompt",162                info="Prompt that guides the generation towards",163            )164            negative_prompt = gr.Textbox(165                label="Negative Prompt",166                value="ugly, disfigured, low quality, blurry, nsfw",167            )168            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")169 170            with gr.Accordion(label="Init Images (Optional)", open=False, visible=False) as init_image_acc:171                init_image = gr.Image(label="Init Image (Optional). Leave blank to generate image with SD 2.1", type="pil")172 173 174            with gr.Accordion(175                label="Params: The generated QR Code functionality is largely influenced by the parameters detailed below",176                open=True,177            ):178                controlnet_conditioning_scale = gr.Slider(179                    minimum=0.0,180                    maximum=5.0,181                    step=0.01,182                    value=1.1,183                    label="Controlnet Conditioning Scale",184                )185                strength = gr.Slider(186                    minimum=0.0, maximum=1.0, step=0.01, value=0.9, label="Strength"187                )188                guidance_scale = gr.Slider(189                    minimum=0.0,190                    maximum=50.0,191                    step=0.25,192                    value=7.5,193                    label="Guidance Scale",194                )195                sampler = gr.Dropdown(choices=list(SAMPLER_MAP.keys()), value="DPM++ Karras SDE")196                seed = gr.Slider(197                    minimum=-1,198                    maximum=9999999999,199                    step=1,200                    value=2313123,201                    label="Seed",202                    randomize=True,203                )204            with gr.Row():205                run_btn = gr.Button("Run")206        with gr.Column():207            result_image = gr.Image(label="Result Image")208    run_btn.click(209        inference,210        inputs=[211            qr_code_content,212            prompt,213            negative_prompt,214            guidance_scale,215            controlnet_conditioning_scale,216            strength,217            seed,218            init_image,219            qr_code_image,220            use_qr_code_as_init_image,221            sampler,222        ],223        outputs=[result_image],224    )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                None,237                None,238                True,239                "DPM++ Karras SDE",240            ],241            [242                "https://huggingface.co/",243                "Bright sunshine coming through the cracks of a wet, cave wall of big rocks",244                "ugly, disfigured, low quality, blurry, nsfw",245                7.5,246                1.11,247                0.9,248                2523992465,249                None,250                None,251                True,252                "DPM++ Karras SDE",253            ],254            [255                "https://huggingface.co/",256                "Sky view of highly aesthetic, ancient greek thermal baths  in beautiful nature",257                "ugly, disfigured, low quality, blurry, nsfw",258                7.5,259                1.5,260                0.9,261                2523992465,262                None,263                None,264                True,265                "DPM++ Karras SDE",266            ],267        ],268        fn=inference,269        inputs=[270            qr_code_content,271            prompt,272            negative_prompt,273            guidance_scale,274            controlnet_conditioning_scale,275            strength,276            seed,277            init_image,278            qr_code_image,279            use_qr_code_as_init_image,280            sampler,281        ],282        outputs=[result_image],283        cache_examples=True,284    )285 286blocks.queue(concurrency_count=1, max_size=20)287blocks.launch(share=bool(os.environ.get("SHARE", False)))288