From a610618938c8efae644cbfa8f272b3f54145db29 Mon Sep 17 00:00:00 2001 From: yada Date: Fri, 31 Mar 2023 21:54:18 -0400 Subject: [PATCH] nodes: add model router --- README.md | 8 ++++++ __init__.py | 3 +++ color_correction.py | 3 ++- image_compress.py | 64 +++++++++++++++++++++++++++++++++++++++++++++ image_layering.py | 2 +- model_router.py | 51 ++++++++++++++++++++++++++++++++++++ 6 files changed, 129 insertions(+), 2 deletions(-) create mode 100644 image_compress.py create mode 100644 model_router.py diff --git a/README.md b/README.md index 758fac6..88b0cda 100644 --- a/README.md +++ b/README.md @@ -22,3 +22,11 @@ image_layering: color_correction: - Adjusts the color of the target image according to another image; ported from stable diffusion WebUI + + +## External Nodes + +[WAS Node suite](https://civitai.com/models/20793/was-node-suites-comfyui) +- Image Blend by Mask: Blend two images by a mask (but all nodes are very good) + +and even look this: https://civitai.com/models/24869/comfyui-custom-nodes-by-xss \ No newline at end of file diff --git a/__init__.py b/__init__.py index e925529..e2cbabe 100644 --- a/__init__.py +++ b/__init__.py @@ -1,9 +1,12 @@ import custom_nodes.comfy_nodes_trojblue.image_layering as image_layering import custom_nodes.comfy_nodes_trojblue.color_correction as color_correction +import custom_nodes.comfy_nodes_trojblue.model_router as model_router + NODE_CLASS_MAPPINGS = { "layering": image_layering.Layering, # Layering "color_correction": color_correction.ColorCorrectionNode, # ColorCorrectionNode + "trRouter": model_router.ModelRouterPlugin, # ModelRouterPlugin } diff --git a/color_correction.py b/color_correction.py index d0ab8b0..d79613e 100644 --- a/color_correction.py +++ b/color_correction.py @@ -13,11 +13,12 @@ class ColorCorrectionNode: "original_image": ("IMAGE",), "target_image": ("IMAGE",), }, + } RETURN_TYPES = ("IMAGE",) FUNCTION = "color_correct" - CATEGORY = "trojblue_folder" + CATEGORY = "trNodes" def tensor_to_pil(self, img): if img is not None: diff --git a/image_compress.py b/image_compress.py new file mode 100644 index 0000000..48a38ae --- /dev/null +++ b/image_compress.py @@ -0,0 +1,64 @@ +# import torch +# from PIL import Image +# import numpy as np +# +# +# jpg_quality_input = ("INT", {"default": 95, +# "min": 50, +# "max": 100, +# "step": 1}) +# class JpgConvertNode: +# @classmethod +# def INPUT_TYPES(s): +# return { +# "required": { +# "original_image": ("IMAGE",), +# "jpg_quality": jpg_quality_input +# }, +# +# } +# +# RETURN_TYPES = ("IMAGE",) +# FUNCTION = "to_jpg" +# CATEGORY = "trNodes" +# +# def tensor_to_pil(self, img): +# if img is not None: +# i = 255. * img.cpu().numpy().squeeze() +# img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) +# return img +# +# def apply_color_correction(self, correction, original_image): +# +# # https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/22bcc7be428c94e9408f589966c2040187245d81/modules/processing.py#L44 +# +# correction_target = cv2.cvtColor(np.asarray(correction.copy()), cv2.COLOR_RGB2LAB) +# +# image = Image.fromarray(cv2.cvtColor(exposure.match_histograms( +# cv2.cvtColor( +# np.asarray(original_image), +# cv2.COLOR_RGB2LAB +# ), +# correction_target, +# channel_axis=2 +# ), cv2.COLOR_LAB2RGB).astype("uint8")) +# +# image = blendLayers(image, original_image, BlendType.LUMINOSITY) +# return image +# +# def png_to_jpg(self, png_file, jpg_file, quality=75): +# with Image.open(png_file) as img: +# img = img.convert('RGB') +# img.save(jpg_file, format='JPEG', quality=quality) +# def color_correct(self, original_image, jpg_quality): +# original_image = self.tensor_to_pil(original_image) +# +# +# target_image = self.tensor_to_pil(target_image) +# +# +# return (target_image,) +# +# NODE_CLASS_MAPPINGS = { +# "JpgConvertNode": JpgConvertNode +# } diff --git a/image_layering.py b/image_layering.py index ee8a065..0d7784d 100644 --- a/image_layering.py +++ b/image_layering.py @@ -37,7 +37,7 @@ class Layering: RETURN_TYPES = ("IMAGE",) FUNCTION = "apply_blend" - CATEGORY = "trojblue_folder" + CATEGORY = "trNodes" def tensor_to_pil(self, img): if img is not None: diff --git a/model_router.py b/model_router.py new file mode 100644 index 0000000..27b552e --- /dev/null +++ b/model_router.py @@ -0,0 +1,51 @@ +class ModelRouterPlugin: + """ + An example node + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tulple. + 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. + 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`. + """ + + FUNCTION = "execute" + CATEGORY = "trNodes" + RETURN_TYPES = ("MODEL", "CLIP", "VAE", "CONDITIONING", "CONDITIONING") + + @classmethod + def INPUT_TYPES(s): + return { + "optional": { + "model": ("MODEL",), + "clip": ("CLIP",), + "vae": ("VAE",), + "conditioning1": ("CONDITIONING",), + "conditioning2": ("CONDITIONING",), + } + } + + def execute(self, model=None, clip=None, vae=None, conditioning1=None, conditioning2=None): + return model, clip, vae, conditioning1, conditioning2 + + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "trRouter": ModelRouterPlugin +}