From 58a1564e4e1f03d4619f6779d5eca40eaaaebe2a Mon Sep 17 00:00:00 2001 From: Gin3601 Date: Mon, 25 May 2026 17:29:26 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=8B=BC=E6=8E=A5=E8=8A=82?= =?UTF-8?q?=E7=82=B9=E7=9A=84=E7=AB=AF=E5=8F=A3=EF=BC=8C=E6=94=AF=E6=8C=81?= =?UTF-8?q?=E5=8F=8C=E6=A8=A1=E5=BC=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/test_concat_by_target_node.py | 58 +++++++-------------- tests/test_concat_by_target_registration.py | 10 ++-- utils/concat_by_target.py | 44 +++++++++------- 3 files changed, 52 insertions(+), 60 deletions(-) diff --git a/tests/test_concat_by_target_node.py b/tests/test_concat_by_target_node.py index def854c..744ee9f 100644 --- a/tests/test_concat_by_target_node.py +++ b/tests/test_concat_by_target_node.py @@ -47,10 +47,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=0, allow_batch_broadcast=True, - image_a=image_a, - image_b=image_b, - mask_a=None, - mask_b=None, + a=image_a, + b=image_b, ) self.assert_image_output(image, mask, width, height, (1, 2, 6, 3)) @@ -74,10 +72,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=4, allow_batch_broadcast=True, - image_a=image_a, - image_b=image_b, - mask_a=None, - mask_b=None, + a=image_a, + b=image_b, ) self.assert_image_output(image, mask, width, height, (1, 4, 12, 3)) @@ -102,10 +98,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.25, multiple_of=0, allow_batch_broadcast=True, - image_a=None, - image_b=None, - mask_a=mask_a, - mask_b=mask_b, + a=mask_a, + b=mask_b, ) self.assertEqual((width, height), (2, 5)) @@ -134,10 +128,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=0, allow_batch_broadcast=True, - image_a=image_a, - image_b=image_b, - mask_a=None, - mask_b=None, + a=image_a, + b=image_b, ) self.assert_image_output(image, mask, width, height, (2, 2, 4, 3)) @@ -160,10 +152,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=0, allow_batch_broadcast=True, - image_a=image_a, - image_b=image_b, - mask_a=None, - mask_b=None, + a=image_a, + b=image_b, ) self.assert_image_output(image, mask, width, height, (1, 2, 5, 3)) @@ -171,7 +161,7 @@ class TestConcatByTargetNode(unittest.TestCase): self.assertTrue(torch.allclose(image[:, :, 2:5, :], torch.full((1, 2, 3, 3), 0.9))) def test_missing_required_inputs_for_current_mode_raise_clear_errors(self): - with self.assertRaisesRegex(RuntimeError, "image_a and image_b"): + with self.assertRaisesRegex(RuntimeError, "inputs a and b"): self.node.concat( data_type="IMAGE", target="A", @@ -184,13 +174,11 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=0, allow_batch_broadcast=True, - image_a=None, - image_b=None, - mask_a=self.make_mask(1, 1, 0.0), - mask_b=self.make_mask(1, 1, 0.0), + a=None, + b=None, ) - with self.assertRaisesRegex(RuntimeError, "mask_a and mask_b"): + with self.assertRaisesRegex(RuntimeError, "inputs a and b"): self.node.concat( data_type="MASK", target="A", @@ -203,10 +191,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=0, allow_batch_broadcast=True, - image_a=self.make_image(1, 1, 0.0), - image_b=self.make_image(1, 1, 0.0), - mask_a=None, - mask_b=None, + a=None, + b=None, ) def test_invalid_background_color_and_batch_mismatch_raise_clear_errors(self): @@ -226,10 +212,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=0, allow_batch_broadcast=True, - image_a=self.make_image(2, 2, 0.1), - image_b=self.make_image(2, 2, 0.9), - mask_a=None, - mask_b=None, + a=self.make_image(2, 2, 0.1), + b=self.make_image(2, 2, 0.9), ) with self.assertRaisesRegex(RuntimeError, "Batch size mismatch"): @@ -245,10 +229,8 @@ class TestConcatByTargetNode(unittest.TestCase): background_value=0.0, multiple_of=0, allow_batch_broadcast=False, - image_a=image_a, - image_b=image_b, - mask_a=None, - mask_b=None, + a=image_a, + b=image_b, ) diff --git a/tests/test_concat_by_target_registration.py b/tests/test_concat_by_target_registration.py index c2f7444..39f0f23 100644 --- a/tests/test_concat_by_target_registration.py +++ b/tests/test_concat_by_target_registration.py @@ -41,10 +41,12 @@ class TestConcatByTargetRegistration(unittest.TestCase): self.assertEqual(required["background_value"][1]["default"], 0.0) self.assertEqual(required["multiple_of"][0], "INT") self.assertEqual(required["allow_batch_broadcast"][0], "BOOLEAN") - self.assertEqual(optional["image_a"], ("IMAGE",)) - self.assertEqual(optional["image_b"], ("IMAGE",)) - self.assertEqual(optional["mask_a"], ("MASK",)) - self.assertEqual(optional["mask_b"], ("MASK",)) + self.assertEqual(str(optional["a"][0]), "*") + self.assertEqual(str(optional["b"][0]), "*") + self.assertNotIn("image_a", optional) + self.assertNotIn("image_b", optional) + self.assertNotIn("mask_a", optional) + self.assertNotIn("mask_b", optional) if __name__ == "__main__": diff --git a/utils/concat_by_target.py b/utils/concat_by_target.py index cd37132..b9e426c 100644 --- a/utils/concat_by_target.py +++ b/utils/concat_by_target.py @@ -18,6 +18,14 @@ RESIZE_MODE_CHOICES = ["keep_proportion", "stretch", "none"] ALIGN_CHOICES = ["start", "center", "end"] +class AnyType(str): + def __ne__(self, other): # pragma: no cover - exercised by ComfyUI type checks + return False + + +ANY_TYPE = AnyType("*") + + def _require_torch(): if torch is None or torch_nn_functional is None: raise RuntimeError("[LLS] torch is required for LLS Concat By Target.") from _TORCH_ERR @@ -50,16 +58,20 @@ def parse_hex_color(color_str: str): def ensure_image_tensor(x): _require_torch() if x is None: - raise RuntimeError("[LLS] IMAGE mode requires both image_a and image_b inputs.") + raise RuntimeError("[LLS] IMAGE mode requires both inputs a and b to be connected.") if not isinstance(x, torch.Tensor): x = torch.as_tensor(x) if x.ndim == 3: + if int(x.shape[-1]) not in {1, 3, 4}: + raise RuntimeError("[LLS] IMAGE input must have shape [B,H,W,C] or [H,W,C].") x = x.unsqueeze(0) if x.ndim != 4: raise RuntimeError("[LLS] IMAGE input must have shape [B,H,W,C] or [H,W,C].") batch, height, width, channels = [int(dim) for dim in x.shape] if batch <= 0 or height <= 0 or width <= 0 or channels <= 0: raise RuntimeError("[LLS] IMAGE input must have positive batch, height, width, and channels.") + if channels not in {1, 3, 4}: + raise RuntimeError("[LLS] IMAGE input channel count must be 1, 3, or 4.") if not torch.is_floating_point(x): x = x.to(dtype=torch.float32) return x @@ -68,7 +80,7 @@ def ensure_image_tensor(x): def ensure_mask_tensor(x): _require_torch() if x is None: - raise RuntimeError("[LLS] MASK mode requires both mask_a and mask_b inputs.") + raise RuntimeError("[LLS] MASK mode requires both inputs a and b to be connected.") if not isinstance(x, torch.Tensor): x = torch.as_tensor(x) if x.ndim == 2: @@ -409,10 +421,8 @@ class LLSConcatByTarget: "allow_batch_broadcast": ("BOOLEAN", {"default": True}), }, "optional": { - "image_a": ("IMAGE",), - "image_b": ("IMAGE",), - "mask_a": ("MASK",), - "mask_b": ("MASK",), + "a": (ANY_TYPE,), + "b": (ANY_TYPE,), }, } @@ -429,21 +439,19 @@ class LLSConcatByTarget: background_value, multiple_of, allow_batch_broadcast, - image_a=None, - image_b=None, - mask_a=None, - mask_b=None, + a=None, + b=None, ): _require_torch() data_type = str(data_type or "IMAGE") background_value = max(0.0, min(1.0, float(background_value))) if data_type == "IMAGE": - if image_a is None or image_b is None: - raise RuntimeError("[LLS] IMAGE mode requires both image_a and image_b inputs.") + if a is None or b is None: + raise RuntimeError("[LLS] IMAGE mode requires both inputs a and b to be connected.") color = parse_hex_color(background_color) - tensor_a = ensure_image_tensor(image_a) - tensor_b = ensure_image_tensor(image_b) + tensor_a = ensure_image_tensor(a) + tensor_b = ensure_image_tensor(b) output_image, width, height = concat_by_target( data_type="IMAGE", tensor_a=tensor_a, @@ -468,10 +476,10 @@ class LLSConcatByTarget: return output_image, output_mask, int(width), int(height) if data_type == "MASK": - if mask_a is None or mask_b is None: - raise RuntimeError("[LLS] MASK mode requires both mask_a and mask_b inputs.") - tensor_a = ensure_mask_tensor(mask_a) - tensor_b = ensure_mask_tensor(mask_b) + if a is None or b is None: + raise RuntimeError("[LLS] MASK mode requires both inputs a and b to be connected.") + tensor_a = ensure_mask_tensor(a) + tensor_b = ensure_mask_tensor(b) output_mask, width, height = concat_by_target( data_type="MASK", tensor_a=tensor_a,