| """ComfyUI node for the P-Q-Image-2.1-Remove-Object LoRAs: draws the red outline around the mask | |
| exactly as expected by the model, and appends the outline sentence to the instruction. | |
| Install: copy this folder into ComfyUI/custom_nodes/ and install requirements.txt (opencv) into | |
| ComfyUI's Python. The node appears as "Precise Removal Outline" (category Pruna). | |
| The mask is optional. It can be painted (Load Image's MaskEditor: `mask`), come from a mask file | |
| (e.g. Load Image (as Mask), channel red, white = object: `mask_file`), or both: they are combined | |
| (union). Without one, or with empty ones (LoadImage gives an empty mask when nothing was painted), | |
| the photo and the instruction pass through unchanged: the LoRA also removes objects from the | |
| instruction alone. The outline itself comes from draw_outline.py, the same file as in | |
| PrunaAI/P-Q-Image-2.1-Remove-Object, so it is the same ring, pixel for pixel, as in the diffusers usage. | |
| """ | |
| from __future__ import annotations | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from .draw_outline import OUTLINE_SUFFIX, RED, outline_ring | |
| def _mask_at(mask: torch.Tensor, index: int, size: tuple[int, int]) -> np.ndarray: | |
| """Boolean mask of batch item `index` at the photo's size (w, h), resized and thresholded as | |
| draw_outline.draw_outline does (8-bit grey, nearest neighbour, > 127).""" | |
| m = mask if mask.dim() == 2 else mask[min(index, mask.shape[0] - 1)] | |
| grey = Image.fromarray(np.clip(np.rint(m.detach().cpu().float().numpy() * 255), 0, 255).astype(np.uint8), "L") | |
| return np.array(grey.resize(size, Image.NEAREST)) > 127 | |
| class PreciseRemovalOutline: | |
| def INPUT_TYPES(cls): | |
| return { | |
| "required": { | |
| "image": ("IMAGE", {"tooltip": "The photo to edit."}), | |
| "instruction": ("STRING", {"multiline": True, "default": "Remove the ", | |
| "tooltip": "What to remove, e.g. \"Remove the man on the left.\""}), | |
| }, | |
| "optional": { | |
| "mask": ("MASK", {"tooltip": "The object to remove (paint it with the MaskEditor). Optional: " | |
| "without a mask, the instruction alone says what to remove."}), | |
| "mask_file": ("MASK", {"tooltip": "A mask from a file (white = the object to remove), any size: " | |
| "resized to the photo. Combined with `mask` if both are given."}), | |
| }, | |
| } | |
| RETURN_TYPES = ("IMAGE", "STRING") | |
| RETURN_NAMES = ("image", "prompt") | |
| OUTPUT_TOOLTIPS = ("The photo with the red outline (unchanged without a mask).", | |
| "The instruction, followed by the outline sentence when an outline was drawn.") | |
| FUNCTION = "run" | |
| CATEGORY = "Pruna" | |
| DESCRIPTION = ("Draws the red outline expected by P-Q-Image-2.1-Remove-Object (2 px outside the " | |
| "mask, 3 px wide, 4 px on photos over 900 px) and adds the matching sentence to the " | |
| "instruction. Without a mask, passes both through.") | |
| def run(self, image: torch.Tensor, instruction: str, mask: torch.Tensor | None = None, | |
| mask_file: torch.Tensor | None = None): | |
| out, drawn = image.clone(), False | |
| masks = [x for x in (mask, mask_file) if x is not None and x.numel() and float(x.max()) > 0] | |
| if masks: | |
| red = torch.tensor([c / 255 for c in RED], dtype=out.dtype, device=out.device) | |
| for i in range(out.shape[0]): | |
| h, w = out.shape[1], out.shape[2] | |
| m = np.logical_or.reduce([_mask_at(x, i, (w, h)) for x in masks]) | |
| if m.any(): | |
| ring = torch.from_numpy(outline_ring(m)).to(out.device) | |
| out[i][ring] = red | |
| drawn = True | |
| prompt = f"{instruction.rstrip()} {OUTLINE_SUFFIX}" if drawn else instruction | |
| return (out, prompt) | |
| NODE_CLASS_MAPPINGS = {"PreciseRemovalOutline": PreciseRemovalOutline} | |
| NODE_DISPLAY_NAME_MAPPINGS = {"PreciseRemovalOutline": "Precise Removal Outline"} | |