From 34e24e27a2d79757b055108cde47dec3d87bc76a Mon Sep 17 00:00:00 2001 From: BXYMartin Date: Mon, 4 Mar 2024 16:39:41 +0000 Subject: [PATCH] Added initial modules --- .gitignore | 155 ++++++++++++++++++++++++++++++++++ __init__.py | 3 + nodes.py | 213 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + 4 files changed, 372 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 nodes.py create mode 100644 requirements.txt diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..be2f233 --- /dev/null +++ b/.gitignore @@ -0,0 +1,155 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +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 + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow +__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 maintainted 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/ +*.pyc + +.idea +/node_modules diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d721463 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..f6c99e9 --- /dev/null +++ b/nodes.py @@ -0,0 +1,213 @@ +from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel +import numpy as np +from PIL import Image +import torch + + +class AnyType(str): + """A special class that is always equal in not equal comparisons. Credit to pythongosssss""" + + def __eq__(self, _) -> bool: + return True + + def __ne__(self, __value: object) -> bool: + return False + +any = AnyType("*") + + +class MultiControlNetConverter: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "controlnets": (any,), + }, + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("multicontrolnet",) + CATEGORY = "ControlNetUtils" + FUNCTION = "execute" + + def execute(self, controlnets): + model = MultiControlNetModel(controlnets) + return [model] + + +class NHWC2NCHWTensor: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + "optional": { + } + } + + CATEGORY = "ControlNetUtils" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "execute" + def execute(self, image): + print(image.shape) + image.permute(0, 3, 1, 2) + print(image.shape) + return [image] + +class NHWCTensor2PIL: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + "optional": { + } + } + + CATEGORY = "ControlNetUtils" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "execute" + + def execute(self, image): + # convert NHWC Tensor to PIL Image + img_array = image.squeeze(0).cpu().numpy() * 255.0 + img_pil = Image.fromarray(np.clip(img_array, 0, 255).astype(np.uint8)) + return [img_pil] + + +class PIL2NHWCTensor: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + "optional": { + } + } + + CATEGORY = "ControlNetUtils" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + + FUNCTION = "execute" + def execute(self, image): + # 将 PIL.Image 转换为 numpy.ndarray + img_array = np.array(image) + # 转换 numpy.ndarray 为 torch.Tensor + img_tensor = torch.from_numpy(img_array).float() / 255. + # 转换图像格式为 CHW (如果需要) + if img_tensor.ndim == 3 and img_tensor.shape[-1] == 3: + img_tensor = img_tensor.permute(2, 0, 1) + # 添加批次维度并转换为 NHWC + img_tensor = img_tensor.unsqueeze(0).permute(0, 2, 3, 1) + return [img_tensor] + + +class ListOfImages: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + }, + "optional": { + "any_1": ("IMAGE",), + "any_2": ("IMAGE",), + "any_3": ("IMAGE",), + "any_4": ("IMAGE",), + "any_5": ("IMAGE",), + "any_6": ("IMAGE",), + "any_7": ("IMAGE",), + "any_8": ("IMAGE",), + } + } + + CATEGORY = "ControlNetUtils" + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + + FUNCTION = "execute" + + def execute(self, + any_1=None, + any_2=None, + any_3=None, + any_4=None, + any_5=None, + any_6=None, + any_7=None, + any_8=None): + + list_any = [] + + if any_1 is not None: + try: + list_any.append(any_1) + except Exception as e: + raise e + if any_2 is not None: + try: + list_any.append(any_2) + except Exception as e: + raise e + if any_3 is not None: + try: + list_any.append(any_3) + except Exception as e: + raise e + if any_4 is not None: + try: + list_any.append(any_4) + except Exception as e: + raise e + if any_5 is not None: + try: + list_any.append(any_5) + except Exception as e: + raise e + if any_6 is not None: + try: + list_any.append(any_6) + except Exception as e: + raise e + if any_7 is not None: + try: + list_any.append(any_7) + except Exception as e: + raise e + if any_8 is not None: + try: + list_any.append(any_8) + except Exception as e: + raise e + + # yes, double brackets are needed because of the OUTPUT_IS_LIST... ¯\_(ツ)_/¯ + return [list_any] + +NODE_CLASS_MAPPINGS = { + "MultiControlNetConverter": MultiControlNetConverter, + "ListOfImages": ListOfImages, + "PIL2NHWCTensor": PIL2NHWCTensor, + "NHWC2NCHWTensor": NHWC2NCHWTensor, + "NHWCTensor2PIL": NHWCTensor2PIL +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MultiControlNetConverter": "Multi-ControlNet Converter", + "ListOfImages": "List of Images", + "PIL2NHWCTensor": "Convert PIL to Tensor (NHWC)", + "NHWC2NCHWTensor": "Convert Tensor (NHWC) to (NCHW)", + "NHWCTensor2PIL": "Convert Tensor (NHWC) to PIL" +} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..81caff4 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +diffusers \ No newline at end of file