From 01596d6a340532afc4e1f7b6fdb3293abe748211 Mon Sep 17 00:00:00 2001 From: Vincent at Jupiter Date: Fri, 24 Jan 2025 00:17:24 +0800 Subject: [PATCH] init image features --- .gitignore | 6 ++ __init__.py | 9 ++- categories.py | 10 +++ example_node.py.example | 155 ---------------------------------------- node_mappings.py | 17 +++++ py/__init__.py | 0 py/log.py | 55 ++++++++++++++ py/nodes_LSD.py | 0 py/nodes_image.py | 151 ++++++++++++++++++++++++++++++++++++++ pyproject.toml | 2 +- requirements.txt | 2 + 11 files changed, 248 insertions(+), 159 deletions(-) create mode 100644 .gitignore create mode 100644 categories.py delete mode 100644 example_node.py.example create mode 100644 node_mappings.py create mode 100644 py/__init__.py create mode 100644 py/log.py create mode 100644 py/nodes_LSD.py create mode 100644 py/nodes_image.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a209b4b --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +_tests +__pycache__ +.venv +.idea +*.pth +*.ini diff --git a/__init__.py b/__init__.py index 570722e..c7366dd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,9 @@ +from .py.log import log +from .node_mappings import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +# WEB_DIRECTORY is the comfyui nodes directory that ComfyUI will link and auto-load. WEB_DIRECTORY = "./js/v1" -NODE_CLASS_MAPPINGS = {} -NODE_DISPLAY_NAME_MAPPINGS = {} +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] -__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file +log('Successfully Initiated 💥🚀') \ No newline at end of file diff --git a/categories.py b/categories.py new file mode 100644 index 0000000..fa7a640 --- /dev/null +++ b/categories.py @@ -0,0 +1,10 @@ +icons = { + ### Image Nodes + "AG Image Resize To": "🐦‍🔥 ArchiGraph/Image", + "AG Image to BW": "🐦‍🔥 ArchiGraph/Image", + + ### Specific Nodes - TODO: not implemented yet + "ArchiGraph/CV": "🐦‍🔥 ArchiGraph/CV", + "DeepLSD": "🐦‍🔥 ArchiGraph/CV", + "Shrink": "image", +} \ No newline at end of file diff --git a/example_node.py.example b/example_node.py.example deleted file mode 100644 index 29ab2aa..0000000 --- a/example_node.py.example +++ /dev/null @@ -1,155 +0,0 @@ -class Example: - """ - A example node - - Class methods - ------------- - INPUT_TYPES (dict): - Tell the main program input parameters of nodes. - IS_CHANGED: - optional method to control when the node is re executed. - - Attributes - ---------- - RETURN_TYPES (`tuple`): - The type of each element in the output tuple. - RETURN_NAMES (`tuple`): - Optional: The name of each output in the output tuple. - FUNCTION (`str`): - The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() - OUTPUT_NODE ([`bool`]): - If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. - The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. - Assumed to be False if not present. - CATEGORY (`str`): - The category the node should appear in the UI. - DEPRECATED (`bool`): - Indicates whether the node is deprecated. Deprecated nodes are hidden by default in the UI, but remain - functional in existing workflows that use them. - EXPERIMENTAL (`bool`): - Indicates whether the node is experimental. Experimental nodes are marked as such in the UI and may be subject to - significant changes or removal in future versions. Use with caution in production workflows. - execute(s) -> tuple || None: - The entry point method. The name of this method must be the same as the value of property `FUNCTION`. - For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. - """ - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - """ - Return a dictionary which contains config for all input fields. - Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". - Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. - The type can be a list for selection. - - Returns: `dict`: - - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` - - Value input_fields (`dict`): Contains input fields config: - * Key field_name (`string`): Name of a entry-point method's argument - * Value field_config (`tuple`): - + First value is a string indicate the type of field or a list for selection. - + Second value is a config for type "INT", "STRING" or "FLOAT". - """ - return { - "required": { - "image": ("IMAGE",), - "int_field": ("INT", { - "default": 0, - "min": 0, #Minimum value - "max": 4096, #Maximum value - "step": 64, #Slider's step - "display": "number", # Cosmetic only: display as "number" or "slider" - "lazy": True # Will only be evaluated if check_lazy_status requires it - }), - "float_field": ("FLOAT", { - "default": 1.0, - "min": 0.0, - "max": 10.0, - "step": 0.01, - "round": 0.001, #The value representing the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. - "display": "number", - "lazy": True - }), - "print_to_screen": (["enable", "disable"],), - "string_field": ("STRING", { - "multiline": False, #True if you want the field to look like the one on the ClipTextEncode node - "default": "Hello World!", - "lazy": True - }), - }, - } - - RETURN_TYPES = ("IMAGE",) - #RETURN_NAMES = ("image_output_name",) - - FUNCTION = "test" - - #OUTPUT_NODE = False - - CATEGORY = "Example" - - def check_lazy_status(self, image, string_field, int_field, float_field, print_to_screen): - """ - Return a list of input names that need to be evaluated. - - This function will be called if there are any lazy inputs which have not yet been - evaluated. As long as you return at least one field which has not yet been evaluated - (and more exist), this function will be called again once the value of the requested - field is available. - - Any evaluated inputs will be passed as arguments to this function. Any unevaluated - inputs will have the value None. - """ - if print_to_screen == "enable": - return ["int_field", "float_field", "string_field"] - else: - return [] - - def test(self, image, string_field, int_field, float_field, print_to_screen): - if print_to_screen == "enable": - print(f"""Your input contains: - string_field aka input text: {string_field} - int_field: {int_field} - float_field: {float_field} - """) - #do some processing on the image, in this example I just invert it - image = 1.0 - image - return (image,) - - """ - The node will always be re executed if any of the inputs change but - this method can be used to force the node to execute again even when the inputs don't change. - You can make this node return a number or a string. This value will be compared to the one returned the last time the node was - executed, if it is different the node will be executed again. - This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash - changes between executions the LoadImage node is executed again. - """ - #@classmethod - #def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen): - # return "" - -# Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension -# WEB_DIRECTORY = "./somejs" - - -# Add custom API routes, using router -from aiohttp import web -from server import PromptServer - -@PromptServer.instance.routes.get("/hello") -async def get_hello(request): - return web.json_response("hello") - - -# A dictionary that contains all nodes you want to export with their names -# NOTE: names should be globally unique -NODE_CLASS_MAPPINGS = { - "Example": Example -} - -# A dictionary that contains the friendly/humanly readable titles for the nodes -NODE_DISPLAY_NAME_MAPPINGS = { - "Example": "Example Node" -} diff --git a/node_mappings.py b/node_mappings.py new file mode 100644 index 0000000..c55cb41 --- /dev/null +++ b/node_mappings.py @@ -0,0 +1,17 @@ +from .py.log import log +# import importlib.util # for auto import + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +def AddMapping(node_class, display_name): + NODE_CLASS_MAPPINGS[node_class.NAME] = node_class + NODE_DISPLAY_NAME_MAPPINGS[node_class.NAME] = display_name + +# Load Image Nodes +try: + from .py.nodes_image import ImageResizeTo, Image2BW + AddMapping(ImageResizeTo, "Image Resize To 🐦‍🔥") + AddMapping(Image2BW, "Image To Black and White 🐦‍🔥") +except ImportError: + log("Failed to load image nodes", msg_color="BRIGHT_GREEN") diff --git a/py/__init__.py b/py/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/py/log.py b/py/log.py new file mode 100644 index 0000000..bf94042 --- /dev/null +++ b/py/log.py @@ -0,0 +1,55 @@ +# https://stackoverflow.com/questions/4842424/list-of-ansi-color-escape-sequences +# https://en.wikipedia.org/wiki/ANSI_escape_code#3-bit_and_4-bit +COLORS = { + 'BLACK': '\33[30m', + 'RED': '\33[31m', + 'GREEN': '\33[32m', + 'YELLOW': '\33[33m', + 'BLUE': '\33[34m', + 'MAGENTA': '\33[35m', + 'CYAN': '\33[36m', + 'WHITE': '\33[37m', + 'GREY': '\33[90m', + 'BRIGHT_RED': '\33[91m', + 'BRIGHT_GREEN': '\33[92m', + 'BRIGHT_YELLOW': '\33[93m', + 'BRIGHT_BLUE': '\33[94m', + 'BRIGHT_MAGENTA': '\33[95m', + 'BRIGHT_CYAN': '\33[96m', + 'BRIGHT_WHITE': '\33[97m', + # Styles. + 'RESET': '\33[00m', + 'BOLD': '\33[01m', + 'NORMAL': '\33[22m', + 'ITALIC': '\33[03m', + 'UNDERLINE': '\33[04m', + 'BLINK': '\33[05m', + 'BLINK2': '\33[06m', + 'SELECTED': '\33[07m', + # Backgrounds + 'BG_BLACK': '\33[40m', + 'BG_RED': '\33[41m', + 'BG_GREEN': '\33[42m', + 'BG_YELLOW': '\33[43m', + 'BG_BLUE': '\33[44m', + 'BG_MAGENTA': '\33[45m', + 'BG_CYAN': '\33[46m', + 'BG_WHITE': '\33[47m', + 'BG_GREY': '\33[100m', + 'BG_BRIGHT_RED': '\33[101m', + 'BG_BRIGHT_GREEN': '\33[102m', + 'BG_BRIGHT_YELLOW': '\33[103m', + 'BG_BRIGHT_BLUE': '\33[104m', + 'BG_BRIGHT_MAGENTA': '\33[105m', + 'BG_BRIGHT_CYAN': '\33[106m', + 'BG_BRIGHT_WHITE': '\33[107m', +} + +def log(message, color=None, msg_color=None, prefix=None): + """Basic logging.""" + color = COLORS[color] if color is not None and color in COLORS else COLORS["BRIGHT_MAGENTA"] + msg_color = COLORS[msg_color] if msg_color is not None and msg_color in COLORS else '' + prefix = f'[{prefix}]' if prefix is not None else '' + msg = f'{color}[ComfyUI-ArchiGraph]{prefix} ' + msg += f'{msg_color}{message}{COLORS["RESET"]}' + print(msg) \ No newline at end of file diff --git a/py/nodes_LSD.py b/py/nodes_LSD.py new file mode 100644 index 0000000..e69de29 diff --git a/py/nodes_image.py b/py/nodes_image.py new file mode 100644 index 0000000..c78f0a6 --- /dev/null +++ b/py/nodes_image.py @@ -0,0 +1,151 @@ +import torch + +from ..categories import icons +from comfy.utils import common_upscale + +#---------------------------------------------------------------------------------------------------------------------# +class ImageResizeTo: + upscale_methods = ["lanczos", "nearest-exact", "bilinear", "area", "bicubic"] + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Second value is a config for type "INT", "STRING" or "FLOAT". + """ + return { + "required": { + "image": ("IMAGE",{ + "lazy": False # Will only be evaluated if check_lazy_status requires it + }), + "edge_size": ("INT", { + "default": 1024, + "min": 16, #Minimum value + "max": 8192, #Maximum value + "step": 2, #Slider's step + "display": "number", # Cosmetic only: display as "number" or "slider" + }), + "edge_type": ("BOOLEAN", { + "default": True, # Default value + "label_on": "By Short Edge", # Label for the first option + "label_off": "By Long Edge", # Label for the second option + }), + # "is_by_short": ("BOOL", { + + # }), + "upscale_method": (cls.upscale_methods,), + }, + } + + RETURN_TYPES = ("IMAGE", "INT", "INT", "STRING") + RETURN_NAMES = ("IMAGE", "prev_width", "prev_height", "log") + FUNCTION = "resize" + NAME = "AG Image Resize To" + CATEGORY = icons.get(NAME) + DESCRIPTION = "Resizes an image by specify one edge size. Proportion is kept by default." + + def resize(self, image, edge_size, edge_type, upscale_method): + + # get original image size + B, H, W, C = image.shape + prev_width = W + prev_height = H + + image = image.movedim(-1, 1) + + if edge_type: + if H < W: + new_height = edge_size + new_width = int(W * edge_size / H) + else: + new_width = edge_size + new_height = int(H * edge_size / W) + else: + if H > W: + new_height = edge_size + new_width = int(W * edge_size / H) + else: + new_width = edge_size + new_height = int(H * edge_size / W) + + image = common_upscale(image, new_width, new_height, upscale_method, "disabled") + image = image.movedim(1, -1) + + log = f"Resized from W{prev_width}xH{prev_height} to W{new_width}xH{new_height}." + + return (image, prev_width, prev_height,log,) + +#---------------------------------------------------------------------------------------------------------------------# +class Image2BW: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",{ + "lazy": False + }), + "reds":("INT", {"default": 40, "min": -200, "max": 300, "display": "slider"}), + "yellows":("INT", {"default": 60, "min": -200, "max": 300, "display": "slider"}), + "greens":("INT", {"default": 40, "min": -200, "max": 300, "display": "slider"}), + "cyans":("INT", {"default": 60, "min": -200, "max": 300, "display": "slider"}), + "blues":("INT", {"default": 20, "min": -200, "max": 300, "display": "slider"}), + "magentas":("INT", {"default": 80, "min": -200, "max": 300, "display": "slider"}), + }, + } + + RETURN_TYPES = ("IMAGE", ) + RETURN_NAMES = ("IMAGE", ) + FUNCTION = "bw" + NAME = "AG Image to BW" + CATEGORY = icons.get(NAME) + DESCRIPTION = "Converts an image to black & white like the adjustments in Photoshop." + + def bw(self, image, reds, yellows, greens, cyans, blues, magentas): + + datatype = torch.float32 + # get original image size + B, H, W, C = image.shape + + color_ratios = torch.tensor([reds, greens, blues, cyans, magentas, yellows], dtype=datatype) + color_ratios = (color_ratios + 200) / 500 + + sorted_vals, sorted_idx = torch.sort(image, dim=3, descending=True) # BxHxWxC + + # Get the color ratio value of the maximum channel + max_idx = sorted_idx[:, :, :, 0].unsqueeze(-1) # BxHxWx1 + ratios_reshaped = color_ratios.unsqueeze(0).unsqueeze(0).unsqueeze(0).expand(B,H,W,-1) # BxHxWx6 + ratio_max = torch.gather(ratios_reshaped, 3, max_idx) # BxHxWx1 + + combo_idx = (sorted_idx[:, :, :, -1] + 3).unsqueeze(-1) # BxHxWx1 + ratio_combo = torch.gather(ratios_reshaped, 3, combo_idx) # BxHxWx1 + + # duplicate sorted_vals to a new tensor + sorted_vals_2 = sorted_vals.clone() + sorted_vals_2[:, :, :, 0] = 0 + + # move the first value to the last in sorted_vals_2 + sorted_vals_2 = torch.cat((sorted_vals_2[:, :, :, 1:], sorted_vals_2[:, :, :, 0].unsqueeze(-1)), dim=-1) # BxHxWxC + + diff = sorted_vals - sorted_vals_2 # BxHxWxC + ratio = torch.cat((ratio_max, ratio_combo, torch.ones_like(ratio_max, dtype=datatype)), dim=-1) # BxHxWx3 + + # calculate the sum product of diff and ratio + final = torch.sum(diff * ratio,dim=-1) # BxHxWx1 + + return (final, ) \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 1b88d74..00dbf43 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,4 +11,4 @@ Repository = "https://github.com/vincentfs/ComfyUI-ArchiGraph" [tool.comfy] PublisherId = "vincentfs" DisplayName = "ComfyUI-ArchiGraph" -Icon = "" +Icon = "https://github.com/vincentfs/ComfyUI-ArchiGraph/blob/main/ArchiGraphLogo.png" diff --git a/requirements.txt b/requirements.txt index e69de29..c36c6c3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -0,0 +1,2 @@ +opencv-python>=4.10.0.84 +opencv-contrib-python>=4.10.0.84 \ No newline at end of file