From e2fb70d4a197a59846bc99440da645afc4c47da3 Mon Sep 17 00:00:00 2001 From: vladmandic Date: Wed, 17 Dec 2025 10:03:17 +0100 Subject: [PATCH] detailer draw segmentation overlays Signed-off-by: vladmandic --- modules/devices.py | 2 +- modules/loader.py | 5 +++-- modules/postprocess/yolo.py | 34 +++++++++++++++++++++++----------- 3 files changed, 27 insertions(+), 14 deletions(-) diff --git a/modules/devices.py b/modules/devices.py index 76d9a8652..70bfdfefc 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -390,7 +390,7 @@ def test_triton(early: bool = False): def test_triton_func(a,b,c): return a * b + c test_triton_func = torch.compile(test_triton_func, fullgraph=True) - test_triton_func(torch.randn(32, device=device), torch.randn(32, device=device), torch.randn(32, device=device)) + test_triton_func(torch.randn(16, device=device), torch.randn(16, device=device), torch.randn(16, device=device)) triton_ok = True else: triton_ok = False diff --git a/modules/loader.py b/modules/loader.py index 43824fc70..5814baabb 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -94,9 +94,9 @@ timer.startup.record("torch") try: import bitsandbytes # pylint: disable=W0611,C0411 + _bnb = True except Exception: - from diffusers.utils import import_utils - import_utils._bitsandbytes_available = False # pylint: disable=protected-access + _bnb = False timer.startup.record("bnb") import transformers # pylint: disable=W0611,C0411 @@ -134,6 +134,7 @@ try: import diffusers.utils.import_utils # pylint: disable=W0611,C0411 diffusers.utils.import_utils._k_diffusion_available = True # pylint: disable=protected-access # monkey-patch since we use k-diffusion from git diffusers.utils.import_utils._k_diffusion_version = '0.0.12' # pylint: disable=protected-access + diffusers.utils.import_utils._bitsandbytes_available = _bnb # pylint: disable=protected-access import diffusers # pylint: disable=W0611,C0411 import diffusers.loaders.single_file # pylint: disable=W0611,C0411 diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py index 3189cb86c..c0f55d465 100644 --- a/modules/postprocess/yolo.py +++ b/modules/postprocess/yolo.py @@ -240,18 +240,30 @@ class YoloRestorer(Detailer): ) return [merged] - def draw_boxes(self, image: Image.Image, items: list[YoloResult]) -> Image.Image: - if isinstance(image, Image.Image): - draw = ImageDraw.Draw(image) - else: + def draw_masks(self, image: Image.Image, items: list[YoloResult]) -> Image.Image: + if not isinstance(image, Image.Image): image = Image.fromarray(image) - draw = ImageDraw.Draw(image) - font = images.get_font(16) + image = image.convert('RGBA') + size = min(image.width, image.height) // 32 + font = images.get_font(size) + color = (0, 190, 190) shared.log.debug(f'Detailer: draw={items}') for i, item in enumerate(items): - draw.rectangle(item.box, outline="#00C8C8", width=3) - draw.text((item.box[0]+4, item.box[1]+4), f'{i+1} {item.label} {item.score:.2f}', fill="black", font=font) - draw.text((item.box[0]+2, item.box[1]+2), f'{i+1} {item.label} {item.score:.2f}', fill="white", font=font) + if shared.opts.detailer_seg and item.mask is not None: + mask = item.mask.convert('L') + else: + mask = Image.new('L', image.size, 0) + draw_mask = ImageDraw.Draw(mask) + draw_mask.rectangle(item.box, fill="white", outline=None, width=0) + alpha = mask.point(lambda p: int(p * 0.5)) + overlay = Image.new("RGBA", image.size, color + (0,)) + overlay.putalpha(alpha) + image = Image.alpha_composite(image, overlay) + + draw_text = ImageDraw.Draw(image) + draw_text.text((item.box[0] + 2, item.box[1] - size - 2), f'{i+1} {item.label} {item.score:.2f}', fill="black", font=font) + draw_text.text((item.box[0] + 0, item.box[1] - size - 4), f'{i+1} {item.label} {item.score:.2f}', fill="white", font=font) + image = image.convert("RGB") return np.array(image) def restore(self, np_image, p: processing.StableDiffusionProcessing = None): @@ -392,7 +404,7 @@ class YoloRestorer(Detailer): if shared.opts.detailer_sort: items = sorted(items, key=lambda x: x.box[0]) # sort items left-to-right to improve consistency if shared.opts.detailer_save: - annotated = self.draw_boxes(annotated, items) + annotated = self.draw_masks(annotated, items) for j, item in enumerate(items): if item.mask is None: @@ -405,7 +417,7 @@ class YoloRestorer(Detailer): pc.negative_prompts = [pc.negative_prompt] pc.prompts, pc.network_data = extra_networks.parse_prompts(pc.prompts) extra_networks.activate(pc, pc.network_data) - shared.log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label} score={item.score:.2f} seg={shared.opts.detailer_seg} prompt="{pc.prompt}"') + shared.log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={shared.opts.detailer_seg} prompt="{pc.prompt}"') pc.init_images = [image] pc.image_mask = [item.mask] pc.overlay_images = []