nodes/Gin3601-LLS-node
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

修改拼接节点的端口，支持双模式

Browse Source
This commit is contained in:
Gin3601
2026-05-25 17:29:26 +08:00
parent 84e02629e6
commit 58a1564e4e
3 changed files with 52 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_concat_by_target_node.py
+20 -38
View File
@@ -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,
)
tests/test_concat_by_target_registration.py
+6 -4
View File
@@ -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__":
utils/concat_by_target.py
+26 -18
View File
@@ -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,