From baeb8496e512b9d4bfd8fc964e7a11cb98774094 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Thu, 25 Jun 2026 03:42:17 +0200 Subject: [PATCH] Mask context padding + downscale-only vision cap - mask_padding: grow the mask crop box by a fraction of image size per side for surrounding context (0 = tight crop, default) - vision resize is now downscale-only (a small mask crop keeps native size instead of being magnified to the megapixel target) - README: document mask_padding and the megapixel cap Co-Authored-By: Claude Opus 4.8 --- README.md | 8 ++++++-- nodes.py | 34 +++++++++++++++++++++++++++------- 2 files changed, 33 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 016ebe6..c9ca1b3 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,10 @@ image+mask pairs** (connect `image1`, a fresh `image2`/`mask2` pair appears, and Each reference image has an optional companion mask. When connected, the image is **cropped to the mask's bounding box** before the vision encoder, so the VLM only sees the -masked region. This is *reference-image* masking — **not inpainting**: Krea2 has no +masked region. Use **`mask_padding`** to keep surrounding context: it grows the crop box by +that fraction of the image size on each side (`0` = tight crop, `0.1` ≈ 10% margin, high +values ≈ the whole image). The mask is used only to compute the crop — it is not itself +sent to the VLM. This is *reference-image* masking — **not inpainting**: Krea2 has no concat/inpaint pathway to regenerate a masked region of the output. (To spatially restrict where the *prompt* applies in the output, use ComfyUI's standard `ConditioningSetMask` downstream of this node — that works generically at the sampler level.) @@ -44,7 +47,8 @@ downstream of this node — that works generically at the sampler level.) | `prompt` | STRING | Your text prompt. | | `image1…N` | IMAGE | Optional reference images; slots grow as you connect them. | | `mask1…N` | MASK | Optional per-image mask; crops `imageN` to the masked bounding box. | -| `vision_megapixels` | FLOAT | Area each image is resized to before the vision encoder (default `1.0`). | +| `mask_padding` | FLOAT | Context kept around the mask, as a fraction of image size per side (`0` = tight, default). | +| `vision_megapixels` | FLOAT | Max size before the vision encoder; references are downscaled to this cap, never upscaled (default `1.0`). | Output is a standard `CONDITIONING` for the Krea2 sampler. With no images connected it works as a plain Krea2 text encoder. diff --git a/nodes.py b/nodes.py index 03968f9..96d78f7 100644 --- a/nodes.py +++ b/nodes.py @@ -50,8 +50,15 @@ class TextEncodeKrea2: "mask1": ("MASK",), "vision_megapixels": ("FLOAT", { "default": 1.0, "min": 0.1, "max": 8.0, "step": 0.1, - "tooltip": "Target size (in megapixels) each reference image is area-resized to " - "before the Qwen3-VL vision encoder. Higher = more visual detail, more VRAM.", + "tooltip": "Maximum size (in megapixels) for each reference before the Qwen3-VL " + "vision encoder. References larger than this are downscaled; smaller " + "ones (e.g. a tight mask crop) are kept at native size, never upscaled.", + }), + "mask_padding": ("FLOAT", { + "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.02, + "tooltip": "Context kept around the mask before cropping, as a fraction of the " + "image size added on EACH side. 0 = tight crop to the mask; 0.1 = ~10% " + "margin of surroundings. Only applies when a mask is connected.", }), }, } @@ -74,8 +81,9 @@ class TextEncodeKrea2: return out @staticmethod - def _crop_to_mask(image, mask): - """Crop image (B,H,W,C) to the bounding box of mask. No-op if mask empty/None.""" + def _crop_to_mask(image, mask, padding=0.0): + """Crop image (B,H,W,C) to the mask bounding box, expanded by `padding` (a + fraction of the image size) on each side. No-op if mask empty/None.""" if mask is None: return image @@ -97,9 +105,18 @@ class TextEncodeKrea2: cols = torch.where(torch.any(presence, dim=0))[0] y0, y1 = int(rows[0]), int(rows[-1]) x0, x1 = int(cols[0]), int(cols[-1]) + + if padding > 0.0: # grow the box outward for surrounding context, clamped to the image + pad_x = round(padding * w) + pad_y = round(padding * h) + x0 = max(0, x0 - pad_x) + x1 = min(w - 1, x1 + pad_x) + y0 = max(0, y0 - pad_y) + y1 = min(h - 1, y1 + pad_y) + return image[:, y0:y1 + 1, x0:x1 + 1, :] - def encode(self, clip, prompt, vision_megapixels=1.0, **kwargs): + def encode(self, clip, prompt, vision_megapixels=1.0, mask_padding=0.0, **kwargs): images = self._collect_indexed(kwargs, "image") masks = self._collect_indexed(kwargs, "mask") ordered = sorted(images.keys()) @@ -109,10 +126,13 @@ class TextEncodeKrea2: total = int(vision_megapixels * 1024 * 1024) for slot, n in enumerate(ordered): - image = self._crop_to_mask(images[n], masks.get(n)) + image = self._crop_to_mask(images[n], masks.get(n), padding=mask_padding) samples = image.movedim(-1, 1) - scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2])) + # vision_megapixels is an upper CAP, not a fixed target: only downscale + # oversized references, never upscale. Otherwise a small mask crop gets + # magnified to fill the VLM frame and the subject reads as huge/zoomed. + scale_by = min(1.0, math.sqrt(total / (samples.shape[3] * samples.shape[2]))) width = round(samples.shape[3] * scale_by) height = round(samples.shape[2] * scale_by)