From d70859a5daa9682ebe5b52b8bb96ca7c19a03cf2 Mon Sep 17 00:00:00 2001 From: Tasha Upchurch Date: Thu, 25 May 2023 14:41:25 -0400 Subject: [PATCH] restructure and rename to be more generic --- __init__.py | 85 +++++++++++++++++++++++++------------ image.py | 119 +++++++++++++++++++++++++--------------------------- latent.py | 2 +- nlp.py | 6 +-- 4 files changed, 120 insertions(+), 92 deletions(-) diff --git a/__init__.py b/__init__.py index 6b46710..042c4cf 100644 --- a/__init__.py +++ b/__init__.py @@ -1,12 +1,21 @@ +NODE_CLASS_MAPPINGS = {} +try: + import comfy.samplers + from custom_nodes.EternalKernelLiteGraphNodes.image import PreviewImageTest + from custom_nodes.EternalKernelLiteGraphNodes.image import TinyTxtToImg + use_generic_nodes = False +except: + use_generic_nodes = True + print("Failed to import comfy.samplers, assuming no comfyui skipping SD related nodes") + from custom_nodes.EternalKernelLiteGraphNodes.latent import LatentInterpolation -from custom_nodes.EternalKernelLiteGraphNodes.image import PadToMatch +from custom_nodes.EternalKernelLiteGraphNodes.image import PadToMatch from custom_nodes.EternalKernelLiteGraphNodes.image import StackImages from custom_nodes.EternalKernelLiteGraphNodes.image import ImageBC from custom_nodes.EternalKernelLiteGraphNodes.image import RGBA_MOD from custom_nodes.EternalKernelLiteGraphNodes.image import rgba_lower_clip from custom_nodes.EternalKernelLiteGraphNodes.image import rgba_upper_clip - from custom_nodes.EternalKernelLiteGraphNodes.image import rgba_merge from custom_nodes.EternalKernelLiteGraphNodes.image import rgba_split from custom_nodes.EternalKernelLiteGraphNodes.image import select_from_batch @@ -15,34 +24,56 @@ from custom_nodes.EternalKernelLiteGraphNodes.image import SelectFromRGBSimilari from custom_nodes.EternalKernelLiteGraphNodes.image import LoadImage from custom_nodes.EternalKernelLiteGraphNodes.image import PromptTemplate from custom_nodes.EternalKernelLiteGraphNodes.image import ExecWidget -from custom_nodes.EternalKernelLiteGraphNodes.image import TinyTxtToImg -from custom_nodes.EternalKernelLiteGraphNodes.image import ImageDistanceMask -from custom_nodes.EternalKernelLiteGraphNodes.image import PreviewImageTest +from custom_nodes.EternalKernelLiteGraphNodes.image import ImageDistanceMask from custom_nodes.EternalKernelLiteGraphNodes.nlp import NODE_CLASS_MAPPINGS as NLP_NODE_CLASS_MAPPINGS -NODE_CLASS_MAPPINGS = { - "Interpolation": LatentInterpolation, - "Image Pad To Match": PadToMatch, - "Image Stack": StackImages, - "Image B/C": ImageBC, - "Image RGBA Mod": RGBA_MOD, - "Image RGBA Lower Clip": rgba_lower_clip, - "Image RGBA Upper Clip": rgba_upper_clip, - "Image RGBA Merge": rgba_merge, - "Image RGBA Split": rgba_split, - "Select From Batch": select_from_batch, - "Quantize": Quantize, - "Select From RGB Similarity": SelectFromRGBSimilarity, - "Load Layered Image": LoadImage, - "Prompt Template": PromptTemplate, - "Eval Widget": ExecWidget, - "Tiny Txt 2 Img":TinyTxtToImg, - "Image Distance Mask":ImageDistanceMask, - "Preview Image Test":PreviewImageTest, -} -for k,v in NLP_NODE_CLASS_MAPPINGS.items(): - NODE_CLASS_MAPPINGS[k] = v \ No newline at end of file + +if use_generic_nodes: + NODE_CLASS_MAPPINGS_GENERIC = { + "Interpolation": LatentInterpolation, + "Image Pad To Match": PadToMatch, + "Image Stack": StackImages, + "Image B/C": ImageBC, + "Image RGBA Mod": RGBA_MOD, + "Image RGBA Lower Clip": rgba_lower_clip, + "Image RGBA Upper Clip": rgba_upper_clip, + "Image RGBA Merge": rgba_merge, + "Image RGBA Split": rgba_split, + "Select From Batch": select_from_batch, + "Quantize": Quantize, + "Select From RGB Similarity": SelectFromRGBSimilarity, + "Load Layered Image": LoadImage, + "Eval Widget": ExecWidget, + "Image Distance Mask": ImageDistanceMask, + } + + NODE_CLASS_MAPPINGS.update(NODE_CLASS_MAPPINGS_GENERIC) +else: + NODE_CLASS_MAPPINGS_COMFY = { + "Interpolation": LatentInterpolation, + "Image Pad To Match": PadToMatch, + "Image Stack": StackImages, + "Image B/C": ImageBC, + "Image RGBA Mod": RGBA_MOD, + "Image RGBA Lower Clip": rgba_lower_clip, + "Image RGBA Upper Clip": rgba_upper_clip, + "Image RGBA Merge": rgba_merge, + "Image RGBA Split": rgba_split, + "Select From Batch": select_from_batch, + "Quantize": Quantize, + "Select From RGB Similarity": SelectFromRGBSimilarity, + "Load Layered Image": LoadImage, + "Prompt Template": PromptTemplate, + "Eval Widget": ExecWidget, + "Tiny Txt 2 Img": TinyTxtToImg, + "Image Distance Mask": ImageDistanceMask, + "Preview Image Test": PreviewImageTest, + } + NODE_CLASS_MAPPINGS.update(NODE_CLASS_MAPPINGS_COMFY) + +for k, v in NLP_NODE_CLASS_MAPPINGS.items(): + NODE_CLASS_MAPPINGS[k] = v diff --git a/image.py b/image.py index 3c2b0da..7408c3d 100644 --- a/image.py +++ b/image.py @@ -1,12 +1,17 @@ -import comfy.samplers -import custom_nodes.Derfuu_ComfyUI_ModdedNodes.components.fields as field -import PIL.Image +try: + import comfy.samplers + import deactivated_custom_nodes.Derfuu_ComfyUI_ModdedNodes.components.fields as field + from nodes import CLIPTextEncode, VAEEncode, VAEDecode, KSampler, CheckpointLoaderSimple, EmptyLatentImage, \ + SaveImage + from custom_nodes.ComfyUI_ADV_CLIP_emb.nodes import AdvancedCLIPTextEncode as CLIPTextEncodeAdvanced +except ImportError as e: + print("comfy.samplers not found, skipping comfyui") + print(e) + SaveImage = None + import torch import torchvision import PIL.Image as Image -from nodes import CLIPTextEncode, VAEEncode, VAEDecode, KSampler, CheckpointLoaderSimple, EmptyLatentImage, SaveImage -from custom_nodes.ComfyUI_ADV_CLIP_emb.nodes import AdvancedCLIPTextEncode as CLIPTextEncodeAdvanced -import doctest try: import folder_paths @@ -107,7 +112,7 @@ class TinyTxtToImg: } } - CATEGORY = "EternalKernelNodes" + CATEGORY = "ETK" RETURN_TYPES = ("IMAGE",) FUNCTION = "tinytxt2img" @@ -186,36 +191,37 @@ class TinyTxtToImg: return (image,) -class PreviewImageTest(SaveImage): - def __init__(self): - self.output_dir = folder_paths.get_temp_directory() - self.type = "temp" +if SaveImage != None: + class PreviewImageTest(SaveImage): + def __init__(self): + self.output_dir = folder_paths.get_temp_directory() + self.type = "temp" - @classmethod - def INPUT_TYPES(s): - return { - "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, - "required": { - "images": ("IMAGE",), - }, - "optional": { - "prompt": ("STRING", {"multiline": True}), - "neg_prompt": ("STRING", {"multiline": True}), - "clip_encoder": (["comfy -ignore below", "advanced"],), - "token_normalization": (["none", "mean", "length", "length+mean"],), - "weight_interpretation": (["comfy", "A1111", "compel", "comfy++"],) - }, - } + @classmethod + def INPUT_TYPES(s): + return { + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + "required": { + "images": ("IMAGE",), + }, + "optional": { + "prompt": ("STRING", {"multiline": True}), + "neg_prompt": ("STRING", {"multiline": True}), + "clip_encoder": (["comfy -ignore below", "advanced"],), + "token_normalization": (["none", "mean", "length", "length+mean"],), + "weight_interpretation": (["comfy", "A1111", "compel", "comfy++"],) + }, + } - def myfunc(self, images, prompt=None, extra_pnginfo=None, ret=None): - return ret + def myfunc(self, images, prompt=None, extra_pnginfo=None, ret=None): + return ret - CATEGORY = "image" + CATEGORY = "ETK/image" - def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, **kwargs): - ret = super().save_images(images, filename_prefix, prompt, extra_pnginfo) - my_ret = self.myfunc(images, prompt, extra_pnginfo, ret) - return my_ret + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None, **kwargs): + ret = super().save_images(images, filename_prefix, prompt, extra_pnginfo) + my_ret = self.myfunc(images, prompt, extra_pnginfo, ret) + return my_ret class ExecWidget: @@ -296,7 +302,7 @@ class LoadImage: {"image": (sorted(os.listdir(input_dir)),)}, } - CATEGORY = "image" + CATEGORY = "ETK/image" RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") FUNCTION = "load_image_handler" @@ -427,7 +433,7 @@ class SelectFromRGBSimilarity: } } - CATEGORY = "image" + CATEGORY = "ETK/image" RETURN_TYPES = ("IMAGE",) @@ -511,7 +517,7 @@ class Quantize: RETURN_TYPES = ("IMAGE",) FUNCTION = "quantize" - CATEGORY = "image" + CATEGORY = "ETK/image" def quantize(self, image: torch.Tensor, colors: int = 256): """ @@ -576,8 +582,7 @@ class select_from_batch: "min": 0, "max": 100, "step": 1 - }),} - + }), } , "optional": { @@ -586,14 +591,13 @@ class select_from_batch: }, } + CATEGORY = "ETK/image" - CATEGORY = "image" - - RETURN_TYPES = ("IMAGE","LATENT",) + RETURN_TYPES = ("IMAGE", "LATENT",) FUNCTION = "select_from_batch" - def select_from_batch(self, index=0, image=None,latent=None): + def select_from_batch(self, index=0, image=None, latent=None): """ >>> image = torch.rand((1, 512, 512,3)) >>> index = 0 @@ -611,16 +615,15 @@ class select_from_batch: o_lat_dict = {} if image is None: print(latent["samples"].shape) - smpls=latent["samples"][index:index+1,:,:,:].expand(1, -1, -1, -1).clone() + smpls = latent["samples"][index:index + 1, :, :, :].expand(1, -1, -1, -1).clone() o_lat_dict["samples"] = smpls olat = o_lat_dict print(olat["samples"].shape) else: - oimg = image[index:index+1,:,:,:].clone() + oimg = image[index:index + 1, :, :, :].clone() print(oimg.shape) - - return (oimg,olat,) + return (oimg, olat,) class ImageDistanceMask: @@ -663,14 +666,13 @@ class ImageDistanceMask: } } - CATEGORY = "image" + CATEGORY = "ETK/image" RETURN_TYPES = ("IMAGE",) FUNCTION = "distance_mask_image" def distance_mask_image(self, image, x, y, falloff_method=None, Z=1, aspect_ratio=1.0): - import torch.nn.functional as F B, H, W, C = image.shape # Create a meshgrid for the coordinates @@ -760,7 +762,7 @@ class rgba_lower_clip: RETURN_TYPES = ("IMAGE",) FUNCTION = "rgba_lower_clip" - CATEGORY = "EternalKernelNodes/postprocessing" + CATEGORY = "ETK/image" def rgba_lower_clip(self, image: torch.Tensor, nrm_after, r, g, b, a): # image is (B,W,H,C) @@ -831,7 +833,7 @@ class rgba_upper_clip: RETURN_TYPES = ("IMAGE",) FUNCTION = "rgba_upper_clip" - CATEGORY = "EternalKernelNodes/postprocessing" + CATEGORY = "ETK/image" def rgba_upper_clip(self, image: torch.Tensor, nrm_after, r, g, b, a): # image is (B,W,H,C) @@ -894,7 +896,7 @@ class RGBA_MOD: RETURN_TYPES = ("IMAGE",) FUNCTION = "rgba_mod" - CATEGORY = "EternalKernelNodes/postprocessing" + CATEGORY = "ETK/image" def rgba_mod(self, image: torch.Tensor, r, g, b, a): """simply multiply the RGBA values by the sliders @@ -942,12 +944,7 @@ class ImageBC: RETURN_TYPES = ("IMAGE",) FUNCTION = "image_b_c" - CATEGORY = "EternalKernelNodes/postprocessing" - - import torch - - import torch - import torchvision + CATEGORY = "ETK/image" def image_b_c(self, image, contrast, brightness): """ @@ -1019,7 +1016,7 @@ class PadToMatch: RETURN_TYPES = ("IMAGE", "IMAGE") FUNCTION = "padme" - CATEGORY = "EternalKernelNodes/postprocessing" + CATEGORY = "ETK/image" def padme(self, image1, image2, im2_x_center, im2_y_center): """ @@ -1087,7 +1084,7 @@ class StackImages: RETURN_TYPES = ("IMAGE",) FUNCTION = "stackme" - CATEGORY = "EternalKernelNodes/postprocessing" + CATEGORY = "ETK/image" def stackme(self, image1, image2=None, number_of_images=2): """ @@ -1128,7 +1125,7 @@ class rgba_merge: } } - CATEGORY = "image" + CATEGORY = "ETK/image" RETURN_TYPES = ("IMAGE",) @@ -1149,7 +1146,7 @@ class rgba_split: } } - CATEGORY = "image" + CATEGORY = "ETK/image" RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE") diff --git a/latent.py b/latent.py index f335bbd..4affafd 100644 --- a/latent.py +++ b/latent.py @@ -20,7 +20,7 @@ class LatentInterpolation: RETURN_TYPES = ("LATENT",) FUNCTION = "interpolate" - CATEGORY = "EternalKernelNodes" + CATEGORY = "ETK/Latent" def interpolate(self, latent1, latent2, t): # Ensure that the latents are dictionaries diff --git a/nlp.py b/nlp.py index 1e8014c..f3719cd 100644 --- a/nlp.py +++ b/nlp.py @@ -1,10 +1,10 @@ -import comfy.samplers -import custom_nodes.Derfuu_ComfyUI_ModdedNodes.components.fields as field + +#import custom_nodes.Derfuu_ComfyUI_ModdedNodes.components.fields as field import PIL.Image import torch import torchvision import PIL.Image as Image -from nodes import CLIPTextEncode, VAEEncode, VAEDecode, KSampler, CheckpointLoaderSimple, EmptyLatentImage + import doctest # make sure to import ABC from abc import ABC, abstractmethod