245 lines
7.6 KiB
Python
245 lines
7.6 KiB
Python
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
MODULE_PATH = Path(__file__).resolve().parents[1] / "nodes.py"
|
|
SPEC = importlib.util.spec_from_file_location("universal_remove_nodes", MODULE_PATH)
|
|
MODULE = importlib.util.module_from_spec(SPEC)
|
|
assert SPEC is not None and SPEC.loader is not None
|
|
SPEC.loader.exec_module(MODULE)
|
|
|
|
|
|
def test_extension_entrypoint_exports_nodes_and_web_directory():
|
|
package_path = MODULE_PATH.parent / "__init__.py"
|
|
spec = importlib.util.spec_from_file_location(
|
|
"ComfyUI_UniversalRemove_test",
|
|
package_path,
|
|
submodule_search_locations=[str(MODULE_PATH.parent)],
|
|
)
|
|
package = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = package
|
|
assert spec is not None and spec.loader is not None
|
|
spec.loader.exec_module(package)
|
|
|
|
assert package.WEB_DIRECTORY == "./web"
|
|
assert "UR Region Editor" in package.NODE_CLASS_MAPPINGS
|
|
assert "UR Regions To Mask Batch" in package.NODE_CLASS_MAPPINGS
|
|
|
|
|
|
def test_prompt_builder_outputs_expected_strings():
|
|
node = MODULE.URRemovePromptBuilder()
|
|
segmentation_prompt, positive_prompt, negative_prompt = node.build(
|
|
target_prompt="watermark",
|
|
scene_hint="night city street",
|
|
preserve_prompt="keep the lighting and camera framing",
|
|
extra_negative_prompt="text remnants",
|
|
fill_mode="background_reconstruct",
|
|
masked_edit_only=True,
|
|
)
|
|
|
|
assert segmentation_prompt == "watermark"
|
|
assert "Remove watermark" in positive_prompt
|
|
assert "night city street" in positive_prompt
|
|
assert "keep the lighting and camera framing" in positive_prompt
|
|
assert "text remnants" in negative_prompt
|
|
|
|
|
|
def test_demo_image_is_self_contained_and_batched():
|
|
node = MODULE.URDemoImage()
|
|
(image,) = node.generate(width=64, height=32, batch_size=2)
|
|
|
|
assert image.shape == (2, 32, 64, 3)
|
|
assert image.min().item() >= 0.0
|
|
assert image.max().item() <= 1.0
|
|
|
|
|
|
def test_region_editor_normalizes_regions_and_counts_them():
|
|
node = MODULE.URRegionEditor()
|
|
regions_json, canvas_width, canvas_height, region_count = node.emit(
|
|
regions_json='[{"type":"rect","x":1,"y":2,"w":3,"h":4}]',
|
|
canvas_width=64,
|
|
canvas_height=32,
|
|
default_frame_index=5,
|
|
)
|
|
|
|
assert '"frame_index": 5' in regions_json
|
|
assert canvas_width == 64
|
|
assert canvas_height == 32
|
|
assert region_count == 1
|
|
|
|
|
|
def test_boxes_to_mask_batch_repeats_mask_to_batch():
|
|
node = MODULE.URBoxesToMaskBatch()
|
|
images = torch.zeros((2, 4, 4, 3), dtype=torch.float32)
|
|
|
|
mask, selected_frame, preview = node.build_masks(
|
|
images=images,
|
|
boxes_json="[[1, 1, 3, 3]]",
|
|
frame_index=0,
|
|
coordinate_format="xyxy",
|
|
broadcast_mode="repeat_to_batch",
|
|
normalized=False,
|
|
preview_color="#FF0000",
|
|
preview_opacity=0.5,
|
|
)
|
|
|
|
expected = torch.tensor(
|
|
[
|
|
[0, 0, 0, 0],
|
|
[0, 1, 1, 0],
|
|
[0, 1, 1, 0],
|
|
[0, 0, 0, 0],
|
|
],
|
|
dtype=torch.float32,
|
|
)
|
|
assert mask.shape == (2, 4, 4)
|
|
assert torch.equal(mask[0], expected)
|
|
assert torch.equal(mask[1], expected)
|
|
assert selected_frame.shape == (1, 4, 4, 3)
|
|
assert preview.shape == (1, 4, 4, 3)
|
|
|
|
|
|
def test_boxes_to_mask_batch_handles_reversed_out_of_bounds_box():
|
|
node = MODULE.URBoxesToMaskBatch()
|
|
images = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
|
|
|
mask, _, _ = node.build_masks(
|
|
images=images,
|
|
boxes_json="[[5, 5, -1, -1]]",
|
|
frame_index=0,
|
|
coordinate_format="xyxy",
|
|
broadcast_mode="single_frame_only",
|
|
normalized=False,
|
|
preview_color="#FF0000",
|
|
preview_opacity=0.5,
|
|
)
|
|
|
|
assert mask.shape == (1, 4, 4)
|
|
assert mask.sum().item() == 16.0
|
|
|
|
|
|
def test_regions_to_mask_batch_supports_rect_circle_and_polygon():
|
|
node = MODULE.URRegionsToMaskBatch()
|
|
images = torch.zeros((1, 10, 10, 3), dtype=torch.float32)
|
|
regions_json = """
|
|
[
|
|
{"type": "rect", "frame_index": 0, "x": 1, "y": 1, "w": 2, "h": 2},
|
|
{"type": "circle", "frame_index": 0, "cx": 7, "cy": 7, "r": 1.5},
|
|
{"type": "polygon", "frame_index": 0, "points": [[5, 1], [8, 1], [8, 3], [5, 3]]}
|
|
]
|
|
"""
|
|
|
|
mask, selected_frame, preview = node.build_masks(
|
|
images=images,
|
|
regions_json=regions_json,
|
|
default_frame_index=0,
|
|
broadcast_mode="single_frame_only",
|
|
normalized=False,
|
|
preview_frame=0,
|
|
preview_color="#FF0000",
|
|
preview_opacity=0.5,
|
|
)
|
|
|
|
assert mask.shape == (1, 10, 10)
|
|
assert mask[0, 1:3, 1:3].sum().item() == 4.0
|
|
assert mask[0, 7, 7].item() == 1.0
|
|
assert mask[0, 1:3, 5:8].sum().item() > 0.0
|
|
assert selected_frame.shape == (1, 10, 10, 3)
|
|
assert preview.shape == (1, 10, 10, 3)
|
|
|
|
|
|
def test_regions_to_mask_batch_holds_between_keyframes():
|
|
node = MODULE.URRegionsToMaskBatch()
|
|
images = torch.zeros((4, 6, 6, 3), dtype=torch.float32)
|
|
regions_json = """
|
|
[
|
|
{"type": "rect", "frame_index": 0, "x": 0, "y": 0, "w": 2, "h": 2},
|
|
{"type": "rect", "frame_index": 2, "x": 4, "y": 4, "w": 2, "h": 2}
|
|
]
|
|
"""
|
|
|
|
mask, _, _ = node.build_masks(
|
|
images=images,
|
|
regions_json=regions_json,
|
|
default_frame_index=0,
|
|
broadcast_mode="hold_until_next_keyframe",
|
|
normalized=False,
|
|
preview_frame=0,
|
|
preview_color="#FF0000",
|
|
preview_opacity=0.5,
|
|
)
|
|
|
|
assert mask.shape == (4, 6, 6)
|
|
assert mask[0, 0:2, 0:2].sum().item() == 4.0
|
|
assert mask[1, 0:2, 0:2].sum().item() == 4.0
|
|
assert mask[2, 4:6, 4:6].sum().item() == 4.0
|
|
assert mask[3, 4:6, 4:6].sum().item() == 4.0
|
|
assert mask[1, 4:6, 4:6].sum().item() == 0.0
|
|
|
|
|
|
def test_mask_processor_can_expand_mask_without_image():
|
|
node = MODULE.URMaskProcessor()
|
|
mask = torch.zeros((1, 5, 5), dtype=torch.float32)
|
|
mask[:, 2, 2] = 1.0
|
|
|
|
processed_mask, mask_image, preview = node.process(
|
|
mask=mask,
|
|
expand_pixels=1,
|
|
blur_sigma=0.0,
|
|
blur_kernel=0,
|
|
preview_color="#00FF00",
|
|
preview_opacity=0.5,
|
|
image=None,
|
|
)
|
|
|
|
assert processed_mask.sum().item() == 9.0
|
|
assert mask_image.shape == (1, 5, 5, 3)
|
|
assert preview.shape == (1, 5, 5, 3)
|
|
|
|
|
|
def test_keyframe_prep_fills_masked_region_and_pads():
|
|
node = MODULE.URKeyframePrep()
|
|
images = torch.zeros((1, 3, 5, 3), dtype=torch.float32)
|
|
mask = torch.zeros((1, 3, 5), dtype=torch.float32)
|
|
mask[:, 1, 2] = 1.0
|
|
|
|
original_frame, prepared_frame, prepared_mask, filled_frame, prepared_width, prepared_height, frame_count = node.prepare(
|
|
images=images,
|
|
mask=mask,
|
|
frame_index=0,
|
|
fill_color="#00FF00",
|
|
resize_longest_side=0,
|
|
pad_to_multiple=4,
|
|
)
|
|
|
|
assert original_frame.shape == (1, 3, 5, 3)
|
|
assert prepared_frame.shape == (1, 4, 8, 3)
|
|
assert prepared_mask.shape == (1, 4, 8)
|
|
assert prepared_width == 8
|
|
assert prepared_height == 4
|
|
assert frame_count == 1
|
|
assert torch.allclose(filled_frame[0, 1, 2], torch.tensor([0.0, 1.0, 0.0]))
|
|
|
|
|
|
def test_soft_composite_uses_mask_for_direct_blend_when_no_blur():
|
|
node = MODULE.URSoftComposite()
|
|
original = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
|
|
inpaint = torch.ones((1, 4, 4, 3), dtype=torch.float32)
|
|
mask = torch.zeros((1, 4, 4), dtype=torch.float32)
|
|
mask[:, 1:3, 1:3] = 1.0
|
|
|
|
composited, soft_mask = node.composite(
|
|
inpaint=inpaint,
|
|
original=original,
|
|
mask=mask,
|
|
kernel=1,
|
|
sigma=0.0,
|
|
)
|
|
|
|
assert torch.equal(soft_mask, mask)
|
|
assert torch.allclose(composited[0, 0, 0], torch.tensor([0.0, 0.0, 0.0]))
|
|
assert torch.allclose(composited[0, 1, 1], torch.tensor([1.0, 1.0, 1.0]))
|