node_upscale_by: v3
This commit is contained in:
Lex Darlog (DRL)
+65
-163
@@ -2,192 +2,93 @@
|
||||
"""
|
||||
"""
|
||||
|
||||
import typing as _t
|
||||
|
||||
from inspect import cleandoc as _cleandoc
|
||||
|
||||
from frozendict import deepfreeze as _deepfreeze
|
||||
|
||||
from comfy.comfy_types.node_typing import IO as _IO
|
||||
# from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel as _ImageUpscaleWithModel
|
||||
# from nodes import ImageScaleBy as _ImageScaleBy
|
||||
from comfy_api.latest import io as _io
|
||||
from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel as _ImageUpscaleWithModel
|
||||
from nodes import ImageScaleBy as _ImageScaleBy
|
||||
|
||||
from . import _io_simple as _i
|
||||
from ._io_simple import _io_override
|
||||
from . import _meta
|
||||
from .docstring_formatter import format_docstring as _format_docstring
|
||||
from .node_scale import _scale_type_dict as __scale_type_dict_base
|
||||
from .funcs import _show_text_on_node
|
||||
|
||||
|
||||
# ==========================================================
|
||||
# Copy of built-in ComfyUI nodes with v1 schema,
|
||||
# as a temporary workaround
|
||||
|
||||
|
||||
import torch
|
||||
|
||||
from comfy import model_management
|
||||
import comfy.utils
|
||||
|
||||
|
||||
class _ImageUpscaleWithModel:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"upscale_model": ("UPSCALE_MODEL",), "image": ("IMAGE",),
|
||||
}}
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "upscale"
|
||||
|
||||
CATEGORY = "image/upscaling"
|
||||
|
||||
def upscale(self, upscale_model, image):
|
||||
device = model_management.get_torch_device()
|
||||
|
||||
memory_required = model_management.module_size(upscale_model.model)
|
||||
memory_required += (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
|
||||
memory_required += image.nelement() * image.element_size()
|
||||
model_management.free_memory(memory_required, device)
|
||||
|
||||
upscale_model.to(device)
|
||||
in_img = image.movedim(-1,-3).to(device)
|
||||
|
||||
tile = 512
|
||||
overlap = 32
|
||||
|
||||
oom = True
|
||||
while oom:
|
||||
try:
|
||||
steps = in_img.shape[0] * comfy.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
|
||||
pbar = comfy.utils.ProgressBar(steps)
|
||||
s = comfy.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
|
||||
oom = False
|
||||
except model_management.OOM_EXCEPTION as e:
|
||||
tile //= 2
|
||||
if tile < 128:
|
||||
raise e
|
||||
|
||||
upscale_model.to("cpu")
|
||||
s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
|
||||
return (s,)
|
||||
|
||||
|
||||
class _ImageScaleBy:
|
||||
upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
|
||||
"scale_by": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.01}),
|
||||
}}
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "upscale"
|
||||
|
||||
CATEGORY = "image/upscaling"
|
||||
|
||||
def upscale(self, image, upscale_method, scale_by):
|
||||
samples = image.movedim(-1,1)
|
||||
width = round(samples.shape[3] * scale_by)
|
||||
height = round(samples.shape[2] * scale_by)
|
||||
s = comfy.utils.common_upscale(samples, width, height, upscale_method, "disabled")
|
||||
s = s.movedim(1,-1)
|
||||
return (s,)
|
||||
|
||||
|
||||
# ==========================================================
|
||||
|
||||
|
||||
_ImageUpscaleWithModel_instance = _ImageUpscaleWithModel()
|
||||
__ImageUpscaleWithModel_input_types: _t.Dict[str, dict] = _ImageUpscaleWithModel_instance.INPUT_TYPES()
|
||||
__ImageUpscaleWithModel_input_types_required = __ImageUpscaleWithModel_input_types.get('required', dict())
|
||||
|
||||
_ImageScaleBy_instance = _ImageScaleBy()
|
||||
|
||||
__scale_type_dict: _t.Dict[str, _t.Any] = dict(__scale_type_dict_base, round=0.00001)
|
||||
|
||||
_input_types = _deepfreeze({
|
||||
'required': {
|
||||
'upscale_model': __ImageUpscaleWithModel_input_types_required.get('upscale_model', (_IO.UPSCALE_MODEL, )),
|
||||
'image': __ImageUpscaleWithModel_input_types_required.get('image', (_IO.IMAGE, )),
|
||||
'model_scale': (
|
||||
_IO.FLOAT,
|
||||
dict(__scale_type_dict, default=2.0, tooltip='The upscale factor a model natively increases image by'),
|
||||
),
|
||||
'scale_method': (_ImageScaleBy.upscale_methods, {'default': 'bicubic'}),
|
||||
'scale': (
|
||||
_IO.FLOAT,
|
||||
dict(__scale_type_dict, default=1.5, tooltip='The actual factor you want to upscale by'),
|
||||
),
|
||||
'show_status': (
|
||||
_IO.BOOLEAN,
|
||||
{
|
||||
'default': False, 'label_on': 'scales performed', 'label_off': 'no',
|
||||
'tooltip': "Show the extracted string on the node itself?"
|
||||
},
|
||||
),
|
||||
},
|
||||
# 'optional': {},
|
||||
'hidden': {
|
||||
'unique_id': 'UNIQUE_ID', # used for text display at the bottom of the node
|
||||
},
|
||||
})
|
||||
|
||||
# Threshold for "almost equal to 1.0".
|
||||
# Supposedly, no image would go beyond 250k pixels in either side (even polygraphy),
|
||||
# and the double of that equals change of less than a half-pixel.
|
||||
_epsilon = 1.0 / 500_000
|
||||
_half_upper_threshold = 0.5 + _epsilon * 0.5
|
||||
|
||||
|
||||
def _status_message(no_model_scale: bool, do_downscale: bool, model_scale: float, scale: float, second_downscale: float):
|
||||
if no_model_scale:
|
||||
return f"🤔❔ x{scale:.3f} (no model)"
|
||||
if not do_downscale:
|
||||
return f"x{scale:.3f} (no second scale)"
|
||||
|
||||
base_msg = f"x{scale:.3f} = x{model_scale:.3f} → <strong>x{second_downscale:.3f}</strong>"
|
||||
if scale > model_scale:
|
||||
return f"‼️ <strong>VERY</strong> blurry output\n{base_msg}"
|
||||
if scale > (model_scale - 0.5 + _epsilon):
|
||||
return f"⚠️ Blurry output\n{base_msg}"
|
||||
return base_msg
|
||||
|
||||
|
||||
class ImageUpscaleByWithModel:
|
||||
class ImageUpscaleByWithModel(_io.ComfyNode):
|
||||
"""
|
||||
A simple wrapper over "Upscale Image (using Model)" and "Upscale Image By".
|
||||
|
||||
First, up-scales with model. Then, (down)scales to get to the desired scale factor.
|
||||
"""
|
||||
NODE_NAME = 'ImageUpscaleByWithModel'
|
||||
CATEGORY = _meta.category
|
||||
DESCRIPTION = _format_docstring(_cleandoc(__doc__))
|
||||
_schema = _io.Schema(
|
||||
'ImageUpscaleByWithModel',
|
||||
display_name="Upscale Image By (with Model)",
|
||||
category=_meta.category,
|
||||
description=_format_docstring(_cleandoc(__doc__)),
|
||||
|
||||
FUNCTION = 'main'
|
||||
RETURN_TYPES = (_IO.IMAGE, )
|
||||
RETURN_NAMES = ('image', )
|
||||
# OUTPUT_TOOLTIPS = tuple()
|
||||
inputs=[
|
||||
_io.UpscaleModel.Input('upscale_model'),
|
||||
_io.Image.Input('image'),
|
||||
_io_override(_i.scale, 'model_scale', default=2.0, round=0.00001, tooltip=(
|
||||
"The upscale factor a model natively increases image by"
|
||||
)),
|
||||
_io.Combo.Input('scale_method', options=_ImageScaleBy.upscale_methods, default='bicubic'),
|
||||
_io_override(_i.scale, 'scale', default=1.5, round=0.00001, tooltip=(
|
||||
"The actual factor you want to upscale by"
|
||||
)),
|
||||
_io.Boolean.Input('show_status', default=False, label_on='scales performed', label_off='no', tooltip=(
|
||||
"Show the extracted string on the node itself?"
|
||||
)),
|
||||
],
|
||||
|
||||
hidden=[_io.Hidden.unique_id],
|
||||
|
||||
outputs=[_io.Image.Output('IMAGE', display_name='image')],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return _input_types
|
||||
def define_schema(cls) -> _io.Schema:
|
||||
return cls._schema
|
||||
|
||||
@staticmethod
|
||||
def main(
|
||||
# Threshold for "almost equal to 1.0".
|
||||
# Supposedly, no image would go beyond 250k pixels in either side (even polygraphy),
|
||||
# and the double of that equals change of less than a half-pixel.
|
||||
_epsilon = 1.0 / 500_000
|
||||
_half_upper_threshold = 0.5 + _epsilon * 0.5
|
||||
|
||||
@classmethod
|
||||
def _status_message(
|
||||
cls,
|
||||
no_model_scale: bool, do_downscale: bool, model_scale: float, scale: float, second_downscale: float
|
||||
):
|
||||
if no_model_scale:
|
||||
return f"🤔❔ x{scale:.3f} (no model)"
|
||||
if not do_downscale:
|
||||
return f"x{scale:.3f} (no second scale)"
|
||||
|
||||
base_msg = f"x{scale:.3f} = x{model_scale:.3f} → <strong>x{second_downscale:.3f}</strong>"
|
||||
if scale > model_scale:
|
||||
return f"‼️ <strong>VERY</strong> blurry output\n{base_msg}"
|
||||
if scale > (model_scale - 0.5 + cls._epsilon):
|
||||
return f"⚠️ Blurry output\n{base_msg}"
|
||||
return base_msg
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
upscale_model, image, model_scale: float, scale_method, scale: float,
|
||||
show_status: bool = False,
|
||||
unique_id: str = None,
|
||||
) -> _t.Tuple[str]:
|
||||
"""
|
||||
A simple wrapper over ``ImageUpscaleWithModel`` and ``ImageScaleBy``.
|
||||
|
||||
First, up-scales with model. Then, (down)scales to get to the desired scale factor.
|
||||
"""
|
||||
) -> _io.NodeOutput:
|
||||
second_downscale = scale / model_scale
|
||||
do_downscale = abs(second_downscale - 1.0) > _epsilon
|
||||
do_downscale = abs(second_downscale - 1.0) > cls._epsilon
|
||||
no_model_scale = (
|
||||
scale <= _half_upper_threshold
|
||||
or abs(model_scale - 1.0) <= _epsilon
|
||||
scale <= cls._half_upper_threshold
|
||||
or abs(model_scale - 1.0) <= cls._epsilon
|
||||
)
|
||||
if no_model_scale:
|
||||
out_image = _ImageScaleBy_instance.upscale(image, scale_method, scale)[0]
|
||||
@@ -196,8 +97,9 @@ class ImageUpscaleByWithModel:
|
||||
if do_downscale:
|
||||
out_image = _ImageScaleBy_instance.upscale(out_image, scale_method, second_downscale)[0]
|
||||
|
||||
unique_id = cls.hidden.unique_id
|
||||
if show_status and unique_id:
|
||||
msg = _status_message(no_model_scale, do_downscale, model_scale, scale, second_downscale)
|
||||
msg = cls._status_message(no_model_scale, do_downscale, model_scale, scale, second_downscale)
|
||||
_show_text_on_node(msg, unique_id)
|
||||
|
||||
return (out_image, )
|
||||
return _io.NodeOutput(out_image, )
|
||||
|
||||
Reference in New Issue
Block a user