844 lines
31 KiB
Python
844 lines
31 KiB
Python
import json
|
|
import math
|
|
from typing import List, Sequence, Tuple
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
CATEGORY = "Universal Remove"
|
|
|
|
|
|
def _ensure_image_batch(images: torch.Tensor) -> torch.Tensor:
|
|
if not isinstance(images, torch.Tensor):
|
|
raise TypeError("Expected IMAGE input as torch.Tensor")
|
|
if images.ndim == 3:
|
|
images = images.unsqueeze(0)
|
|
if images.ndim != 4:
|
|
raise ValueError(f"Expected IMAGE tensor with 4 dims, got {images.ndim}")
|
|
return images.clamp(0.0, 1.0)
|
|
|
|
|
|
def _ensure_mask_batch(mask: torch.Tensor) -> torch.Tensor:
|
|
if not isinstance(mask, torch.Tensor):
|
|
raise TypeError("Expected MASK input as torch.Tensor")
|
|
if mask.ndim == 2:
|
|
mask = mask.unsqueeze(0)
|
|
elif mask.ndim == 4 and mask.shape[1] == 1:
|
|
mask = mask[:, 0]
|
|
elif mask.ndim == 4 and mask.shape[-1] == 1:
|
|
mask = mask[..., 0]
|
|
if mask.ndim != 3:
|
|
raise ValueError(f"Expected MASK tensor with 3 dims, got {mask.ndim}")
|
|
return mask.float().clamp(0.0, 1.0)
|
|
|
|
|
|
def _tile_to_batch(tensor: torch.Tensor, target_batch: int) -> torch.Tensor:
|
|
if tensor.shape[0] == target_batch:
|
|
return tensor
|
|
if tensor.shape[0] <= 0:
|
|
raise ValueError("Cannot tile an empty batch")
|
|
repeats = math.ceil(target_batch / tensor.shape[0])
|
|
repeat_shape = [repeats] + [1] * (tensor.ndim - 1)
|
|
return tensor.repeat(*repeat_shape)[:target_batch]
|
|
|
|
|
|
def _match_mask_to_images(mask: torch.Tensor, images: torch.Tensor) -> torch.Tensor:
|
|
mask = _ensure_mask_batch(mask)
|
|
images = _ensure_image_batch(images)
|
|
if mask.shape[0] != images.shape[0]:
|
|
mask = _tile_to_batch(mask, images.shape[0])
|
|
if mask.shape[1] != images.shape[1] or mask.shape[2] != images.shape[2]:
|
|
mask = F.interpolate(
|
|
mask.unsqueeze(1),
|
|
size=(images.shape[1], images.shape[2]),
|
|
mode="nearest",
|
|
).squeeze(1)
|
|
return mask.clamp(0.0, 1.0)
|
|
|
|
|
|
def _to_bchw(images: torch.Tensor) -> torch.Tensor:
|
|
return images.permute(0, 3, 1, 2)
|
|
|
|
|
|
def _to_bhwc(images: torch.Tensor) -> torch.Tensor:
|
|
return images.permute(0, 2, 3, 1)
|
|
|
|
|
|
def _resize_images(images: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
|
|
bchw = _to_bchw(images)
|
|
resized = F.interpolate(bchw, size=size, mode="bilinear", align_corners=False)
|
|
return _to_bhwc(resized).clamp(0.0, 1.0)
|
|
|
|
|
|
def _resize_masks(mask: torch.Tensor, size: Tuple[int, int]) -> torch.Tensor:
|
|
resized = F.interpolate(mask.unsqueeze(1), size=size, mode="nearest")
|
|
return resized.squeeze(1).clamp(0.0, 1.0)
|
|
|
|
|
|
def _pad_images(images: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
|
|
pad_h = max(0, target_h - images.shape[1])
|
|
pad_w = max(0, target_w - images.shape[2])
|
|
if pad_h == 0 and pad_w == 0:
|
|
return images
|
|
bchw = _to_bchw(images)
|
|
padded = F.pad(bchw, (0, pad_w, 0, pad_h), mode="constant", value=0.0)
|
|
return _to_bhwc(padded).clamp(0.0, 1.0)
|
|
|
|
|
|
def _pad_masks(mask: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
|
|
pad_h = max(0, target_h - mask.shape[1])
|
|
pad_w = max(0, target_w - mask.shape[2])
|
|
if pad_h == 0 and pad_w == 0:
|
|
return mask
|
|
padded = F.pad(mask.unsqueeze(1), (0, pad_w, 0, pad_h), mode="constant", value=0.0)
|
|
return padded.squeeze(1).clamp(0.0, 1.0)
|
|
|
|
|
|
def _parse_hex_color(color: str) -> torch.Tensor:
|
|
text = color.strip()
|
|
if text.startswith("#"):
|
|
text = text[1:]
|
|
if len(text) != 6:
|
|
raise ValueError("Color must be a 6-digit hex string, e.g. #00FF00")
|
|
rgb = [int(text[i : i + 2], 16) / 255.0 for i in range(0, 6, 2)]
|
|
return torch.tensor(rgb, dtype=torch.float32)
|
|
|
|
|
|
def _overlay_mask(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor, opacity: float) -> torch.Tensor:
|
|
images = _ensure_image_batch(images)
|
|
mask = _match_mask_to_images(mask, images)
|
|
color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3)
|
|
alpha = mask.unsqueeze(-1).clamp(0.0, 1.0) * float(opacity)
|
|
return (images * (1.0 - alpha) + color * alpha).clamp(0.0, 1.0)
|
|
|
|
|
|
def _gaussian_kernel_1d(kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
radius = kernel_size // 2
|
|
if sigma <= 0:
|
|
sigma = max(float(kernel_size) / 6.0, 1e-6)
|
|
coords = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
|
|
kernel = torch.exp(-0.5 * (coords / sigma) ** 2)
|
|
kernel = kernel / kernel.sum()
|
|
return kernel
|
|
|
|
|
|
def _gaussian_blur_mask(mask: torch.Tensor, kernel_size: int, sigma: float) -> torch.Tensor:
|
|
mask = _ensure_mask_batch(mask)
|
|
if kernel_size <= 1 and sigma <= 0:
|
|
return mask
|
|
if kernel_size <= 1:
|
|
kernel_size = int(max(3, math.ceil(sigma * 6)))
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
|
|
kernel_1d = _gaussian_kernel_1d(kernel_size, sigma, mask.device, mask.dtype)
|
|
kernel_x = kernel_1d.view(1, 1, 1, -1)
|
|
kernel_y = kernel_1d.view(1, 1, -1, 1)
|
|
|
|
blurred = mask.unsqueeze(1)
|
|
pad = kernel_size // 2
|
|
blurred = F.pad(blurred, (pad, pad, 0, 0), mode="replicate")
|
|
blurred = F.conv2d(blurred, kernel_x)
|
|
blurred = F.pad(blurred, (0, 0, pad, pad), mode="replicate")
|
|
blurred = F.conv2d(blurred, kernel_y)
|
|
return blurred.squeeze(1).clamp(0.0, 1.0)
|
|
|
|
|
|
def _expand_mask(mask: torch.Tensor, expand_pixels: int) -> torch.Tensor:
|
|
mask = _ensure_mask_batch(mask)
|
|
if expand_pixels <= 0:
|
|
return mask
|
|
kernel = expand_pixels * 2 + 1
|
|
expanded = F.max_pool2d(mask.unsqueeze(1), kernel_size=kernel, stride=1, padding=expand_pixels)
|
|
return expanded.squeeze(1).clamp(0.0, 1.0)
|
|
|
|
|
|
def _fill_masked_region(images: torch.Tensor, mask: torch.Tensor, color: torch.Tensor) -> torch.Tensor:
|
|
images = _ensure_image_batch(images)
|
|
mask = _match_mask_to_images(mask, images)
|
|
color = color.to(device=images.device, dtype=images.dtype).view(1, 1, 1, 3)
|
|
return (images * (1.0 - mask.unsqueeze(-1)) + color * mask.unsqueeze(-1)).clamp(0.0, 1.0)
|
|
|
|
|
|
def _normalize_box_values(values: Sequence[float], width: int, height: int, normalized: bool) -> Tuple[float, float, float, float]:
|
|
x1, y1, x2, y2 = values
|
|
if normalized:
|
|
x1 *= width
|
|
x2 *= width
|
|
y1 *= height
|
|
y2 *= height
|
|
return x1, y1, x2, y2
|
|
|
|
|
|
def _parse_single_box(item, coordinate_format: str, width: int, height: int, normalized: bool) -> Tuple[int, int, int, int]:
|
|
if isinstance(item, dict):
|
|
if {"x", "y", "w", "h"}.issubset(item):
|
|
x1 = float(item["x"])
|
|
y1 = float(item["y"])
|
|
w = float(item["w"])
|
|
h = float(item["h"])
|
|
if normalized:
|
|
x1 *= width
|
|
y1 *= height
|
|
w *= width
|
|
h *= height
|
|
x2 = x1 + w
|
|
y2 = y1 + h
|
|
elif {"x1", "y1", "x2", "y2"}.issubset(item):
|
|
x1, y1, x2, y2 = _normalize_box_values(
|
|
[float(item["x1"]), float(item["y1"]), float(item["x2"]), float(item["y2"])],
|
|
width,
|
|
height,
|
|
normalized,
|
|
)
|
|
elif {"left", "top", "right", "bottom"}.issubset(item):
|
|
x1, y1, x2, y2 = _normalize_box_values(
|
|
[float(item["left"]), float(item["top"]), float(item["right"]), float(item["bottom"])],
|
|
width,
|
|
height,
|
|
normalized,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported box object: {item}")
|
|
elif isinstance(item, (list, tuple)) and len(item) == 4:
|
|
values = [float(v) for v in item]
|
|
if coordinate_format == "xywh":
|
|
x1, y1, w, h = values
|
|
if normalized:
|
|
x1 *= width
|
|
y1 *= height
|
|
w *= width
|
|
h *= height
|
|
x2 = x1 + w
|
|
y2 = y1 + h
|
|
else:
|
|
x1, y1, x2, y2 = _normalize_box_values(values, width, height, normalized)
|
|
else:
|
|
raise ValueError(f"Unsupported box entry: {item}")
|
|
|
|
left = int(math.floor(min(x1, x2)))
|
|
right = int(math.ceil(max(x1, x2)))
|
|
top = int(math.floor(min(y1, y2)))
|
|
bottom = int(math.ceil(max(y1, y2)))
|
|
|
|
left = max(0, min(width, left))
|
|
right = max(0, min(width, right))
|
|
top = max(0, min(height, top))
|
|
bottom = max(0, min(height, bottom))
|
|
return left, top, right, bottom
|
|
|
|
|
|
def _parse_boxes_json(boxes_json: str, coordinate_format: str, width: int, height: int, normalized: bool) -> List[Tuple[int, int, int, int]]:
|
|
text = boxes_json.strip()
|
|
if not text:
|
|
return []
|
|
data = json.loads(text)
|
|
if isinstance(data, dict):
|
|
if "boxes" in data:
|
|
data = data["boxes"]
|
|
else:
|
|
data = [data]
|
|
if not isinstance(data, list):
|
|
raise ValueError("boxes_json must be a JSON array or object")
|
|
boxes = []
|
|
for item in data:
|
|
box = _parse_single_box(item, coordinate_format, width, height, normalized)
|
|
if box[2] > box[0] and box[3] > box[1]:
|
|
boxes.append(box)
|
|
return boxes
|
|
|
|
|
|
def _scale_point(x: float, y: float, width: int, height: int, normalized: bool) -> Tuple[float, float]:
|
|
if normalized:
|
|
x *= width
|
|
y *= height
|
|
return x, y
|
|
|
|
|
|
def _region_frame_index(region: dict, default_frame_index: int, batch_size: int) -> int:
|
|
index = int(region.get("frame_index", default_frame_index))
|
|
return max(0, min(index, batch_size - 1))
|
|
|
|
|
|
def _region_type(region: dict) -> str:
|
|
value = str(region.get("type", region.get("shape", "rect"))).strip().lower()
|
|
aliases = {
|
|
"rectangle": "rect",
|
|
"box": "rect",
|
|
"bbox": "rect",
|
|
"ellipse": "circle",
|
|
"poly": "polygon",
|
|
}
|
|
return aliases.get(value, value)
|
|
|
|
|
|
def _polygon_to_mask(points: Sequence[Sequence[float]], width: int, height: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
|
|
if len(points) < 3:
|
|
return torch.zeros((height, width), dtype=dtype, device=device)
|
|
|
|
px = torch.arange(width, dtype=dtype, device=device).view(1, width) + 0.5
|
|
py = torch.arange(height, dtype=dtype, device=device).view(height, 1) + 0.5
|
|
inside = torch.zeros((height, width), dtype=torch.bool, device=device)
|
|
|
|
wrapped = list(points) + [points[0]]
|
|
eps = torch.tensor(1e-6, dtype=dtype, device=device)
|
|
for current, nxt in zip(wrapped[:-1], wrapped[1:]):
|
|
x1 = torch.tensor(float(current[0]), dtype=dtype, device=device)
|
|
y1 = torch.tensor(float(current[1]), dtype=dtype, device=device)
|
|
x2 = torch.tensor(float(nxt[0]), dtype=dtype, device=device)
|
|
y2 = torch.tensor(float(nxt[1]), dtype=dtype, device=device)
|
|
crosses_y = torch.logical_xor(y1 > py, y2 > py)
|
|
x_intersection = (x2 - x1) * (py - y1) / (y2 - y1 + eps) + x1
|
|
intersects = crosses_y & (px < x_intersection)
|
|
inside = torch.logical_xor(inside, intersects)
|
|
|
|
return inside.to(dtype=dtype)
|
|
|
|
|
|
def _draw_region(mask: torch.Tensor, region: dict, width: int, height: int, normalized: bool) -> None:
|
|
shape = _region_type(region)
|
|
if shape == "rect":
|
|
box = _parse_single_box(region, "xyxy", width, height, normalized)
|
|
x1, y1, x2, y2 = box
|
|
if x2 > x1 and y2 > y1:
|
|
mask[y1:y2, x1:x2] = 1.0
|
|
return
|
|
|
|
if shape == "circle":
|
|
if "center" in region:
|
|
cx, cy = region["center"]
|
|
else:
|
|
cx = region.get("cx", region.get("x", 0.0))
|
|
cy = region.get("cy", region.get("y", 0.0))
|
|
cx, cy = _scale_point(float(cx), float(cy), width, height, normalized)
|
|
radius = float(region.get("r", region.get("radius", 0.0)))
|
|
if normalized:
|
|
radius *= min(width, height)
|
|
if radius <= 0:
|
|
return
|
|
yy, xx = torch.meshgrid(
|
|
torch.arange(height, dtype=mask.dtype, device=mask.device),
|
|
torch.arange(width, dtype=mask.dtype, device=mask.device),
|
|
indexing="ij",
|
|
)
|
|
circle = (xx + 0.5 - cx) ** 2 + (yy + 0.5 - cy) ** 2 <= radius ** 2
|
|
mask[circle] = 1.0
|
|
return
|
|
|
|
if shape == "polygon":
|
|
raw_points = region.get("points", [])
|
|
points = []
|
|
for point in raw_points:
|
|
if isinstance(point, dict):
|
|
x = float(point.get("x", 0.0))
|
|
y = float(point.get("y", 0.0))
|
|
else:
|
|
x = float(point[0])
|
|
y = float(point[1])
|
|
points.append(_scale_point(x, y, width, height, normalized))
|
|
mask.copy_(torch.maximum(mask, _polygon_to_mask(points, width, height, mask.device, mask.dtype)))
|
|
return
|
|
|
|
raise ValueError(f"Unsupported region type: {shape}")
|
|
|
|
|
|
def _parse_regions_json(regions_json: str) -> List[dict]:
|
|
text = regions_json.strip()
|
|
if not text:
|
|
return []
|
|
data = json.loads(text)
|
|
if isinstance(data, dict):
|
|
if "regions" in data:
|
|
data = data["regions"]
|
|
else:
|
|
data = [data]
|
|
if not isinstance(data, list):
|
|
raise ValueError("regions_json must be a JSON array or object")
|
|
for region in data:
|
|
if not isinstance(region, dict):
|
|
raise ValueError(f"Region entries must be objects, got: {region}")
|
|
return data
|
|
|
|
|
|
def _apply_keyframe_masks(keyframe_masks: dict, batch: int, height: int, width: int, dtype: torch.dtype, device: torch.device, broadcast_mode: str) -> torch.Tensor:
|
|
mask = torch.zeros((batch, height, width), dtype=dtype, device=device)
|
|
if not keyframe_masks:
|
|
return mask
|
|
|
|
sorted_items = sorted(keyframe_masks.items())
|
|
if broadcast_mode == "repeat_to_batch":
|
|
combined = torch.zeros((height, width), dtype=dtype, device=device)
|
|
for _, key_mask in sorted_items:
|
|
combined = torch.maximum(combined, key_mask)
|
|
return combined.unsqueeze(0).repeat(batch, 1, 1)
|
|
|
|
if broadcast_mode == "hold_until_next_keyframe":
|
|
for idx, (start_frame, key_mask) in enumerate(sorted_items):
|
|
end_frame = sorted_items[idx + 1][0] if idx + 1 < len(sorted_items) else batch
|
|
if end_frame > start_frame:
|
|
mask[start_frame:end_frame] = key_mask
|
|
return mask
|
|
|
|
for frame, key_mask in sorted_items:
|
|
if broadcast_mode == "hold_after_frame":
|
|
mask[frame:] = key_mask
|
|
else:
|
|
mask[frame] = key_mask
|
|
return mask
|
|
|
|
|
|
class URDemoImage:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = CATEGORY
|
|
|
|
def generate(self, width: int, height: int, batch_size: int):
|
|
x = torch.linspace(0.0, 1.0, int(width), dtype=torch.float32).view(1, 1, width)
|
|
y = torch.linspace(0.0, 1.0, int(height), dtype=torch.float32).view(1, height, 1)
|
|
red = 0.15 + 0.35 * x.expand(1, height, width)
|
|
green = 0.18 + 0.45 * y.expand(1, height, width)
|
|
blue = 0.25 + 0.25 * (1.0 - x.expand(1, height, width))
|
|
image = torch.stack([red, green, blue], dim=-1).repeat(int(batch_size), 1, 1, 1)
|
|
|
|
x1 = max(0, min(int(width), int(width * 0.22)))
|
|
x2 = max(0, min(int(width), int(width * 0.55)))
|
|
y1 = max(0, min(int(height), int(height * 0.24)))
|
|
y2 = max(0, min(int(height), int(height * 0.50)))
|
|
image[:, y1:y2, x1:x2, 0] = 0.92
|
|
image[:, y1:y2, x1:x2, 1] = 0.16
|
|
image[:, y1:y2, x1:x2, 2] = 0.12
|
|
return (image.clamp(0.0, 1.0),)
|
|
|
|
|
|
class URRemovePromptBuilder:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"target_prompt": ("STRING", {"multiline": False, "default": "object to remove"}),
|
|
"scene_hint": ("STRING", {"multiline": True, "default": ""}),
|
|
"preserve_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"extra_negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"fill_mode": (["background_reconstruct", "clean_plate", "direct_replace"],),
|
|
"masked_edit_only": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("segmentation_prompt", "remove_positive_prompt", "remove_negative_prompt")
|
|
FUNCTION = "build"
|
|
CATEGORY = CATEGORY
|
|
|
|
def build(
|
|
self,
|
|
target_prompt: str,
|
|
scene_hint: str,
|
|
preserve_prompt: str,
|
|
extra_negative_prompt: str,
|
|
fill_mode: str,
|
|
masked_edit_only: bool,
|
|
):
|
|
target = target_prompt.strip() or "target object"
|
|
segmentation_prompt = target
|
|
|
|
positive_parts = [f"Remove {target} from the masked region."]
|
|
if fill_mode == "background_reconstruct":
|
|
positive_parts.append("Reconstruct a plausible background behind the removed content.")
|
|
elif fill_mode == "clean_plate":
|
|
positive_parts.append("Create a clean plate result that matches the surrounding scene.")
|
|
else:
|
|
positive_parts.append("Replace the masked region with visually coherent scene content.")
|
|
|
|
if masked_edit_only:
|
|
positive_parts.append("Only modify the masked region and keep unmasked content unchanged.")
|
|
if scene_hint.strip():
|
|
positive_parts.append(f"Scene context: {scene_hint.strip()}.")
|
|
if preserve_prompt.strip():
|
|
positive_parts.append(f"Preserve: {preserve_prompt.strip()}.")
|
|
|
|
negative_parts = [
|
|
"artifacts",
|
|
"ghosting",
|
|
"duplicate subject",
|
|
"warped anatomy",
|
|
"broken perspective",
|
|
"extra text",
|
|
"watermark",
|
|
"subtitle",
|
|
"logo",
|
|
"blur",
|
|
"low quality",
|
|
]
|
|
if extra_negative_prompt.strip():
|
|
negative_parts.append(extra_negative_prompt.strip())
|
|
|
|
return (
|
|
segmentation_prompt,
|
|
" ".join(positive_parts),
|
|
", ".join(negative_parts),
|
|
)
|
|
|
|
|
|
class URRegionEditor:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"regions_json": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "[]",
|
|
},
|
|
),
|
|
"canvas_width": ("INT", {"default": 768, "min": 64, "max": 8192, "step": 8}),
|
|
"canvas_height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
|
|
"default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "INT", "INT", "INT")
|
|
RETURN_NAMES = ("regions_json", "canvas_width", "canvas_height", "region_count")
|
|
FUNCTION = "emit"
|
|
CATEGORY = CATEGORY
|
|
|
|
def emit(self, regions_json: str, canvas_width: int, canvas_height: int, default_frame_index: int):
|
|
regions = _parse_regions_json(regions_json)
|
|
normalized_regions = []
|
|
for region in regions:
|
|
copied = dict(region)
|
|
copied["frame_index"] = int(copied.get("frame_index", default_frame_index))
|
|
normalized_regions.append(copied)
|
|
text = json.dumps(normalized_regions, ensure_ascii=False)
|
|
return (text, int(canvas_width), int(canvas_height), len(normalized_regions))
|
|
|
|
|
|
class URBoxesToMaskBatch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"boxes_json": ("STRING", {"multiline": True, "default": "[]"}),
|
|
"frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
|
"coordinate_format": (["xyxy", "xywh"],),
|
|
"broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame"],),
|
|
"normalized": ("BOOLEAN", {"default": False}),
|
|
"preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}),
|
|
"preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("mask", "selected_frame", "preview")
|
|
FUNCTION = "build_masks"
|
|
CATEGORY = CATEGORY
|
|
|
|
def build_masks(
|
|
self,
|
|
images: torch.Tensor,
|
|
boxes_json: str,
|
|
frame_index: int,
|
|
coordinate_format: str,
|
|
broadcast_mode: str,
|
|
normalized: bool,
|
|
preview_color: str,
|
|
preview_opacity: float,
|
|
):
|
|
images = _ensure_image_batch(images)
|
|
batch, height, width, _ = images.shape
|
|
index = max(0, min(int(frame_index), batch - 1))
|
|
|
|
boxes = _parse_boxes_json(boxes_json, coordinate_format, width, height, normalized)
|
|
selected_mask = torch.zeros((1, height, width), dtype=images.dtype, device=images.device)
|
|
for x1, y1, x2, y2 in boxes:
|
|
selected_mask[:, y1:y2, x1:x2] = 1.0
|
|
|
|
if broadcast_mode == "repeat_to_batch":
|
|
mask = selected_mask.repeat(batch, 1, 1)
|
|
elif broadcast_mode == "hold_after_frame":
|
|
mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device)
|
|
mask[index:] = selected_mask
|
|
else:
|
|
mask = torch.zeros((batch, height, width), dtype=images.dtype, device=images.device)
|
|
mask[index] = selected_mask[0]
|
|
|
|
preview = _overlay_mask(
|
|
images[index : index + 1],
|
|
selected_mask,
|
|
_parse_hex_color(preview_color),
|
|
preview_opacity,
|
|
)
|
|
return (mask, images[index : index + 1], preview)
|
|
|
|
|
|
class URRegionsToMaskBatch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"regions_json": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": '[{"type":"rect","frame_index":0,"x":160,"y":120,"w":260,"h":140}]',
|
|
},
|
|
),
|
|
"default_frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
|
"broadcast_mode": (["repeat_to_batch", "single_frame_only", "hold_after_frame", "hold_until_next_keyframe"],),
|
|
"normalized": ("BOOLEAN", {"default": False}),
|
|
"preview_frame": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
|
"preview_color": ("STRING", {"multiline": False, "default": "#FF3B30"}),
|
|
"preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("mask", "selected_frame", "preview")
|
|
FUNCTION = "build_masks"
|
|
CATEGORY = CATEGORY
|
|
|
|
def build_masks(
|
|
self,
|
|
images: torch.Tensor,
|
|
regions_json: str,
|
|
default_frame_index: int,
|
|
broadcast_mode: str,
|
|
normalized: bool,
|
|
preview_frame: int,
|
|
preview_color: str,
|
|
preview_opacity: float,
|
|
):
|
|
images = _ensure_image_batch(images)
|
|
batch, height, width, _ = images.shape
|
|
default_index = max(0, min(int(default_frame_index), batch - 1))
|
|
preview_index = max(0, min(int(preview_frame), batch - 1))
|
|
|
|
keyframe_masks = {}
|
|
for region in _parse_regions_json(regions_json):
|
|
frame = _region_frame_index(region, default_index, batch)
|
|
if frame not in keyframe_masks:
|
|
keyframe_masks[frame] = torch.zeros((height, width), dtype=images.dtype, device=images.device)
|
|
_draw_region(keyframe_masks[frame], region, width, height, normalized)
|
|
|
|
mask = _apply_keyframe_masks(
|
|
keyframe_masks,
|
|
batch,
|
|
height,
|
|
width,
|
|
images.dtype,
|
|
images.device,
|
|
broadcast_mode,
|
|
)
|
|
preview = _overlay_mask(
|
|
images[preview_index : preview_index + 1],
|
|
mask[preview_index : preview_index + 1],
|
|
_parse_hex_color(preview_color),
|
|
preview_opacity,
|
|
)
|
|
return (mask, images[preview_index : preview_index + 1], preview)
|
|
|
|
|
|
class URMaskProcessor:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"mask": ("MASK",),
|
|
"expand_pixels": ("INT", {"default": 0, "min": 0, "max": 2048}),
|
|
"blur_sigma": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 512.0, "step": 0.01}),
|
|
"blur_kernel": ("INT", {"default": 0, "min": 0, "max": 2049}),
|
|
"preview_color": ("STRING", {"multiline": False, "default": "#00FF00"}),
|
|
"preview_opacity": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK", "IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("mask", "mask_image", "preview")
|
|
FUNCTION = "process"
|
|
CATEGORY = CATEGORY
|
|
|
|
def process(
|
|
self,
|
|
mask: torch.Tensor,
|
|
expand_pixels: int,
|
|
blur_sigma: float,
|
|
blur_kernel: int,
|
|
preview_color: str,
|
|
preview_opacity: float,
|
|
image: torch.Tensor = None,
|
|
):
|
|
mask = _ensure_mask_batch(mask)
|
|
if expand_pixels > 0:
|
|
mask = _expand_mask(mask, expand_pixels)
|
|
if blur_sigma > 0.0 or blur_kernel > 1:
|
|
mask = _gaussian_blur_mask(mask, blur_kernel, blur_sigma)
|
|
mask = mask.clamp(0.0, 1.0)
|
|
|
|
mask_image = mask.unsqueeze(-1).repeat(1, 1, 1, 3)
|
|
if image is not None:
|
|
image = _ensure_image_batch(image)
|
|
preview = _overlay_mask(image, mask, _parse_hex_color(preview_color), preview_opacity)
|
|
else:
|
|
preview = _overlay_mask(mask_image, mask, _parse_hex_color(preview_color), preview_opacity)
|
|
return (mask, mask_image, preview)
|
|
|
|
|
|
class URKeyframePrep:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
"frame_index": ("INT", {"default": 0, "min": 0, "max": 1000000}),
|
|
"fill_color": ("STRING", {"multiline": False, "default": "#00FF00"}),
|
|
"resize_longest_side": ("INT", {"default": 0, "min": 0, "max": 8192, "step": 8}),
|
|
"pad_to_multiple": ("INT", {"default": 16, "min": 1, "max": 256}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE", "MASK", "IMAGE", "INT", "INT", "INT")
|
|
RETURN_NAMES = (
|
|
"original_frame",
|
|
"prepared_frame",
|
|
"prepared_mask",
|
|
"filled_frame",
|
|
"prepared_width",
|
|
"prepared_height",
|
|
"frame_count",
|
|
)
|
|
FUNCTION = "prepare"
|
|
CATEGORY = CATEGORY
|
|
|
|
def prepare(
|
|
self,
|
|
images: torch.Tensor,
|
|
mask: torch.Tensor,
|
|
frame_index: int,
|
|
fill_color: str,
|
|
resize_longest_side: int,
|
|
pad_to_multiple: int,
|
|
):
|
|
images = _ensure_image_batch(images)
|
|
mask = _match_mask_to_images(mask, images)
|
|
batch, height, width, _ = images.shape
|
|
index = max(0, min(int(frame_index), batch - 1))
|
|
|
|
original_frame = images[index : index + 1]
|
|
prepared_frame = original_frame.clone()
|
|
prepared_mask = mask[index : index + 1].clone()
|
|
|
|
if resize_longest_side > 0:
|
|
longest = max(height, width)
|
|
scale = float(resize_longest_side) / float(longest)
|
|
target_h = max(1, int(round(height * scale)))
|
|
target_w = max(1, int(round(width * scale)))
|
|
prepared_frame = _resize_images(prepared_frame, (target_h, target_w))
|
|
prepared_mask = _resize_masks(prepared_mask, (target_h, target_w))
|
|
|
|
prepared_h = prepared_frame.shape[1]
|
|
prepared_w = prepared_frame.shape[2]
|
|
if pad_to_multiple > 1:
|
|
target_h = int(math.ceil(prepared_h / pad_to_multiple) * pad_to_multiple)
|
|
target_w = int(math.ceil(prepared_w / pad_to_multiple) * pad_to_multiple)
|
|
prepared_frame = _pad_images(prepared_frame, target_h, target_w)
|
|
prepared_mask = _pad_masks(prepared_mask, target_h, target_w)
|
|
|
|
color = _parse_hex_color(fill_color)
|
|
filled_frame = _fill_masked_region(prepared_frame, prepared_mask, color)
|
|
return (
|
|
original_frame,
|
|
prepared_frame,
|
|
prepared_mask,
|
|
filled_frame,
|
|
prepared_frame.shape[2],
|
|
prepared_frame.shape[1],
|
|
batch,
|
|
)
|
|
|
|
|
|
class URSoftComposite:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"inpaint": ("IMAGE",),
|
|
"original": ("IMAGE",),
|
|
"mask": ("MASK",),
|
|
"kernel": ("INT", {"default": 15, "min": 1, "max": 2049}),
|
|
"sigma": ("FLOAT", {"default": 7.5, "min": 0.0, "max": 512.0, "step": 0.01}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "soft_mask")
|
|
FUNCTION = "composite"
|
|
CATEGORY = CATEGORY
|
|
|
|
def composite(
|
|
self,
|
|
inpaint: torch.Tensor,
|
|
original: torch.Tensor,
|
|
mask: torch.Tensor,
|
|
kernel: int,
|
|
sigma: float,
|
|
):
|
|
original = _ensure_image_batch(original)
|
|
inpaint = _ensure_image_batch(inpaint)
|
|
target_batch = max(original.shape[0], inpaint.shape[0], _ensure_mask_batch(mask).shape[0])
|
|
|
|
original = _tile_to_batch(original, target_batch)
|
|
inpaint = _tile_to_batch(inpaint, target_batch)
|
|
mask = _tile_to_batch(_ensure_mask_batch(mask), target_batch)
|
|
mask = _match_mask_to_images(mask, original)
|
|
|
|
if inpaint.shape[1] != original.shape[1] or inpaint.shape[2] != original.shape[2]:
|
|
inpaint = _resize_images(inpaint, (original.shape[1], original.shape[2]))
|
|
|
|
soft_mask = _gaussian_blur_mask(mask, kernel, sigma).clamp(0.0, 1.0)
|
|
composited = original * (1.0 - soft_mask.unsqueeze(-1)) + inpaint * soft_mask.unsqueeze(-1)
|
|
return (composited.clamp(0.0, 1.0), soft_mask)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"UR Demo Image": URDemoImage,
|
|
"UR Remove Prompt Builder": URRemovePromptBuilder,
|
|
"UR Region Editor": URRegionEditor,
|
|
"UR Boxes To Mask Batch": URBoxesToMaskBatch,
|
|
"UR Regions To Mask Batch": URRegionsToMaskBatch,
|
|
"UR Mask Processor": URMaskProcessor,
|
|
"UR Keyframe Prep": URKeyframePrep,
|
|
"UR Soft Composite": URSoftComposite,
|
|
}
|
|
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"UR Demo Image": "UR Demo Image",
|
|
"UR Remove Prompt Builder": "UR Remove Prompt Builder",
|
|
"UR Region Editor": "UR Region Editor",
|
|
"UR Boxes To Mask Batch": "UR Boxes To Mask Batch",
|
|
"UR Regions To Mask Batch": "UR Regions To Mask Batch",
|
|
"UR Mask Processor": "UR Mask Processor",
|
|
"UR Keyframe Prep": "UR Keyframe Prep",
|
|
"UR Soft Composite": "UR Soft Composite",
|
|
}
|