From 811baef7991880bb7c62feccc55846b2eea1cbf3 Mon Sep 17 00:00:00 2001
From: Kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 3 May 2024 13:15:35 +0300
Subject: [PATCH] fix

---
 __init__.py         |  1 +
 nodes/mask_nodes.py |  9 +++++----
 web/js/jsnodes.js   | 21 +++++++++++++++++++++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/__init__.py b/__init__.py
index adcf11e..b568427 100644
--- a/__init__.py
+++ b/__init__.py
@@ -40,6 +40,7 @@ NODE_CONFIG = {
     #images
     "ColorMatch": {"class": ColorMatch, "name": "Color Match"},
     "AddLabel": {"class": AddLabel, "name": "Add Label"},
+    "ImageAndMaskPreview": {"class": ImageAndMaskPreview},
     "ImageBatchMulti": {"class": ImageBatchMulti, "name": "Image Batch Multi"},
     "ImageBatchTestPattern": {"class": ImageBatchTestPattern, "name": "Image Batch Test Pattern"},
     "ImageConcanate": {"class": ImageConcanate, "name": "Image Concanate"},
diff --git a/nodes/mask_nodes.py b/nodes/mask_nodes.py
index 9ea1f65..6889279 100644
--- a/nodes/mask_nodes.py
+++ b/nodes/mask_nodes.py
@@ -830,8 +830,8 @@ class GetMaskSize:
             "mask": ("MASK",),
         }}
 
-    RETURN_TYPES = ("MASK","INT", "INT", )
-    RETURN_NAMES = ("mask", "width", "height",)
+    RETURN_TYPES = ("MASK","INT", "INT", "INT",)
+    RETURN_NAMES = ("mask", "width", "height", "count",)
     FUNCTION = "getsize"
     CATEGORY = "KJNodes/masking"
     DESCRIPTION = """
@@ -843,9 +843,10 @@ and passes through the mask unchanged.
     def getsize(self, mask):
         width = mask.shape[2]
         height = mask.shape[1]
+        count = mask.shape[0]
         return {"ui": {
-            "text": [f"{width}x{height}"]}, 
-            "result": (mask, width, height) 
+            "text": [f"{count}x{width}x{height}"]}, 
+            "result": (mask, width, height, count) 
         }
 
 class GrowMaskWithBlur:
diff --git a/web/js/jsnodes.js b/web/js/jsnodes.js
index 9a9483e..da34cd3 100644
--- a/web/js/jsnodes.js
+++ b/web/js/jsnodes.js
@@ -73,6 +73,27 @@ app.registerExtension({
 						});
 					}
 					break;
+
+			case "GetMaskSize":
+				const onConnectInput = nodeType.prototype.onConnectInput;
+				nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) {
+					const v = onConnectInput?.(this, arguments);
+					targetSlot.outputs[1]["name"] = "width"
+					targetSlot.outputs[2]["name"] = "height" 
+					targetSlot.outputs[3]["name"] = "count"
+					return v;
+				}
+				const onExecuted = nodeType.prototype.onExecuted;
+				nodeType.prototype.onExecuted = function(message) {
+					const r = onExecuted? onExecuted.apply(this,arguments): undefined
+					let values = message["text"].toString().split('x').map(Number);
+					this.outputs[1]["name"] = values[1] + " width"
+					this.outputs[2]["name"] = values[2] + " height" 
+					this.outputs[3]["name"] = values[0] + " count" 
+					return r
+				}
+				break;
+
 			case "JoinStringMulti":
 				nodeType.prototype.onNodeCreated = function () {
 				this._type = "STRING"