SparknightLLC-ComfyUI-ImageAutosize/__init__.py

import comfy.utils
from comfy_api.latest import ComfyExtension, io


SCALE_METHODS = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
CONSTRAINT_PRIORITIES = ["min_size", "max_size"]
CROP_MODES = [
	"none",
	"center",
	"top",
	"bottom",
	"left",
	"right",
	"top_left",
	"top_right",
	"bottom_left",
	"bottom_right",
]


def _calculate_target_dimensions(
	width: int,
	height: int,
	max_size: int,
	min_size: int,
	constraint_priority: str,
	divisible_by: int,
) -> tuple[int, int]:
	max_scale = max_size / max(width, height)
	min_scale = min_size / min(width, height)
	if constraint_priority == "min_size":
		scale = max(max_scale, min_scale)
	else:
		scale = min(min_scale, max_scale)

	target_width = max(divisible_by, round(width * scale / divisible_by) * divisible_by)
	target_height = max(divisible_by, round(height * scale / divisible_by) * divisible_by)
	return target_width, target_height


def _get_crop_origin(
	width: int,
	height: int,
	target_width: int,
	target_height: int,
	crop_mode: str,
) -> tuple[int, int]:
	if crop_mode in ("left", "top_left", "bottom_left"):
		x = 0
	elif crop_mode in ("right", "top_right", "bottom_right"):
		x = width - target_width
	else:
		x = (width - target_width) // 2

	if crop_mode in ("top", "top_left", "top_right"):
		y = 0
	elif crop_mode in ("bottom", "bottom_left", "bottom_right"):
		y = height - target_height
	else:
		y = (height - target_height) // 2

	return x, y


class ImageAutosize(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="ImageAutosize",
			display_name="Image/Mask Autosize",
			category="image",
			description="Automatically resizes an image or mask for diffusion workflows.",
			search_aliases=[
				"autosize",
				"image mask autosize",
				"auto resize",
				"resize to multiple",
				"resize image",
				"resize mask",
			],
			inputs=[
				io.MatchType.Input(
					"image",
					template=image_type,
					tooltip="The image or mask to resize.",
				),
				io.Int.Input(
					"max_size",
					default=1280,
					min=1,
					max=8192,
					step=1,
					tooltip="Longer-dimension target used to calculate one candidate resize scale.",
				),
				io.Int.Input(
					"min_size",
					default=512,
					min=1,
					max=4096,
					step=1,
					tooltip="Shorter-dimension target used to calculate one candidate resize scale.",
				),
				io.Int.Input(
					"divisible_by",
					default=32,
					min=1,
					max=8192,
					step=1,
					tooltip="Rounds both output dimensions to the nearest multiple of this value.",
				),
				io.Combo.Input(
					"interpolation_mode",
					options=SCALE_METHODS,
					default="lanczos",
					tooltip="Interpolation algorithm used for resizing.",
				),
				io.Combo.Input(
					"crop_mode",
					options=CROP_MODES,
					default="center",
					tooltip="Preserves aspect ratio by cropping from this position. None stretches to the output dimensions.",
				),
				io.Combo.Input(
					"constraint_priority",
					options=CONSTRAINT_PRIORITIES,
					default="min_size",
					tooltip="Chooses the candidate scale. min_size uses the larger scale; max_size uses the smaller scale. Divisibility rounding runs afterward.",
				),
			],
			outputs=[
				io.MatchType.Output(template=image_type, display_name="resized"),
				io.Int.Output(display_name="width"),
				io.Int.Output(display_name="height"),
				io.Float.Output(display_name="scale_x"),
				io.Float.Output(display_name="scale_y"),
			],
		)

	@classmethod
	def execute(
		cls,
		image: io.Image.Type | io.Mask.Type,
		max_size: int,
		min_size: int,
		constraint_priority: str,
		divisible_by: int,
		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)

		target_width, target_height = _calculate_target_dimensions(
			width,
			height,
			max_size,
			min_size,
			constraint_priority,
			divisible_by,
		)

		resize_width = target_width
		resize_height = target_height
		if 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, y = _get_crop_origin(
				resize_width,
				resize_height,
				target_width,
				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)

		return io.NodeOutput(
			output,
			target_width,
			target_height,
			resize_width / width,
			resize_height / height,
		)


class ImageAutosizeExtension(ComfyExtension):
	async def get_node_list(self) -> list[type[io.ComfyNode]]:
		return [ImageAutosize]


async def comfy_entrypoint() -> ImageAutosizeExtension:
	return ImageAutosizeExtension()