修改图像拼接节点名字

This commit is contained in:
Gin3601
2026-05-25 18:44:25 +08:00
parent 281a1b8289
commit 9d95ac58b1
3 changed files with 51 additions and 17 deletions
+25 -2
View File
@@ -31,6 +31,29 @@ class TestConcatByTargetNode(unittest.TestCase):
self.assertLessEqual(float(image.max().item()), 1.0)
self.assertTrue(torch.allclose(mask, torch.zeros_like(mask)))
def test_new_port_names_are_accepted(self):
image_a = self.make_image(width=2, height=2, value=0.2)
image_b = self.make_image(width=1, height=2, value=0.8)
image, mask, width, height = self.node.concat(
data_type="IMAGE",
target="A",
position="right",
match_target_size=True,
resize_mode="none",
align="center",
gap=0,
background_color="#000000",
background_value=0.0,
multiple_of=0,
allow_batch_broadcast=True,
**{"image/mask_A": image_a, "image/mask_B": image_b},
)
self.assert_image_output(image, mask, width, height, (1, 2, 3, 3))
self.assertTrue(torch.allclose(image[:, :, 0:2, :], torch.full((1, 2, 2, 3), 0.2)))
self.assertTrue(torch.allclose(image[:, :, 2:3, :], torch.full((1, 2, 1, 3), 0.8)))
def test_image_mode_target_b_right_places_a_to_the_right_of_b(self):
image_a = self.make_image(width=2, height=2, value=0.2)
image_b = self.make_image(width=3, height=2, value=0.8)
@@ -229,7 +252,7 @@ class TestConcatByTargetNode(unittest.TestCase):
)
def test_missing_required_inputs_for_current_mode_raise_clear_errors(self):
with self.assertRaisesRegex(RuntimeError, "inputs a and b"):
with self.assertRaisesRegex(RuntimeError, "image/mask_A and image/mask_B"):
self.node.concat(
data_type="IMAGE",
target="A",
@@ -246,7 +269,7 @@ class TestConcatByTargetNode(unittest.TestCase):
b=None,
)
with self.assertRaisesRegex(RuntimeError, "inputs a and b"):
with self.assertRaisesRegex(RuntimeError, "image/mask_A and image/mask_B"):
self.node.concat(
data_type="MASK",
target="A",
+4 -2
View File
@@ -41,8 +41,10 @@ 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["a"][0], "IMAGE,MASK")
self.assertEqual(optional["b"][0], "IMAGE,MASK")
self.assertEqual(optional["image/mask_A"][0], "IMAGE,MASK")
self.assertEqual(optional["image/mask_B"][0], "IMAGE,MASK")
self.assertNotIn("a", optional)
self.assertNotIn("b", optional)
self.assertNotIn("image_a", optional)
self.assertNotIn("image_b", optional)
self.assertNotIn("mask_a", optional)
+22 -13
View File
@@ -17,6 +17,8 @@ POSITION_CHOICES = ["top", "bottom", "left", "right"]
RESIZE_MODE_CHOICES = ["keep_proportion", "stretch", "none"]
ALIGN_CHOICES = ["start", "center", "end"]
INPUT_DATA_CHOICES = "IMAGE,MASK"
PORT_A_NAME = "image/mask_A"
PORT_B_NAME = "image/mask_B"
def _require_torch():
@@ -57,8 +59,8 @@ def resolve_concat_data_type(a, b, requested_data_type: str):
if len(inferred) > 1:
raise RuntimeError(
f"[LLS] Inputs a and b must both resolve to IMAGE or both resolve to MASK. "
f"Detected a={inferred_a or 'unknown'}, b={inferred_b or 'unknown'}."
f"[LLS] Inputs {PORT_A_NAME} and {PORT_B_NAME} must both resolve to IMAGE or both resolve to MASK. "
f"Detected {PORT_A_NAME}={inferred_a or 'unknown'}, {PORT_B_NAME}={inferred_b or 'unknown'}."
)
if len(inferred) == 1:
return inferred.pop()
@@ -450,8 +452,8 @@ class LLSConcatByTarget:
"allow_batch_broadcast": ("BOOLEAN", {"default": True}),
},
"optional": {
"a": (INPUT_DATA_CHOICES,),
"b": (INPUT_DATA_CHOICES,),
PORT_A_NAME: (INPUT_DATA_CHOICES,),
PORT_B_NAME: (INPUT_DATA_CHOICES,),
},
}
@@ -470,18 +472,23 @@ class LLSConcatByTarget:
allow_batch_broadcast,
a=None,
b=None,
**kwargs,
):
_require_torch()
data_type = str(data_type or "IMAGE")
background_value = max(0.0, min(1.0, float(background_value)))
resolved_data_type = resolve_concat_data_type(a, b, data_type)
input_a = kwargs[PORT_A_NAME] if PORT_A_NAME in kwargs else a
input_b = kwargs[PORT_B_NAME] if PORT_B_NAME in kwargs else b
resolved_data_type = resolve_concat_data_type(input_a, input_b, data_type)
if resolved_data_type == "IMAGE":
if a is None or b is None:
raise RuntimeError("[LLS] IMAGE mode requires both inputs a and b to be connected.")
if input_a is None or input_b is None:
raise RuntimeError(
f"[LLS] IMAGE mode requires both inputs {PORT_A_NAME} and {PORT_B_NAME} to be connected."
)
color = parse_hex_color(background_color)
tensor_a = ensure_image_tensor(a)
tensor_b = ensure_image_tensor(b)
tensor_a = ensure_image_tensor(input_a)
tensor_b = ensure_image_tensor(input_b)
output_image, width, height = concat_by_target(
data_type="IMAGE",
tensor_a=tensor_a,
@@ -506,10 +513,12 @@ class LLSConcatByTarget:
return output_image, output_mask, int(width), int(height)
if resolved_data_type == "MASK":
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)
if input_a is None or input_b is None:
raise RuntimeError(
f"[LLS] MASK mode requires both inputs {PORT_A_NAME} and {PORT_B_NAME} to be connected."
)
tensor_a = ensure_mask_tensor(input_a)
tensor_b = ensure_mask_tensor(input_b)
output_mask, width, height = concat_by_target(
data_type="MASK",
tensor_a=tensor_a,