hjconstas/qrcode-diffusion
0
1from typing import Optional2 3import gradio as gr4import qrcode5import torch6from diffusers import (7 ControlNetModel,8 EulerAncestralDiscreteScheduler,9 StableDiffusionControlNetPipeline,10)11from gradio.components import Image, Radio, Slider, Textbox, Number12from PIL import Image as PilImage13from typing_extensions import Literal14 15 16def main():17 device = (18 'cuda' if torch.cuda.is_available() 19 else 'mps' if torch.backends.mps.is_available() 20 else 'cpu'21 )22 23 controlnet_tile = ControlNetModel.from_pretrained(24 "lllyasviel/control_v11f1e_sd15_tile",25 torch_dtype=torch.float16 if device == "cuda" else torch.float32,26 use_safetensors=False,27 cache_dir="./cache"28 ).to(device)29 30 controlnet_brightness = ControlNetModel.from_pretrained(31 "ioclab/control_v1p_sd15_brightness",32 torch_dtype=torch.float16 if device == "cuda" else torch.float32,33 use_safetensors=True,34 cache_dir="./cache"35 ).to(device)36 37 def make_pipe(hf_repo: str, device: str) -> StableDiffusionControlNetPipeline:38 pipe = StableDiffusionControlNetPipeline.from_pretrained(39 hf_repo,40 controlnet=[controlnet_tile, controlnet_brightness],41 torch_dtype=torch.float16 if device == "cuda" else torch.float32,42 cache_dir="./cache",43 )44 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)45 # pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)46 return pipe.to(device)47 48 pipes = {49 "DreamShaper": make_pipe("Lykon/DreamShaper", device),50 # "DreamShaper": make_pipe("Lykon/DreamShaper", "cpu"),51 # "Realistic Vision V1.4": make_pipe("SG161222/Realistic_Vision_V1.4", "cpu"),52 # "OpenJourney": make_pipe("prompthero/openjourney", "cpu"),53 # "Anything V3": make_pipe("Linaqruf/anything-v3.0", "cpu"),54 }55 56 def move_pipe(hf_repo: str):57 for pipe_name, pipe in pipes.items():58 if pipe_name != hf_repo:59 pipe.to("cpu")60 return pipes[hf_repo].to(device)61 62 def predict(63 model: Literal[64 "DreamShaper",65 # "Realistic Vision V1.4",66 # "OpenJourney",67 # "Anything V3"68 ],69 qrcode_data: str,70 prompt: str,71 negative_prompt: Optional[str] = None,72 num_inference_steps: int = 100,73 guidance_scale: int = 9,74 controlnet_conditioning_tile: float = 0.25,75 controlnet_conditioning_brightness: float = 0.45,76 seed: int = 1331,77 ) -> PilImage:78 generator = torch.Generator(device).manual_seed(seed)79 if model == "DreamShaper":80 pipe = pipes["DreamShaper"]81 # pipe = move_pipe("DreamShaper Vision V1.4")82 # elif model == "Realistic Vision V1.4":83 # pipe = move_pipe("Realistic Vision V1.4")84 # elif model == "OpenJourney":85 # pipe = move_pipe("OpenJourney")86 # elif model == "Anything V3":87 # pipe = move_pipe("Anything V3")88 89 90 qr = qrcode.QRCode(91 error_correction=qrcode.constants.ERROR_CORRECT_H,92 box_size=11,93 border=9,94 )95 qr.add_data(qrcode_data)96 qr.make(fit=True)97 qrcode_image = qr.make_image(98 fill_color="black",99 back_color="white"100 ).convert("RGB")101 qrcode_image = qrcode_image.resize((512, 512), PilImage.LANCZOS)102 103 image = pipe(104 prompt,105 [qrcode_image, qrcode_image],106 num_inference_steps=num_inference_steps,107 generator=generator,108 negative_prompt=negative_prompt,109 guidance_scale=guidance_scale,110 controlnet_conditioning_scale=[111 controlnet_conditioning_tile,112 controlnet_conditioning_brightness113 ]114 ).images[0]115 116 return image117 118 119 ui = gr.Interface(120 fn=predict,121 inputs=[122 Radio(123 value="DreamShaper",124 label="Model",125 choices=[126 "DreamShaper",127 # "Realistic Vision V1.4",128 # "OpenJourney",129 # "Anything V3"130 ],131 ),132 Textbox(133 value="https://twitter.com/JulienBlanchon",134 label="QR Code Data",135 ),136 Textbox(137 value="Japanese ramen with chopsticks, egg and steam, ultra detailed 8k",138 label="Prompt",139 ),140 Textbox(141 value="logo, watermark, signature, text, BadDream, UnrealisticDream",142 label="Negative Prompt",143 optional=True144 ),145 Slider(146 value=100,147 label="Number of Inference Steps",148 minimum=10,149 maximum=400,150 step=1,151 ),152 Slider(153 value=9,154 label="Guidance Scale",155 minimum=1,156 maximum=20,157 step=1,158 ),159 Slider(160 value=0.25,161 label="Controlnet Conditioning Tile",162 minimum=0.0,163 maximum=1.0,164 step=0.05,165 166 ),167 Slider(168 value=0.45,169 label="Controlnet Conditioning Brightness",170 minimum=0.0,171 maximum=1.0,172 step=0.05,173 ),174 Number(175 value=1,176 label="Seed",177 precision=0,178 ),179 180 ],181 outputs=Image(182 label="Generated Image",183 type="pil",184 ),185 examples=[186 [187 "DreamShaper",188 "https://twitter.com/JulienBlanchon",189 "rock, mountain",190 "",191 100,192 9,193 0.25,194 0.45,195 1,196 ],197 [198 "DreamShaper",199 "https://twitter.com/JulienBlanchon",200 "Japanese ramen with chopsticks, egg and steam, ultra detailed 8k",201 "logo, watermark, signature, text, BadDream, UnrealisticDream",202 100,203 9,204 0.25,205 0.45,206 1,207 ],208 # [209 # "Anything V3",210 # "https://twitter.com/JulienBlanchon",211 # "Japanese ramen with chopsticks, egg and steam, ultra detailed 8k",212 # "logo, watermark, signature, text, BadDream, UnrealisticDream",213 # 100,214 # 9,215 # 0.25,216 # 0.60,217 # 1,218 # ],219 [220 "DreamShaper",221 "https://twitter.com/JulienBlanchon",222 "processor, chipset, electricity, black and white board",223 "logo, watermark, signature, text, BadDream, UnrealisticDream",224 300,225 9,226 0.50,227 0.30,228 1,229 ],230 ],231 cache_examples=True,232 title="Stable Diffusion QR Code Controlnet",233 description="Generate QR Code with Stable Diffusion and Controlnet",234 allow_flagging="never",235 max_batch_size=1,236 )237 238 ui.queue(concurrency_count=10).launch()239 240if __name__ == "__main__":241 main()