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 6ebf1dbd..58b29f9a 100644
--- a/comfy_extras/nodes_images.py
+++ b/comfy_extras/nodes_images.py
@@ -16,7 +16,8 @@ from inspect import cleandoc
 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 @@ class SaveSVGNode:
             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 @@ NODE_CLASS_MAPPINGS = {
     "SaveAnimatedPNG": SaveAnimatedPNG,
     "SaveSVGNode": SaveSVGNode,
     "ImageStitch": ImageStitch,
+    "GetImageSize": GetImageSize,
 }
diff --git a/nodes.py b/nodes.py
index 67360e7d..637279ff 100644
--- a/nodes.py
+++ b/nodes.py
@@ -2067,6 +2067,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "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 fbaef756..b5a0f022 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 @@ from unittest.mock import patch, MagicMock
 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