add app

app.py

@@ -0,0 +1,261 @@
+import torch
+import gradio as gr
+from multiprocessing import cpu_count
+import os
+from generate_code import create_code, backgrounds, correction_map
+
+from diffusers import (
+    StableDiffusionControlNetPipeline,
+    ControlNetModel,
+    EulerAncestralDiscreteScheduler,
+)
+
+main_generator = torch.Generator()
+
+# MONSTER_V2 = "/home/ubuntu/training/diffusers/examples/controlnet/out_model_2023-06-18_17-27-06"
+# LANDMARKS = "/home/ubuntu/training/diffusers/examples/controlnet/out_model_2023-06-19_23-43-50/"
+MONSTER_V2 = "monster-labs/V2"
+LANDMARKS = "monster-labs/V2"
+
+controlnet = [
+    ControlNetModel.from_pretrained(MONSTER_V2, torch_dtype=torch.float16, subfolder="step1"),
+    ControlNetModel.from_pretrained(LANDMARKS, torch_dtype=torch.float16, subfolder="step2"),
+]
+
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    controlnet=controlnet,
+    safety_checker=None,
+    torch_dtype=torch.float16,
+).to("cuda")
+pipe.enable_xformers_memory_efficient_attention()
+pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
+def inference_map(
+    qr_code_content: str,
+    prompt: str,
+    negative_prompt: str,
+    guidance_scale: float,
+    controlnet_conditioning_scale_0: float,
+    seed: int,
+    controlnet_start_0: float,
+    controlnet_start_1: float,
+    controlnet_end_0: float,
+    controlnet_end_1: float,
+    background: str,
+    error_correction: str,
+    margin: int,
+    module_size: int,
+    width: int,
+    height: int,
+):
+    return inference(
+        qr_code_content,
+        prompt,
+        negative_prompt,
+        guidance_scale,
+        (controlnet_conditioning_scale_0, 1),
+        seed,
+        (controlnet_start_0, controlnet_start_1),
+        (controlnet_end_0, controlnet_end_1),
+        background,
+        error_correction,
+        margin,
+        module_size,
+        width,
+        height,
+    )
+
+def inference(
+    qr_code_content: str,
+    prompt: str,
+    negative_prompt: str,
+    guidance_scale: float = 10.0,
+    controlnet_conditioning_scale: tuple[float, float] = (1.0, 1.0),
+    seed: int = -1,
+    controlnet_start: tuple[float, float] = (0.2, 0.0),
+    controlnet_end: tuple[float, float] = (0.95, 1.0),
+    background: str = "gray",
+    error_correction: str = "H",
+    margin: int = 1,
+    module_size: int = 16,
+    width: int = None,
+    height: int = None,
+):
+    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")
+
+    if background not in backgrounds:
+        raise gr.Error("Invalid background")
+    
+    if error_correction not in correction_map:
+        raise gr.Error("Invalid error correction")
+
+
+    generator = torch.manual_seed(seed) if seed != -1 else main_generator
+
+    # print("Generating QR Code from content")
+    qrcode_images = create_code(qr_code_content, module_size, margin, background, error_correction, False, 1, True)
+
+    out = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=list(qrcode_images),
+        width=qrcode_images[0].width,
+        height=qrcode_images[0].height,
+        guidance_scale=float(guidance_scale),
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        # controlnet_start=controlnet_start,
+        # controlnet_end=controlnet_end,
+        controlnet_guidance=[(controlnet_start[0], controlnet_end[0]), (controlnet_start[1], controlnet_end[1])],
+        generator=generator,
+        num_inference_steps=40,
+    )
+    return out.images[0]
+
+with gr.Blocks() as blocks:
+    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, nsfw",
+            )
+
+            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",
+                )
+                guidance_scale = gr.Slider(
+                    minimum=0.0,
+                    maximum=25.0,
+                    step=0.25,
+                    value=7,
+                    label="Guidance Scale",
+                )
+                seed = gr.Number(
+                    minimum=-1,
+                    maximum=9999999999,
+                    step=1,
+                    value=2313123,
+                    label="Seed",
+                    randomize=True,
+                )
+                controlnet_start_0 = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.2,
+                    label="Controlnet Start 0",
+                )
+                controlnet_start_1 = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.0,
+                    label="Controlnet Start 1",
+                )
+                controlnet_end_0 = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.95,
+                    label="Controlnet End 0",
+                )
+                controlnet_end_1 = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=1.0,
+                    label="Controlnet End 1",
+                )
+                background = gr.Dropdown(
+                    label="Background",
+                    choices=backgrounds,
+                    value="gray",
+                )
+                error_correction = gr.Dropdown(
+                    label="Error Correction",
+                    choices=correction_map.keys(),
+                    value="H",
+                )
+                margin = gr.Slider(
+                    minimum=0,
+                    maximum=10,
+                    step=1,
+                    value=1,
+                    label="Margin",
+                )
+                module_size = gr.Slider(
+                    minimum=1,
+                    maximum=100,
+                    step=1,
+                    value=16,
+                    label="Module Size",
+                )
+                width = gr.Slider(
+                    minimum=512,
+                    maximum=1024,
+                    step=256,
+                    value=512,
+                    label="Width",
+                )
+                height = gr.Slider(
+                    minimum=512,
+                    maximum=1024,
+                    step=256,
+                    value=512,
+                    label="Height",
+                )
+            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_map,
+        inputs=[
+            qr_code_content,
+            prompt,
+            negative_prompt,
+            guidance_scale,
+            controlnet_conditioning_scale,
+            seed,
+            controlnet_start_0,
+            controlnet_start_1,
+            controlnet_end_0,
+            controlnet_end_1,
+            background,
+            error_correction,
+            margin,
+            module_size,
+            width,
+            height,
+        ],
+        outputs=[result_image],
+    )
+
+# login = os.environ.get("LOGIN", "admin")
+# password = os.environ.get("PASSWORD", "1234")
+
+blocks.queue(concurrency_count=1, max_size=40)
+blocks.launch(share=False)
+# blocks.launch(share=False, auth=(login, password))

