Skip to content

Commit 8564480

Browse files
[feat] Add GetImageSize node (Comfy-Org#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.
1 parent 312d511 commit 8564480

File tree

3 files changed

+38
-2
lines changed

3 files changed

+38
-2
lines changed

comfy_extras/nodes_images.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@
1616
import torch
1717
import comfy.utils
1818

19-
from comfy.comfy_types import FileLocator
19+
from comfy.comfy_types import FileLocator, IO
20+
from server import PromptServer
2021

2122
MAX_RESOLUTION = nodes.MAX_RESOLUTION
2223

@@ -491,6 +492,36 @@ def replacement(match):
491492
counter += 1
492493
return { "ui": { "images": results } }
493494

495+
class GetImageSize:
496+
497+
@classmethod
498+
def INPUT_TYPES(s):
499+
return {
500+
"required": {
501+
"image": (IO.IMAGE,),
502+
},
503+
"hidden": {
504+
"unique_id": "UNIQUE_ID",
505+
}
506+
}
507+
508+
RETURN_TYPES = (IO.INT, IO.INT)
509+
RETURN_NAMES = ("width", "height")
510+
FUNCTION = "get_size"
511+
512+
CATEGORY = "image"
513+
DESCRIPTION = """Returns width and height of the image, and passes it through unchanged."""
514+
515+
def get_size(self, image, unique_id=None) -> tuple[int, int]:
516+
height = image.shape[1]
517+
width = image.shape[2]
518+
519+
# Send progress text to display size on the node
520+
if unique_id:
521+
PromptServer.instance.send_progress_text(f"width: {width}, height: {height}", unique_id)
522+
523+
return width, height
524+
494525
NODE_CLASS_MAPPINGS = {
495526
"ImageCrop": ImageCrop,
496527
"RepeatImageBatch": RepeatImageBatch,
@@ -500,4 +531,5 @@ def replacement(match):
500531
"SaveAnimatedPNG": SaveAnimatedPNG,
501532
"SaveSVGNode": SaveSVGNode,
502533
"ImageStitch": ImageStitch,
534+
"GetImageSize": GetImageSize,
503535
}

nodes.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2067,6 +2067,7 @@ def expand_image(self, image, left, top, right, bottom, feathering):
20672067
"ImageQuantize": "Image Quantize",
20682068
"ImageSharpen": "Image Sharpen",
20692069
"ImageScaleToTotalPixels": "Scale Image to Total Pixels",
2070+
"GetImageSize": "Get Image Size",
20702071
# _for_testing
20712072
"VAEDecodeTiled": "VAE Decode (Tiled)",
20722073
"VAEEncodeTiled": "VAE Encode (Tiled)",

tests-unit/comfy_extras_test/image_stitch_test.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,10 @@
55
mock_nodes = MagicMock()
66
mock_nodes.MAX_RESOLUTION = 16384
77

8-
with patch.dict('sys.modules', {'nodes': mock_nodes}):
8+
# Mock server module for PromptServer
9+
mock_server = MagicMock()
10+
11+
with patch.dict('sys.modules', {'nodes': mock_nodes, 'server': mock_server}):
912
from comfy_extras.nodes_images import ImageStitch
1013

1114

0 commit comments

Comments
 (0)