diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e535eae --- /dev/null +++ b/__init__.py @@ -0,0 +1,7 @@ +from .node import Image_Blending_Mode_Mask + +NODE_CLASS_MAPPINGS = { + "Image Blending Mode Mask" : Image_Blending_Mode_Mask, +} + +__all__ = ['NODE_CLASS_MAPPINGS'] diff --git a/node.py b/node.py new file mode 100644 index 0000000..ece4e80 --- /dev/null +++ b/node.py @@ -0,0 +1,204 @@ +from PIL import Image, ImageOps +from PIL.PngImagePlugin import PngInfo +from typing import Optional, Union, List +from urllib.request import urlopen +import comfy.diffusers_convert +import comfy.samplers +import comfy.sd +import comfy.utils +import comfy.clip_vision +import comfy.model_management +import numpy as np +import os +import subprocess +import sys +import torch + +class cstr(str): + class color: + END = '\33[0m' + BOLD = '\33[1m' + ITALIC = '\33[3m' + UNDERLINE = '\33[4m' + BLINK = '\33[5m' + BLINK2 = '\33[6m' + SELECTED = '\33[7m' + + BLACK = '\33[30m' + RED = '\33[31m' + GREEN = '\33[32m' + YELLOW = '\33[33m' + BLUE = '\33[34m' + VIOLET = '\33[35m' + BEIGE = '\33[36m' + WHITE = '\33[37m' + + BLACKBG = '\33[40m' + REDBG = '\33[41m' + GREENBG = '\33[42m' + YELLOWBG = '\33[43m' + BLUEBG = '\33[44m' + VIOLETBG = '\33[45m' + BEIGEBG = '\33[46m' + WHITEBG = '\33[47m' + + GREY = '\33[90m' + LIGHTRED = '\33[91m' + LIGHTGREEN = '\33[92m' + LIGHTYELLOW = '\33[93m' + LIGHTBLUE = '\33[94m' + LIGHTVIOLET = '\33[95m' + LIGHTBEIGE = '\33[96m' + LIGHTWHITE = '\33[97m' + + GREYBG = '\33[100m' + LIGHTREDBG = '\33[101m' + LIGHTGREENBG = '\33[102m' + LIGHTYELLOWBG = '\33[103m' + LIGHTBLUEBG = '\33[104m' + LIGHTVIOLETBG = '\33[105m' + LIGHTBEIGEBG = '\33[106m' + LIGHTWHITEBG = '\33[107m' + + @staticmethod + def add_code(name, code): + if not hasattr(cstr.color, name.upper()): + setattr(cstr.color, name.upper(), code) + else: + raise ValueError(f"'cstr' object already contains a code with the name '{name}'.") + + def __new__(cls, text): + return super().__new__(cls, text) + + def __getattr__(self, attr): + if attr.lower().startswith("_cstr"): + code = getattr(self.color, attr.upper().lstrip("_cstr")) + modified_text = self.replace(f"__{attr[1:]}__", f"{code}") + return cstr(modified_text) + elif attr.upper() in dir(self.color): + code = getattr(self.color, attr.upper()) + modified_text = f"{code}{self}{self.color.END}" + return cstr(modified_text) + elif attr.lower() in dir(cstr): + return getattr(cstr, attr.lower()) + else: + raise AttributeError(f"'cstr' object has no attribute '{attr}'") + + def print(self, **kwargs): + print(self, **kwargs) + +# Tensor to PIL +def tensor2pil(image): + return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) +# PIL to Tensor +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + +# Freeze PIP modules +def packages(versions=False): + import sys + import subprocess + return [( r.decode().split('==')[0] if not versions else r.decode() ) for r in subprocess.check_output([sys.executable, '-s', '-m', 'pip', 'freeze']).split()] +def install_package(package, uninstall_first: Union[List[str], str] = None): + if os.getenv("WAS_BLOCK_AUTO_INSTALL", 'False').lower() in ('true', '1', 't'): + cstr(f"Preventing package install of '{package}' due to WAS_BLOCK_INSTALL env").msg.print() + else: + if uninstall_first is None: + return + + if isinstance(uninstall_first, str): + uninstall_first = [uninstall_first] + + cstr(f"Uninstalling {', '.join(uninstall_first)}..") + subprocess.check_call([sys.executable, '-s', '-m', 'pip', 'uninstall', *uninstall_first]) + cstr("Installing package...").msg.print() + subprocess.check_call([sys.executable, '-s', '-m', 'pip', '-q', 'install', package]) + +class Image_Blending_Mode_Mask: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_a": ("IMAGE",), + "image_b": ("IMAGE",), + "mode": ([ + "add", + "color", + "color_burn", + "color_dodge", + "darken", + "difference", + "exclusion", + "hard_light", + "hue", + "lighten", + "multiply", + "overlay", + "screen", + "soft_light" + ],), + "blend_percentage": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "mask": ("MASK",), # 添加遮罩输入 + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "image_blending_mode_mask" + + CATEGORY = "Image/Image" + + def image_blending_mode_mask(self, image_a, image_b, mode='add', blend_percentage=1.0, mask=None): + + # Install Pilgram + if 'pilgram' not in packages(): + install_package("pilgram") + + # Import Pilgram module + import pilgram + + # Convert images to PIL + img_a = tensor2pil(image_a) + img_b = tensor2pil(image_b) + + # Ensure images are the same size + if img_a.size != img_b.size: + raise ValueError("Input images must have the same dimensions") + + # Apply blending mode + blending_modes = { + "color": pilgram.css.blending.color, + "color_burn": pilgram.css.blending.color_burn, + "color_dodge": pilgram.css.blending.color_dodge, + "darken": pilgram.css.blending.darken, + "difference": pilgram.css.blending.difference, + "exclusion": pilgram.css.blending.exclusion, + "hard_light": pilgram.css.blending.hard_light, + "hue": pilgram.css.blending.hue, + "lighten": pilgram.css.blending.lighten, + "multiply": pilgram.css.blending.multiply, + "add": pilgram.css.blending.normal, + "overlay": pilgram.css.blending.overlay, + "screen": pilgram.css.blending.screen, + "soft_light": pilgram.css.blending.soft_light + } + + out_image = blending_modes.get(mode, pilgram.css.blending.normal)(img_a, img_b) + + out_image = out_image.convert("RGB") + + # Apply mask if provided + if mask is not None: + mask = ImageOps.invert(tensor2pil(mask).convert('L')) + out_image = Image.composite(img_a, out_image, mask.resize(img_a.size)) + + # Blend image based on blend percentage + blend_mask = Image.new(mode="L", size=img_a.size, color=(round(blend_percentage * 255))) + blend_mask = ImageOps.invert(blend_mask) + out_image = Image.composite(img_a, out_image, blend_mask) + + return (pil2tensor(out_image), ) + \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c243da3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +numpy +pilgram \ No newline at end of file