generate_code.py

@@ -0,0 +1,103 @@
+import qrcode
+import numpy as np
+from PIL import Image
+import base64
+
+correction_map = {
+  'L': qrcode.constants.ERROR_CORRECT_L,
+  'M': qrcode.constants.ERROR_CORRECT_M,
+  'Q': qrcode.constants.ERROR_CORRECT_Q,
+  'H': qrcode.constants.ERROR_CORRECT_H,
+}
+
+backgrounds = ['white', 'black', 'gray', 'noise']
+
+bg_map = {
+    'white': 255,
+    'black': 0,
+    'gray': 128,
+}
+
+def create_code(text, module_size, margin, background, error_correction, centered, submodule_prop, split=False):
+  qr = qrcode.QRCode(
+    version=1,
+    error_correction=correction_map[error_correction],
+    box_size=module_size,
+    border=margin if margin > 0 else 1,
+  )
+  qr.add_data(text)
+  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 * module_size
+  w, h = img.size
+
+  mask = None
+
+  if margin == 0:
+    corner_size = 9 * module_size
+    # make a mask that hides the margin when not in the corners
+    mask = np.ones((h, w)).astype(bool)
+    mask[corner_size:-corner_size, :] = False
+    mask[:, corner_size:-corner_size] = False
+    mask[-corner_size:, -corner_size:] = False
+    mask[module_size:-module_size, module_size:-module_size] = True
+
+  if submodule_prop != 1:
+    submodule_size = round(module_size * submodule_prop)
+    k = (module_size - submodule_size) // 2
+    # qr_array = np.array(img).reshape((h // module_size, w // module_size, module_size, module_size))
+    new_mask = np.zeros((h // module_size, module_size, w // module_size, module_size)).astype(bool)
+    new_mask[:, k:k+submodule_size, :, k:k+submodule_size] = True
+    mask = np.logical_and(mask, new_mask.reshape((h, w)))
+
+  w = (w + 255 + offset_min) // 256 * 256
+  h = (h + 255 + offset_min) // 256 * 256
+
+  # create new image with background chosen by user
+  if background == 'noise':
+    # noise = np.random.randint(0, 256, (h // module_size, w // module_size), dtype=np.uint8)
+    noise = np.random.normal(128, 64, (h // module_size, w // module_size)).astype(np.uint8)
+    # noise = (noise // 128) * 255
+    noise = np.round(noise / 128) * 128
+    # clamp values
+    noise = np.clip(noise, 0, 255)
+    
+    bg = Image.fromarray(noise.astype(np.uint8)).resize((w, h), Image.NEAREST)
+  else:
+    bg = Image.new('L', (w, h), bg_map[background])
+
+  # paste qr code in the center of the image
+  if centered:
+    coords = ((w - img.size[0]) // 2, (h - img.size[1]) // 2)
+  else:
+    # paste it aligned on module size, closest to center
+    coords = ((w - img.size[0]) // 2 // module_size * module_size, (h - img.size[1]) // 2 // module_size * module_size)
+
+  if mask is not None:
+    new_img = Image.new('L', img.size, bg_map['gray'] if background == 'noise' else bg_map[background])
+    new_img.paste(img, mask=Image.fromarray(mask))
+    bg.paste(new_img, coords)
+  else:
+    bg.paste(img, coords)
+
+  if split:
+    # isolate the 3 qr markers, paste them onto a new gray image
+    m_size = module_size * 9
+
+    new_bg = Image.new('L', (w, h), bg_map['gray'])
+
+    # create mask for the 3 markers from a numpy array
+    mask = np.zeros((h, w), dtype=bool)
+    mask[coords[1]:coords[1] + m_size, coords[0]:coords[0] + m_size] = True
+    mask[coords[1]:coords[1] + m_size, coords[0] + img.size[0] - m_size:coords[0] + img.size[0]] = True
+    mask[coords[1] + img.size[1] - m_size:coords[1] + img.size[1], coords[0]:coords[0] + m_size] = True
+
+    # paste the 3 markers on the new image
+    new_bg.paste(bg, mask=Image.fromarray(mask))
+
+    return bg, new_bg
+  else:
+    return bg,

requirements.txt

@@ -0,0 +1,9 @@
+diffusers
+transformers
+accelerate
+torch
+xformers
+gradio
+Pillow
+qrcode
+numpy