import torch import gradio as gr from PIL import Image import qrcode from pathlib import Path import requests import io import os import spaces import random from diffusers import ( StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, DDIMScheduler, DPMSolverMultistepScheduler, DEISMultistepScheduler, HeunDiscreteScheduler, EulerDiscreteScheduler, ) MAX_SEED = 2**32 - 1 # QR Code generation setup qrcode_generator = qrcode.QRCode( version=1, error_correction=qrcode.ERROR_CORRECT_H, box_size=16, border=4, ) # SDXL and ControlNet setup device = torch.device("cuda" if torch.cuda.is_available() else "cpu") vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) controlnet = ControlNetModel.from_pretrained( "AGCobra/1", torch_dtype=torch.float16 ).to(device) pipe = StableDiffusionXLControlNetPipeline.from_pretrained( "stabilityai/stable-diffusion-xl-base-1.0", vae=vae, controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True, variant="fp16", ).to(device) # Sampler setup 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 @spaces.GPU() def inference( qr_code_content: str, prompt: str, negative_prompt: str, guidance_scale: float = 7.5, controlnet_conditioning_scale: float = 1.1, strength: float = 0.9, seed: int = -1, sampler: str = "DPM++ Karras SDE", ): if prompt is None or prompt == "": raise gr.Error("Prompt is required") if qr_code_content == "": raise gr.Error("QR Code Content is required") pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config) if seed == -1: seed = random.randint(0, MAX_SEED) # Use a sub-seed for additional randomness subseed = random.randint(0, MAX_SEED) generator = torch.Generator(device=device).manual_seed(seed + subseed) print("Generating QR Code from content") qr = qrcode.QRCode( version=1, error_correction=qrcode.constants.ERROR_CORRECT_H, box_size=16, border=4, ) qr.add_data(qr_code_content) qr.make(fit=True) qrcode_image = qr.make_image(fill_color="black", back_color="white") qrcode_image = resize_for_condition_image(qrcode_image, 1024) init_image = qrcode_image out = pipe( prompt=prompt, negative_prompt=negative_prompt, image=init_image, control_image=qrcode_image, controlnet_conditioning_scale=float(controlnet_conditioning_scale), guidance_scale=float(guidance_scale), generator=generator, strength=float(strength), num_inference_steps=30, ) return out.images[0] with gr.Blocks() as demo: with gr.Row(): with gr.Column(): qr_code_content = gr.Textbox( label="QR Code Content", info="QR Code Content or URL", value="", ) 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", ) with gr.Accordion( label="Advanced Parameters", open=True, ): controlnet_conditioning_scale = gr.Slider( minimum=0.0, maximum=2.0, step=0.01, value=1.1, label="Controlnet Conditioning Scale", ) strength = gr.Slider( minimum=0.0, maximum=1.0, step=0.01, value=0.9, label="Strength" ) guidance_scale = gr.Slider( minimum=0.0, maximum=50.0, step=0.25, value=7.5, label="Guidance Scale", ) sampler = gr.Dropdown(choices=list(SAMPLER_MAP.keys()), value="DPM++ Karras SDE", label="Sampler") seed = gr.Slider( minimum=-1, maximum=MAX_SEED, step=1, value=-1, label="Seed", randomize=True, ) with gr.Row(): run_btn = gr.Button("Run") with gr.Column(): result_image = gr.Image(label="Result Image") run_btn.click( inference, inputs=[ qr_code_content, prompt, negative_prompt, guidance_scale, controlnet_conditioning_scale, strength, seed, sampler, ], outputs=[result_image], ) demo.queue(max_size=20).launch(share=bool(os.environ.get("SHARE", False)))