Edit QR stylizer

app.py

@@ -5,6 +5,7 @@ from PIL import Image
 import qrcode
 from qrcode.constants import ERROR_CORRECT_H
 
+
 # ========= device/dtype =========
 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
@@ -44,58 +45,64 @@ def make_qr(url: str = "http://www.mybirdfire.com", size: int = 512, border: int
     return img.resize((size, size), resample=Image.NEAREST)
 
 # ========= ControlNet Stylizer (SD1.5 + sd15-canny) =========
-_cn = {"pipe": None}
-def _load_controlnet_dual():
-    if _cn["pipe"] is None:
-        from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
-        from diffusers.schedulers.scheduling_euler_discrete import EulerDiscreteScheduler
-        from controlnet_aux import CannyDetector
-
-        canny = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=dtype)
-        tile  = ControlNetModel.from_pretrained("lllyasviel/control_v11f1e_sd15_tile", torch_dtype=dtype)
-
-        pipe = StableDiffusionControlNetPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5",
-            controlnet=[canny, tile],            # <— dual CN
-            torch_dtype=dtype,
-            safety_checker=None
+# --- SDXL dual ControlNet stylizer (canny + tile) ---
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from diffusers.schedulers.scheduling_euler_discrete import EulerDiscreteScheduler
+from controlnet_aux import CannyDetector
+
+SDXL_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"   # swap to your SDXL anime checkpoint if you have one
+CN_CANNY   = "diffusers/controlnet-canny-sdxl-1.0"
+CN_TILE    = "diffusers/controlnet-tile-sdxl-1.0"
+
+_sdxl = {"pipe": None}
+def _load_sdxl_dual():
+    if _sdxl["pipe"] is None:
+        cn1 = ControlNetModel.from_pretrained(CN_CANNY, torch_dtype=dtype)
+        cn2 = ControlNetModel.from_pretrained(CN_TILE,   torch_dtype=dtype)
+        pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+            SDXL_MODEL, controlnet=[cn1, cn2], torch_dtype=dtype, safety_checker=None
         ).to(device)
-
         pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
-        pipe.enable_attention_slicing()
         pipe.enable_vae_slicing()
+        _sdxl["pipe"] = pipe
+        _sdxl["canny"] = CannyDetector()
+    return _sdxl["pipe"], _sdxl["canny"]
 
-        _cn["pipe"]  = pipe
-        _cn["cannyx"] = CannyDetector()
-    return _cn["pipe"], _cn["cannyx"]
-
-NEG_DEFAULT = "lowres, low contrast, blurry, washed out, jpeg artifacts, worst quality"
-
+NEG = "lowres, low contrast, blurry, jpeg artifacts, worst quality, extra digits, bad anatomy"
 
-def stylize_qr(prompt: str, negative: str, steps: int, guidance: float, seed: int,
-               canny_low: int, canny_high: int):
-    qr_image = make_qr("http://www.mybirdfire.com", size=512, border=6)
-    pipe, canny = _load_controlnet_dual()
+def stylize_qr_sdxl(prompt: str, steps: int=28, guidance: float=7.0, seed: int=1470713301,
+                    canny_low: int=80, canny_high: int=160):
+    # 1) make a strong QR @1024
+    qr = make_qr("http://www.mybirdfire.com", size=1024, border=6)
+    # 2) edges for canny CN
+    pipe, canny = _load_sdxl_dual()
+    edges = canny(qr, low_threshold=int(canny_low), high_threshold=int(canny_high))
 
