From 9807e12198b6c7af47de2c708d999220cfd90ffd Mon Sep 17 00:00:00 2001 From: Tasha Upchurch Date: Thu, 25 May 2023 12:55:40 -0400 Subject: [PATCH] initial commit variety of nodes to support many things including a start on alpha channel support for comfyui --- __init__.py | 48 ++ components/__init__.py | 0 components/fields.py | 13 + components/sizes.py | 22 + components/tree.py | 22 + image.py | 1192 ++++++++++++++++++++++++++++++++++++++++ latent.py | 77 +++ nlp.py | 105 ++++ 8 files changed, 1479 insertions(+) create mode 100644 __init__.py create mode 100644 components/__init__.py create mode 100644 components/fields.py create mode 100644 components/sizes.py create mode 100644 components/tree.py create mode 100644 image.py create mode 100644 latent.py create mode 100644 nlp.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..6b46710 --- /dev/null +++ b/__init__.py @@ -0,0 +1,48 @@ +from custom_nodes.EternalKernelLiteGraphNodes.latent import LatentInterpolation +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 +from custom_nodes.EternalKernelLiteGraphNodes.image import Quantize +from custom_nodes.EternalKernelLiteGraphNodes.image import SelectFromRGBSimilarity +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.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 diff --git a/components/__init__.py b/components/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/components/fields.py b/components/fields.py new file mode 100644 index 0000000..fb0dd31 --- /dev/null +++ b/components/fields.py @@ -0,0 +1,13 @@ +import sys + +FLOAT = ("FLOAT", {"default": 1, + "min": -sys.float_info.max, + "max": sys.float_info.max, + "step": 0.01}) +BOOL = ("BOOL", {"default": False}) +INT = ("INT", {"default": 1, + "min": -sys.maxsize, + "max": sys.maxsize, + "step": 1}) +STRING = ("STRING", {"default": ""}) +STRING_ML = ("STRING", {"multiline": True}) \ No newline at end of file diff --git a/components/sizes.py b/components/sizes.py new file mode 100644 index 0000000..00cf725 --- /dev/null +++ b/components/sizes.py @@ -0,0 +1,22 @@ +def get_latent_size(LATENT, ORIGINAL_VALUES=False): + lc = LATENT.copy() + size = lc["samples"].shape[3], lc["samples"].shape[2] + if ORIGINAL_VALUES == False: + size = size[0] * 8, size[1] * 8 + return size + + +def get_image_size(IMAGE): + samples = IMAGE.movedim(-1, 1) + size = samples.shape[3], samples.shape[2] + # size = size.movedim(1, -1) + return size + + +def get_conditioning_size(CONDITIONING): + size = CONDITIONING["area"] + width = size[1] + height = size[0] + x_offs = size[3] + y_offs = size[2] + return ({"w": width, "h": height}, {"x": x_offs, "y": y_offs}) diff --git a/components/tree.py b/components/tree.py new file mode 100644 index 0000000..acbfcfc --- /dev/null +++ b/components/tree.py @@ -0,0 +1,22 @@ +TREE_MAIN = "EternalKernelNodes" + +TREE_VARIABLE = TREE_MAIN + "/Variables" +TREE_TUPLES = TREE_MAIN + "/Tuples" +TREE_MATH = TREE_MAIN + "/Math" +TREE_FUNCTIONS = TREE_MAIN + "/Functions" +TREE_MODDED = TREE_MAIN + "/Modded nodes" + +TREE_CONVERTERS = TREE_FUNCTIONS + "/Converters" + +TREE_TRIGONOMETRY = TREE_MATH + "/Trigonometry" + +TREE_TUPLE_MODDED = TREE_TUPLES + "/Modded nodes" + +TREE_TUPLE_LATENTS = TREE_TUPLE_MODDED + "/Latents" +TREE_TUPLE_CONDITIONING = TREE_TUPLE_MODDED + "/Conditioning" + +TREE_COND = TREE_MODDED + "/Conditions" +TREE_IMAGES = TREE_MODDED + "/Image" +TREE_LATENTS = TREE_MODDED + "/Latent" + +TREE_DEBUG = TREE_MAIN + "/Debug" diff --git a/image.py b/image.py new file mode 100644 index 0000000..3c2b0da --- /dev/null +++ b/image.py @@ -0,0 +1,1192 @@ +import comfy.samplers +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, SaveImage +from custom_nodes.ComfyUI_ADV_CLIP_emb.nodes import AdvancedCLIPTextEncode as CLIPTextEncodeAdvanced +import doctest + +try: + import folder_paths +except: + pass + +import os +import numpy as np +import hashlib + +### info for code completion AI ### +""" +all of these classes are plugins for comfyui and follow the same pattern +all of the images are torch tensors and it is unknown and unimportant if they are on the cpu or gpu +all image inputs are (B,W,H,C) + +avoid numpy and PIL as much as possible +""" + + +def rgba_per_channel_norm(image): + """ normalize per channel """ + print("normalizing per channel") + + if image[:, :, :, 0].min() != image[:, :, :, 0].max(): + image[:, :, :, 0] = image[:, :, :, 0] - image[:, :, :, 0].min() + image[:, :, :, 0] = image[:, :, :, 0] / image[:, :, :, 0].max() + + if image[:, :, :, 1].min() != image[:, :, :, 1].max(): + image[:, :, :, 1] = image[:, :, :, 1] - image[:, :, :, 1].min() + image[:, :, :, 1] = image[:, :, :, 1] / image[:, :, :, 1].max() + + if image[:, :, :, 2].min() != image[:, :, :, 2].max(): + image[:, :, :, 2] = image[:, :, :, 2] - image[:, :, :, 2].min() + image[:, :, :, 2] = image[:, :, :, 2] / image[:, :, :, 2].max() + + if image[:, :, :, 3].min() != image[:, :, :, 3].max(): + image[:, :, :, 3] = image[:, :, :, 3] - image[:, :, :, 3].min() + image[:, :, :, 3] = image[:, :, :, 3] / image[:, :, :, 3].max() + return image + + +def torch_image_show(image): + """Show a torch image""" + from PIL import Image + image = torch.mul(image[0], 255, ) + image = image.to(torch.uint8) # Convert tensor to uint8 data type + + # must have at least 3 channels + if image.shape[-1] == 1: + image = torch.cat((image, image, image), dim=-1) + + image = image.numpy() + image = Image.fromarray(image) # Create a PIL image + image.show() + + +class TinyTxtToImg: + """small text to image generator""" + + def __init__(self): + import random + self.mdl = None + self.clp = None + self.vae = None + self.vision = None + self.seed = random.randint(0, 2 ** 32 - 1) + self.steps = 10 + self.cfg = 8 + self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0] + self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0] + self.positive = None + self.negative = None + self.width = 512 + self.height = 512 + self.batch_size = 1 + self.denoise = 1.0 + self.latent_image = None + self.samples = None + + @classmethod + def INPUT_TYPES(cls): + return {"required": + { + "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++"],), + + }, + "optional": { + "name": ("STRING", {"multiline": False, + "default": "tinytxt2img"} + ), + "overrides": ("STRING", {"multiline": True, + "default": ""}) + } + } + + CATEGORY = "EternalKernelNodes" + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "tinytxt2img" + + def tinytxt2img(self, prompt, neg_prompt, name="tinytxt2img", overrides: str = "", + clip_encoder="comfy", + token_normalization="length+mean", + weight_interpretation="comfy++"): + """ use the imports from nodes to generate an image from text """ + + import random + import json + if CheckpointLoaderSimple.share_mdl is None: + self.mdl, self.clp, self.vae, self.vision = CheckpointLoaderSimple.load_checkpoint(None, + ckpt_name="v1-5-pruned-emaonly.safetensors" + ) + CSL = CheckpointLoaderSimple + CSL.share_mdl = self.mdl + CSL.share_clip = self.clp + CSL.share_vae = self.vae + CSL.share_vision = self.vision + + else: + self.mdl, self.clp, self.vae = (CheckpointLoaderSimple.share_mdl, + CheckpointLoaderSimple.share_clip, + CheckpointLoaderSimple.share_vae + ) + self.seed = random.randint(0, 2 ** 32 - 1) + self.steps = 10 + self.cfg = 8 + self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0] + self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0] + self.positive = prompt + self.negative = neg_prompt + self.width = 512 + self.height = 512 + self.batch_size = 1 + self.denoise = 1.0 + if overrides: + if overrides != "": + overrides = overrides.replace("'", "\"") + if overrides[0] != "{": + overrides = "{" + overrides + "}" + valid_dict = json.loads(overrides) + for k, v in valid_dict.items(): + if k in self.__dict__: + self.__setattr__(k, v) + else: + print(f"invalid override key: {k}") + if clip_encoder == "comfy -ignore below": + self.pos_cond = CLIPTextEncode.encode(None, self.clp, self.positive)[0] + self.neg_cond = CLIPTextEncode.encode(None, self.clp, self.negative)[0] + elif clip_encoder == "advanced": + self.pos_cond = CLIPTextEncodeAdvanced.encode(None, + self.clp, + self.positive, + token_normalization, + weight_interpretation)[0] + self.neg_cond = CLIPTextEncodeAdvanced.encode(None, + self.clp, + self.negative, + token_normalization, + weight_interpretation)[0] + + self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0] + + samples = KSampler.sample(None, + self.mdl, + self.seed, self.steps, self.cfg, self.sampler_name, self.scheduler, + self.pos_cond, self.neg_cond, + self.latent_image, denoise=self.denoise)[0] + + image = VAEDecode.decode(None, self.vae, samples)[0] + image = image.detach().cpu() + + return (image,) + + +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++"],) + }, + } + + def myfunc(self, images, prompt=None, extra_pnginfo=None, ret=None): + return ret + + CATEGORY = "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 + + +class ExecWidget: + """runs eval on the given text""" + + @classmethod + def INPUT_TYPES(cls): + return {"optional": + { + "text_to_eval": ("STRING", {"multiline": True}), + "image1": ("IMAGE", {"multiline": False}), + "float1": ("FLOAT", {"multiline": False}), + + }, + "required": { + "name": ("STRING", {"multiline": False, + "default": "exec_func"} + ), + } + } + + CATEGORY = "text" + RETURN_TYPES = ("STRING", "IMAGE", "FLOAT") + FUNCTION = "exec_handler" + + def exec_handler(self, text_to_eval, image1: torch.Tensor = None, float1: float = 0.0, name: str = "exec_func"): + """ + >>> ExecWidget().exec_handler("2 + 3") + '5' + """ + result = None + if image1 is not None: + image_in = image1.clone() + else: + image_in = None + new_locals = {"float_out": float1, "image_in": image_in} + exec(text_to_eval, globals(), new_locals) + # merge new_locals into globals + # globals().update(new_locals) + # locals().update(new_locals) + # print(globals()) + out_str = new_locals.get("out_str", None) + out_image = new_locals.get("out_image", None) + out_float = new_locals.get("out_float", None) + return (out_str, out_image, out_float,) + + +class PromptTemplate: + """replaces the text in the given text string with a given other text at some key positions""" + + @classmethod + def INPUT_TYPES(s): + return {"required": + { + "text": ("STRING", {"multiline": True}), + "replacement": ("STRING", {"multiline": True}), + "key1": ("STRING", {"multiline": False}), + } + } + + CATEGORY = "text" + RETURN_TYPES = ("STRING",) + FUNCTION = "prompt_template_handler" + + def prompt_template_handler(self, text: str, replacement, key1): + """ + >>> PromptTemplate().prompt_template_handler("hello world", "universe", "world") + 'hello universe' + """ + return (text.replace(key1, replacement),) + + +class LoadImage: + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + return {"required": + {"image": (sorted(os.listdir(input_dir)),)}, + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE", "MASK", "IMAGE") + FUNCTION = "load_image_handler" + + def load_image_handler(self, image): + """ + >>> fp = './mouths.tif' + >>> results = LoadImage().load_image_handler(fp) + >>> torch_image_show(results[0][0]) + :param image: + :return: + """ + if image.endswith(".xcf"): + return self.load_xcf(image) + elif image.endswith(".tif"): + return self.load_tif(image) + else: + return self.load_image(image) + + def load_xcf(self, image): + pass + + def load_tif(self, image): + input_dir = folder_paths.get_input_directory() + image_path = os.path.join(input_dir, image) + img = Image.open(image_path) + out = [] + + # first find the maximum width and height + max_width = 0; + max_height = 0 + for i in range(img.n_frames): + img.seek(i) + max_width = max(max_width, img.width) + max_height = max(max_height, img.height) + + # now create a new tensor with the max width and height + out = torch.zeros((img.n_frames, max_height, max_width, 4), dtype=torch.float) + + for i in range(img.n_frames): + img.seek(i) + np_arr = np.array(img).astype(np.float32) / 255.0 + + padx = max_width - img.width + pady = max_height - img.height + np_arr = np.pad(np_arr, ((0, pady), (0, padx), (0, 0)), mode='constant') + + t_arr = torch.tensor(np_arr) + out[i] = t_arr + + mask = out[:, :, :, 2:3] + out_rgb = out[:, :, :, :3] + out_rgba = out + + return (out_rgb, mask, out_rgba,) + + def load_image(self, image): + input_dir = folder_paths.get_input_directory() + image_path = os.path.join(input_dir, image) + i = Image.open(image_path) + + img_rgba = Image.open(image_path) + img_rgba = img_rgba.convert("RGBA") + img_rgba = np.array(img_rgba).astype(np.float32) / 255.0 + img_rgba = torch.from_numpy(img_rgba)[None,] + + image = i.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + if 'A' in i.getbands(): + mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0 + mask = 1. - torch.from_numpy(mask) + else: + mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu") + return (image, mask, img_rgba) + + @classmethod + def IS_CHANGED(s, image): + input_dir = folder_paths.get_input_directory() + image_path = os.path.join(input_dir, image) + m = hashlib.sha256() + try: + with open(image_path, 'rb') as f: + m.update(f.read()) + except FileNotFoundError: + return True + return m.digest().hex() + + +class SelectFromRGBSimilarity: + """ + selects pixel positions based on similarity to a given color + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + + "r": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.1 + }), + "g": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.1 + }), + "b": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.1 + }), + "similarity": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }) + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "select_from_rgb_similarity" + + def select_from_rgb_similarity(self, image: torch.Tensor, r, g, b, similarity: float = 0.5): + """ + >>> image = torch.rand((1, 64, 64,3)) + >>> color = [0.5, 0.5, 0.5] + >>> similarity = 0.5 + >>> mask = SelectFromRGBSimilarity.select_from_rgb_similarity(None,image,color[0],color[1],color[2],similarity)[0] + #>>> torch_image_show(mask) + >>> mask.shape + torch.Size([1, 64, 64, 3]) + >>> mask.mean() + tensor(1.) + + >>> image = torch.rand((2, 64, 64,3)) + >>> color = [0.5, 0.5, 0.5] + >>> similarity = 0.25 + >>> mask = SelectFromRGBSimilarity.select_from_rgb_similarity(None,image,color[0],color[1],color[2],similarity)[0] + #>>> torch_image_show(mask) + >>> mask.shape + torch.Size([1, 64, 64, 3]) + >>> mask.mean() + tensor(0.34) + + :param image: torch tensor of shape (B,W,H,C) + + :param similarity: array of 3 floats + :return: image mask of the positions where the color is similar white=selected black=not selected + """ + + batch_size, height, width, _ = image.shape + color = [r, g, b] + + result = torch.zeros_like(image) + # remember that the input image is BWHC + similarity = torch.tensor(similarity) + out = [] + for b in range(batch_size): + # just use the first 3 channels + use_image = image[b:b + 1, :, :, :3] + white = torch.ones_like(use_image) + + max_distance = torch.norm(torch.tensor([1.0, 1.0, 1.0])) + image_vector = torch.norm(use_image - torch.tensor([0, 0, 0]), dim=-1, keepdim=True) + + color_tensor = torch.tensor(color).to(image.device) + color_tensor = color_tensor[None, None, None, :] # Expand dimensions to match image shape for broadcasting + color_distance = torch.linalg.vector_norm(use_image - color_tensor, dim=-1, + keepdim=True) # Euclidean distance between color and pixels + color_distance = color_distance / max_distance + + mask_bool = color_distance < similarity + mask_white = white * mask_bool + out.append(mask_white) + + out = torch.cat(out, dim=0) + return (out,) + + +class Quantize: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image": ("IMAGE",), + "colors": ("INT", { + "default": 256, + "min": 1, + "max": 256, + "step": 1 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "quantize" + + CATEGORY = "image" + + def quantize(self, image: torch.Tensor, colors: int = 256): + """ + >>> image = torch.rand((2, 512, 512,3)) + >>> colors = 256 + >>> quantized = Quantize.quantize(None,image,colors)[0] + >>> quantized.shape + torch.Size([2, 512, 512, 3]) + >>> image = torch.rand((2, 512, 512,4)) + >>> colors = 256 + >>> quantized = Quantize.quantize(None,image,colors)[0] + >>> quantized.shape + torch.Size([2, 512, 512, 4]) + + """ + batch_size, height, width, _ = image.shape + + result = torch.zeros_like(image) + + for b in range(batch_size): + tensor_image_start = image[b] + if tensor_image_start.shape[2] == 4: + tensor_image = image[b, :, :, :3] + else: + tensor_image = image[b] + + img = (tensor_image * 255).to(torch.uint8).numpy() + pil_image = Image.fromarray(img, mode='RGB') + + palette = pil_image.quantize( + colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836 + quantized_image = pil_image.quantize(colors=colors, palette=palette, dither=0) + + rgb = quantized_image.convert("RGB") + quantized_array = torchvision.transforms.PILToTensor()(rgb).float() / 255 + + # fix the indices to b, w, h, c + quantized_array = quantized_array.permute(1, 2, 0) + + if tensor_image_start.shape[2] == 4: + result[b][:, :, :3] = quantized_array.to(image.dtype) + result[b][:, :, 3] = image[b, :, :, 3] + else: + result[b] = quantized_array.to(image.dtype) + + return (result,) + + +class select_from_batch: + """selects a single image from a batch""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + + "index": ("INT", { + "default": 0, + "min": 0, + "max": 100, + "step": 1 + }),} + + + , + "optional": { + "image": ("IMAGE",), + "latent": ("LATENT",), + }, + } + + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE","LATENT",) + + FUNCTION = "select_from_batch" + + def select_from_batch(self, index=0, image=None,latent=None): + """ + >>> image = torch.rand((1, 512, 512,3)) + >>> index = 0 + >>> selected = select_from_batch.select_from_batch(None,image,index) + >>> selected.shape + torch.Size([1, 512, 512, 3]) + >>> image = torch.rand((2, 512, 512,3)) + >>> selected = select_from_batch.select_from_batch(None,image,1) + >>> selected.shape + torch.Size([1, 512, 512, 3]) + + """ + olat = None + oimg = None + o_lat_dict = {} + if image is None: + print(latent["samples"].shape) + 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() + print(oimg.shape) + + + return (oimg,olat,) + + +class ImageDistanceMask: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "x": ("INT", { + "default": 0, + "min": 0, + "max": 10240, + "step": 1 + }), + "y": ("INT", { + "default": 0, + "min": 0, + "max": 10240, + "step": 1 + }), + + "Z": ("FLOAT", { + "default": 1, + "min": -2, + "max": 2, + "step": .05 + }), + "aspect_ratio": ("FLOAT", { + "default": 1, + "min": 0, + "max": 2, + "step": .05 + }), + + "falloff_method": (["linear", "spherical"],), + + } + } + + CATEGORY = "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 + grid_y, grid_x = torch.meshgrid(torch.arange(H), torch.arange(W)) + + # Adjust aspect ratio + grid_x = grid_x * aspect_ratio + x = x * aspect_ratio + + # Ensure grid_x and grid_y are on the same device and dtype as the image + grid_x, grid_y = grid_x.to(image.device).type(image.dtype), grid_y.to(image.device).type(image.dtype) + + # Calculate distance from (x, y) to each pixel + dist = torch.sqrt((grid_x - x) ** 2 + (grid_y - y) ** 2) + + # Normalize the distance to be between 0 and 1 + dist = (dist - dist.min()) / (dist.max() - dist.min()) + + if falloff_method: + if falloff_method == "linear": + dist = 1 - dist + elif falloff_method == "spherical": + dist = torch.sqrt(1 - dist ** 2) + + if Z < 0: + dist = 1 - dist + Z = abs(Z) + + dist = dist - (1 - Z) + dist = torch.clamp(dist, 0, 1) + + # Normalize again after the modifications + dist = (dist - dist.min()) / (dist.max() - dist.min()) + + # Extend dimensionality to match the input tensor (B, H, W, C) + dist = dist[None, ..., None].expand(B, H, W, C) + + return (dist,) + + +class rgba_lower_clip: + """Clip RGBA values below a threshold""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "nrm_after": ("INT", { + "default": 0, + "min": 0, + "max": 1, + "step": 1 + + }), + "r": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "g": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "b": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "a": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "rgba_lower_clip" + + CATEGORY = "EternalKernelNodes/postprocessing" + + def rgba_lower_clip(self, image: torch.Tensor, nrm_after, r, g, b, a): + # image is (B,W,H,C) + image = image.clone() + # check for alpha channel, add if not present + + if image.shape[3] == 3: + alpha = torch.ones_like(image[:, :, :, 0]) + image = torch.cat((image, alpha.unsqueeze(3)), dim=3) + + image[:, :, :, 0] = torch.clamp(image[:, :, :, 0], min=r, max=1.0) + image[:, :, :, 1] = torch.clamp(image[:, :, :, 1], min=g, max=1.0) + image[:, :, :, 2] = torch.clamp(image[:, :, :, 2], min=b, max=1.0) + image[:, :, :, 3] = torch.clamp(image[:, :, :, 3], min=a, max=1.0) + + if int(nrm_after) == 1: + image = rgba_per_channel_norm(image) + + return (image,) + + +class rgba_upper_clip: + """Clip RGBA values above a threshold""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "nrm_after": ("INT", { + "default": 0, + "min": 0, + "max": 1, + "step": 1 + + }), + "r": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "g": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "b": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "a": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "rgba_upper_clip" + + CATEGORY = "EternalKernelNodes/postprocessing" + + def rgba_upper_clip(self, image: torch.Tensor, nrm_after, r, g, b, a): + # image is (B,W,H,C) + image = image.clone() + # check for alpha channel, add if not present + + if image.shape[3] == 3: + alpha = torch.ones_like(image[:, :, :, 0]) + image = torch.cat((image, alpha.unsqueeze(3)), dim=3) + + image[:, :, :, 0] = torch.clamp(image[:, :, :, 0], min=0, max=r) + image[:, :, :, 1] = torch.clamp(image[:, :, :, 1], min=0, max=g) + image[:, :, :, 2] = torch.clamp(image[:, :, :, 2], min=0, max=b) + image[:, :, :, 3] = torch.clamp(image[:, :, :, 3], min=0, max=a) + + if int(nrm_after) == 1: + image = rgba_per_channel_norm(image) + + return (image,) + + +class RGBA_MOD: + """provide sliders to modify RGBA values""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "r": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 2.0, + "step": 0.01 + }), + "g": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 2.0, + "step": 0.01 + }), + "b": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 2.0, + "step": 0.01 + }), + "a": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 2.0, + "step": 0.01 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "rgba_mod" + + CATEGORY = "EternalKernelNodes/postprocessing" + + def rgba_mod(self, image: torch.Tensor, r, g, b, a): + """simply multiply the RGBA values by the sliders + images are (B, W, H, C) + """ + # check for alpha channel + if image.shape[3] == 3: + image = torch.cat((image, torch.ones_like(image[:, :, :, 0:1])), dim=3) + + image = image.clone() + image[:, :, :, 0] = image[:, :, :, 0] * r + image[:, :, :, 1] = image[:, :, :, 1] * g + image[:, :, :, 2] = image[:, :, :, 2] * b + image[:, :, :, 3] = image[:, :, :, 3] * a + + return (image,) + + +class ImageBC: + """Normalize an images brightness,droping lower and higher values of n perctile""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "contrast": ("FLOAT", { + "default": 0.0, + "min": -10.0, + "max": 10.0, + "step": 0.01 + }), + "brightness": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 10.0, + "step": 0.01 + }), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "image_b_c" + + CATEGORY = "EternalKernelNodes/postprocessing" + + import torch + + import torch + import torchvision + + def image_b_c(self, image, contrast, brightness): + """ + Adjusts the contrast of an image, while preserving the alpha channel + + >>> image = (torch.rand(2, 256, 256, 4)/10)+.45 + >>> image[:,0,0,:] = torch.tensor([0.0,0.0,0.0,0.0]) + >>> image[:,1,1,:] = torch.tensor([1.0,1.0,1.0,1.0]) + >>> torch_image_show(image) + >>> normalized_image = ImageBC.image_b_c(None,image,2,2) + >>> torch_image_show(normalized_image) + >>> normalized_image.shape + torch.Size([2, 256, 256, 4]) + """ + # check for alpha channel + if image.shape[3] == 3: + image = torch.cat((image, torch.ones_like(image[:, :, :, 0:1])), dim=3) + + # get the image size + size = image.shape + + # Separate the alpha channel + img_a = image[..., -1:] + image = image[..., :-1] + + # Permute the dimensions to [..., 1 or 3, H, W] + image = image.permute(0, 3, 1, 2) + + # Adjust the contrast + image = torchvision.transforms.functional.adjust_brightness(image, brightness) + image = torchvision.transforms.functional.adjust_contrast(image, contrast) + + # Permute the dimensions back to [Batch, H, W, Channels] + image = image.permute(0, 2, 3, 1) + + # Concatenate the alpha channel back + image = torch.cat((image, img_a), dim=-1) + + return (image,) + + +class PadToMatch: + """ a node that will pad the smaller of two images to match the size of the larger one """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image1": ("IMAGE",), + "image2": ("IMAGE",), + "im2_x_center": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + "im2_y_center": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01 + }), + }, + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + FUNCTION = "padme" + + CATEGORY = "EternalKernelNodes/postprocessing" + + def padme(self, image1, image2, im2_x_center, im2_y_center): + """ + pads the smaller of two images to match the size of the larger one + x_center and y_center are how far to the right and down the smaller image should be placed + images are (B, W, H, C) + >>> PadToMatch.padme(None,torch.ones((1, 2, 2, 3)), torch.ones((1, 3, 3, 3)), 0.5, 0.5)[0].mean(axis=(3)) + tensor([[[1., 1., 0.], + [1., 1., 0.], + [0., 0., 0.]]]) + + + """ + target_width = max(image1.shape[1], image2.shape[1]) + target_height = max(image1.shape[2], image2.shape[2]) + + padded1 = torch.zeros((image1.shape[0], target_width, target_height, image1.shape[3])) + padded2 = torch.zeros((image2.shape[0], target_width, target_height, image2.shape[3])) + + pad_w_1 = target_width - image1.shape[1] + pad_h_1 = target_height - image1.shape[2] + + # img 1 is taken as the centered one + xs = int(pad_w_1 / 2) + xe = xs + image1.shape[1] + ys = int(pad_h_1 / 2) + ye = ys + image1.shape[2] + padded1[:, xs:xe, ys:ye, :] = image1 + + # img 2 is placed according to the centering parameters + pad_w_2 = target_width - image2.shape[1] + pad_h_2 = target_height - image2.shape[2] + + xs = int(pad_w_2 * im2_x_center) + ys = int(pad_h_2 * im2_y_center) + xe = xs + image2.shape[1] + ye = ys + image2.shape[2] + + if xe == 0: + xe = None + if ye == 0: + ye = None + padded2[:, xs:xe, ys:ye, :] = image2 + + return (padded1, padded2) + + +class StackImages: + """ a node that will stack two images on top of each other """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image1": ("IMAGE",), + }, + "optional": { + "image2": ("IMAGE",), + "number_of_images": ("INT", {"default": 2, "min": 2, "max": 100}), + }, } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "stackme" + + CATEGORY = "EternalKernelNodes/postprocessing" + + def stackme(self, image1, image2=None, number_of_images=2): + """ + stacks two images on top of each other + images are (B, W, H, C) torch tensors + """ + image1 = image1.clone() + + if image2 is not None: + image2 = image2.clone() + stacked = torch.cat((image1, image2), dim=0) + else: + stacked = torch.cat((image1, image1), dim=0) + + if number_of_images > 2: + while stacked.shape[0] < number_of_images: + stacked = torch.cat((stacked, stacked), dim=0) + print(stacked.shape) + out = stacked[:number_of_images, :, :, :] + print(out.shape) + return (out,) + + +class rgba_merge: + """merges the channels of an RGBA image""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "r": ("IMAGE",), + "g": ("IMAGE",), + "b": ("IMAGE",), + "a": ("IMAGE",), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE",) + + FUNCTION = "merge" + + +class rgba_split: + """splits the channels of an RGBA image""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + } + } + + CATEGORY = "image" + + RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE") + + FUNCTION = "split" + + def split(self, image): + """ + >>> image = torch.rand((1, 512, 512,3)) + >>> r, g, b, a = rgba_split.split(None,image) + >>> r.shape + torch.Size([1, 1, 512, 512]) + :param image: + :return: + """ + # check for alhpa channel + if image.shape[3] == 3: + # add alpha channel + alpha = torch.ones((image.shape[0], image.shape[1], image.shape[2], 1)) + image = torch.cat((image, alpha), 3) + # for now each channel must be (b,w,h,3) so repeat the channel 3 times + r = image[:, :, :, 0].unsqueeze(3).repeat(1, 1, 1, 3) + g = image[:, :, :, 1].unsqueeze(3).repeat(1, 1, 1, 3) + b = image[:, :, :, 2].unsqueeze(3).repeat(1, 1, 1, 3) + a = image[:, :, :, 3].unsqueeze(3).repeat(1, 1, 1, 3) + + return (r, g, b, a,) + + +def get_image_size(image): + """ returns the width and height of an image """ + size = image.shape + width = int(size[0]) + height = int(size[1]) + return (width * height) + + +if __name__ == "__main__": + import doctest + + doctest.testmod() diff --git a/latent.py b/latent.py new file mode 100644 index 0000000..f335bbd --- /dev/null +++ b/latent.py @@ -0,0 +1,77 @@ +#from nodes import MAX_RESOLUTION +import torch + +from custom_nodes.EternalKernelLiteGraphNodes.components import fields as field + + +class LatentInterpolation: + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent1": ("LATENT",), + "latent2": ("LATENT",), + "t": field.FLOAT, + } + } + + RETURN_TYPES = ("LATENT",) + FUNCTION = "interpolate" + CATEGORY = "EternalKernelNodes" + + def interpolate(self, latent1, latent2, t): + # Ensure that the latents are dictionaries + if not (isinstance(latent1, dict) and isinstance(latent2, dict)): + raise ValueError("Latent values must be dictionaries.") + + # Ensure that both latents have the same keys + if set(latent1.keys()) != set(latent2.keys()): + raise ValueError("Latent dictionaries must have the same keys.") + + inter_latent = {} + for key in latent1: + if isinstance(latent1[key], (list, tuple)): + inter_latent[key] = [(1 - t) * l1_val + t * l2_val for l1_val, l2_val in + zip(latent1[key], latent2[key])] + elif isinstance(latent1[key], torch.Tensor): + inter_latent[key] = (1 - t) * latent1[key] + t * latent2[key] + else: + raise ValueError(f"Unsupported value type for key '{key}' in latent dictionaries.") + + return (inter_latent,) + + +def test_latent_interpolation(): + import torch + # Create an instance of the LatentInterpolation class + li = LatentInterpolation() + + # Define test cases for the interpolate method + # latents are at least 4 diminsional + # latents are of the same length + test_cases = [ + { + "input": { + "latent1": torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]), + "latent2": torch.tensor([[[[5.0, 6.0], [7.0, 8.0]]]]), + "t": 0.5, + + }, + "expected": torch.tensor([[[[3.0, 4.0], [5.0, 6.0]]]]), + }, + ] + + # Run the test cases + for test_case in test_cases: + result = li.interpolate(**test_case["input"]) + if not (result[0][0] == test_case["expected"]).all().item(): + print(f"Expected {test_case['expected']} but got {result}") + + + + +if __name__ == "__main__": + test_latent_interpolation() \ No newline at end of file diff --git a/nlp.py b/nlp.py new file mode 100644 index 0000000..1e8014c --- /dev/null +++ b/nlp.py @@ -0,0 +1,105 @@ +import comfy.samplers +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 + +import os +import openai + + +class Config: + api_key = os.environ.get('OPENAI_API_KEY') + organization = os.environ.get('OPENAI_ORG') + + +NODE_CLASS_MAPPINGS = {} + +openai.organization = Config.organization +openai.api_key = Config.api_key + + +class SimpleTextWidget: + """base class for simple construction""" + + def __init__(self, func): + self.func = func + + + @classmethod + def register(self, name): + global NODE_CLASS_MAPPINGS + + # register the widget + NODE_CLASS_MAPPINGS[name] = self + + @classmethod + def INPUT_TYPES(cls): + return {"required": + { + "text": ("STRING", {"multiline": True}), + "Display": ("SERVER_STRING", {"multiline": True}), + }, + + } + + CATEGORY = "text" + RETURN_TYPES = ("STRING",) + FUNCTION = "handler" + + def handler(self, text,Display): + ret = self.func(text) + return (ret,) + + +try: + import folder_paths +except: + pass + +import os +import numpy as np +import hashlib + +### info for code completion AI ### +""" +all of these classes are plugins for comfyui and follow the same pattern +all of the images are torch tensors and it is unknown and unimportant if they are on the cpu or gpu +all image inputs are (B,W,H,C) + +avoid numpy and PIL as much as possible +""" + + +class LLMCompletion(SimpleTextWidget): + """uses the input text to call the specified LLM model and returns the output string""" + + def OAI_completion(self, text): + out_message = [{"role": "user", "content": text}] + + response = openai.ChatCompletion.create( + model="gpt-3.5-turbo", + messages=out_message + ) + gpt_message = response["choices"][0]["message"]["content"].strip() + self.server_strings[self.SSID] = gpt_message + return gpt_message + + def __init__(self): + self.func = self.OAI_completion + super().__init__(self.func) + import uuid + self.SSID = str(uuid.uuid4()) + from main import server_obj_holder + self.server_strings = server_obj_holder[0].server_strings + self.server_strings["self.SSID"] = "empty" + + + + +LLMCompletion.register("LLM_Completion")