From fa3db9751d84a254985a07c91e58cee4d25a9779 Mon Sep 17 00:00:00 2001 From: ThereforeGames Date: Thu, 6 Aug 2026 07:09:17 -0400 Subject: [PATCH] Add pad mode & autosize transform nodes Introduce reversible padding support and transform metadata for Image/Mask Autosize. - Add a new 'pad' crop_mode that preserves aspect ratio with reversible padding. - Introduce AutosizeTransform dataclass and AUTOSIZE_TRANSFORM IO type. - Add helper functions (_to_samples, _from_samples, _apply_transform, _resize_samples) and use them in ImageAutosize. - Add two nodes: ImageAutosizeApplyTransform and ImageAutosizeRestore to apply transforms to aligned masks and restore original dimensions. - Update tests, README, and CHANGELOG and bump package version to 0.3.0. --- CHANGELOG.md | 12 +++ README.md | 14 ++- __init__.py | 202 ++++++++++++++++++++++++++++++----- pyproject.toml | 2 +- tests/test_image_autosize.py | 102 ++++++++++++++++-- 5 files changed, 294 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d587d4b..5befeef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,17 @@ # Changelog +## 0.3.0 - 2026 August 5 + +### Added + +- Add reversible `pad` sizing that preserves pixel aspect ratio while producing + divisible diffusion dimensions. +- Add `Apply Autosize Transform` for applying the recorded image geometry to + aligned masks. +- Add `Restore Autosized Image/Mask` for removing diffusion padding and + returning to the exact source dimensions. +- Append reusable transform metadata to Image/Mask Autosize outputs. + ## 0.2.0 - 2026 July 31 ### Changed diff --git a/README.md b/README.md index d8b5971..b8fcd48 100644 --- a/README.md +++ b/README.md @@ -5,8 +5,8 @@ A set-and-forget image and mask resizer for diffusion workflows in Image/Mask Autosize combines a longer-dimension target, a shorter-dimension target, configurable constraint priority, and divisible output dimensions in -one operation. It also provides anchored cropping and reports the exact -horizontal and vertical resize scales. +one operation. It also provides anchored cropping, reversible diffusion +padding, and reports the exact horizontal and vertical resize scales. ![workflow](example_workflows/workflow.png) @@ -38,6 +38,12 @@ When `crop_mode` is `none`, the input is resized directly to the calculated dimensions. Divisibility rounding can therefore produce slightly different horizontal and vertical scales. +When `crop_mode` is `pad`, the input is proportionally resized to fit inside +the calculated dimensions. Images use replicated edge padding and masks use +zero padding. The returned transform can be applied to aligned masks with +`Apply Autosize Transform`, then reversed after diffusion with +`Restore Autosized Image/Mask`. + ## Inputs - `image`: The `IMAGE` or `MASK` to resize. @@ -49,7 +55,7 @@ horizontal and vertical scales. value. - `interpolation_mode`: `nearest-exact`, `bilinear`, `area`, `bicubic`, or `lanczos`. -- `crop_mode`: `none`, center, an edge, or a corner. +- `crop_mode`: `none`, `pad`, center, an edge, or a corner. ## Outputs @@ -58,6 +64,8 @@ horizontal and vertical scales. - `height`: Final output height. - `scale_x`: Actual intermediate resize width divided by the original width. - `scale_y`: Actual intermediate resize height divided by the original height. +- `transform`: Recorded geometry for aligning related inputs and removing pad + regions after diffusion. For anchored cropping, `scale_x` and `scale_y` describe the resize performed before the crop. A crop also introduces an offset, so apply the same settings diff --git a/__init__.py b/__init__.py index 2eba9e8..e0010e4 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,7 @@ +from dataclasses import dataclass + +import torch.nn.functional as F + import comfy.utils from comfy_api.latest import ComfyExtension, io @@ -6,6 +10,7 @@ SCALE_METHODS = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] CONSTRAINT_PRIORITIES = ["min_size", "max_size"] CROP_MODES = [ "none", + "pad", "center", "top", "bottom", @@ -16,6 +21,20 @@ CROP_MODES = [ "bottom_left", "bottom_right", ] +AUTOSIZE_TRANSFORM = io.Custom("AUTOSIZE_TRANSFORM") + + +@dataclass(frozen=True) +class AutosizeTransform: + original_width: int + original_height: int + target_width: int + target_height: int + resize_width: int + resize_height: int + offset_x: int + offset_y: int + crop_mode: str def _calculate_target_dimensions( @@ -62,6 +81,65 @@ def _get_crop_origin( return x, y +def _to_samples(image): + is_image = len(image.shape) == 4 + if is_image: + return image.movedim(-1, 1), True + return image.unsqueeze(1), False + + +def _from_samples(samples, is_image): + if is_image: + return samples.movedim(1, -1) + return samples.squeeze(1) + + +def _resize_samples(samples, width, height, interpolation_mode): + if samples.shape[-1] == width and samples.shape[-2] == height: + return samples + samples = comfy.utils.common_upscale( + samples, + width, + height, + interpolation_mode, + "disabled", + ) + if len(samples.shape) == 3: + samples = samples.unsqueeze(1) + return samples + + +def _apply_transform(image, transform, interpolation_mode): + samples, is_image = _to_samples(image) + if samples.shape[-1] != transform.original_width or samples.shape[-2] != transform.original_height: + raise ValueError( + "Autosize transform input dimensions must match its source dimensions. " + f"Expected {(transform.original_width, transform.original_height)}, got " + f"{(samples.shape[-1], samples.shape[-2])}." + ) + + samples = _resize_samples( + samples, + transform.resize_width, + transform.resize_height, + interpolation_mode, + ) + if transform.crop_mode == "pad": + pad_right = transform.target_width - transform.resize_width - transform.offset_x + pad_bottom = transform.target_height - transform.resize_height - transform.offset_y + padding = (transform.offset_x, pad_right, transform.offset_y, pad_bottom) + if any(padding): + samples = F.pad(samples, padding, mode="replicate" if is_image else "constant") + elif transform.crop_mode != "none": + samples = samples[ + :, + :, + transform.offset_y:transform.offset_y + transform.target_height, + transform.offset_x:transform.offset_x + transform.target_width, + ] + return _from_samples(samples, is_image) + + class ImageAutosize(io.ComfyNode): @classmethod def define_schema(cls) -> io.Schema: @@ -120,7 +198,10 @@ class ImageAutosize(io.ComfyNode): "crop_mode", options=CROP_MODES, default="center", - tooltip="Preserves aspect ratio by cropping from this position. None stretches to the output dimensions.", + tooltip=( + "Anchored modes preserve aspect ratio by cropping. Pad preserves aspect ratio " + "with reversible padding. None stretches to the output dimensions." + ), ), io.Combo.Input( "constraint_priority", @@ -135,6 +216,7 @@ class ImageAutosize(io.ComfyNode): io.Int.Output(display_name="height"), io.Float.Output(display_name="scale_x"), io.Float.Output(display_name="scale_y"), + AUTOSIZE_TRANSFORM.Output(display_name="transform"), ], ) @@ -149,13 +231,8 @@ class ImageAutosize(io.ComfyNode): interpolation_mode: str, crop_mode: str, ) -> io.NodeOutput: - is_image = len(image.shape) == 4 - if is_image: - _, height, width, _ = image.shape - samples = image.movedim(-1, 1) - else: - _, height, width = image.shape - samples = image.unsqueeze(1) + samples, _is_image = _to_samples(image) + height, width = samples.shape[-2:] target_width, target_height = _calculate_target_dimensions( width, @@ -168,23 +245,21 @@ class ImageAutosize(io.ComfyNode): resize_width = target_width resize_height = target_height - if crop_mode != "none": + if crop_mode == "pad": + scale = min(target_width / width, target_height / height) + resize_width = min(target_width, max(1, round(width * scale))) + resize_height = min(target_height, max(1, round(height * scale))) + elif crop_mode != "none": scale = max(target_width / width, target_height / height) resize_width = max(target_width, round(width * scale)) resize_height = max(target_height, round(height * scale)) - if resize_width != width or resize_height != height: - samples = comfy.utils.common_upscale( - samples, - resize_width, - resize_height, - interpolation_mode, - "disabled", - ) - if len(samples.shape) == 3: - samples = samples.unsqueeze(1) - - if crop_mode != "none": + x = 0 + y = 0 + if crop_mode == "pad": + x = (target_width - resize_width) // 2 + y = (target_height - resize_height) // 2 + elif crop_mode != "none": x, y = _get_crop_origin( resize_width, resize_height, @@ -192,12 +267,18 @@ class ImageAutosize(io.ComfyNode): target_height, crop_mode, ) - samples = samples[:, :, y:y + target_height, x:x + target_width] - - if is_image: - output = samples.movedim(1, -1) - else: - output = samples.squeeze(1) + transform = AutosizeTransform( + original_width=width, + original_height=height, + target_width=target_width, + target_height=target_height, + resize_width=resize_width, + resize_height=resize_height, + offset_x=x, + offset_y=y, + crop_mode=crop_mode, + ) + output = _apply_transform(image, transform, interpolation_mode) return io.NodeOutput( output, @@ -205,12 +286,77 @@ class ImageAutosize(io.ComfyNode): target_height, resize_width / width, resize_height / height, + transform, ) +class ImageAutosizeApplyTransform(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + image_type = io.MatchType.Template("image_type", [io.Image, io.Mask]) + return io.Schema( + node_id="ImageAutosizeApplyTransform", + display_name="Apply Autosize Transform", + category="image", + inputs=[ + io.MatchType.Input("image", template=image_type), + AUTOSIZE_TRANSFORM.Input("transform"), + io.Combo.Input("interpolation_mode", options=SCALE_METHODS, default="nearest-exact"), + ], + outputs=[io.MatchType.Output(template=image_type, display_name="resized")], + ) + + @classmethod + def execute(cls, image, transform, interpolation_mode) -> io.NodeOutput: + return io.NodeOutput(_apply_transform(image, transform, interpolation_mode)) + + +class ImageAutosizeRestore(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + image_type = io.MatchType.Template("image_type", [io.Image, io.Mask]) + return io.Schema( + node_id="ImageAutosizeRestore", + display_name="Restore Autosized Image/Mask", + category="image", + inputs=[ + io.MatchType.Input("image", template=image_type), + AUTOSIZE_TRANSFORM.Input("transform"), + io.Combo.Input("interpolation_mode", options=SCALE_METHODS, default="lanczos"), + ], + outputs=[io.MatchType.Output(template=image_type, display_name="restored")], + ) + + @classmethod + def execute(cls, image, transform, interpolation_mode) -> io.NodeOutput: + if transform.crop_mode != "pad": + raise ValueError("Restore Autosized Image/Mask requires an Autosize transform using pad mode.") + + samples, is_image = _to_samples(image) + if samples.shape[-1] != transform.target_width or samples.shape[-2] != transform.target_height: + raise ValueError( + "Autosized input dimensions must match the transform output dimensions. " + f"Expected {(transform.target_width, transform.target_height)}, got " + f"{(samples.shape[-1], samples.shape[-2])}." + ) + samples = samples[ + :, + :, + transform.offset_y:transform.offset_y + transform.resize_height, + transform.offset_x:transform.offset_x + transform.resize_width, + ] + samples = _resize_samples( + samples, + transform.original_width, + transform.original_height, + interpolation_mode, + ) + return io.NodeOutput(_from_samples(samples, is_image)) + + class ImageAutosizeExtension(ComfyExtension): async def get_node_list(self) -> list[type[io.ComfyNode]]: - return [ImageAutosize] + return [ImageAutosize, ImageAutosizeApplyTransform, ImageAutosizeRestore] async def comfy_entrypoint() -> ImageAutosizeExtension: diff --git a/pyproject.toml b/pyproject.toml index 52d61cf..ab2a6b8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,7 +2,7 @@ [project] name = "comfyui-imageautosize" # Unique identifier for your node. Immutable after creation. description = "A node for ComfyUI that provides a convenient way of resizing or cropping an image for diffusion tasks." -version = "0.2.0" # Custom node version. Must be semantically versioned. +version = "0.3.0" # Custom node version. Must be semantically versioned. license = { file = "LICENSE.txt" } dependencies = [] # Filled in from requirements.txt diff --git a/tests/test_image_autosize.py b/tests/test_image_autosize.py index 5c019d4..bb43c0f 100644 --- a/tests/test_image_autosize.py +++ b/tests/test_image_autosize.py @@ -20,7 +20,14 @@ class ImageAutosizeTests(unittest.TestCase): extension = asyncio.run(MODULE.comfy_entrypoint()) nodes = asyncio.run(extension.get_node_list()) - self.assertEqual(nodes, [MODULE.ImageAutosize]) + self.assertEqual( + nodes, + [ + MODULE.ImageAutosize, + MODULE.ImageAutosizeApplyTransform, + MODULE.ImageAutosizeRestore, + ], + ) def test_schema_accepts_images_and_masks(self): inputs = MODULE.ImageAutosize.INPUT_TYPES() @@ -45,7 +52,7 @@ class ImageAutosizeTests(unittest.TestCase): self.assertEqual(MODULE.ImageAutosize.define_schema().display_name, "Image/Mask Autosize") self.assertEqual( MODULE.ImageAutosize.RETURN_TYPES, - ["COMFY_MATCHTYPE_V3", "INT", "INT", "FLOAT", "FLOAT"], + ["COMFY_MATCHTYPE_V3", "INT", "INT", "FLOAT", "FLOAT", "AUTOSIZE_TRANSFORM"], ) def test_autosizes_image_and_reports_pre_crop_scales(self): @@ -61,7 +68,7 @@ class ImageAutosizeTests(unittest.TestCase): crop_mode="center", ) - resized, width, height, scale_x, scale_y = output.args + resized, width, height, scale_x, scale_y = output.args[:5] self.assertEqual(resized.shape, (1, 8, 4, 3)) self.assertEqual((width, height), (4, 8)) self.assertAlmostEqual(scale_x, 5 / 3) @@ -80,7 +87,7 @@ class ImageAutosizeTests(unittest.TestCase): crop_mode="none", ) - resized, width, height, scale_x, scale_y = output.args + resized, width, height, scale_x, scale_y = output.args[:5] self.assertEqual(resized.shape, (1, 8, 4, 3)) self.assertEqual((width, height), (4, 8)) self.assertAlmostEqual(scale_x, 4 / 3) @@ -141,7 +148,7 @@ class ImageAutosizeTests(unittest.TestCase): ) self.assertTrue(torch.equal(output.args[0], image)) - self.assertEqual(output.args[3:], (1.0, 1.0)) + self.assertEqual(output.args[3:5], (1.0, 1.0)) def test_minimum_shorter_dimension_overrides_longer_target(self): self.assertEqual( @@ -195,11 +202,94 @@ class ImageAutosizeTests(unittest.TestCase): crop_mode="none", ) - resized, width, height, scale_x, scale_y = output.args + resized, width, height, scale_x, scale_y = output.args[:5] self.assertEqual(resized.shape, (1, 50, 5, 3)) self.assertEqual((width, height), (5, 50)) self.assertEqual((scale_x, scale_y), (0.5, 0.5)) + def test_pad_mode_preserves_content_aspect_and_records_transform(self): + image = torch.arange(5 * 3 * 3, dtype=torch.float32).reshape(1, 5, 3, 3) + + output = MODULE.ImageAutosize.execute( + image=image, + max_size=8, + min_size=1, + constraint_priority="min_size", + divisible_by=4, + interpolation_mode="nearest-exact", + crop_mode="pad", + ) + + resized, width, height, scale_x, scale_y, transform = output.args + self.assertEqual(resized.shape, (1, 8, 4, 3)) + self.assertEqual((width, height), (4, 8)) + self.assertEqual((transform.resize_width, transform.resize_height), (4, 7)) + self.assertEqual((transform.offset_x, transform.offset_y), (0, 0)) + self.assertAlmostEqual(scale_x, 4 / 3) + self.assertAlmostEqual(scale_y, 7 / 5) + + def test_transform_applies_identical_padding_to_mask(self): + image = torch.zeros((1, 5, 3, 3)) + mask = torch.ones((1, 5, 3)) + transform = MODULE.ImageAutosize.execute( + image=image, + max_size=8, + min_size=1, + constraint_priority="min_size", + divisible_by=4, + interpolation_mode="nearest-exact", + crop_mode="pad", + ).args[-1] + + resized_mask, = MODULE.ImageAutosizeApplyTransform.execute( + mask, + transform, + "nearest-exact", + ).args + + self.assertEqual(resized_mask.shape, (1, 8, 4)) + self.assertTrue(torch.all(resized_mask[:, :7] == 1)) + self.assertTrue(torch.all(resized_mask[:, 7:] == 0)) + + def test_restore_removes_padding_and_returns_original_dimensions(self): + image = torch.rand((1, 5, 3, 3)) + resized, *_metadata, transform = MODULE.ImageAutosize.execute( + image=image, + max_size=8, + min_size=1, + constraint_priority="min_size", + divisible_by=4, + interpolation_mode="nearest-exact", + crop_mode="pad", + ).args + + restored, = MODULE.ImageAutosizeRestore.execute( + resized, + transform, + "nearest-exact", + ).args + + self.assertEqual(restored.shape, image.shape) + + def test_restore_rejects_wrong_canvas_dimensions(self): + image = torch.zeros((1, 5, 3, 3)) + transform = MODULE.ImageAutosize.execute( + image=image, + max_size=8, + min_size=1, + constraint_priority="min_size", + divisible_by=4, + interpolation_mode="nearest-exact", + crop_mode="pad", + ).args[-1] + + with self.assertRaisesRegex(ValueError, "transform output dimensions"): + MODULE.ImageAutosizeRestore.execute( + torch.zeros((1, 4, 4, 3)), + transform, + "nearest-exact", + ) + def test_crop_origins_follow_selected_anchor(self): self.assertEqual(MODULE._get_crop_origin(7, 9, 4, 6, "top_left"), (0, 0)) self.assertEqual(MODULE._get_crop_origin(7, 9, 4, 6, "bottom_right"), (3, 3))