diff --git a/__init__.py b/__init__.py index a298851..d80379d 100644 --- a/__init__.py +++ b/__init__.py @@ -5,12 +5,14 @@ from .nodes import * NODE_CLASS_MAPPINGS = { - "BestResolutionSimple": BestResolutionSimple, + "BestResolutionFromArea": BestResolutionFromArea, "BestResolutionFromAspectRatio": BestResolutionFromAspectRatio, + "BestResolutionSimple": BestResolutionSimple, } NODE_DISPLAY_NAME_MAPPINGS = { - "BestResolutionSimple": "Best-Res (simple)", + "BestResolutionFromArea": "Best-Res (area)", "BestResolutionFromAspectRatio": "Best-Res (ratio)", + "BestResolutionSimple": "Best-Res (simple)", } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/_funcs.py b/_funcs.py new file mode 100644 index 0000000..ff10eb0 --- /dev/null +++ b/_funcs.py @@ -0,0 +1,86 @@ +# encoding: utf-8 +""" +""" + +import typing as _t + + +_t_number = _t.Union[int, float] + + +def aspect_ratios_sorted(aspect_a: float, aspect_b: float, min_clamp: float = 1.0): + """Return aspect ratio in a standard form: 2 floats, both 1+, in descending order (16:9, not 9:16).""" + aspect_a = float(max(abs(aspect_a), min_clamp)) + aspect_b = float(max(abs(aspect_b), min_clamp)) + return (aspect_b, aspect_a) if aspect_b > aspect_a else (aspect_a, aspect_b) + + +def number_to_int(value: _t.Union[int, float], min: int = 1) -> int: + if not isinstance(value, int): + if isinstance(value, float): + sign = -1 if value < 0.0 else 1 + value = int(value * sign + 0.5) * sign + else: + value = int(value) + return max(value, min) + + +def round_abs(abs_value: _t_number, step: int): + """ + Assuming both args are positive and ``step`` is already an int, detect the closest positive (non-zero) value + which is also divisible by step. + """ + n_steps = int(float(abs_value) / step + 0.5) + n_steps = max(n_steps, 1) + return step * n_steps, n_steps + + +def round_width_and_height_closest_to_the_ratio(width_f: _t_number, height_f: _t_number, step: int): + """3-pass detection of the best rounded resolution. Best = closest to the desired ratio.""" + desired_width_to_height_ratio = float(width_f) / height_f + + # First pass: directly from width_f and height_f + width, n_steps_x = round_abs(width_f, step) + height, n_steps_y = round_abs(height_f, step) + + # Second pass: try calculating one side from already rounded another one: + height_from_width, n_steps_y_from_x = round_abs(float(width) / desired_width_to_height_ratio, step) + width_from_height, n_steps_x_from_y = round_abs(float(height) * desired_width_to_height_ratio, step) + + # ... and select one of three options, closest to the perfect ratio: + closest_delta = abs((float(width) / height) - desired_width_to_height_ratio) + for w, n_x, h, n_y in [ + (width, n_steps_x, height_from_width, n_steps_y_from_x), + (width_from_height, n_steps_x_from_y, height, n_steps_y), + ]: + cur_delta = abs((float(w) / h) - desired_width_to_height_ratio) + if cur_delta < closest_delta: + closest_delta = cur_delta + width, n_steps_x = w, n_x + height, n_steps_y = h, n_y + + return width, n_steps_x, height, n_steps_y + + +def simple_result_from_approx_wh(width_f: float, height_f: _t_number, step: int): + """Final part of the main func for simple (non-upscale) nodes - when desired float/height are already calculated""" + step = number_to_int(step) + width, n_steps_x, height, n_steps_y = round_width_and_height_closest_to_the_ratio(width_f, height_f, step) + + ar_desired = width_f / height_f + ar_real = float(width) / height + ar_text = ( + "AR: perfect match" + if abs(ar_desired - ar_real) < 0.0005 + else f"AR desired/real: {ar_desired:.3f}/{ar_real:.3f}" + ) + + text = ( + f"{width} x {height}\n" + f"{n_steps_x} * {step} x {n_steps_y} * {step}\n" + f"{ar_text}" + ) + return { + # "ui": {"text": [text]}, # TODO: for some reason, it isn't displayed + "result": (width, height, text) + } diff --git a/node_types.py b/node_types.py index f358e39..7d5e89c 100644 --- a/node_types.py +++ b/node_types.py @@ -19,16 +19,6 @@ def number_type_dict(default: int = 512, min: int = 1, max: int = _sys.maxsize, } -def number_to_int(value: _t.Union[int, float], min: int = 1) -> int: - if not isinstance(value, int): - if isinstance(value, float): - sign = -1 if value < 0.0 else 1 - value = int(value * sign + 0.5) * sign - else: - value = int(value) - return max(value, min) - - type_dict_step_init = number_type_dict(48) type_dict_step_upscale1 = number_type_dict(128) type_dict_res = number_type_dict(1024) diff --git a/nodes.py b/nodes.py index 2cd0c2d..400a794 100644 --- a/nodes.py +++ b/nodes.py @@ -5,6 +5,7 @@ import typing as _t from inspect import cleandoc as _cleandoc +from math import sqrt as _sqrt import sys as _sys from frozendict import deepfreeze @@ -12,46 +13,11 @@ from frozendict import deepfreeze from comfy.comfy_types.node_typing import IO as _IO from .node_types import * - - -_t_number = _t.Union[int, float] - - -def _round_abs(abs_value: _t_number, step: int): - """ - Assuming both args are positive and ``step`` is already an int, detect the closest positive (non-zero) value - which is also divisible by step. - """ - n_steps = int(float(abs_value) / step + 0.5) - n_steps = max(n_steps, 1) - return step * n_steps, n_steps - - -def _round_width_and_height_closest_to_the_ratio(width_f: _t_number, height_f: _t_number, step: int): - """3-pass detection of the best rounded resolution. Best = closest to the desired ratio.""" - desired_width_to_height_ratio = float(width_f) / height_f - - # First pass: directly from width_f and height_f - width, n_steps_x = _round_abs(width_f, step) - height, n_steps_y = _round_abs(height_f, step) - - # Second pass: try calculating one side from already rounded another one: - height_from_width, n_steps_y_from_x = _round_abs(float(width) / desired_width_to_height_ratio, step) - width_from_height, n_steps_x_from_y = _round_abs(float(height) * desired_width_to_height_ratio, step) - - # ... and select one of three options, closest to the perfect ratio: - closest_delta = abs((float(width) / height) - desired_width_to_height_ratio) - for w, n_x, h, n_y in [ - (width, n_steps_x, height_from_width, n_steps_y_from_x), - (width_from_height, n_steps_x_from_y, height, n_steps_y), - ]: - cur_delta = abs((float(w) / h) - desired_width_to_height_ratio) - if cur_delta < closest_delta: - closest_delta = cur_delta - width, n_steps_x = w, n_x - height, n_steps_y = h, n_y - - return width, n_steps_x, height, n_steps_y +from ._funcs import ( + aspect_ratios_sorted as _aspect_ratios_sorted, + number_to_int as _number_to_int, + simple_result_from_approx_wh as _simple_result_from_approx_wh +) _48_tooltip = ( @@ -77,6 +43,9 @@ _input_types_simple = deepfreeze({ # "optional": {}, }) +_return_types_simple = (_IO.INT, _IO.INT, _IO.STRING) +_return_names_simple = ("width", "height", "report") + class BestResolutionSimple: """ @@ -90,28 +59,16 @@ class BestResolutionSimple: # OUTPUT_NODE = True # TODO FUNCTION = "main" - RETURN_TYPES = (_IO.INT, _IO.INT) + RETURN_TYPES = _return_types_simple # RETURN_TYPES_TOOLTIPS = {} - RETURN_NAMES = ("width", "height") + RETURN_NAMES = _return_names_simple @classmethod def INPUT_TYPES(cls): return _input_types_simple def main(self, width: int, height: int, step: int): - step = number_to_int(step) - width, n_steps_x, height, n_steps_y = _round_width_and_height_closest_to_the_ratio(width, height, step) - - text = f"{width} x {height}\n{n_steps_x} * {step} x {n_steps_y} * {step}" - # print(f"text: {text}") - - # The following return format is snatched from built-in PreviewAny - # and easy_nodes.show_text from (https://github.com/andrewharp/ComfyUI-EasyNodes); - # doesn't seem to be documented anywhere 🤷🏻‍♂️ - return { - # "ui": {"text": [text]}, # TODO: for some reason, it isn't displayed - "result": (width, height) - } + return _simple_result_from_approx_wh(float(width), height, step) _tooltip_aspect = ( @@ -155,7 +112,7 @@ _input_types_orient = deepfreeze({ class BestResolutionFromAspectRatio: """ A more advanced node for easier selection of an optimal resolution: - image size specified indirectly - by one of the sides + aspect ratio. + image size selected indirectly - by one of the sides + aspect ratio. """ NODE_NAME = "BestResolutionFromAspectRatio" CATEGORY = "utils/resolution" @@ -164,35 +121,89 @@ class BestResolutionFromAspectRatio: # OUTPUT_NODE = True # TODO FUNCTION = "main" - RETURN_TYPES = (_IO.INT, _IO.INT) - RETURN_NAMES = ("width", "height") + RETURN_TYPES = _return_types_simple + RETURN_NAMES = _return_names_simple @classmethod def INPUT_TYPES(cls): return _input_types_orient def main(self, size: int, step: int, size_is_big: bool, landscape: bool, aspect_a: float, aspect_b: float): - aspect_a = max(float(aspect_a), 1.0) - aspect_b = max(float(aspect_b), 1.0) - aspect_big = max(aspect_a, aspect_b) - aspect_small = min(aspect_a, aspect_b) + aspect_big, aspect_small = _aspect_ratios_sorted(aspect_a, aspect_b) side_main_f = float(size) - side_secondary_f = side_main_f * ( + side_other_f = side_main_f * ( (aspect_small / aspect_big) if size_is_big else (aspect_big / aspect_small) ) - width_f = max(side_main_f, side_secondary_f) - height_f = min(side_main_f, side_secondary_f) + width_f, height_f = (side_other_f, side_main_f) if side_other_f > side_main_f else (side_main_f, side_other_f) if not landscape: # The opposite: height is bigger width_f, height_f = height_f, width_f - step = number_to_int(step) - width, n_steps_x, height, n_steps_y = _round_width_and_height_closest_to_the_ratio(width_f, height_f, step) + return _simple_result_from_approx_wh(width_f, height_f, step) - text = f"{width} x {height}\n{n_steps_x} * {step} x {n_steps_y} * {step}" - return { - # "ui": {"text": [text]}, # TODO: for some reason, it isn't displayed - "result": (width, height) - } + +_input_types_area = deepfreeze({ + "required": { + "square_size": (_IO.INT, dict(type_dict_res, **{"tooltip": ( + "The total resolution of the image would be the same as of a square with this side.\n" + "The width and height would be such to respect aspect ratio, but also be as close as possible to the " + "total number of pixels as in this square image.\n\n" + "- 512x512 square (SD 1.5): ~0.25 megapixels\n" + "- 1024x1024 square (SDXL): ~1 megapixel" + )})), + "step": _input_types_simple["required"]["step"], + "landscape": _input_types_orient["required"]["landscape"], + "aspect_a": _input_types_orient["required"]["aspect_a"], + "aspect_b": _input_types_orient["required"]["aspect_b"], + }, + # "hidden": {}, + # "optional": {}, +}) + + +class BestResolutionFromArea: + """ + The most efficient way of selecting an optimal resolution: + image size selected indirectly - by the total desired resolution (area) + aspect ratio. + + Desired resolution (aka image area/megapixels/pixel count) is specified with a side of a square image. This isn't + accidental: most models disclose what image resolution they're trained on, and usually they're square: + + - SD 1.5 - 512x512 pixels + - SDXL - 1024x1024 pixels + + By simply providing this single number and setting your aspect ratio/orientation, you get the width and height to + produce the closest total resolution to the training set, while also respecting image proportions and step-rounding. + """ + NODE_NAME = "BestResolutionFromArea" + CATEGORY = "utils/resolution" + DESCRIPTION = _cleandoc(__doc__) + + # OUTPUT_NODE = True # TODO + + FUNCTION = "main" + RETURN_TYPES = _return_types_simple + RETURN_NAMES = _return_names_simple + + @classmethod + def INPUT_TYPES(cls): + return _input_types_area + + def main(self, square_size: int, step: int, landscape: bool, aspect_a: float, aspect_b: float): + # square_size = 1024; step = 48; landscape = True; aspect_a = 9.0; aspect_b = 16.0 + aspect_big, aspect_small = _aspect_ratios_sorted(aspect_a, aspect_b) + aspect_x, aspect_y = (aspect_big, aspect_small) if landscape else (aspect_small, aspect_big) + + aspect_area = aspect_x * aspect_y + aspect_norm_scale = 1.0 / _sqrt(aspect_area) + aspect_x *= aspect_norm_scale + aspect_y *= aspect_norm_scale + # Now the two aspects produce a normalized rectangle - i.e., it's area is 1 + + square_size = _number_to_int(square_size) + width_f = aspect_x * square_size + height_f = aspect_y * square_size + + return _simple_result_from_approx_wh(width_f, height_f, step)