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]))