开发图像拼接节点
This commit is contained in:
Gin3601
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,230 @@
|
||||
import unittest
|
||||
|
||||
try:
|
||||
from .test_mask_draw_helpers import load_plugin_package, torch
|
||||
except ImportError:
|
||||
from test_mask_draw_helpers import load_plugin_package, torch
|
||||
|
||||
|
||||
HAS_CONCAT_RUNTIME_DEPS = torch is not None
|
||||
CONCAT_RUNTIME_DEPS_MESSAGE = "concat node tests require torch"
|
||||
|
||||
|
||||
@unittest.skipUnless(HAS_CONCAT_RUNTIME_DEPS, CONCAT_RUNTIME_DEPS_MESSAGE)
|
||||
class TestConcatByTargetNode(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
plugin = load_plugin_package()
|
||||
cls.node = plugin.NODE_CLASS_MAPPINGS["LLSConcatByTarget"]()
|
||||
|
||||
def make_image(self, width, height, value, batch=1):
|
||||
return torch.full((batch, height, width, 3), float(value), dtype=torch.float32)
|
||||
|
||||
def make_mask(self, width, height, value, batch=1):
|
||||
return torch.full((batch, height, width), float(value), dtype=torch.float32)
|
||||
|
||||
def assert_image_output(self, image, mask, width, height, shape):
|
||||
self.assertEqual((width, height), (shape[2], shape[1]))
|
||||
self.assertEqual(tuple(image.shape), shape)
|
||||
self.assertEqual(tuple(mask.shape), (shape[0], shape[1], shape[2]))
|
||||
self.assertGreaterEqual(float(image.min().item()), 0.0)
|
||||
self.assertLessEqual(float(image.max().item()), 1.0)
|
||||
self.assertTrue(torch.allclose(mask, torch.zeros_like(mask)))
|
||||
|
||||
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)
|
||||
|
||||
image, mask, width, height = self.node.concat(
|
||||
data_type="IMAGE",
|
||||
target="B",
|
||||
position="right",
|
||||
match_target_size=True,
|
||||
resize_mode="none",
|
||||
align="center",
|
||||
gap=1,
|
||||
background_color="#000000",
|
||||
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,
|
||||
)
|
||||
|
||||
self.assert_image_output(image, mask, width, height, (1, 2, 6, 3))
|
||||
self.assertTrue(torch.allclose(image[:, :, 0:3, :], torch.full((1, 2, 3, 3), 0.8)))
|
||||
self.assertTrue(torch.allclose(image[:, :, 3:4, :], torch.zeros((1, 2, 1, 3))))
|
||||
self.assertTrue(torch.allclose(image[:, :, 4:6, :], torch.full((1, 2, 2, 3), 0.2)))
|
||||
|
||||
def test_image_mode_keep_proportion_gap_and_multiple_of(self):
|
||||
image_a = self.make_image(width=4, height=4, value=0.25)
|
||||
image_b = self.make_image(width=2, height=2, value=0.75)
|
||||
|
||||
image, mask, width, height = self.node.concat(
|
||||
data_type="IMAGE",
|
||||
target="A",
|
||||
position="right",
|
||||
match_target_size=True,
|
||||
resize_mode="keep_proportion",
|
||||
align="center",
|
||||
gap=1,
|
||||
background_color="#808080",
|
||||
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,
|
||||
)
|
||||
|
||||
self.assert_image_output(image, mask, width, height, (1, 4, 12, 3))
|
||||
self.assertTrue(torch.allclose(image[:, :, 0:4, :], torch.full((1, 4, 4, 3), 0.25)))
|
||||
self.assertTrue(torch.allclose(image[:, :, 4:5, :], torch.full((1, 4, 1, 3), 128.0 / 255.0)))
|
||||
self.assertTrue(torch.allclose(image[:, :, 5:9, :], torch.full((1, 4, 4, 3), 0.75)))
|
||||
self.assertTrue(torch.allclose(image[:, :, 9:12, :], torch.full((1, 4, 3, 3), 128.0 / 255.0)))
|
||||
|
||||
def test_mask_mode_outputs_mask_and_preview_image(self):
|
||||
mask_a = self.make_mask(width=2, height=2, value=1.0)
|
||||
mask_b = self.make_mask(width=2, height=2, value=0.0)
|
||||
|
||||
image, mask, width, height = self.node.concat(
|
||||
data_type="MASK",
|
||||
target="A",
|
||||
position="bottom",
|
||||
match_target_size=True,
|
||||
resize_mode="none",
|
||||
align="center",
|
||||
gap=1,
|
||||
background_color="#000000",
|
||||
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,
|
||||
)
|
||||
|
||||
self.assertEqual((width, height), (2, 5))
|
||||
self.assertEqual(tuple(mask.shape), (1, 5, 2))
|
||||
self.assertEqual(tuple(image.shape), (1, 5, 2, 3))
|
||||
self.assertTrue(torch.allclose(mask[:, 0:2, :], torch.full((1, 2, 2), 1.0)))
|
||||
self.assertTrue(torch.allclose(mask[:, 2:3, :], torch.full((1, 1, 2), 0.25)))
|
||||
self.assertTrue(torch.allclose(mask[:, 3:5, :], torch.zeros((1, 2, 2))))
|
||||
self.assertTrue(torch.allclose(image[..., 0], mask))
|
||||
self.assertTrue(torch.allclose(image[..., 1], mask))
|
||||
self.assertTrue(torch.allclose(image[..., 2], mask))
|
||||
|
||||
def test_batch_broadcast_works_when_enabled(self):
|
||||
image_a = self.make_image(width=2, height=2, value=0.1, batch=1)
|
||||
image_b = self.make_image(width=2, height=2, value=0.9, batch=2)
|
||||
|
||||
image, mask, width, height = self.node.concat(
|
||||
data_type="IMAGE",
|
||||
target="A",
|
||||
position="left",
|
||||
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_a=image_a,
|
||||
image_b=image_b,
|
||||
mask_a=None,
|
||||
mask_b=None,
|
||||
)
|
||||
|
||||
self.assert_image_output(image, mask, width, height, (2, 2, 4, 3))
|
||||
self.assertTrue(torch.allclose(image[:, :, 0:2, :], torch.full((2, 2, 2, 3), 0.9)))
|
||||
self.assertTrue(torch.allclose(image[:, :, 2:4, :], torch.full((2, 2, 2, 3), 0.1)))
|
||||
|
||||
def test_missing_required_inputs_for_current_mode_raise_clear_errors(self):
|
||||
with self.assertRaisesRegex(RuntimeError, "image_a and image_b"):
|
||||
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_a=None,
|
||||
image_b=None,
|
||||
mask_a=self.make_mask(1, 1, 0.0),
|
||||
mask_b=self.make_mask(1, 1, 0.0),
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "mask_a and mask_b"):
|
||||
self.node.concat(
|
||||
data_type="MASK",
|
||||
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_a=self.make_image(1, 1, 0.0),
|
||||
image_b=self.make_image(1, 1, 0.0),
|
||||
mask_a=None,
|
||||
mask_b=None,
|
||||
)
|
||||
|
||||
def test_invalid_background_color_and_batch_mismatch_raise_clear_errors(self):
|
||||
image_a = self.make_image(width=2, height=2, value=0.1, batch=2)
|
||||
image_b = self.make_image(width=2, height=2, value=0.9, batch=3)
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "background_color"):
|
||||
self.node.concat(
|
||||
data_type="IMAGE",
|
||||
target="A",
|
||||
position="right",
|
||||
match_target_size=True,
|
||||
resize_mode="none",
|
||||
align="center",
|
||||
gap=0,
|
||||
background_color="bad-color",
|
||||
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,
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "Batch size mismatch"):
|
||||
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=False,
|
||||
image_a=image_a,
|
||||
image_b=image_b,
|
||||
mask_a=None,
|
||||
mask_b=None,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,51 @@
|
||||
import unittest
|
||||
|
||||
try:
|
||||
from .test_mask_draw_helpers import load_plugin_package
|
||||
except ImportError:
|
||||
from test_mask_draw_helpers import load_plugin_package
|
||||
|
||||
|
||||
class TestConcatByTargetRegistration(unittest.TestCase):
|
||||
def test_plugin_registers_concat_by_target_node(self):
|
||||
plugin = load_plugin_package()
|
||||
|
||||
self.assertIn("LLSConcatByTarget", plugin.NODE_CLASS_MAPPINGS)
|
||||
self.assertEqual(
|
||||
plugin.NODE_DISPLAY_NAME_MAPPINGS["LLSConcatByTarget"],
|
||||
"LLS Concat By Target",
|
||||
)
|
||||
|
||||
def test_concat_by_target_schema_matches_contract(self):
|
||||
plugin = load_plugin_package()
|
||||
node_cls = plugin.NODE_CLASS_MAPPINGS["LLSConcatByTarget"]
|
||||
schema = node_cls.INPUT_TYPES()
|
||||
required = schema["required"]
|
||||
optional = schema["optional"]
|
||||
|
||||
self.assertEqual(node_cls.CATEGORY, "LLS/Utils")
|
||||
self.assertEqual(node_cls.FUNCTION, "concat")
|
||||
self.assertEqual(node_cls.RETURN_TYPES, ("IMAGE", "MASK", "INT", "INT"))
|
||||
self.assertEqual(node_cls.RETURN_NAMES, ("image", "mask", "width", "height"))
|
||||
self.assertEqual(required["data_type"][0], ["IMAGE", "MASK"])
|
||||
self.assertEqual(required["target"][0], ["A", "B"])
|
||||
self.assertEqual(required["position"][0], ["top", "bottom", "left", "right"])
|
||||
self.assertEqual(required["resize_mode"][0], ["keep_proportion", "stretch", "none"])
|
||||
self.assertEqual(required["align"][0], ["start", "center", "end"])
|
||||
self.assertEqual(required["match_target_size"][0], "BOOLEAN")
|
||||
self.assertEqual(required["gap"][0], "INT")
|
||||
self.assertEqual(required["gap"][1]["default"], 0)
|
||||
self.assertEqual(required["background_color"][0], "STRING")
|
||||
self.assertEqual(required["background_color"][1]["default"], "#000000")
|
||||
self.assertEqual(required["background_value"][0], "FLOAT")
|
||||
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",))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,488 @@
|
||||
from __future__ import annotations
|
||||
|
||||
try:
|
||||
import torch
|
||||
import torch.nn.functional as torch_nn_functional
|
||||
except Exception as exc: # pragma: no cover - optional runtime dependency
|
||||
torch = None
|
||||
torch_nn_functional = None
|
||||
_TORCH_ERR = exc
|
||||
else: # pragma: no cover - exercised in runtime/tests
|
||||
_TORCH_ERR = None
|
||||
|
||||
|
||||
DATA_TYPE_CHOICES = ["IMAGE", "MASK"]
|
||||
TARGET_CHOICES = ["A", "B"]
|
||||
POSITION_CHOICES = ["top", "bottom", "left", "right"]
|
||||
RESIZE_MODE_CHOICES = ["keep_proportion", "stretch", "none"]
|
||||
ALIGN_CHOICES = ["start", "center", "end"]
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def _validate_choice(name: str, value: str, choices):
|
||||
if value not in choices:
|
||||
raise RuntimeError(f"[LLS] Invalid {name} '{value}'. Expected one of: {', '.join(choices)}.")
|
||||
|
||||
|
||||
def parse_hex_color(color_str: str):
|
||||
value = str(color_str or "").strip()
|
||||
if value.startswith("#"):
|
||||
value = value[1:]
|
||||
if len(value) == 3:
|
||||
value = "".join(part * 2 for part in value)
|
||||
if len(value) != 6:
|
||||
raise RuntimeError("[LLS] background_color must be a valid HEX color like '#000000' or '#ffffff'.")
|
||||
try:
|
||||
red = int(value[0:2], 16) / 255.0
|
||||
green = int(value[2:4], 16) / 255.0
|
||||
blue = int(value[4:6], 16) / 255.0
|
||||
except ValueError as exc:
|
||||
raise RuntimeError(
|
||||
"[LLS] background_color must be a valid HEX color like '#000000' or '#ffffff'."
|
||||
) from exc
|
||||
return red, green, blue
|
||||
|
||||
|
||||
def ensure_image_tensor(x):
|
||||
_require_torch()
|
||||
if x is None:
|
||||
raise RuntimeError("[LLS] IMAGE mode requires both image_a and image_b inputs.")
|
||||
if not isinstance(x, torch.Tensor):
|
||||
x = torch.as_tensor(x)
|
||||
if x.ndim == 3:
|
||||
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 not torch.is_floating_point(x):
|
||||
x = x.to(dtype=torch.float32)
|
||||
return 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.")
|
||||
if not isinstance(x, torch.Tensor):
|
||||
x = torch.as_tensor(x)
|
||||
if x.ndim == 2:
|
||||
x = x.unsqueeze(0)
|
||||
elif x.ndim == 3:
|
||||
pass
|
||||
elif x.ndim == 4 and int(x.shape[-1]) == 1:
|
||||
x = x.squeeze(-1)
|
||||
elif x.ndim == 4 and int(x.shape[1]) == 1:
|
||||
x = x.squeeze(1)
|
||||
else:
|
||||
raise RuntimeError("[LLS] MASK input must have shape [B,H,W], [H,W], [B,H,W,1], or [B,1,H,W].")
|
||||
if not torch.is_floating_point(x):
|
||||
x = x.to(dtype=torch.float32)
|
||||
return x.to(dtype=torch.float32).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _repeat_batch(tensor, target_batch: int):
|
||||
repeat_shape = [1] * tensor.ndim
|
||||
repeat_shape[0] = target_batch
|
||||
return tensor.repeat(*repeat_shape)
|
||||
|
||||
|
||||
def broadcast_batch(a, b, allow_batch_broadcast: bool):
|
||||
batch_a = int(a.shape[0])
|
||||
batch_b = int(b.shape[0])
|
||||
if batch_a == batch_b:
|
||||
return a, b
|
||||
if allow_batch_broadcast and batch_a == 1 and batch_b > 1:
|
||||
return _repeat_batch(a, batch_b), b
|
||||
if allow_batch_broadcast and batch_b == 1 and batch_a > 1:
|
||||
return a, _repeat_batch(b, batch_a)
|
||||
raise RuntimeError(
|
||||
f"[LLS] Batch size mismatch: A has batch {batch_a}, B has batch {batch_b}, and broadcasting is disabled or impossible."
|
||||
)
|
||||
|
||||
|
||||
def resize_tensor_to_match(tensor, *, data_type: str, match_axis: str, target_size: int, resize_mode: str):
|
||||
_validate_choice("resize_mode", resize_mode, RESIZE_MODE_CHOICES)
|
||||
if resize_mode == "none":
|
||||
return tensor
|
||||
|
||||
if data_type == "IMAGE":
|
||||
current_height = int(tensor.shape[1])
|
||||
current_width = int(tensor.shape[2])
|
||||
else:
|
||||
current_height = int(tensor.shape[1])
|
||||
current_width = int(tensor.shape[2])
|
||||
|
||||
if match_axis == "height":
|
||||
if current_height == target_size:
|
||||
return tensor
|
||||
new_height = max(1, int(target_size))
|
||||
if resize_mode == "keep_proportion":
|
||||
scale = float(new_height) / float(current_height)
|
||||
new_width = max(1, int(round(current_width * scale)))
|
||||
else:
|
||||
new_width = current_width
|
||||
elif match_axis == "width":
|
||||
if current_width == target_size:
|
||||
return tensor
|
||||
new_width = max(1, int(target_size))
|
||||
if resize_mode == "keep_proportion":
|
||||
scale = float(new_width) / float(current_width)
|
||||
new_height = max(1, int(round(current_height * scale)))
|
||||
else:
|
||||
new_height = current_height
|
||||
else:
|
||||
raise RuntimeError(f"[LLS] Unsupported match axis '{match_axis}'.")
|
||||
|
||||
if data_type == "IMAGE":
|
||||
resized = torch_nn_functional.interpolate(
|
||||
tensor.movedim(-1, 1),
|
||||
size=(new_height, new_width),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
return resized.movedim(1, -1)
|
||||
|
||||
resized = torch_nn_functional.interpolate(
|
||||
tensor.unsqueeze(1),
|
||||
size=(new_height, new_width),
|
||||
mode="nearest",
|
||||
)
|
||||
return resized.squeeze(1)
|
||||
|
||||
|
||||
def create_canvas(
|
||||
*,
|
||||
data_type: str,
|
||||
batch: int,
|
||||
height: int,
|
||||
width: int,
|
||||
device,
|
||||
dtype,
|
||||
channels: int | None = None,
|
||||
background_color=None,
|
||||
background_value: float = 0.0,
|
||||
):
|
||||
if data_type == "IMAGE":
|
||||
if channels is None:
|
||||
raise RuntimeError("[LLS] IMAGE canvas creation requires a channels value.")
|
||||
canvas = torch.zeros((batch, height, width, channels), device=device, dtype=dtype)
|
||||
red, green, blue = background_color
|
||||
if channels == 1:
|
||||
canvas.fill_((red + green + blue) / 3.0)
|
||||
return canvas
|
||||
fill_values = [0.0] * channels
|
||||
fill_values[0] = red
|
||||
if channels > 1:
|
||||
fill_values[1] = green
|
||||
if channels > 2:
|
||||
fill_values[2] = blue
|
||||
if channels == 4:
|
||||
fill_values[3] = 1.0
|
||||
fill_tensor = torch.tensor(fill_values, device=device, dtype=dtype).view(1, 1, 1, channels)
|
||||
canvas[:] = fill_tensor
|
||||
return canvas
|
||||
|
||||
return torch.full(
|
||||
(batch, height, width),
|
||||
float(background_value),
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
|
||||
def paste_tensor(canvas, tensor, top: int, left: int):
|
||||
if canvas.ndim == 4:
|
||||
canvas[:, top : top + tensor.shape[1], left : left + tensor.shape[2], :] = tensor
|
||||
else:
|
||||
canvas[:, top : top + tensor.shape[1], left : left + tensor.shape[2]] = tensor
|
||||
return canvas
|
||||
|
||||
|
||||
def pad_to_multiple(
|
||||
tensor,
|
||||
*,
|
||||
data_type: str,
|
||||
multiple_of: int,
|
||||
background_color=None,
|
||||
background_value: float = 0.0,
|
||||
):
|
||||
if multiple_of <= 0:
|
||||
return tensor
|
||||
|
||||
if data_type == "IMAGE":
|
||||
batch, height, width, channels = [int(dim) for dim in tensor.shape]
|
||||
dtype = tensor.dtype
|
||||
device = tensor.device
|
||||
else:
|
||||
batch, height, width = [int(dim) for dim in tensor.shape]
|
||||
channels = None
|
||||
dtype = tensor.dtype
|
||||
device = tensor.device
|
||||
|
||||
padded_height = ((height + multiple_of - 1) // multiple_of) * multiple_of
|
||||
padded_width = ((width + multiple_of - 1) // multiple_of) * multiple_of
|
||||
if padded_height == height and padded_width == width:
|
||||
return tensor
|
||||
|
||||
canvas = create_canvas(
|
||||
data_type=data_type,
|
||||
batch=batch,
|
||||
height=padded_height,
|
||||
width=padded_width,
|
||||
channels=channels,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
background_color=background_color,
|
||||
background_value=background_value,
|
||||
)
|
||||
return paste_tensor(canvas, tensor, 0, 0)
|
||||
|
||||
|
||||
def _alignment_offset(container_size: int, item_size: int, align: str) -> int:
|
||||
_validate_choice("align", align, ALIGN_CHOICES)
|
||||
if align == "start":
|
||||
return 0
|
||||
if align == "center":
|
||||
return max(0, (container_size - item_size) // 2)
|
||||
return max(0, container_size - item_size)
|
||||
|
||||
|
||||
def concat_by_target(
|
||||
*,
|
||||
data_type: str,
|
||||
tensor_a,
|
||||
tensor_b,
|
||||
target: str,
|
||||
position: str,
|
||||
match_target_size: bool,
|
||||
resize_mode: str,
|
||||
align: str,
|
||||
gap: int,
|
||||
multiple_of: int,
|
||||
allow_batch_broadcast: bool,
|
||||
background_color=None,
|
||||
background_value: float = 0.0,
|
||||
):
|
||||
_validate_choice("data_type", data_type, DATA_TYPE_CHOICES)
|
||||
_validate_choice("target", target, TARGET_CHOICES)
|
||||
_validate_choice("position", position, POSITION_CHOICES)
|
||||
_validate_choice("resize_mode", resize_mode, RESIZE_MODE_CHOICES)
|
||||
_validate_choice("align", align, ALIGN_CHOICES)
|
||||
if gap < 0:
|
||||
raise RuntimeError("[LLS] gap must be >= 0.")
|
||||
if multiple_of < 0:
|
||||
raise RuntimeError("[LLS] multiple_of must be >= 0.")
|
||||
|
||||
tensor_a, tensor_b = broadcast_batch(tensor_a, tensor_b, allow_batch_broadcast)
|
||||
target_tensor = tensor_a if target == "A" else tensor_b
|
||||
other_tensor = tensor_b if target == "A" else tensor_a
|
||||
|
||||
if data_type == "IMAGE":
|
||||
work_dtype = target_tensor.dtype if torch.is_floating_point(target_tensor) else torch.float32
|
||||
target_tensor = target_tensor.to(dtype=work_dtype)
|
||||
other_tensor = other_tensor.to(device=target_tensor.device, dtype=work_dtype)
|
||||
else:
|
||||
target_tensor = target_tensor.to(dtype=torch.float32)
|
||||
other_tensor = other_tensor.to(device=target_tensor.device, dtype=torch.float32)
|
||||
|
||||
if match_target_size:
|
||||
match_axis = "height" if position in {"left", "right"} else "width"
|
||||
target_size = int(target_tensor.shape[1] if match_axis == "height" else target_tensor.shape[2])
|
||||
other_tensor = resize_tensor_to_match(
|
||||
other_tensor,
|
||||
data_type=data_type,
|
||||
match_axis=match_axis,
|
||||
target_size=target_size,
|
||||
resize_mode=resize_mode,
|
||||
)
|
||||
|
||||
if data_type == "IMAGE":
|
||||
batch = int(target_tensor.shape[0])
|
||||
target_height = int(target_tensor.shape[1])
|
||||
target_width = int(target_tensor.shape[2])
|
||||
other_height = int(other_tensor.shape[1])
|
||||
other_width = int(other_tensor.shape[2])
|
||||
channels = int(target_tensor.shape[3])
|
||||
else:
|
||||
batch = int(target_tensor.shape[0])
|
||||
target_height = int(target_tensor.shape[1])
|
||||
target_width = int(target_tensor.shape[2])
|
||||
other_height = int(other_tensor.shape[1])
|
||||
other_width = int(other_tensor.shape[2])
|
||||
channels = None
|
||||
|
||||
if position in {"left", "right"}:
|
||||
canvas_height = max(target_height, other_height)
|
||||
canvas_width = target_width + int(gap) + other_width
|
||||
target_top = _alignment_offset(canvas_height, target_height, align)
|
||||
other_top = _alignment_offset(canvas_height, other_height, align)
|
||||
if position == "right":
|
||||
target_left = 0
|
||||
other_left = target_width + int(gap)
|
||||
else:
|
||||
other_left = 0
|
||||
target_left = other_width + int(gap)
|
||||
else:
|
||||
canvas_height = target_height + int(gap) + other_height
|
||||
canvas_width = max(target_width, other_width)
|
||||
target_left = _alignment_offset(canvas_width, target_width, align)
|
||||
other_left = _alignment_offset(canvas_width, other_width, align)
|
||||
if position == "bottom":
|
||||
target_top = 0
|
||||
other_top = target_height + int(gap)
|
||||
else:
|
||||
other_top = 0
|
||||
target_top = other_height + int(gap)
|
||||
|
||||
canvas = create_canvas(
|
||||
data_type=data_type,
|
||||
batch=batch,
|
||||
height=canvas_height,
|
||||
width=canvas_width,
|
||||
channels=channels,
|
||||
device=target_tensor.device,
|
||||
dtype=target_tensor.dtype,
|
||||
background_color=background_color,
|
||||
background_value=background_value,
|
||||
)
|
||||
paste_tensor(canvas, target_tensor, target_top, target_left)
|
||||
paste_tensor(canvas, other_tensor, other_top, other_left)
|
||||
canvas = pad_to_multiple(
|
||||
canvas,
|
||||
data_type=data_type,
|
||||
multiple_of=int(multiple_of),
|
||||
background_color=background_color,
|
||||
background_value=background_value,
|
||||
)
|
||||
|
||||
if data_type == "IMAGE":
|
||||
final_height = int(canvas.shape[1])
|
||||
final_width = int(canvas.shape[2])
|
||||
else:
|
||||
final_height = int(canvas.shape[1])
|
||||
final_width = int(canvas.shape[2])
|
||||
return canvas, final_width, final_height
|
||||
|
||||
|
||||
def mask_to_preview_image(mask):
|
||||
normalized = ensure_mask_tensor(mask)
|
||||
return normalized.unsqueeze(-1).repeat(1, 1, 1, 3).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
class LLSConcatByTarget:
|
||||
CATEGORY = "LLS/Utils"
|
||||
FUNCTION = "concat"
|
||||
RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT")
|
||||
RETURN_NAMES = ("image", "mask", "width", "height")
|
||||
DESCRIPTION = "Concatenate IMAGE or MASK inputs around a chosen target canvas."
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"data_type": (DATA_TYPE_CHOICES, {"default": "IMAGE"}),
|
||||
"target": (TARGET_CHOICES, {"default": "A"}),
|
||||
"position": (POSITION_CHOICES, {"default": "right"}),
|
||||
"match_target_size": ("BOOLEAN", {"default": True}),
|
||||
"resize_mode": (RESIZE_MODE_CHOICES, {"default": "keep_proportion"}),
|
||||
"align": (ALIGN_CHOICES, {"default": "center"}),
|
||||
"gap": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}),
|
||||
"background_color": ("STRING", {"default": "#000000", "multiline": False}),
|
||||
"background_value": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
||||
"multiple_of": ("INT", {"default": 0, "min": 0, "max": 512, "step": 1}),
|
||||
"allow_batch_broadcast": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
"optional": {
|
||||
"image_a": ("IMAGE",),
|
||||
"image_b": ("IMAGE",),
|
||||
"mask_a": ("MASK",),
|
||||
"mask_b": ("MASK",),
|
||||
},
|
||||
}
|
||||
|
||||
def concat(
|
||||
self,
|
||||
data_type,
|
||||
target,
|
||||
position,
|
||||
match_target_size,
|
||||
resize_mode,
|
||||
align,
|
||||
gap,
|
||||
background_color,
|
||||
background_value,
|
||||
multiple_of,
|
||||
allow_batch_broadcast,
|
||||
image_a=None,
|
||||
image_b=None,
|
||||
mask_a=None,
|
||||
mask_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.")
|
||||
color = parse_hex_color(background_color)
|
||||
tensor_a = ensure_image_tensor(image_a)
|
||||
tensor_b = ensure_image_tensor(image_b)
|
||||
output_image, width, height = concat_by_target(
|
||||
data_type="IMAGE",
|
||||
tensor_a=tensor_a,
|
||||
tensor_b=tensor_b,
|
||||
target=str(target),
|
||||
position=str(position),
|
||||
match_target_size=bool(match_target_size),
|
||||
resize_mode=str(resize_mode),
|
||||
align=str(align),
|
||||
gap=int(gap),
|
||||
multiple_of=int(multiple_of),
|
||||
allow_batch_broadcast=bool(allow_batch_broadcast),
|
||||
background_color=color,
|
||||
background_value=background_value,
|
||||
)
|
||||
output_image = output_image.clamp(0.0, 1.0)
|
||||
output_mask = torch.zeros(
|
||||
(int(output_image.shape[0]), int(output_image.shape[1]), int(output_image.shape[2])),
|
||||
device=output_image.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
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)
|
||||
output_mask, width, height = concat_by_target(
|
||||
data_type="MASK",
|
||||
tensor_a=tensor_a,
|
||||
tensor_b=tensor_b,
|
||||
target=str(target),
|
||||
position=str(position),
|
||||
match_target_size=bool(match_target_size),
|
||||
resize_mode=str(resize_mode),
|
||||
align=str(align),
|
||||
gap=int(gap),
|
||||
multiple_of=int(multiple_of),
|
||||
allow_batch_broadcast=bool(allow_batch_broadcast),
|
||||
background_color=None,
|
||||
background_value=background_value,
|
||||
)
|
||||
output_mask = output_mask.clamp(0.0, 1.0).to(dtype=torch.float32)
|
||||
output_image = mask_to_preview_image(output_mask)
|
||||
return output_image, output_mask, int(width), int(height)
|
||||
|
||||
raise RuntimeError(f"[LLS] Invalid data_type '{data_type}'. Expected IMAGE or MASK.")
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"LLSConcatByTarget": LLSConcatByTarget}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"LLSConcatByTarget": "LLS Concat By Target"}
|
||||
@@ -32,6 +32,10 @@ LLS / Utils
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from .concat_by_target import (
|
||||
NODE_CLASS_MAPPINGS as CONCAT_NODE_CLASS_MAPPINGS,
|
||||
NODE_DISPLAY_NAME_MAPPINGS as CONCAT_NODE_DISPLAY_NAME_MAPPINGS,
|
||||
)
|
||||
from .model_info import (
|
||||
FAMILY_DEFAULT_PRESET,
|
||||
MODEL_FAMILY_CHOICES,
|
||||
@@ -238,6 +242,7 @@ NODE_CLASS_MAPPINGS: dict[str, type] = {
|
||||
"LLSResolutionSelector": LLSResolutionSelector,
|
||||
"LLSGenerationConfig": LLSGenerationConfig,
|
||||
}
|
||||
NODE_CLASS_MAPPINGS.update(CONCAT_NODE_CLASS_MAPPINGS)
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {
|
||||
"LLSStringLiteral": "LLS String Literal",
|
||||
@@ -246,3 +251,4 @@ NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = {
|
||||
"LLSResolutionSelector": "LLS Resolution Selector",
|
||||
"LLSGenerationConfig": "LLS Generation Config",
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CONCAT_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
Reference in New Issue
Block a user