-    # edge map for canny CN; tile CN uses the raw QR
-    edges = canny(qr_image, 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
 
-    gen = torch.Generator(device=device).manual_seed(int(seed)) if int(seed) != 0 else None
-
-    # weights per controlnet: [canny_weight, tile_weight]
-    cn_scales = [1.20, 0.60]   # stronger structure, lighter texture
+    # Control weights + schedule (canny, tile)
+    cn_scales  = [1.1, 0.6]
+    cn_start   = [0.25, 0.00]
+    cn_end     = [0.95, 1.00]
 
     def run():
-        return pipe(
-            prompt=str(prompt),
-            negative_prompt=negative or NEG_DEFAULT,
-            image=[edges, qr_image],
-            controlnet_conditioning_scale=cn_scales,
+        img = pipe(
+            prompt=prompt,
+            negative_prompt=NEG,
+            image=[edges, qr],                               # canny first, tile second
+            controlnet_conditioning_scale=cn_scales,         # strengths
+            control_guidance_start=cn_start,                 # when they start
+            control_guidance_end=cn_end,                     # when they end
             num_inference_steps=int(steps),
             guidance_scale=float(guidance),
             generator=gen
         ).images[0]
 
+        # optional: re‑overlay razor‑sharp finder squares to boost scanning
+        # (uncomment if scans are borderline)
+        # img = overlay_finders(img, qr)
+        return img
+
     if device in ("cuda", "mps"):
         with torch.autocast(device):
             return run()
@@ -107,7 +114,7 @@ with gr.Blocks() as demo:
     gr.Markdown("## Stable Diffusion + QR Code + ControlNet")
 
     with gr.Tab("Stable Diffusion (prompt → image)"):
-        prompt = gr.Textbox(label="Prompt", value="A fantasy castle at sunset")
+        prompt = gr.Textbox(label="Prompt", value="Sky, Moon, Bird, Blue, In the dark, Goddess, Sweet, Beautiful, Fantasy, Art, Anime")
         negative = gr.Textbox(label="Negative Prompt", value="lowres, bad anatomy, worst quality")
         steps  = gr.Slider(10, 50, value=30, label="Steps", step=1)
         cfg    = gr.Slider(1, 12, value=7.0, label="Guidance Scale", step=0.1)
@@ -122,16 +129,27 @@ 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 sd15-canny + Euler)"):
-        s_prompt = gr.Textbox(label="Style Prompt", value="Sky, Moon, Bird, Blue, In the dark, Goddess, Sweet")
-        s_negative = gr.Textbox(label="Negative Prompt", value="lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, worst quality, low quality, very displeasing, (bad)")
-        s_steps  = gr.Slider(10, 50, value=28, label="Steps", step=1)
-        s_cfg    = gr.Slider(1, 12, value=7.0, label="CFG", step=0.1)
-        s_seed   = gr.Number(value=1470713301, label="Seed", precision=0)
-        canny_l  = gr.Slider(0, 255, value=80, step=1, label="Canny low")
-        canny_h  = gr.Slider(0, 255, value=160, step=1, label="Canny high")
-        out_styl = gr.Image(label="Stylized QR")
-        gr.Button("Stylize").click(stylize_qr, [s_prompt, s_negative, s_steps, s_cfg, s_seed, canny_l, canny_h], out_styl)
+    # with gr.Tab("QR Stylizer (ControlNet sd15-canny + Euler)"):
+    #     s_prompt = gr.Textbox(label="Style Prompt", value="Sky, Moon, Bird, Blue, In the dark, Goddess, Sweet, Beautiful, Fantasy, Art, Anime")
+    #     s_negative = gr.Textbox(label="Negative Prompt", value="lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, worst quality, low quality, very displeasing, (bad)")
+    #     s_steps  = gr.Slider(10, 50, value=28, label="Steps", step=1)
+    #     s_cfg    = gr.Slider(1, 12, value=7.0, label="CFG", step=0.1)
+    #     s_seed   = gr.Number(value=1470713301, label="Seed", precision=0)
+    #     canny_l  = gr.Slider(0, 255, value=80, step=1, label="Canny low")
+    #     canny_h  = gr.Slider(0, 255, value=160, step=1, label="Canny high")
+    #     out_styl = gr.Image(label="Stylized QR")
+    #     gr.Button("Stylize").click(stylize_qr, [s_prompt, s_negative, s_steps, s_cfg, s_seed, canny_l, canny_h], out_styl)
+
+    with gr.Tab("QR Stylizer (SDXL dual ControlNet)"):
+        p = gr.Textbox(label="Prompt", value="Sky, Moon, Bird, Blue, In the dark, Goddess, Sweet, Beautiful, Fantasy, Art, Anime")
+        st = gr.Slider(20, 40, 28, step=1, label="Steps")
+        cfg = gr.Slider(4.5, 9.0, 7.0, step=0.1, label="CFG")
+        sd = gr.Number(value=1470713301, label="Seed", precision=0)
+        cl = gr.Slider(0, 255, 80, step=1, label="Canny low")
+        ch = gr.Slider(0, 255, 160, step=1, label="Canny high")
+        out = gr.Image(label="Stylized QR (SDXL)")
+        gr.Button("Stylize").click(stylize_qr_sdxl, [p, st, cfg, sd, cl, ch], out)
+
 
 if __name__ == "__main__":
     demo.launch()