diff --git a/__.code-workspace b/__.code-workspace new file mode 100644 index 0000000..27f7fbb --- /dev/null +++ b/__.code-workspace @@ -0,0 +1,13 @@ +{ + "folders": [ + { + "path": ".", + } + ], + "settings": { + "files.exclude": { + "__pycache__":true, + } + + } +} \ No newline at end of file diff --git a/__init__.py b/__init__.py index c9237e4..54766a8 100644 --- a/__init__.py +++ b/__init__.py @@ -5,17 +5,22 @@ NODE_CLASS_MAPPINGS = { **NCM, } -if os.path.exists(f"{os.path.abspath('.')}\ComfyUI\custom_nodes\ComfyUI-TeaNodes\_misc.py"): - print('[TeaNodes]: Misc. Nodes Loaded.') +if os.path.exists( + f"{os.path.abspath('.')}\ComfyUI\custom_nodes\ComfyUI-TeaNodes\_misc.py" +): + print("[TeaNodes]: Misc. Nodes Loaded.") from ._misc import NODE_CLASS_MAPPINGS as MISC_NCM + NODE_CLASS_MAPPINGS = { **NCM, **MISC_NCM, } + def remove_cm_prefix(node_mapping: str) -> str: if node_mapping.startswith("TC_"): return node_mapping[3:] return node_mapping + NODE_DISPLAY_NAME_MAPPINGS = {key: remove_cm_prefix(key) for key in NODE_CLASS_MAPPINGS} diff --git a/_nodes.py b/_nodes.py index f5b503d..5ba3945 100644 --- a/_nodes.py +++ b/_nodes.py @@ -1,23 +1,28 @@ -import torch, random -from torchvision.transforms.functional import center_crop -import numpy as np -import comfy.utils - +import random from colorsys import hsv_to_rgb -from kornia.enhance import equalize_clahe, adjust_gamma, add_weighted -from ._func import pixel_approx, po2, Color, byte +import numpy as np +import torch +from kornia.enhance import adjust_gamma, equalize_clahe from PIL import Image +from torchvision.transforms.functional import center_crop + +import comfy.utils +import folder_paths + +from ._func import Color, byte, pixel_approx, po2 + class CropTo: @classmethod def INPUT_TYPES(s): return { "required": { - "image_src": ("IMAGE", ), - "image_ref": ("IMAGE", ), + "image_src": ("IMAGE",), + "image_ref": ("IMAGE",), }, } + RETURN_TYPES = ("IMAGE",) FUNCTION = "weightedadd" CATEGORY = "TeaNodes/Image" @@ -31,16 +36,24 @@ class CropTo: return (result,) + class KorniaGamma: @classmethod def INPUT_TYPES(s): return { "required": { - "image": ("IMAGE", ), - "gamma": ("FLOAT", {"default": 1, "min": 0.0, "max": 100, "step": 0.01}), - "gain": ("FLOAT", {"default": 1.0, "min": -100, "max": 100, "step": 0.01}), + "image": ("IMAGE",), + "gamma": ( + "FLOAT", + {"default": 1, "min": 0.0, "max": 100, "step": 0.01}, + ), + "gain": ( + "FLOAT", + {"default": 1.0, "min": -100, "max": 100, "step": 0.01}, + ), }, } + RETURN_TYPES = ("IMAGE",) FUNCTION = "gamma" CATEGORY = "TeaNodes/Image" @@ -52,17 +65,22 @@ class KorniaGamma: return (result,) + class EqualizeCLAHE: @classmethod def INPUT_TYPES(s): return { "required": { - "image": ("IMAGE", ), + "image": ("IMAGE",), "size": ("TUPLE", {"default": (1024, 1024)}), - "clip_limit": ("FLOAT", {"default": 64, "min": 0.0, "max": 255, "step": 0.1}), + "clip_limit": ( + "FLOAT", + {"default": 64, "min": 0.0, "max": 255, "step": 0.1}, + ), "grid_size": ("INT", {"default": 8, "min": 1, "max": 64, "step": 1}), }, } + RETURN_TYPES = ("IMAGE",) FUNCTION = "equalize" CATEGORY = "TeaNodes/Image" @@ -80,15 +98,19 @@ class EqualizeCLAHE: return (result,) + class SizeApproximation: @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "square": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 32}), + "square": ( + "INT", + {"default": 1024, "min": 512, "max": 4096, "step": 32}, + ), + } } - } RETURN_TYPES = ("TUPLE", "INT", "INT") FUNCTION = "calculate" @@ -106,6 +128,7 @@ class SizeApproximation: height, width = pixel_approx(height, width, square) return ((width, height), width, height) + class ImageResize: @classmethod def INPUT_TYPES(cls): @@ -113,10 +136,16 @@ class ImageResize: "required": { "image": ("IMAGE",), "size": ("TUPLE", {"default": (1024, 1024)}), - "width": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 32}), - "height": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 32}), + "width": ( + "INT", + {"default": 1024, "min": 512, "max": 4096, "step": 32}, + ), + "height": ( + "INT", + {"default": 1024, "min": 512, "max": 4096, "step": 32}, + ), }, - } + } RETURN_TYPES = ("IMAGE",) FUNCTION = "resize" @@ -125,22 +154,29 @@ class ImageResize: def resize(self, image, size, width, height): _image = image.movedim(-1, 1) - if size != (1024, 1024): width, height = size + if size != (1024, 1024): + width, height = size - result = comfy.utils.common_upscale(_image, int( width ), int( height ), 'lanczos', 'center') + result = comfy.utils.common_upscale( + _image, int(width), int(height), "lanczos", "center" + ) result = result.movedim(1, -1) return (result,) + class ImageScale: @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "factor": ("FLOAT", {"default": 1.0, "min": 0.2, "max": 5.0, "step": 0.01}), + "factor": ( + "FLOAT", + {"default": 1.0, "min": 0.2, "max": 5.0, "step": 0.01}, + ), # "step" of 0.00 breaks the node graph engine... + } } - } RETURN_TYPES = ("IMAGE",) FUNCTION = "resize" @@ -151,21 +187,24 @@ class ImageScale: height, width = image.shape[1:3] - result = comfy.utils.common_upscale(_image, int(width * factor), int(height * factor), 'lanczos', 'center') + result = comfy.utils.common_upscale( + _image, int(width * factor), int(height * factor), "lanczos", "center" + ) result = result.movedim(1, -1) return (result,) -class RandomColorFill(): + +class RandomColorFill: @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "color": ("STRING", {"default": '#7f7f7fff'}), + "color": ("STRING", {"default": "#7f7f7fff"}), "hue_range": ("FLOAT", {"default": 0.005}), "sat_range": ("FLOAT", {"default": 0.005}), - "val_range": ("FLOAT", {"default": 0.005}) + "val_range": ("FLOAT", {"default": 0.005}), }, } @@ -187,19 +226,20 @@ class RandomColorFill(): height, width = image.shape[1:3] _color = tuple(byte(e) for e in color_rgba.RGBA) - _image = Image.new('RGBA', (width, height), _color) - _image = np.array(_image.convert('RGBA')).astype(np.float32) / 255.0 + _image = Image.new("RGBA", (width, height), _color) + _image = np.array(_image.convert("RGBA")).astype(np.float32) / 255.0 result = torch.from_numpy(_image).unsqueeze(0) print(type(result)) return (result,) -class ColorFill(): + +class ColorFill: @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "color": ("STRING", {"default": '#7f7f7fff'}), + "color": ("STRING", {"default": "#7f7f7fff"}), }, } @@ -208,26 +248,236 @@ class ColorFill(): CATEGORY = "TeaNodes/Input" def fill(self, image, color): - if color.startswith('#'): - _color = color.lstrip('#') - try: color_rgba = tuple(int(_color[i:i+2], 16) for i in (0, 2, 4, 6)) - except: color_rgba = tuple(int(_color[i:i+2], 16) for i in (0, 2, 4)) - else: print(f"Something went wrong here: {color}") + if color.startswith("#"): + _color = color.lstrip("#") + try: + color_rgba = tuple(int(_color[i : i + 2], 16) for i in (0, 2, 4, 6)) + except: + color_rgba = tuple(int(_color[i : i + 2], 16) for i in (0, 2, 4)) + else: + print(f"Something went wrong here: {color}") else: - _color = color.split(',') - try: _color = tuple(int(e) for e in _color if 255>int(e)>0) - except: print(f"Something went wrong here: {color}") + _color = color.split(",") + try: + _color = tuple(int(e) for e in _color if 255 > int(e) > 0) + except: + print(f"Something went wrong here: {color}") color_rgba = _color - color_mode = 'RGBA' if len(color_rgba) > 3 else 'RGB' + color_mode = "RGBA" if len(color_rgba) > 3 else "RGB" height, width = image.shape[1:3] - _image = Image.new('RGBA', (width, height), color_rgba) + _image = Image.new("RGBA", (width, height), color_rgba) _image = np.array(_image.convert(color_mode)).astype(np.float32) / 255.0 result = torch.from_numpy(_image).unsqueeze(0) return (result,) +import os + +def history_order(data, func=os.path.basename, predicate=lambda e: bool(e), reverse=True): + match data: + case dict(): return '\n'.join(f"{v:>4d}: {func(k)}" for k, v in sorted((t for t in data.items() if predicate(t)), key=lambda e: e[-1], reverse=reverse)) + case list()|set()|tuple(): return '\n'.join(f"{func(e)}" for e in data) + case _: return '' + +import re + +RETYPE = type(re.compile('')) + +class Pool: + def __init__(self, content): + self.regex = None + self.content = content + self.filtered = [] + + def sieve(self, pattern): + if isinstance(self.regex, RETYPE) and self.regex.pattern == pattern and self.filtered: return None + else: self.regex = re.compile(pattern, re.I); self.filtered = [name for name in self.content if self.regex.search(name)] + return True + +class LoraPool(Pool): + name = 'Lora' + loaded = dict() + history = dict() + def __init__(self): + super().__init__([os.path.join(p, e) for p, _, f in os.walk('ComfyUI\\models\\loras', followlinks=True) for e in f if e.endswith('.safetensors')]) + +class ModelPool(Pool): + name = "Model" + loaded = dict() + history = dict() + def __init__(self): + super().__init__([os.path.join(p, e) for p, _, f in os.walk('ComfyUI\\models\\checkpoints', followlinks=True) for e in f if e.endswith('.safetensors')]) + +class Randomizer(): + def __init__(self, pool): + self.path = '' + self.pool : Pool = pool + self.loop = 0 + self.every = 1 + self.pause = False + + def out(self, skip=False): + if not skip: self.loop += 1 + warn = f"Warning: No {self.pool.name} Found! Using the last available.\n" if not self.pool.filtered else '' + return '\n'.join(( + f"{warn}Loop-<{self.loop}/{self.every}>: {self.path}", + "\nHistory (has generated with):", + history_order(self.pool.history, predicate=lambda e: e in self.pool.filtered), + "\nPool (hasn't generated with):", + history_order(self.pool.filtered), + )) + + def yet(self): + self.loop %= self.every + if not self.path: self.path = random.choice(self.pool.content) + return self.pause or not self.loop + + def run(self, pattern): + self.loop %= self.every + # Execute the randomization only after every # generations + if self.pause or self.loop != 0: return self.path + if self.pool.sieve(pattern): + if self.path not in self.pool.history: self.pool.history[self.path] = 0 + self.pool.history[self.path] += 1 + else: + self.pool.history = dict() + if self.pool.filtered: self.path = random.choice(self.pool.filtered) + + def get(self): + return self.pool.loaded[self.path] + + def put(self, content): + self.pool.loaded[self.path] = content + + +class RandomLora: + def __init__(self): + self.randomizer = Randomizer(LoraPool()) + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("MODEL", {"tooltip": "The diffusion model the LoRA will be applied to."}), + "clip": ("CLIP", {"tooltip": "The CLIP model the LoRA will be applied to."}), + "pattern": ("STRING", {"default": "", "multiline": True, "tooltip": "Regular Expression"}), + "every": ("INT", {"default": 7, "min": 1, "max": 99, "tooltip": "Change only takes effect every N generations."}), + "pause": ("BOOLEAN", {"default": False, "tooltip": "Pause the randomization and counting, keep generating with current lora."}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the diffusion model. This value can be negative."}), + "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}), + } + } + + RETURN_TYPES = ("MODEL", "CLIP", "STRING") + RETURN_NAMES = ("MODEL", "CLIP", "RESULT") + OUTPUT_TOOLTIPS = ("The modified diffusion model.", "The modified CLIP model.", "List of lora names that matches the pattern.") + FUNCTION = "load_lora" + + CATEGORY = "loaders" + DESCRIPTION = "Randomly apply a different Lora that matches the search pattern after every N generations." + SEARCH_ALIASES = ["lora", "load lora", "apply lora", "lora loader", "lora model"] + + def try_get(self, model, clip, strength_model, strength_clip): + try: + model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, self.randomizer.get(), strength_model, strength_clip) + return (model_lora, clip_lora, self.randomizer.out()) + except: + self.randomizer.put(comfy.utils.load_torch_file(self.randomizer.path)) + model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, self.randomizer.get(), strength_model, strength_clip) + return (model_lora, clip_lora, self.randomizer.out()) + + def load_lora(self, model, clip, pattern, every, pause, strength_model, strength_clip, **kwargs): + self.randomizer.every = every + self.randomizer.pause = pause + if strength_model == 0 and strength_clip == 0: + return (model, clip) + if not self.randomizer.yet(): + return self.try_get(model, clip, strength_clip, strength_model) + else: + self.randomizer.run(pattern) + return self.try_get(model, clip, strength_clip, strength_model) + +class RandomModel: + def __init__(self): + self.randomizer = Randomizer(ModelPool()) + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "pattern": ("STRING", {"default": "", "multiline": True, "tooltip": "Regular Expression"}), + "every": ("INT", {"default": 7, "min": 1, "max": 99, "tooltip": "Change only takes effect every N generations."}), + "pause": ("BOOLEAN", {"default": False, "tooltip": "Pause the randomization and counting, keep generating with current model."}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + # "ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}), + } + } + RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING") + RETURN_NAMES = ("MODEL", "CLIP", "VAE", "RESULT") + OUTPUT_TOOLTIPS = ("The model used for denoising latents.", + "The CLIP model used for encoding text prompts.", + "The VAE model used for encoding and decoding images to and from latent space.", + "List of model names that matches the pattern.") + FUNCTION = "load_checkpoint" + + CATEGORY = "loaders" + DESCRIPTION = "Randomly switch to a different Model which matches the search pattern after every N generations." + SEARCH_ALIASES = ["load model", "checkpoint", "model loader", "load checkpoint", "ckpt", "model"] + + def try_get(self): + try: + return *self.randomizer.get(), self.randomizer.out() + except: + self.randomizer.put(comfy.sd.load_checkpoint_guess_config(self.randomizer.path, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))[:3]) + return *self.randomizer.get(), self.randomizer.out() + + def load_checkpoint(self, pattern, every, pause, **kwargs): + self.randomizer.every = every + self.randomizer.pause = pause + if not self.randomizer.yet(): + return self.try_get() + else: + self.randomizer.run(pattern) + return self.try_get() + +# class RegexCkptList: +# @classmethod +# def INPUT_TYPES(s): +# return { +# "required": { +# "regex": ("STRING", {"default": ""}), +# "ckpt_name": ( +# folder_paths.get_filename_list("checkpoints"), +# {"tooltip": "The name of the checkpoint (model) to load."}, +# ), +# } +# } + +# RETURN_TYPES = ("MODEL", "CLIP", "VAE") +# OUTPUT_TOOLTIPS = ( +# "The model used for denoising latents.", +# "The CLIP model used for encoding text prompts.", +# "The VAE model used for encoding and decoding images to and from latent space.", +# ) +# FUNCTION = "load_checkpoint" + +# CATEGORY = "loaders" +# DESCRIPTION = "Loads a diffusion model checkpoint, diffusion models are used to denoise latents." + +# def load_checkpoint(self, ckpt_name): +# ckpt_path = folder_paths.get_full_path_or_raise("checkpoints", ckpt_name) +# out = comfy.sd.load_checkpoint_guess_config( +# ckpt_path, +# output_vae=True, +# output_clip=True, +# embedding_directory=folder_paths.get_folder_paths("embeddings"), +# ) +# return out[:3] + + NODE_CLASS_MAPPINGS = { "TC_CropTo": CropTo, "TC_KorniaGamma": KorniaGamma, @@ -237,4 +487,6 @@ NODE_CLASS_MAPPINGS = { "TC_ImageScale": ImageScale, "TC_ColorFill": ColorFill, "TC_RandomColorFill": RandomColorFill, -} \ No newline at end of file + "TC_RandomLora": RandomLora, + "TC_RandomModel": RandomModel, +}