ControlNet Blend the image to QRCode

app.py

@@ -5,11 +5,11 @@ from PIL import Image
 import qrcode
 from qrcode.constants import ERROR_CORRECT_H
 
-# -------- Stable Diffusion setup --------
+# ========= Stable Diffusion (prompt-only) =========
 device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 dtype = torch.float16 if device != "cpu" else torch.float32
 
-pipe = StableDiffusionPipeline.from_pretrained(
+sd_pipe = StableDiffusionPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     torch_dtype=dtype
 ).to(device)
@@ -17,19 +17,17 @@ pipe = StableDiffusionPipeline.from_pretrained(
 def sd_generate(prompt, steps, guidance, seed):
     gen = torch.Generator(device=device).manual_seed(int(seed)) if int(seed) != 0 else None
     def run():
-        return pipe(prompt, num_inference_steps=int(steps), guidance_scale=float(guidance), generator=gen).images[0]
-    # autocast only where supported
+        return sd_pipe(prompt, num_inference_steps=int(steps), guidance_scale=float(guidance), generator=gen).images[0]
     if device in ("cuda", "mps"):
         with torch.autocast(device):
             return run()
-    else:
-        return run()
+    return run()
 
-# -------- QR code helper --------
+# ========= QR Maker =========
 def make_qr(url: str = "http://www.mybirdfire.com", size: int = 512, border: int = 2) -> Image.Image:
     qr = qrcode.QRCode(
         version=None,
-        error_correction=ERROR_CORRECT_H,  # high EC to survive stylization later
+        error_correction=ERROR_CORRECT_H,
         box_size=10,
         border=border
     )
@@ -38,9 +36,51 @@ def make_qr(url: str = "http://www.mybirdfire.com", size: int = 512, border: int
     img = qr.make_image(fill_color="black", back_color="white").convert("RGB")
     return img.resize((size, size), resample=Image.NEAREST)
 
-# -------- UI --------
+# ========= ControlNet Stylizer (prompt + QR) =========
+# lazy-load to speed initial startup
+_cn_loaded = {"pipe": None}
+def _load_controlnet():
+    if _cn_loaded["pipe"] is None:
+        from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+        from controlnet_aux import CannyDetector
+        controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=dtype)
+        pipe = StableDiffusionControlNetPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            controlnet=controlnet,
+            torch_dtype=dtype,
+            safety_checker=None
+        ).to(device)
+        pipe.enable_attention_slicing()
+        pipe.enable_vae_slicing()
+        _cn_loaded["pipe"] = pipe
+        _cn_loaded["canny"] = CannyDetector()
+    return _cn_loaded["pipe"], _cn_loaded["canny"]
+
+def stylize_qr(prompt: str, qr_image: Image.Image, steps: int, guidance: float, seed: int, canny_low: int, canny_high: int):
+    if qr_image is None:
+        raise gr.Error("Please provide a QR image (use the QR Maker tab first or upload one).")
+    pipe, canny = _load_controlnet()
+    # ensure 512x512 for speed/quality
+    qr_img = qr_image.convert("RGB").resize((512, 512), Image.NEAREST)
+    edges = canny(qr_img, low_threshold=int(canny_low), high_threshold=int(canny_high))
+
+    gen = torch.Generator(device=device).manual_seed(int(seed)) if int(seed) != 0 else None
+    def run():
+        return pipe(
+            prompt=str(prompt),
+            image=edges,
+            num_inference_steps=int(steps),
+            guidance_scale=float(guidance),
+            generator=gen
+        ).images[0]
+    if device in ("cuda", "mps"):
+        with torch.autocast(device):
+            return run()
+    return run()
+
+# ========= UI =========
 with gr.Blocks() as demo:
-    gr.Markdown("## Stable Diffusion + QR (step by step)")
+    gr.Markdown("## Stable Diffusion + QR Code + ControlNet (step by step)")
 
     with gr.Tab("Stable Diffusion (prompt → image)"):
         prompt = gr.Textbox(label="Prompt", value="A fantasy castle at sunset")
@@ -57,5 +97,16 @@ with gr.Blocks() as demo:
         out_qr = gr.Image(label="QR Code", type="pil")
         gr.Button("Generate QR").click(make_qr, [url, size, quiet], out_qr)
 
+    with gr.Tab("QR Stylizer (ControlNet)"):
+        s_prompt = gr.Textbox(label="Style Prompt", value="floral papercut style, high contrast, preserve sharp squares")
+        s_steps  = gr.Slider(10, 50, value=30, label="Steps", step=1)
+        s_cfg    = gr.Slider(1, 12, value=7.5, label="Guidance Scale", step=0.1)
+        s_seed   = gr.Number(value=0, label="Seed (0 = random)", precision=0)
+        canny_l  = gr.Slider(0, 255, value=100, step=1, label="Canny low")
+        canny_h  = gr.Slider(0, 255, value=200, step=1, label="Canny high")
+        qr_in    = gr.Image(label="QR Input (use output from QR Maker or upload)", type="pil")
+        out_styl = gr.Image(label="Stylized QR")
+        gr.Button("Stylize").click(stylize_qr, [s_prompt, qr_in, s_steps, s_cfg, s_seed, canny_l, canny_h], out_styl)
+
 if __name__ == "__main__":
     demo.launch()

requirements.txt

@@ -4,4 +4,5 @@ transformers
 accelerate
 safetensors
 gradio
-qrcode[pil]
+qrcode[pil]
+controlnet-aux