diff --git a/.gitignore b/.gitignore index 88f2525..ddd2b42 100644 --- a/.gitignore +++ b/.gitignore @@ -162,3 +162,12 @@ __pypackages__/ # =========================================================================== # example_node.py +qqq* +www* +eee* +aaa* +sss* +ddd* +zzz* +xxx* +ccc* diff --git a/README.md b/README.md index 3918c19..652f388 100644 --- a/README.md +++ b/README.md @@ -2,9 +2,9 @@ > "Simple is better than complex." _— Zen of Python_ -# `Best Resolution` nodes for ComfyUI +# `Best Resolution` ComfyUI nodes [UX/QoL] -A small modular pack — QoL nodes for semi-automatic calculation of the best (most optimal) sampling resolution: +A small modular pack — nodes for semi-automatic calculation of the best (most optimal) sampling resolution: - ... model-agnostic — i.e., compatible with **ANY** model (from now **or** the future), - ... accounting for upscale, - ... and for (the essential) pixel-step. diff --git a/__init__.py b/__init__.py index 7cf92ef..0ba59ab 100644 --- a/__init__.py +++ b/__init__.py @@ -4,9 +4,10 @@ import typing as _t +from .node_crop_pad import BestResolutionUpscaledCropPad +from .nodes_prims import * from .nodes_simple import * from .nodes_upscale import * -from .nodes_prims import * NODE_CLASS_MAPPINGS: _t.Dict[str, type] = { "BestResolutionFromArea": BestResolutionFromArea, @@ -15,6 +16,9 @@ NODE_CLASS_MAPPINGS: _t.Dict[str, type] = { "BestResolutionSimple": BestResolutionSimple, "BestResolutionPrimResPriority": BestResolutionPrimResPriority, + "BestResolutionPrimUpStrategy": BestResolutionPrimUpStrategy, + + "BestResolutionUpscaledCropPad": BestResolutionUpscaledCropPad, } NODE_DISPLAY_NAME_MAPPINGS: _t.Dict[str, str] = { "BestResolutionFromArea": "Best-Res (area)", @@ -23,6 +27,9 @@ NODE_DISPLAY_NAME_MAPPINGS: _t.Dict[str, str] = { "BestResolutionSimple": "Best-Res (simple)", "BestResolutionPrimResPriority": "Priority (Best-Res)", + "BestResolutionPrimUpStrategy": "Upscale Strategy (Best-Res)", + + "BestResolutionUpscaledCropPad": "Upscaled Crop/Pad (Best-Res)", } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/_dataclass.py b/_dataclass.py new file mode 100644 index 0000000..bfedbb9 --- /dev/null +++ b/_dataclass.py @@ -0,0 +1,15 @@ +# encoding: utf-8 +""" +""" + +import typing as _t + +from dataclasses import dataclass as dataclass_with_slots_if_possible +from sys import version_info as _version_info + +if _version_info > (3, 10): + _dataclass_raw = dataclass_with_slots_if_possible + + def dataclass_with_slots_if_possible(*args, **kwargs): + kwargs['slots'] = True + return _dataclass_raw(*args, **kwargs) diff --git a/_funcs.py b/_funcs.py index 74a96a4..13c75d9 100644 --- a/_funcs.py +++ b/_funcs.py @@ -11,6 +11,7 @@ from math import sqrt as _sqrt from server import PromptServer as _PromptServer from .enums import * +from .return_tuples import * _t_number = _t.Union[int, float] @@ -159,12 +160,12 @@ def simple_result_from_approx_wh( width_f: float, height_f: _t_number, step: int, show: bool = True, unique_id: str = None, target_square_size: _t_number = None -): +) -> ResultSimple: """Final part of the main func for simple (non-upscale) nodes - when desired width/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) - result = (width, height) + result = ResultSimple(width, height) if not unique_id: return result @@ -199,7 +200,7 @@ def upscale_result_from_approx_wh( priority: _t.Union[RoundingPriority, str], upscale: float, up_step:int, show: bool = True, unique_id: str = None, target_square_size: _t_number = None -): +) -> ResultUpscaled: """Primary part of the main func for nodes with upscaling - when desired initial-width/height are already calculated.""" upscale = max(float(upscale), 1.0) up_width_f: float = upscale * width_f @@ -227,10 +228,7 @@ def upscale_result_from_approx_wh( float(up_width) / upscale, float(up_height) / upscale, step ) - needs_resize, real_upscale_avg, real_upscale_x, real_upscale_y = _need_post_resize(width, height, up_width, up_height) - out_upscale: float = upscale if needs_resize else real_upscale_avg - - result = (width, height, out_upscale, up_width, up_height, needs_resize) + result = ResultUpscaled(upscale, width, height, up_width, up_height) if not unique_id: return result @@ -249,12 +247,14 @@ def upscale_result_from_approx_wh( f"{prefix}{x}" for x in cur_report.split('\n') )) + needs_resize, real_upscale_avg, real_upscale_x, real_upscale_y = _need_post_resize(width, height, up_width, up_height) + separator_line: str = ( # Real upscale factor in X and Y differ enough so we can't get the up-res from init-res with ANY uniform value: f"\n--- ⚠️ x{real_upscale_x:.3f} / x{real_upscale_y:.3f} ---\n" if needs_resize # The upscale is uniform: - else f"\n--- x{real_upscale_x:.3f} ---\n" + else f"\n--- x{real_upscale_avg:.3f} ---\n" ) text = separator_line.join(report_parts) diff --git a/_funcs_crop_pad.py b/_funcs_crop_pad.py new file mode 100644 index 0000000..efc59bb --- /dev/null +++ b/_funcs_crop_pad.py @@ -0,0 +1,184 @@ +# encoding: utf-8 +""" +The actual behavior of "Upscaled Crop/Pad" node. +""" + +import typing as _t + +from ._dataclass import dataclass_with_slots_if_possible as _dataclass_with_slots_if_possible +from ._funcs import round_pos_int as _round_pos_int, _need_post_resize, _show_text_on_node +from .enums import * +from .return_tuples import * + + +@_dataclass_with_slots_if_possible +class _CropPadInput: + """ + For the ease of passing between the functions, the entire set of inputs for "Upscaled Crop/Pad" node + is internally turned into a dataclass. + """ + upscale: float + + orig_w: int + orig_h: int + up_w: int + up_h: int + + strategy: str + align_x: float + align_y: float + + +def _upscaled_crop_xy_offset(extra_w: int, extra_h: int, align_x: float, align_y: float): + crop_x_origin = _round_pos_int(align_x * extra_w) + crop_y_origin = _round_pos_int((1.0 - align_y) * extra_h) + return crop_x_origin, crop_y_origin + + +def _upscaled_crop(_in: _CropPadInput, real_upscale: float) -> ResultUpscaledCropPad: + raw_upscaled_w = _round_pos_int(real_upscale * _in.orig_w) + raw_upscaled_h = _round_pos_int(real_upscale * _in.orig_h) + extra_w = max(raw_upscaled_w - _in.up_w, 0) + extra_h = max(raw_upscaled_h - _in.up_h, 0) + crop_x_offset, crop_y_offset = _upscaled_crop_xy_offset(extra_w, extra_h, _in.align_x, _in.align_y) + return ResultUpscaledCropPad( + real_upscale, + True, _in.up_w, _in.up_h, crop_x_offset, crop_y_offset, + False, 0, 0, 0, 0, + ) + + +def _upscaled_pad_side_values(extra_w: int, extra_h: int, align_x: float, align_y: float): + pad_left = _round_pos_int(align_x * extra_w) + pad_right = extra_w - pad_left + pad_bottom = _round_pos_int(align_y * extra_h) + pad_top = extra_h - pad_bottom + return pad_left, pad_top, pad_right, pad_bottom + + +def _upscaled_pad(_in: _CropPadInput, real_upscale: float) -> ResultUpscaledCropPad: + raw_upscaled_w = _round_pos_int(real_upscale * _in.orig_w) + raw_upscaled_h = _round_pos_int(real_upscale * _in.orig_h) + pad_w = max(_in.up_w - raw_upscaled_w, 0) + pad_h = max(_in.up_h - raw_upscaled_h, 0) + return ResultUpscaledCropPad( + real_upscale, + False, raw_upscaled_w, raw_upscaled_h, 0, 0, + True, *_upscaled_pad_side_values(pad_w, pad_h, _in.align_x, _in.align_y), + ) + + +def _upscaled_crop_delta_pixels(_in: _CropPadInput, result: ResultUpscaledCropPad) -> int: + """Calculate the area of total cropped patch for crop-only mode.""" + raw_upscaled_w = _round_pos_int(result.upscale * _in.orig_w) + raw_upscaled_h = _round_pos_int(result.upscale * _in.orig_h) + cropped_w = min(result.crop_width, raw_upscaled_w) + cropped_h = min(result.crop_height, raw_upscaled_h) + return abs(raw_upscaled_w * raw_upscaled_h - cropped_w * cropped_h) + + +def _upscaled_pad_delta_pixels(_in: _CropPadInput, result: ResultUpscaledCropPad) -> int: + """Calculate the area of total out-painted patch for pad-only mode.""" + raw_upscaled_w = _round_pos_int(result.upscale * _in.orig_w) + raw_upscaled_h = _round_pos_int(result.upscale * _in.orig_h) + padded_w = raw_upscaled_w + max(result.pad_left, 0) + max(result.pad_right, 0) + padded_h = raw_upscaled_h + max(result.pad_top, 0) + max(result.pad_bottom, 0) + return abs(padded_w * padded_h - raw_upscaled_w * raw_upscaled_h) + + +def _upscaled_crop_pad(_in: _CropPadInput) -> ResultUpscaledCropPad: + """The actual function for "Upscaled Crop/Pad" node - extracted to wrap it with displaying the message.""" + needs_resize, real_upscale_avg, real_upscale_x, real_upscale_y = _need_post_resize(_in.orig_w, _in.orig_h, _in.up_w, _in.up_h) + if not needs_resize: + return ResultUpscaledCropPad( + real_upscale_avg, + False, _in.up_w, _in.up_h, 0, 0, + False, 0, 0, 0, 0, + ) + + _in.align_x = min(max(_in.align_x, 0.0), 1.0) + _in.align_y = min(max(_in.align_y, 0.0), 1.0) + + if _in.strategy == UpscaledCropPadStrategy.CROP: + return _upscaled_crop(_in, max(real_upscale_x, real_upscale_y)) + + if _in.strategy == UpscaledCropPadStrategy.PAD: + return _upscaled_pad(_in, min(real_upscale_x, real_upscale_y)) + + if _in.strategy == UpscaledCropPadStrategy.NEAREST: + crop_result = _upscaled_crop(_in, max(real_upscale_x, real_upscale_y)) + pad_result = _upscaled_pad(_in, min(real_upscale_x, real_upscale_y)) + return ( + pad_result + if _upscaled_pad_delta_pixels(_in, pad_result) < _upscaled_crop_delta_pixels(_in, crop_result) + else crop_result + ) + + assert _in.strategy == UpscaledCropPadStrategy.EXACT_UPSCALE + + upscale = _in.upscale + raw_upscaled_w = _round_pos_int(upscale * _in.orig_w) + raw_upscaled_h = _round_pos_int(upscale * _in.orig_h) + # We're in the most complex case. These ^ can be both below and above the target up-res. + # So, we need to follow the whole upscale->crop->pad chain: we'll just get zeroes where there will be nothing to do. + + cropped_w = min(_in.up_w, raw_upscaled_w) + cropped_h = min(_in.up_h, raw_upscaled_h) + crop_extra_w = raw_upscaled_w - cropped_w + crop_extra_h = raw_upscaled_h - cropped_h + crop_x_offset, crop_y_offset = _upscaled_crop_xy_offset(crop_extra_w, crop_extra_h, _in.align_x, _in.align_y) + + pad_extra_w = _in.up_w - cropped_w + pad_extra_h = _in.up_h - cropped_h + pad_side_values = _upscaled_pad_side_values(pad_extra_w, pad_extra_h, _in.align_x, _in.align_y) + return ResultUpscaledCropPad( + upscale, + crop_extra_w > 0 or crop_extra_h > 0, cropped_w, cropped_h, crop_x_offset, crop_y_offset, + pad_extra_w > 0 or pad_extra_h > 0, *pad_side_values, + ) + + +def upscaled_crop_pad( + upscale: float, + orig_w: int, orig_h: int, up_w: int, up_h: int, + strategy: str, + align_x: float, align_y: float, + show: bool = True, + unique_id: str = None +) -> ResultUpscaledCropPad: + """ + Detect, whether some cropping/out-painting needs to be done after upscaling the original image. + And if so - what are the values for them. + """ + # noinspection PyArgumentList + result = _upscaled_crop_pad(_CropPadInput( + upscale, + orig_w, orig_h, up_w, up_h, + strategy, + align_x, align_y, + )) + + if not unique_id: + return result + + if not show: + _show_text_on_node(None, unique_id) + return result + + if not(result.do_crop or result.do_padding): + _show_text_on_node(f"x{result.upscale:.3f}✅", unique_id) + return result + + # Either crop, or pad, or both + text_parts: _t.List[str] = [f"x{result.upscale:.3f}", ] + + if result.do_crop: + crop_r = up_w - (result.crop_x_origin + result.crop_width) + crop_b = up_h - (result.crop_y_origin + result.crop_height) + text_parts.append(f"Crop:\nL {result.crop_x_origin}, R {crop_r}, B {crop_b}, T {result.crop_y_origin}") + + if result.do_padding: + text_parts.append(f"Padding:\nL {result.pad_left}, R {result.pad_right}, B {result.pad_bottom}, T {result.pad_top}") + + _show_text_on_node('\n\n'.join(text_parts), unique_id) + return result diff --git a/_meta.py b/_meta.py new file mode 100644 index 0000000..d7d59fa --- /dev/null +++ b/_meta.py @@ -0,0 +1,6 @@ +# encoding: utf-8 +""" +Metadata-related module. +""" + +category: str = "Best Resolution" diff --git a/docstring_formatter.py b/docstring_formatter.py new file mode 100644 index 0000000..b8a7ede --- /dev/null +++ b/docstring_formatter.py @@ -0,0 +1,146 @@ +# encoding: utf-8 +""" +""" + +import typing as _t + +from inspect import getdoc as _getdoc +import re as _re + + +_re_indent_match = _re.compile("(\t*)( +)(\t*)(.*?)$").match +_re_tab_indent_match = _re.compile("(\t+)(.*?)$").match +_re_list_line_match = _re.compile( + "(\s*)(" + "[-*•]+" + "|" + "[a-zA-Z]\s*[.)]" + "|" + "[0-9+]\s*[.)]" + ")\s+" +).match + + +def _recover_tab_indents(line: str, tab_size: int): + """Turn indenting spaces back to tabs using regexp. Half-tab indents are rounded.""" + assert bool(line) and isinstance(line, str) + + n_tabs = 0.0 + + match = _re_indent_match(line) + while match: + pre_tabs, spaces, post_tabs, line = match.groups() + n_tabs_from_spaces = float(len(spaces)) / tab_size + 0.00001 + n_post_tabs = len(post_tabs) + if n_post_tabs > 0: + # There are tabs after spaces. Don't preserve the fractional spaces-indent, truncate it: + n_tabs_from_spaces = int(n_tabs_from_spaces) + n_tabs += len(pre_tabs) + n_tabs_from_spaces + n_post_tabs + match = _re_indent_match(line) + + if n_tabs < 0.5: + return line + + tabs_prefix = '\t' * int(n_tabs + 0.50001) + return f"{tabs_prefix}{line}" + + +def _join_paragraph_and_format_tabs(paragraph: _t.List[str], tab_size: int): + """ + Given "continuous" paragraph (i.e., with no empty newlines between chunks), recover tabs for each chunk + and join them together into a single actual line. + Works as a generator to account for blocks with different indents - to make each its own line. + """ + pending_indent = 0 + pending_chunks: _t.List[str] = list() + + def join_pending_chunks() -> str: + return "{}{}".format('\t' * pending_indent, ' '.join(pending_chunks)) + + for chunk in paragraph: + chunk = _recover_tab_indents(chunk, tab_size) + + cur_indent = 0 + match = _re_tab_indent_match(chunk) + if match: + tab_indent, chunk = match.groups() # We've detected indent. Now, get rid of it. + cur_indent = len(tab_indent) + + match_list_line = _re_list_line_match(chunk) + # In case of a bulleted/numbered list, we'll need to start a new block, too. + if cur_indent == pending_indent and not match_list_line: + pending_chunks.append(chunk) + continue + + # Indent mismatch or a list line: + # we're either ended one block or entered another. Either way, the previous block ends. + if pending_chunks: + yield join_pending_chunks() + pending_chunks = list() + assert not pending_chunks + pending_chunks.append(chunk) + pending_indent = cur_indent + + if pending_chunks: + yield join_pending_chunks() + + +def _formatted_paragraphs_gen(doc: str, tab_size: int): + """ + Generator, which splits docstring into lines and transforms them into an actual printable output: + - From each bulk of empty lines, the first one is skipped... + - ... thus, non-empty lines are joined into continuous paragraphs. + - Recover tabs in the beginning oh lines (``inspect.cleandoc()`` converts them into spaces). + """ + if not doc: + return + doc = str(doc) + if not doc.strip(): + return + + tab_size = max(int(tab_size), 1) + + cur_paragraph: _t.List[str] = list() + + for line in doc.splitlines(): + line: str = line.rstrip() + if line: + cur_paragraph.append(line) + continue + + assert not line + if cur_paragraph: + for block in _join_paragraph_and_format_tabs(cur_paragraph, tab_size): + yield block + cur_paragraph = list() + # Just skip the current empty line entirely - do nothing with it. + continue + + # We're in a chain of empty lines, and we've already skipped the first one. Preserve the remaining ones: + yield '' + + # Return the last paragraph post-loop: + if cur_paragraph: + for block in _join_paragraph_and_format_tabs(cur_paragraph, tab_size): + yield block + + +def format_docstring(doc: str, tab_size: int = 8) -> str: + """ + Turn a pre-cleaned-up docstring (with tabs as spaces and newlines mid-sentence) + into an actually printable output: + - mid-paragraph new lines are replaced with spaces... + - ... while still keeping indented blocks separate. + + Remember to pass a pre-cleaned-up docstring - i.e.: ``format_docstring(inspect.cleandoc(__doc__))`` + """ + return '\n'.join(_formatted_paragraphs_gen(doc, tab_size)) + + +def format_object_docstring(_obj, tab_size: int = 8) -> str: + """Find the object's docstring and format it with ``format_docstring()``""" + doc = _getdoc(_obj) + if not doc: + return '' + # noinspection PyArgumentList + return format_docstring(doc, tab_size=tab_size) diff --git a/enums.py b/enums.py index 99b4871..f3fce28 100644 --- a/enums.py +++ b/enums.py @@ -11,11 +11,47 @@ except ImportError: from comfy.comfy_types.node_typing import StrEnum as _StrEnum -class RoundingPriority(_StrEnum): +class __BaseEnum(_StrEnum): + @classmethod + def all_values(cls) -> _t.Tuple[str, ...]: + return tuple(x.value for x in cls) + + +class RoundingPriority(__BaseEnum): + """ + Defines what resolution to prioritize, as well as which order to perform rounding in: + + + • desired - first, approximate resolutions are calculated for both initial and upscaled size; then, both are rounded. + In both cases, the rounded resolution is closest to the desired one, + but aspect ratio might differ the most between sizes. + • original - first, initial resolution is calculated and rounded; then, upscaled one is detected from it. + Upscaled resolution might differ the most from the desired aspect ratio, but it follows the ratio from + initial size as much as possible. + • upscaled - vice versa: first, the rounded upscaled resolution is calculated; + then, the initial one back-tracked from it. + """ DESIRED = 'desired' ORIGINAL = 'original' UPSCALED = 'upscaled' - @classmethod - def all_values(cls) -> _t.Tuple[str, ...]: - return tuple(x.value for x in cls) + +class UpscaledCropPadStrategy(__BaseEnum): + """ + If the upscaled resolution can't be achieved by uniform scaling of the initial-res, + how to tweak the image to get there: + + + • pad - upscale the image to fit it into the desired frame on one side, + then add missing pixels on the other one for outpaint. + • crop - upscale the image to fill the entire frame, then crop extra pixels on one side. + • nearest - automatically choose one of the above, to crop/pad the least number of pixels (by area). + • exact-upscale - follow the provided upscale-value precisely. This is the only option + that uses it. Also, it's the only one that might require both outpainting and crop. + + In case when both are required, the output values assume this order: upscale -> crop -> pad. + """ + PAD = 'pad' + CROP = 'crop' + NEAREST = 'nearest' + EXACT_UPSCALE = 'exact-upscale' diff --git a/node_crop_pad.py b/node_crop_pad.py new file mode 100644 index 0000000..97e3d3b --- /dev/null +++ b/node_crop_pad.py @@ -0,0 +1,119 @@ +# encoding: utf-8 +""" +Nodes to calculate crop / padding values (for out-paint). +""" + +import typing as _t + +from inspect import cleandoc as _cleandoc + +from frozendict import deepfreeze as _deepfreeze, frozendict as _frozendict + +from comfy.comfy_types.node_typing import IO as _IO + +from ._funcs_crop_pad import upscaled_crop_pad as _upscaled_crop_pad +from . import _meta +from .docstring_formatter import format_docstring as _format_docstring +from .enums import * +from .nodes_prims import _up_strategy_in_type, _up_strategy_verify +from .nodes_upscale import _return_ttips_upscale +from .slot_types import ( + type_dict_res as _type_dict_res, + rel_pos_in_type as _rel_pos_in_type, + upscale_in_type as _upscale_in_type +) + +# ---------------------------------------------------------- + +_return_types_crop_pad = ( + _IO.FLOAT, + _IO.BOOLEAN, _IO.INT, _IO.INT, _IO.INT, _IO.INT, + _IO.BOOLEAN, _IO.INT, _IO.INT, _IO.INT, _IO.INT, +) +_return_ttips_crop_pad = _frozendict({ + 'upscale': "The actual uniform upscale value to do.", + + 'do_crop': "Whether post-upscale crop needs to be done to get the specified resolution.", + 'crop_width': "Width for the post-upscale crop.", + 'crop_height': "Height for the post-upscale crop.", + 'crop_x': "X-offset for the post-upscale crop.", + 'crop_y': "Y-offset for the post-upscale crop.", + + 'do_padding': "Whether post-upscale padding and out-paint need to be done to get the specified resolution.", + 'pad_left': "Left padding for the post-upscale out-paint.", + 'pad_top': "Top padding for the post-upscale out-paint.", + 'pad_right': "Right padding for the post-upscale out-paint.", + 'pad_bottom': "Bottom padding for the post-upscale out-paint.", +}) +_input_types_crop_pad = _deepfreeze({ + 'required': { + 'upscale': (_IO.FLOAT, dict(_upscale_in_type[1], **{ + 'tooltip': f"Only used when {UpscaledCropPadStrategy.EXACT_UPSCALE!r} strategy selected." + })), + 'orig_w': (_IO.INT, dict(_type_dict_res, **{'tooltip': _return_ttips_upscale['orig_w']})), + 'orig_h': (_IO.INT, dict(_type_dict_res, **{'tooltip': _return_ttips_upscale['orig_h']})), + 'up_w': (_IO.INT, dict(_type_dict_res, **{'tooltip': _return_ttips_upscale['up_w']})), + 'up_h': (_IO.INT, dict(_type_dict_res, **{'tooltip': _return_ttips_upscale['up_h']})), + 'strategy': _up_strategy_in_type, + 'align_x': (_IO.FLOAT, dict(_rel_pos_in_type[1], **{ + 'tooltip': ( + "Where's the image pivot for horizontal alignment (which side stays in place during crop/out-paint):\n" + "0 - Left\n" + "0.5 - Center\n" + "1 - Right" + ), + })), + 'align_y': (_IO.FLOAT, dict(_rel_pos_in_type[1], **{ + 'tooltip': ( + "Where's the image pivot for vertical alignment (which side stays in place during crop/out-paint):\n" + "0 - Bottom\n" + "0.5 - Center\n" + "1 - Top" + ), + 'default': 0.0 + })), + }, + 'hidden': { + 'unique_id': 'UNIQUE_ID', + }, + # 'optional': {}, +}) + + +class BestResolutionUpscaledCropPad: + """ + If the original resolution can't be scaled to the higher one uniformly (i.e., without stretching), + this node detects the appropriate values for cropping/out-painting. + """ + NODE_NAME = 'BestResolutionUpscaledCropPad' + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) + + OUTPUT_NODE = True + + FUNCTION = 'main' + RETURN_TYPES = _return_types_crop_pad + RETURN_NAMES = tuple(_return_ttips_crop_pad.keys()) + OUTPUT_TOOLTIPS = tuple(_return_ttips_crop_pad.values()) + + @classmethod + def INPUT_TYPES(cls): + return _input_types_crop_pad + + def main( + self, + upscale: float, + orig_w: int, orig_h: int, up_w: int, up_h: int, + strategy: _t.Union[UpscaledCropPadStrategy, str], + align_x: float, align_y: float, + # show: bool, + unique_id: str = None + ): + return _upscaled_crop_pad( + upscale, + orig_w, orig_h, up_w, up_h, + _up_strategy_verify(strategy), + align_x, align_y, + # show, + unique_id=unique_id + ) diff --git a/nodes_prims.py b/nodes_prims.py index 1844027..3d7fe71 100644 --- a/nodes_prims.py +++ b/nodes_prims.py @@ -9,44 +9,27 @@ from inspect import cleandoc as _cleandoc from frozendict import deepfreeze as _deepfreeze +from . import _meta +from .docstring_formatter import ( + format_docstring as _format_docstring, + format_object_docstring as _format_object_docstring +) from .enums import * # ---------------------------------------------------------- _res_priority_data_type = RoundingPriority.all_values() _res_priority_data_type_set = set(_res_priority_data_type) - -# _priority_in_type = (_IO.BOOLEAN, { -# 'default': True, 'label_on': 'upscaled', 'label_off': 'original', -# 'tooltip': ( -# "Which of the resolutions is more important to be as close to the desired as possible:\n\n" -# "When ON, the rounded with/height are first calculated for the UPSCALED resolution - and " -# "only then the original ones are back-tracked from them.\n" -# "When OFF, they're first calculated for the ORIGINAL resolution - and " -# "the upscaled ones are tracked forward from them." -# ) -# }) _res_priority_in_type = ( _res_priority_data_type, { 'default': RoundingPriority.DESIRED.value, - 'tooltip': ( - "Defines what resolution to prioritize, as well as which order to perform rounding in:\n\n" - "• desired - first, approximate resolutions are calculated for both initial and upscaled size; " - "then, both are rounded. In both cases, the rounded resolution is closest to the desired one, " - "but aspect ratio might differ the most between sizes.\n\n" - "• original - first, initial resolution is calculated and rounded; then, upscaled one is detected from it. " - "Upscaled resolution might differ the most from the desired aspect ratio, but it follows the ratio from " - "initial size as much as possible.\n\n" - "• upscaled - vice versa: first, the rounded upscaled resolution is calculated; then, the initial one " - "back-tracked from it." - ) + 'tooltip': _format_object_docstring(RoundingPriority), } ) - -_input_types_prim_priority = _deepfreeze({ +_input_types_res_priority = _deepfreeze({ 'required': { - 'priority': _res_priority_in_type + 'priority': _res_priority_in_type, }, # 'hidden': { # 'unique_id': 'UNIQUE_ID', @@ -55,24 +38,30 @@ _input_types_prim_priority = _deepfreeze({ }) +def _res_priority_verify(priority: _t.Union[RoundingPriority, str]) -> str: + if priority not in _res_priority_data_type_set: + raise ValueError(f"Invalid value for resolution priority: {priority!r}\nExpected one of: {_res_priority_data_type!r}") + return str(priority) + + class BestResolutionPrimResPriority: """ 'priority' selector for "Best Resolution" upscale-nodes. """ NODE_NAME = 'BestResolutionPrimResPriority' - CATEGORY = "utils/resolution" - DESCRIPTION = _cleandoc(__doc__) + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) OUTPUT_NODE = False FUNCTION = 'main' - RETURN_TYPES = (_res_priority_data_type,) + RETURN_TYPES = (_res_priority_data_type, ) RETURN_NAMES = ('priority', ) # OUTPUT_TOOLTIPS = ('', ) @classmethod def INPUT_TYPES(cls): - return _input_types_prim_priority + return _input_types_res_priority def main( self, priority: _t.Union[RoundingPriority, str], @@ -81,8 +70,56 @@ class BestResolutionPrimResPriority: ): return (_res_priority_verify(priority), ) +# ---------------------------------------------------------- -def _res_priority_verify(priority: _t.Union[RoundingPriority, str]) -> str: - if priority not in _res_priority_data_type_set: - raise ValueError(f"Invalid value for resolution priority: {priority!r}\nExpected one of: {_res_priority_data_type!r}") - return str(priority) +_up_strategy_data_type = UpscaledCropPadStrategy.all_values() +_up_strategy_data_type_set = set(_up_strategy_data_type) +_up_strategy_in_type = ( + _up_strategy_data_type, + { + 'default': UpscaledCropPadStrategy.EXACT_UPSCALE.value, + 'tooltip': _format_object_docstring(UpscaledCropPadStrategy), + } +) +_input_types_up_strategy = _deepfreeze({ + 'required': { + 'strategy': _up_strategy_in_type, + }, + # 'hidden': { + # 'unique_id': 'UNIQUE_ID', + # }, + # 'optional': {}, +}) + + +def _up_strategy_verify(strategy: _t.Union[RoundingPriority, str]) -> str: + if strategy not in _up_strategy_data_type_set: + raise ValueError(f"Invalid value for upscale strategy: {strategy!r}\nExpected one of: {_up_strategy_data_type!r}") + return str(strategy) + + +class BestResolutionPrimUpStrategy: + """ + 'strategy' selector for "Upscaled Crop/Pad" node in "Best Resolution" pack. + """ + NODE_NAME = 'BestResolutionPrimUpStrategy' + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) + + OUTPUT_NODE = False + + FUNCTION = 'main' + RETURN_TYPES = (_up_strategy_data_type, ) + RETURN_NAMES = ('strategy', ) + # OUTPUT_TOOLTIPS = ('', ) + + @classmethod + def INPUT_TYPES(cls): + return _input_types_up_strategy + + def main( + self, strategy: _t.Union[RoundingPriority, str], + # show: bool, + # unique_id: str = None + ): + return (_up_strategy_verify(strategy), ) diff --git a/nodes_simple.py b/nodes_simple.py index 1418ae9..1ab39ca 100644 --- a/nodes_simple.py +++ b/nodes_simple.py @@ -19,7 +19,12 @@ from ._funcs import ( float_width_height_from_area as _float_width_height_from_area, simple_result_from_approx_wh as _simple_result_from_approx_wh ) -from .node_types import * +from . import _meta +from .docstring_formatter import format_docstring as _format_docstring +from .slot_types import ( + type_dict_res as _type_dict_res, + type_dict_step_init as _type_dict_step_init +) # ---------------------------------------------------------- @@ -31,9 +36,9 @@ _return_names_simple = ('width', 'height') # A tiny optimization by reusing the same immutable dict: _input_types_simple = _deepfreeze({ 'required': { - 'width': (_IO.INT, dict(type_dict_res, **{'tooltip': "Approximate width"})), - 'height': (_IO.INT, dict(type_dict_res, **{'tooltip': "Approximate height"})), - 'step': (_IO.INT, dict(type_dict_step_init, **{'tooltip': ( + 'width': (_IO.INT, dict(_type_dict_res, **{'tooltip': "Approximate width"})), + 'height': (_IO.INT, dict(_type_dict_res, **{'tooltip': "Approximate height"})), + 'step': (_IO.INT, dict(_type_dict_step_init, **{'tooltip': ( "Both width and height will be divisible by this value - by rounding them " "to the closest appropriate resolution.\n\n" "The default 48 is (8 * 3 * 2), so it's a safe choice because:\n" @@ -62,8 +67,8 @@ class BestResolutionSimple: only rounding performed, desired image size specified directly. """ NODE_NAME = 'BestResolutionSimple' - CATEGORY = "utils/resolution" - DESCRIPTION = _cleandoc(__doc__) + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) OUTPUT_NODE = True @@ -97,7 +102,7 @@ _tooltip_aspect = ( _input_types_orient = _deepfreeze({ 'required': { - 'size': (_IO.INT, dict(type_dict_res, **{'tooltip': ( + 'size': (_IO.INT, dict(_type_dict_res, **{'tooltip': ( "Approximate size of one of the image sides.\nWhich one - see the 'size_is_big' and 'landscape' toggles." )})), 'step': _input_types_simple['required']['step'], @@ -136,8 +141,8 @@ class BestResolutionFromAspectRatio: image size selected indirectly - by one of the sides + aspect ratio. """ NODE_NAME = 'BestResolutionFromAspectRatio' - CATEGORY = "utils/resolution" - DESCRIPTION = _cleandoc(__doc__) + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) OUTPUT_NODE = True @@ -177,7 +182,7 @@ class BestResolutionFromAspectRatio: _input_types_area = _deepfreeze({ 'required': { - 'square_size': (_IO.INT, dict(type_dict_res, **{'tooltip': ( + '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" @@ -202,18 +207,19 @@ 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__) + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) OUTPUT_NODE = True diff --git a/nodes_upscale.py b/nodes_upscale.py index 83c25b6..52d0e37 100644 --- a/nodes_upscale.py +++ b/nodes_upscale.py @@ -18,39 +18,31 @@ from ._funcs import ( float_width_height_from_area as _float_width_height_from_area, upscale_result_from_approx_wh as _upscale_result_from_approx_wh ) +from . import _meta +from .docstring_formatter import format_docstring as _format_docstring from .enums import * -from .node_types import * from .nodes_simple import _input_types_area from .nodes_prims import _res_priority_in_type, _res_priority_verify +from .slot_types import ( + type_dict_step_upscale1 as _type_dict_step_upscale1, + upscale_in_type as _upscale_in_type +) # ---------------------------------------------------------- -_return_types_upscale = (_IO.INT, _IO.INT, _IO.FLOAT, _IO.INT, _IO.INT, _IO.BOOLEAN) +_return_types_upscale = (_IO.FLOAT, _IO.INT, _IO.INT, _IO.INT, _IO.INT) _return_ttips_upscale = _frozendict({ + 'upscale': "Simply outputs the same upscale-value you've set on the node.", 'orig_w': "Width for original/initial image", 'orig_h': "Height for original/initial image", - 'upscale': ( - "If 'needs_resize' is FALSE, this would be the actual uniform value to scale your initial-res in order " - "to get the upscaled-res.\n" - "If 'needs_resize' is TRUE, this slot simply outputs the same upscale-value you've set on the node." - ), 'up_w': "Width for the (main) upscaled image", 'up_h': "Height for the (main) upscaled image", - 'needs_resize': ( - "This will be FALSE if you can get the exact upscaled resolution by simply scaling the initial-res uniformly " - "(by the value output into 'upscale' slot).\n\n" - "Otherwise, it will be TRUE — indicating that you'll need to do some cropping/out-painting right AFTER " - "the actual upscale (increasing the resolution itself) but BEFORE post-upscale sampling (\"HD-fix\")." - ), }) __extra_inputs_for_upscale_only = { 'priority': _res_priority_in_type, - 'upscale': (_IO.FLOAT, { - 'default': 1.5, 'min': 1.0, 'max': _sys.float_info.max, 'step': 0.25, 'round': 0.001, - # 'tooltip': "", # TODO - }), - 'up_step': (_IO.INT, dict(type_dict_step_upscale1, **{ + 'upscale': _upscale_in_type, + 'up_step': (_IO.INT, dict(_type_dict_step_upscale1, **{ 'tooltip': "Same as the main `step`, but for the upscaled resolution.", })), } @@ -75,18 +67,19 @@ class BestResolutionFromAreaUpscale: ... PLUS, account for the immediate upscale right away. + 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 = 'BestResolutionFromAreaUpscale' - CATEGORY = "utils/resolution" - DESCRIPTION = _cleandoc(__doc__) + CATEGORY = _meta.category + DESCRIPTION = _format_docstring(_cleandoc(__doc__)) OUTPUT_NODE = True diff --git a/pyproject.toml b/pyproject.toml index 61cb4db..3d638e3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "best-resolution" description = "QoL nodes for semi-automatic calculation of the best (most optimal) sampling resolution \n• compatible with ANY model (from now or the future), \n• accounting for upscale... \n• ...and pixel-step." -version = "1.0.2" +version = "1.1.0" authors = [ {name = "Lex Darlog"} ] diff --git a/return_tuples.py b/return_tuples.py new file mode 100644 index 0000000..f8d1e3f --- /dev/null +++ b/return_tuples.py @@ -0,0 +1,38 @@ +# encoding: utf-8 +""" +For readability, all nodes in this node pack return NamedTuples instead of regular tuples. +""" + +import typing as _t + + +class ResultSimple(_t.NamedTuple): + """Returned NamedTuple for simple (non-upscale) nodes.""" + width: int + height: int + + +class ResultUpscaled(_t.NamedTuple): + """Returned NamedTuple for nodes with upscaling.""" + upscale: float + orig_width: int + orig_height: int + up_width: int + up_height: int + + +class ResultUpscaledCropPad(_t.NamedTuple): + """Returned NamedTuple for "Upscaled Crop/Pad" node.""" + upscale: float + + do_crop: bool + crop_width: int + crop_height: int + crop_x_origin: int + crop_y_origin: int + + do_padding: bool + pad_left: int + pad_top: int + pad_right: int + pad_bottom: int diff --git a/node_types.py b/slot_types.py similarity index 64% rename from node_types.py rename to slot_types.py index 8c6b2a2..7118d6c 100644 --- a/node_types.py +++ b/slot_types.py @@ -7,6 +7,8 @@ import typing as _t import sys as _sys +from comfy.comfy_types.node_typing import IO as _IO + def number_type_dict(default: int = 512, min: int = 1, max: int = _sys.maxsize, step: int = 1) -> _t.Dict[str, int]: """Prepare type-dict for an input int/float parameter.""" @@ -20,6 +22,14 @@ def number_type_dict(default: int = 512, min: int = 1, max: int = _sys.maxsize, } +type_dict_res = number_type_dict(1024) type_dict_step_init = number_type_dict(48) type_dict_step_upscale1 = number_type_dict(128) -type_dict_res = number_type_dict(1024) + +upscale_in_type = (_IO.FLOAT, { + 'default': 1.5, 'min': 1.0, 'max': _sys.float_info.max, 'step': 0.25, 'round': 0.001, + # 'tooltip': "", # TODO +}) +rel_pos_in_type = (_IO.FLOAT, { + 'default': 0.5, 'min': 0.0, 'max': 1.0, 'step': 0.125, 'round': 0.0001, +})