refine

app.py

@@ -1,7 +1,7 @@
 import os, gc, random, re, inspect
 from contextlib import nullcontext
 
-from PIL import Image, ImageOps
+from PIL import Image, ImageOps, ImageFilter
 
 import gradio as gr
 import torch, spaces
@@ -127,15 +127,6 @@ def normalize_color(c):
         return s
     return "white"
 
-def make_qr(url="https://example.com", size=768, border=12, back_color="#FFFFFF", blur_radius=0.0):
-    qr = qrcode.QRCode(version=None, error_correction=ERROR_CORRECT_H, box_size=10, border=int(border))
-    qr.add_data(url.strip()); qr.make(fit=True)
-    img = qr.make_image(fill_color="black", back_color=normalize_color(back_color)).convert("RGB")
-    img = img.resize((int(size), int(size)), Image.NEAREST)
-    if blur_radius and blur_radius > 0:
-        img = img.filter(ImageFilter.GaussianBlur(radius=float(blur_radius)))
-    return img
-
 def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: float = 0.0, feather: float = 1.0) -> Image.Image:
     if strength <= 0: return stylized
     q = qr_img.convert("L")
@@ -148,29 +139,113 @@ def enforce_qr_contrast(stylized: Image.Image, qr_img: Image.Image, strength: fl
     s = np.clip(s, 0.0, 1.0)
     return Image.fromarray((s * 255.0).astype(np.uint8), mode="RGB")
 
-def prep_qr_upload_image(qr_upload: Image.Image, size: int, strict: bool = False) -> Image.Image:
-    # If QR has alpha transparency → flatten onto white
-    if qr_upload.mode in ("RGBA", "LA") or "transparency" in qr_upload.info:
-        rgba = qr_upload.convert("RGBA")
-        # white background (instead of default black!)
-        white_bg = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
-        im = Image.alpha_composite(white_bg, rgba).convert("RGB")
-    else:
-        im = qr_upload.convert("RGB")
-
-    if strict:
-        g = ImageOps.grayscale(im)
-        bw = g.point(lambda p: 0 if p < 128 else 255, "L")
-        im = Image.merge("RGB", (bw, bw, bw))
-
-    # resize for ControlNet
-    im = im.resize((size, size), Image.NEAREST)
-    im.save("debug_qr_after_prep.png")   # debug output
-    return im
+def make_qr(url="https://example.com", size=768, border=12, back_color="#FFFFFF", blur_radius=0.0):
+    qr = qrcode.QRCode(version=None, error_correction=ERROR_CORRECT_H, box_size=10, border=int(border))
+    qr.add_data(url.strip()); qr.make(fit=True)
+    img = qr.make_image(fill_color="black", back_color=normalize_color(back_color)).convert("RGB")
+    img = img.resize((int(size), int(size)), Image.NEAREST)
+    if blur_radius and blur_radius > 0:
+        img = img.filter(ImageFilter.GaussianBlur(radius=float(blur_radius)))
+    return img
 
+def _binarize(img_gray, thr=128):
+    return img_gray.point(lambda p: 255 if p >= thr else 0, mode="1").convert("L")
+
+def _estimate_module_px(bin_img):
+    """Rough module width estimate via run-lengths of black segments."""
+    a = np.array(bin_img, dtype=np.uint8)  # 0..255
+    row = a[a.shape[0]//2]                 # mid row
+    # find lengths of consecutive black runs
+    runs = []
+    cnt, prev = 0, None
+    for v in row:
+        if prev is None:
+            prev, cnt = v, 1
+            continue
+        if v == prev:
+            cnt += 1
+        else:
+            if prev == 0: runs.append(cnt)  # black run
+            prev, cnt = v, 1
+    if prev == 0: runs.append(cnt)
+    if not runs: return None
+    # use a robust small quantile (modules are among the shorter black runs)
+    return max(1, int(np.percentile(runs, 25)))
+
+def _enforce_quiet_zone(bin_img, module_px, border_modules=4):
+    # pad so that white border >= border_modules * module_px
+    arr = np.array(bin_img)
+    ys, xs = np.where(arr == 0)
+    if len(xs) == 0:
+        return ImageOps.expand(bin_img, border=border_modules*module_px, fill=255)
+    left   = xs.min()
+    right  = arr.shape[1]-1 - xs.max()
+    top    = ys.min()
+    bottom = arr.shape[0]-1 - ys.max()
+    need = border_modules*module_px
+    pad = (
+        max(0, need - left),
+        max(0, need - top),
+        max(0, need - right),
+        max(0, need - bottom),
+    )
+    return ImageOps.expand(bin_img, border=pad, fill=255)
+
+def prep_qr_uploaded_for_controlnet(
+    qr_upload: Image.Image,
+    size: int = 768,
+    border_modules: int = 4,
+    rounded: bool = True,           # set True when input is circle/rounded style
+    inflate_ratio: float = 0.9,     # target dot fill ≈ 90% of cell
+    threshold: int = 128
+) -> Image.Image:
+    # 1) flatten & grayscale
+    if qr_upload.mode in ("RGBA", "LA"):
+        bg = Image.new("RGBA", qr_upload.size, (255, 255, 255, 255))
+        qr_upload = Image.alpha_composite(bg, qr_upload.convert("RGBA"))
+    gray = qr_upload.convert("L")
+
+    # 2) hard binarize (kill anti-aliasing)
+    bin_img = _binarize(gray, thr=threshold)
+
+    # 3) estimate module size (px)
+    module_px = _estimate_module_px(bin_img)
+    if module_px is None:
+        module_px = max(2, size // 120)  # safe fallback
+
+    # 4) for rounded QRs, inflate dots a bit (morphological dilation)
+    if rounded:
+        # choose a small kernel relative to module size (≈ 1/5 of module)
+        k = max(3, int(round(module_px * 0.2)) | 1)  # odd kernel size
+        bin_img = bin_img.filter(ImageFilter.MaxFilter(k))
+
+    # 5) enforce quiet zone ≥ border_modules
+    bin_img = _enforce_quiet_zone(bin_img, module_px, border_modules)
+
+    # 6) align-to-grid scaling (try to keep integer pixels per module)
+    w, h = bin_img.size
+    # content width (without the padded quiet zone we just enforced)
+    arr = np.array(bin_img)
+    ys, xs = np.where(arr == 0)
+    if len(xs) > 0:
+        content_w = (xs.max() - xs.min() + 1)
+    else:
+        content_w = w
 
+    # estimated number of modules across content
+    n_est = max(21, int(round(content_w / max(1, module_px))))  # QR v1 min=21
+    scale = max(1, size // (n_est + 2*border_modules))
+    final_dim = (n_est + 2*border_modules) * scale
 
+    # resize with NEAREST to the computed dim, then center-pad to target size
+    img_scaled = bin_img.resize((final_dim, final_dim), Image.NEAREST)
+    pad_each = max(0, (size - final_dim) // 2)
+    img_out = ImageOps.expand(img_scaled, border=pad_each, fill=255)
+    if img_out.size != (size, size):
+        img_out = img_out.resize((size, size), Image.NEAREST)  # rare off-by-one
 
+    # 7) return RGB (3 channels)
+    return img_out.convert("RGB")
 
 # ----- Brightness map preprocessing & mixing -----
 def prep_brightness_map(img: Image.Image, size: int, source: str,
@@ -301,7 +376,7 @@ def _qr_txt2img_core(model_id: str,
     # --- Build base-size control images (s x s)
     # qr_img = make_qr(url=url, size=s, border=int(border), back_color="#FFFFFF", blur_radius=0.0)
     if qr_upload is not None:
-        qr_img = prep_qr_upload_image(qr_upload, s, strict=False)
+        qr_img = prep_qr_uploaded_for_controlnet(qr_upload, s, border=int(border))
     else:
         qr_img = make_qr(url=url, size=s, border=int(border), back_color="#FFFFFF")
 
@@ -457,7 +532,7 @@ def _qr_img2img_core(model_id: str,
     #     make_qr(url=url, size=s, border=int(border), back_color="#FFFFFF", blur_radius=0.0)
     # )
     if qr_upload is not None:
-        qr_img = prep_qr_upload_image(qr_upload, s, strict=False)
+        qr_img = prep_qr_uploaded_for_controlnet(qr_upload, s, border=int(border))
     else:
         qr_img = make_qr(url=url, size=s, border=int(border), back_color="#FFFFFF")