nodes/TeaCrab-ComfyUI-TeaNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

20260412:

Browse Source 
WIP RandomLora, RandomModel
This commit is contained in:
TeaCrab
2026-04-12 21:38:57 -04:00
parent 3fcf9f3568
commit 90408f02cb
3 changed files with 314 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__.code-workspace
+13
View File
@@ -0,0 +1,13 @@
{
"folders": [
{
"path": ".",
}
],
"settings": {
"files.exclude": {
"__pycache__":true,
}
}
}
__init__.py
+7 -2
View File
@@ -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}
_nodes.py
+294 -42
View File
@@ -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,
}
"TC_RandomLora": RandomLora,
"TC_RandomModel": RandomModel,
}