From 856448060ce42674eea66c835bd754644c322723 Mon Sep 17 00:00:00 2001
From: Christian Byrne <cbyrne@comfy.org>
Date: Mon, 2 Jun 2025 18:57:50 -0700
Subject: [PATCH] [feat] Add GetImageSize node (#8386)

* [feat] Add GetImageSize node to return image dimensions

Added a simple GetImageSize node in comfy_extras/nodes_images.py that returns width and height of input images. The node displays dimensions on the UI via PromptServer and provides width/height as outputs for further processing.

* add display name mapping

* [fix] Add server module mock to unit tests for PromptServer import

Updated test to mock server module preventing import errors from the new PromptServer usage in GetImageSize node. Uses direct import pattern consistent with rest of codebase.
---
 comfy_extras/nodes_images.py                  | 34 ++++++++++++++++++-
 nodes.py                                      |  1 +
 .../comfy_extras_test/image_stitch_test.py    |  5 ++-
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
index 6ebf1dbd8841..58b29f9a911f 100644
--- a/comfy_extras/nodes_images.py
+++ b/comfy_extras/nodes_images.py
@@ -16,7 +16,8 @@
 import torch
 import comfy.utils
 
-from comfy.comfy_types import FileLocator
+from comfy.comfy_types import FileLocator, IO
+from server import PromptServer
 
 MAX_RESOLUTION = nodes.MAX_RESOLUTION
 
@@ -491,6 +492,36 @@ def replacement(match):
             counter += 1
         return { "ui": { "images": results } }
 
+class GetImageSize:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": (IO.IMAGE,),
+            },
+            "hidden": {
+                "unique_id": "UNIQUE_ID",
+            }
+        }
+
+    RETURN_TYPES = (IO.INT, IO.INT)
+    RETURN_NAMES = ("width", "height")
+    FUNCTION = "get_size"
+
+    CATEGORY = "image"
+    DESCRIPTION = """Returns width and height of the image, and passes it through unchanged."""
+
+    def get_size(self, image, unique_id=None) -> tuple[int, int]:
+        height = image.shape[1]
+        width = image.shape[2]
+
+        # Send progress text to display size on the node
+        if unique_id:
+            PromptServer.instance.send_progress_text(f"width: {width}, height: {height}", unique_id)
+
+        return width, height
+
 NODE_CLASS_MAPPINGS = {
     "ImageCrop": ImageCrop,
     "RepeatImageBatch": RepeatImageBatch,
@@ -500,4 +531,5 @@ def replacement(match):
     "SaveAnimatedPNG": SaveAnimatedPNG,
     "SaveSVGNode": SaveSVGNode,
     "ImageStitch": ImageStitch,
+    "GetImageSize": GetImageSize,
 }
diff --git a/nodes.py b/nodes.py
index 67360e7da98e..637279ffbd6e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2067,6 +2067,7 @@ def expand_image(self, image, left, top, right, bottom, feathering):
     "ImageQuantize": "Image Quantize",
     "ImageSharpen": "Image Sharpen",
     "ImageScaleToTotalPixels": "Scale Image to Total Pixels",
+    "GetImageSize": "Get Image Size",
     # _for_testing
     "VAEDecodeTiled": "VAE Decode (Tiled)",
     "VAEEncodeTiled": "VAE Encode (Tiled)",
diff --git a/tests-unit/comfy_extras_test/image_stitch_test.py b/tests-unit/comfy_extras_test/image_stitch_test.py
index fbaef756c375..b5a0f022cd0c 100644
--- a/tests-unit/comfy_extras_test/image_stitch_test.py
+++ b/tests-unit/comfy_extras_test/image_stitch_test.py
@@ -5,7 +5,10 @@
 mock_nodes = MagicMock()
 mock_nodes.MAX_RESOLUTION = 16384
 
-with patch.dict('sys.modules', {'nodes': mock_nodes}):
+# Mock server module for PromptServer
+mock_server = MagicMock()
+
+with patch.dict('sys.modules', {'nodes': mock_nodes, 'server': mock_server}):
     from comfy_extras.nodes_images import ImageStitch