diff --git a/nodes/image_utils.py b/nodes/image_utils.py index a2ccc3f..6632fa4 100644 --- a/nodes/image_utils.py +++ b/nodes/image_utils.py @@ -21,7 +21,11 @@ class MTB_StackImages: "match_method": ( ["error", "smallest", "largest"], {"default": "error"}, - ) + ), + "output_rgb": ( + "BOOLEAN", + {"default": True, "tooltip": "Output RGB instead of RGBA"}, + ), }, } @@ -29,7 +33,7 @@ class MTB_StackImages: FUNCTION = "stack" CATEGORY = "mtb/image utils" - def stack(self, vertical, match_method="error", **kwargs): + def stack(self, vertical, match_method="error", output_rgb=True, **kwargs): if not kwargs: raise ValueError("At least one tensor must be provided.") @@ -98,6 +102,9 @@ class MTB_StackImages: stacked_tensor = torch.cat(normalized_tensors, dim=dim) + if output_rgb: + stacked_tensor = stacked_tensor[:, :, :, :3] + return (stacked_tensor,) def normalize_to_rgba(self, tensor):