[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.
This commit is contained in:
Christian Byrne 2025-06-02 18:57:50 -07:00 committed by GitHub
parent 312d511630
commit 856448060c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 38 additions and 2 deletions

View File

@ -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,
}

View File

@ -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)",

View File

@ -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