336 lines
11 KiB
Python
336 lines
11 KiB
Python
# encoding: utf-8
|
|
"""
|
|
The actual nodes defined by the package.
|
|
"""
|
|
|
|
import typing as _t
|
|
|
|
from inspect import cleandoc as _cleandoc
|
|
from math import sqrt as _sqrt
|
|
import sys as _sys
|
|
|
|
from frozendict import deepfreeze, frozendict
|
|
|
|
from comfy.comfy_types.node_typing import IO as _IO
|
|
|
|
from .node_types import *
|
|
from .enums import *
|
|
from ._funcs import (
|
|
aspect_ratios_sorted as _aspect_ratios_sorted,
|
|
number_to_int as _number_to_int,
|
|
float_width_height_from_area as _float_width_height_from_area,
|
|
simple_result_from_approx_wh as _simple_result_from_approx_wh,
|
|
upscale_result_from_approx_wh as _upscale_result_from_approx_wh
|
|
)
|
|
|
|
|
|
# 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': (
|
|
"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"
|
|
"- it's compatible with SD1.5/XL downsampling factor (divisible by 8),\n"
|
|
"- it can be upscaled by x1.5 or x1.333 at the first iteration, which is optimal for latent-upscale,\n"
|
|
"- after x1.5 upscale, if you only do x2 later (it's OK for already high resolutions) - it will be divisible "
|
|
"by 3 AND 9, which might become handy at that point, where you'll probably use UltimateSDUpscale.\n\n"
|
|
"Other values worth trying first: 64, 96, 128."
|
|
)})),
|
|
# 'show': (_IO.BOOLEAN, {
|
|
# 'default': True,
|
|
# # 'label_on': 'on', 'label_off': 'off',
|
|
# 'tooltip': "Show the result on the node itself?",
|
|
# }),
|
|
},
|
|
'hidden': {
|
|
'unique_id': 'UNIQUE_ID', # used for text display at the bottom of the node
|
|
},
|
|
# 'optional': {},
|
|
})
|
|
|
|
# _return_types_simple = (_IO.INT, _IO.INT, _IO.STRING)
|
|
# _return_names_simple = ('width', 'height', 'report')
|
|
_return_types_simple = (_IO.INT, _IO.INT)
|
|
_return_names_simple = ('width', 'height')
|
|
|
|
|
|
class BestResolutionSimple:
|
|
"""
|
|
The simplest node for easier selection of an optimal resolution:
|
|
only rounding performed, desired image size specified directly.
|
|
"""
|
|
NODE_NAME = 'BestResolutionSimple'
|
|
CATEGORY = "utils/resolution"
|
|
DESCRIPTION = _cleandoc(__doc__)
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = 'main'
|
|
RETURN_TYPES = _return_types_simple
|
|
# RETURN_TYPES_TOOLTIPS = {}
|
|
RETURN_NAMES = _return_names_simple
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return _input_types_simple
|
|
|
|
def main(
|
|
self, width: int, height: int, step: int,
|
|
# show: bool,
|
|
unique_id: str = None
|
|
):
|
|
return _simple_result_from_approx_wh(
|
|
float(width), height, step,
|
|
# show,
|
|
unique_id=unique_id, target_square_size=_sqrt(width * height)
|
|
)
|
|
|
|
|
|
_tooltip_aspect = (
|
|
"Two aspects together define an aspect ratio (16:9, 4:3, etc).\n"
|
|
"Order doesn't matter: image orientation is defined by the 'landscape' toggle.\n\n"
|
|
"The specified aspect ratio is APPROXIMATE: step parameter has priority over the exact image proportions."
|
|
)
|
|
|
|
_input_types_orient = deepfreeze({
|
|
'required': {
|
|
'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'],
|
|
'size_is_big': (_IO.BOOLEAN, {
|
|
'default': True, 'label_on': 'bigger', 'label_off': 'smaller',
|
|
'tooltip': "When ON, size parameter represents the bigger image side.\nWhen OFF, it specifies the smaller one."
|
|
}),
|
|
'landscape': (_IO.BOOLEAN, {
|
|
'default': True, 'label_on': 'landscape', 'label_off': 'portrait',
|
|
'tooltip': (
|
|
"Specifies image orientation:\n\n"
|
|
"When ON, width is bigger (image is horizontal).\n"
|
|
"When OFF, height is bigger (image is vertical)."
|
|
),
|
|
}),
|
|
'aspect_a': (_IO.FLOAT, {
|
|
'default': 16.0, 'min': 1.0, 'max': _sys.float_info.max, 'step': 1.0, 'round': 0.001,
|
|
'tooltip': _tooltip_aspect
|
|
}),
|
|
'aspect_b': (_IO.FLOAT, {
|
|
'default': 9.0, 'min': 1.0, 'max': _sys.float_info.max, 'step': 1.0, 'round': 0.001,
|
|
'tooltip': _tooltip_aspect
|
|
}),
|
|
# 'show': _input_types_simple['required']['show'],
|
|
},
|
|
'hidden': {
|
|
'unique_id': 'UNIQUE_ID',
|
|
},
|
|
# 'optional': {},
|
|
})
|
|
|
|
|
|
class BestResolutionFromAspectRatio:
|
|
"""
|
|
A more advanced node for easier selection of an optimal resolution:
|
|
image size selected indirectly - by one of the sides + aspect ratio.
|
|
"""
|
|
NODE_NAME = 'BestResolutionFromAspectRatio'
|
|
CATEGORY = "utils/resolution"
|
|
DESCRIPTION = _cleandoc(__doc__)
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = 'main'
|
|
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,
|
|
# show: bool,
|
|
unique_id: str = None
|
|
):
|
|
aspect_big, aspect_small = _aspect_ratios_sorted(aspect_a, aspect_b)
|
|
|
|
side_main_f = float(size)
|
|
side_other_f = side_main_f * (
|
|
(aspect_small / aspect_big) if size_is_big else (aspect_big / aspect_small)
|
|
)
|
|
|
|
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
|
|
|
|
return _simple_result_from_approx_wh(
|
|
width_f, height_f, step,
|
|
# show,
|
|
unique_id=unique_id, target_square_size=_sqrt(width_f * height_f)
|
|
)
|
|
|
|
|
|
_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'],
|
|
# 'show': _input_types_simple['required']['show'],
|
|
},
|
|
'hidden': {
|
|
'unique_id': 'UNIQUE_ID',
|
|
},
|
|
# '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
|
|
|
|
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,
|
|
# show: bool,
|
|
unique_id: str = None
|
|
):
|
|
square_size: int = _number_to_int(square_size)
|
|
width_f, height_f = _float_width_height_from_area(square_size, landscape, aspect_a, aspect_b)
|
|
return _simple_result_from_approx_wh(
|
|
width_f, height_f, step,
|
|
# show,
|
|
unique_id=unique_id, target_square_size=square_size
|
|
)
|
|
|
|
|
|
_return_types_upscale = (_IO.INT, _IO.INT, _IO.INT, _IO.INT)
|
|
_return_ttips_upscale = frozendict({
|
|
'orig_w': "Width for original image",
|
|
'orig_h': "Height for original image",
|
|
'up_w': "Width for upscaled image",
|
|
'up_h': "Height for upscaled image"
|
|
})
|
|
_return_names_upscale = tuple(_return_ttips_upscale.keys())
|
|
|
|
|
|
|
|
|
|
|
|
# __priority_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."
|
|
# )
|
|
# })
|
|
__priority_type = (
|
|
RoundingPriority.all_values(),
|
|
{
|
|
'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."
|
|
)
|
|
}
|
|
)
|
|
__extra_inputs_for_upscale_only = {
|
|
'priority': __priority_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, **{
|
|
'tooltip': "Same as the main `step`, but for the upscaled resolution.",
|
|
})),
|
|
}
|
|
|
|
_input_types_area_upscale = deepfreeze({
|
|
'required': dict(
|
|
(k_v for k_v in _input_types_area['required'].items() if k_v[0] != 'show'),
|
|
**__extra_inputs_for_upscale_only,
|
|
# show=_input_types_area['required']['show'],
|
|
),
|
|
'hidden': {
|
|
'unique_id': 'UNIQUE_ID',
|
|
},
|
|
# 'optional': {},
|
|
})
|
|
|
|
|
|
class BestResolutionFromAreaUpscale:
|
|
NODE_NAME = 'BestResolutionFromAreaUpscale'
|
|
CATEGORY = "utils/resolution"
|
|
DESCRIPTION = _cleandoc(__doc__)
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = 'main'
|
|
RETURN_TYPES = _return_types_upscale
|
|
RETURN_NAMES = _return_names_upscale
|
|
RETURN_TYPES_TOOLTIPS = _return_ttips_upscale
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return _input_types_area_upscale
|
|
|
|
def main(
|
|
self, square_size: int, step: int, landscape: bool, aspect_a: float, aspect_b: float,
|
|
priority: _t.Union[RoundingPriority, str], upscale: float, up_step:int,
|
|
# show: bool,
|
|
unique_id: str = None
|
|
):
|
|
square_size: int = _number_to_int(square_size)
|
|
width_f, height_f = _float_width_height_from_area(square_size, landscape, aspect_a, aspect_b)
|
|
return _upscale_result_from_approx_wh(
|
|
width_f, height_f, step,
|
|
priority, upscale, up_step,
|
|
# show,
|
|
unique_id=unique_id, target_square_size=square_size
|
|
)
|