diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..32c957e --- /dev/null +++ b/.gitignore @@ -0,0 +1,177 @@ +isnet/ + +# Other GitIgnore Templates: + +.vscode/* +!.vscode/settings.json +!.vscode/tasks.json +!.vscode/launch.json +!.vscode/extensions.json +!.vscode/*.code-snippets + +# Local History for Visual Studio Code +.history/ + +# Built Visual Studio Code Extensions +*.vsix + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ diff --git a/README.md b/README.md index 17288ff..20939b0 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,11 @@ # ComfyUI-TeaNodes - + +Adds a few new nodes: + +Image Equalization CLAHE style. + +Image Size Approximation based on pixel count that retains Image ratio. + +Image Resize Node that takes size tuple from Size Approximation Node. + +Image Scale Node that simply multiplies image size by a factor. \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..7f95786 --- /dev/null +++ b/__init__.py @@ -0,0 +1,13 @@ +from .nodes import NODE_CLASS_MAPPINGS as NCM + + +NODE_CLASS_MAPPINGS = { + **NCM, +} + +def remove_cm_prefix(node_mapping: str) -> str: + if node_mapping.startswith("TC_"): + return node_mapping[3:] + return node_mapping + +NODE_DISPLAY_NAME_MAPPINGS = {key: remove_cm_prefix(key) for key in NODE_CLASS_MAPPINGS} diff --git a/_func.py b/_func.py new file mode 100644 index 0000000..4e9e6d2 --- /dev/null +++ b/_func.py @@ -0,0 +1,43 @@ + +from math import ceil, floor, log + +debouncer = set() + +def po2(value, fill=False): + func = ceil if fill else floor + return pow(2, func(log(value)/log(2))) + +def pixel_approx(primary, secondary, total=1024, ratio=1.0, + threshold=0.125, _inc=1024, calculate_inc=True) -> tuple: + global debouncer + if calculate_inc: + debouncer.clear() + total *= total + ratio = secondary / primary + primary = po2(primary) + secondary = primary * ratio + _inc = primary / 4 + elif total * 3 > int(primary * secondary) > total * 0.333: + if abs(_inc) >= 64: _inc /= 2 + + # print(f"Action: {_inc}") + + count = round(primary) * round(secondary) + # print(round(primary), round(secondary), count) + + _recurse = False + if count > total: + _inc = abs(_inc) * -1 + elif count < total: + _inc = abs(_inc) + + if not (total * (1+threshold) > count > total * (1-threshold)) and total not in debouncer: + debouncer.add(count) + _recurse = True + + if _recurse: + primary += _inc + secondary = primary * ratio + primary, secondary = pixel_approx(primary, secondary, total, ratio, threshold, _inc, False) + + return (int(primary), int(secondary)) diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..4a1f917 --- /dev/null +++ b/nodes.py @@ -0,0 +1,194 @@ +import os, torch +import numpy as np +import comfy.utils + +from kornia.enhance import equalize_clahe +from ._func import pixel_approx, po2 +# from .isnet import dis_process +from PIL import Image + +class EqualizeCLAHE: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "clip_limit": ("FLOAT", {"default": 64, "min": 0.0, "max": 255, "step": 0.1}), + "grid_size": ("INT", {"default": 8, "min": 1, "max": 64, "step": 1}), + }, + "optional": { + "size": ("TUPLE", {"default": (1024, 1024)}), + } + } + RETURN_TYPES = ("IMAGE",) + FUNCTION = "equalize" + CATEGORY = "TeaNodes/Image" + + def equalize(self, image, size, clip_limit, grid_size): + _image = image.movedim(-1, 1) + if size != (1024, 1024): + grid_ratio = grid_size / clip_limit + size_ratio = min(size) / max(size) + clip_limit = po2(max(size), True) + grid_x = clip_limit * grid_ratio + grid_y = max(2, grid_x * size_ratio // 2 * 2) + if size[0] < size[1]: grid_x, grid_y = grid_y, grid_x + _image = equalize_clahe(_image, clip_limit, (grid_x, grid_y)) + result = _image.movedim(1, -1) + + return (result,) + +class SizeApproximation: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "square": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 32}), + } + } + + RETURN_TYPES = ("TUPLE", "INT", "INT") + FUNCTION = "calculate" + CATEGORY = "TeaNodes/Image" + + def calculate(self, image, square): + _image = image.movedim(-1, 1) + + height, width = image.shape[1:3] + + # print(width, height) + + if width >= height: + width, height = pixel_approx(width, height, square) + else: + height, width = pixel_approx(height, width, square) + + return ((width, height), width, height) + + +class ImageResize: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "width": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 32}), + "height": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 32}), + }, + "optional": { + "size": ("TUPLE", {"default": (1024, 1024)}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "resize" + CATEGORY = "TeaNodes/Image" + + def resize(self, image, size, width, height): + _image = image.movedim(-1, 1) + + if size != (1024, 1024): width, height = size + + result = comfy.utils.common_upscale(_image, int( width ), int( height ), "area", 'center') + result = result.movedim(1, -1) + return (result,) + +class ImageScale: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "factor": ("FLOAT", {"default": 1.0, "min": 0.2, "max": 5.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "resize" + CATEGORY = "TeaNodes/Image" + + def resize(self, image, factor): + _image = image.movedim(-1, 1) + + height, width = image.shape[1:3] + + result = comfy.utils.common_upscale(_image, int( width * factor ), int( height * factor ), "area", 'center') + result = result.movedim(1, -1) + + return (result,) + +class ColorFill(): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "color": ("STRING", {"default": '#7f7f7fff'}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "fill" + CATEGORY = "TeaNodes/Input" + + def fill(self, image, color): + if color.startswith('#'): + _color = color.lstrip('#') + try: color_rgba = tuple(int(_color[i:i+2], 16) for i in (0, 2, 4, 6)) + except: color_rgba = tuple(int(_color[i:i+2], 16) for i in (0, 2, 4)) + else: print(f"Something went wrong here: {color}") + else: + _color = color.split(',') + try: _color = tuple(int(e) for e in _color if 255>int(e)>0) + except: print(f"Something went wrong here: {color}") + color_rgba = _color + color_mode = 'RGBA' if len(color_rgba) > 3 else 'RGB' + + height, width = image.shape[1:3] + + _image = Image.new('RGBA', (width, height), color_rgba) + _image = np.array(_image.convert(color_mode)).astype(np.float32) / 255.0 + result = torch.from_numpy(_image).unsqueeze(0) + + return (result,) + +# class MaskBG_DIS: +# @classmethod +# def INPUT_TYPES(s): +# return { +# "required": { +# "image": ("IMAGE", ), +# } +# } + +# RETURN_TYPES = ("IMAGE",) +# FUNCTION = "process" +# CATEGORY = "TeaNodes/Image" + +# def process(self, image): +# i = 255. * image[-1].numpy() +# img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + +# img.save("..\__temp__.png", format='png', pnginfo=None, compress_level=4) + +# images_transformed = dis_process("..\__temp__.png") + +# os.remove("..\__temp__.png") + +# _image = Image.fromarray(images_transformed) + +# _image = _image.image_to_tensor() +# _image.unsqueeze_(0) +# result = _image.repeat(1,1,1,3) + +# return (result,) + +NODE_CLASS_MAPPINGS = { + "TC_EqualizeCLAHE": EqualizeCLAHE, + "TC_SizeApproximation": SizeApproximation, + "TC_ImageResize": ImageResize, + "TC_ImageScale": ImageScale, + "TC_ColorFill": ColorFill, + # "TC_MaskBG_DIS": MaskBG_DIS, +} \ No newline at end of file