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