259 lines
8.6 KiB
Python
259 lines
8.6 KiB
Python
import colorsys
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
MAX_REGIONS = 8
|
|
|
|
|
|
def _conditioning_set_values(conditioning, values):
|
|
out = []
|
|
for cond in conditioning:
|
|
data = cond[1].copy()
|
|
data.update(values)
|
|
out.append([cond[0], data])
|
|
return out
|
|
|
|
|
|
def _encode_prompt(clip, text):
|
|
tokens = clip.tokenize(text or "")
|
|
return clip.encode_from_tokens_scheduled(tokens)
|
|
|
|
|
|
def _palette_rgb(index):
|
|
"""Regional Prompter-style deterministic palette: HSV(deg, 50%, 50%)."""
|
|
if index < 1:
|
|
raise ValueError("Region index starts at 1.")
|
|
|
|
offset = index - 1
|
|
if offset == 0:
|
|
hue = 0.0
|
|
else:
|
|
level = int(math.floor(math.log2(offset)))
|
|
start = 2**level
|
|
position = offset - start
|
|
denominator = 2 ** (level + 1)
|
|
hue = (2 * position + 1) / denominator
|
|
rgb = colorsys.hsv_to_rgb(hue, 0.5, 0.5)
|
|
return torch.tensor(rgb, dtype=torch.float32)
|
|
|
|
|
|
def _normalise_image(image):
|
|
if image is None:
|
|
raise ValueError("mask_image is required.")
|
|
if image.ndim != 4:
|
|
raise ValueError(f"Expected IMAGE tensor [B,H,W,C], got shape {tuple(image.shape)}.")
|
|
if image.shape[-1] < 3:
|
|
raise ValueError("mask_image must have at least 3 color channels.")
|
|
return image[..., :3].float().clamp(0.0, 1.0)
|
|
|
|
|
|
def _soften_mask(mask, feather):
|
|
if feather <= 0:
|
|
return mask
|
|
|
|
radius = max(1, int(feather))
|
|
kernel_size = radius * 2 + 1
|
|
sigma = max(radius / 2.0, 0.5)
|
|
device = mask.device
|
|
dtype = mask.dtype
|
|
|
|
coords = torch.arange(kernel_size, device=device, dtype=dtype) - radius
|
|
kernel_1d = torch.exp(-(coords * coords) / (2 * sigma * sigma))
|
|
kernel_1d = kernel_1d / kernel_1d.sum()
|
|
|
|
b = mask.shape[0]
|
|
x = mask.unsqueeze(1)
|
|
x = F.pad(x, (radius, radius, 0, 0), mode="replicate")
|
|
x = F.conv2d(x, kernel_1d.view(1, 1, 1, kernel_size))
|
|
x = F.pad(x, (0, 0, radius, radius), mode="replicate")
|
|
x = F.conv2d(x, kernel_1d.view(1, 1, kernel_size, 1))
|
|
return x.reshape(b, mask.shape[-2], mask.shape[-1]).clamp(0.0, 1.0)
|
|
|
|
|
|
def _mask_for_region(image, region_index, tolerance, min_alpha):
|
|
target = _palette_rgb(region_index).to(device=image.device, dtype=image.dtype)
|
|
diff = torch.linalg.vector_norm(image - target.view(1, 1, 1, 3), dim=-1)
|
|
distance_limit = max(float(tolerance), 0.0) / 255.0
|
|
|
|
if distance_limit <= 0:
|
|
mask = (diff == 0).float()
|
|
else:
|
|
mask = (1.0 - (diff / distance_limit)).clamp(0.0, 1.0)
|
|
|
|
if min_alpha > 0:
|
|
mask = torch.where(mask >= float(min_alpha), mask, torch.zeros_like(mask))
|
|
return mask
|
|
|
|
|
|
class RegionalMaskPromptEncode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"mask_image": ("IMAGE",),
|
|
"base_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_1_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_2_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_3_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_4_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_count": ("INT", {"default": 4, "min": 1, "max": MAX_REGIONS, "step": 1}),
|
|
"region_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05}),
|
|
"base_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.05}),
|
|
"mask_tolerance": ("INT", {"default": 48, "min": 0, "max": 255, "step": 1}),
|
|
"mask_feather": ("INT", {"default": 3, "min": 0, "max": 64, "step": 1}),
|
|
"min_mask_alpha": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"set_cond_area": (["default", "mask bounds"], {"default": "default"}),
|
|
},
|
|
"optional": {
|
|
"region_5_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_6_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_7_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"region_8_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"CONDITIONING",
|
|
"MASK",
|
|
"MASK",
|
|
"MASK",
|
|
"MASK",
|
|
"MASK",
|
|
"MASK",
|
|
"MASK",
|
|
"MASK",
|
|
"STRING",
|
|
)
|
|
RETURN_NAMES = (
|
|
"conditioning",
|
|
"region_1_mask",
|
|
"region_2_mask",
|
|
"region_3_mask",
|
|
"region_4_mask",
|
|
"region_5_mask",
|
|
"region_6_mask",
|
|
"region_7_mask",
|
|
"region_8_mask",
|
|
"region_info",
|
|
)
|
|
FUNCTION = "encode"
|
|
CATEGORY = "conditioning/regional"
|
|
|
|
def encode(
|
|
self,
|
|
clip,
|
|
mask_image,
|
|
base_prompt,
|
|
region_1_prompt,
|
|
region_2_prompt,
|
|
region_3_prompt,
|
|
region_4_prompt,
|
|
region_count,
|
|
region_strength,
|
|
base_strength,
|
|
mask_tolerance,
|
|
mask_feather,
|
|
min_mask_alpha,
|
|
set_cond_area,
|
|
region_5_prompt="",
|
|
region_6_prompt="",
|
|
region_7_prompt="",
|
|
region_8_prompt="",
|
|
):
|
|
image = _normalise_image(mask_image)
|
|
region_count = int(max(1, min(MAX_REGIONS, region_count)))
|
|
prompts = [
|
|
region_1_prompt,
|
|
region_2_prompt,
|
|
region_3_prompt,
|
|
region_4_prompt,
|
|
region_5_prompt,
|
|
region_6_prompt,
|
|
region_7_prompt,
|
|
region_8_prompt,
|
|
]
|
|
|
|
masks = []
|
|
active_lines = []
|
|
for idx in range(1, MAX_REGIONS + 1):
|
|
if idx <= region_count:
|
|
mask = _mask_for_region(image, idx, mask_tolerance, min_mask_alpha)
|
|
mask = _soften_mask(mask, mask_feather)
|
|
else:
|
|
mask = torch.zeros(image.shape[0], image.shape[1], image.shape[2], device=image.device, dtype=image.dtype)
|
|
masks.append(mask)
|
|
|
|
coverage = float(mask.mean().detach().cpu()) if idx <= region_count else 0.0
|
|
if idx <= region_count:
|
|
active_lines.append(f"Region {idx}: coverage={coverage:.4f}, prompt={'yes' if prompts[idx - 1].strip() else 'no'}")
|
|
|
|
conditioning = []
|
|
if base_prompt.strip() and base_strength > 0:
|
|
base_cond = _conditioning_set_values(_encode_prompt(clip, base_prompt), {"strength": float(base_strength)})
|
|
conditioning.extend(base_cond)
|
|
|
|
set_area_to_bounds = set_cond_area == "mask bounds"
|
|
for idx in range(region_count):
|
|
prompt = prompts[idx].strip()
|
|
if not prompt or region_strength <= 0:
|
|
continue
|
|
if float(masks[idx].max().detach().cpu()) <= 0:
|
|
continue
|
|
|
|
region_cond = _encode_prompt(clip, prompt)
|
|
region_cond = _conditioning_set_values(
|
|
region_cond,
|
|
{
|
|
"mask": masks[idx],
|
|
"set_area_to_bounds": set_area_to_bounds,
|
|
"mask_strength": float(region_strength),
|
|
"strength": float(region_strength),
|
|
},
|
|
)
|
|
conditioning.extend(region_cond)
|
|
|
|
if not conditioning:
|
|
conditioning = [[_encode_prompt(clip, ""), {}]]
|
|
|
|
h, w = image.shape[1], image.shape[2]
|
|
info = f"Regional Mask Prompt: {w}x{h}, regions={region_count}, tolerance={mask_tolerance}, feather={mask_feather}\n" + "\n".join(active_lines)
|
|
return (conditioning, *masks, info)
|
|
|
|
|
|
class RegionalMaskPromptPalette:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"region_count": ("INT", {"default": 4, "min": 1, "max": MAX_REGIONS, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("palette_info",)
|
|
FUNCTION = "palette"
|
|
CATEGORY = "conditioning/regional"
|
|
|
|
def palette(self, region_count):
|
|
lines = []
|
|
for idx in range(1, int(region_count) + 1):
|
|
rgb = _palette_rgb(idx)
|
|
values = [int(round(float(v) * 255)) for v in rgb]
|
|
lines.append(f"Region {idx}: RGB({values[0]}, {values[1]}, {values[2]}) #{values[0]:02x}{values[1]:02x}{values[2]:02x}")
|
|
return ("\n".join(lines),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RegionalMaskPromptEncode": RegionalMaskPromptEncode,
|
|
"RegionalMaskPromptPalette": RegionalMaskPromptPalette,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RegionalMaskPromptEncode": "Regional Mask Prompt Encode",
|
|
"RegionalMaskPromptPalette": "Regional Mask Prompt Palette",
|
|
}
|