TeaCrab
9c3af9fc8a
* RandomLora & RandomModel
* Walks the `models/checkpoints` and `models/loras` folder for `.safetensors` files and uses Regex patterns to filter them down, then randomly select one to apply.
* If the `pattern` is changed, filtering and randomization will be triggered next time the node is executed.
* `every` parameter allows the selected model to loop for a certain number of times before next randomization.
* `pause` holds off randomization and keep using the same model for generation indefinitely.
* `skip` will randomly select another model ignoring the `every` parameter but the looping will continue to count up.
* `RESULT` output will show in text format the current number of loops, which model is currently selected, which model has ran and how many times, and lists all the models that hasn't ran. Use `Preview as Text` or `Show Text` node to access.
* Models are pooled - loaded to RAM from HDD/SSDs, this causes RAM to be filled to the brim, however, Python seem to be able to automatically release memory for new ones
* If previously loaded models are selected again, it is already in the RAM, so no need to access SSD or HDD for it again.
* This saves the tiny bit of load time when using SSD, but could potentially save a lot of loading time for HDD.
* Eventually this could be used for cycling through multiple designated models to compare their differences.
461 lines
16 KiB
Python
461 lines
16 KiB
Python
import random
|
|
from colorsys import hsv_to_rgb
|
|
|
|
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",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "weightedadd"
|
|
CATEGORY = "TeaNodes/Image"
|
|
|
|
def weightedadd(self, image_src, image_ref):
|
|
R = image_ref.movedim(-1, 1)
|
|
S = image_src.movedim(-1, 1)
|
|
|
|
_image = center_crop(S, R.shape[2:])
|
|
result = _image.movedim(1, -1)
|
|
|
|
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},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "gamma"
|
|
CATEGORY = "TeaNodes/Image"
|
|
|
|
def gamma(self, image, gamma, gain):
|
|
_image = image.movedim(-1, 1)
|
|
_image = adjust_gamma(_image, gamma, gain)
|
|
result = _image.movedim(1, -1)
|
|
|
|
return (result,)
|
|
|
|
|
|
class EqualizeCLAHE:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"size": ("TUPLE", {"default": (1024, 1024)}),
|
|
"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"
|
|
|
|
def equalize(self, image, size, clip_limit, grid_size):
|
|
_image = image.movedim(-1, 1)
|
|
if size != (1024, 1024):
|
|
grid_ratio = grid_size / clip_limit
|
|
clip_limit = int(max(8, po2(max(size)) * (clip_limit / 1024)))
|
|
grid_size = int(max(2, clip_limit * grid_ratio))
|
|
print(clip_limit, grid_size)
|
|
|
|
_image = equalize_clahe(_image, float(clip_limit), (grid_size, grid_size))
|
|
result = _image.movedim(1, -1)
|
|
|
|
return (result,)
|
|
|
|
|
|
class SizeApproximation:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"square": (
|
|
"INT",
|
|
{"default": 1024, "min": 512, "max": 4096, "step": 32},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TUPLE", "INT", "INT")
|
|
FUNCTION = "calculate"
|
|
CATEGORY = "TeaNodes/Image"
|
|
|
|
def calculate(self, image, square):
|
|
_image = image.movedim(-1, 1)
|
|
height, width = image.shape[1:3]
|
|
|
|
# print(width, height)
|
|
|
|
if width >= height:
|
|
width, height = pixel_approx(width, height, square)
|
|
else:
|
|
height, width = pixel_approx(height, width, square)
|
|
return ((width, height), width, height)
|
|
|
|
|
|
class ImageResize:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"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},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "resize"
|
|
CATEGORY = "TeaNodes/Image"
|
|
|
|
def resize(self, image, size, width, height):
|
|
_image = image.movedim(-1, 1)
|
|
|
|
if size != (1024, 1024):
|
|
width, height = size
|
|
|
|
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},
|
|
),
|
|
# "step" of 0.00 breaks the node graph engine...
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "resize"
|
|
CATEGORY = "TeaNodes/Image"
|
|
|
|
def resize(self, image, factor):
|
|
_image = image.movedim(-1, 1)
|
|
|
|
height, width = image.shape[1:3]
|
|
|
|
result = comfy.utils.common_upscale(
|
|
_image, int(width * factor), int(height * factor), "lanczos", "center"
|
|
)
|
|
result = result.movedim(1, -1)
|
|
|
|
return (result,)
|
|
|
|
|
|
class RandomColorFill:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"color": ("STRING", {"default": "#7f7f7fff"}),
|
|
"hue_range": ("FLOAT", {"default": 0.005}),
|
|
"sat_range": ("FLOAT", {"default": 0.005}),
|
|
"val_range": ("FLOAT", {"default": 0.005}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "randomize"
|
|
CATEGORY = "TeaNodes/Input"
|
|
|
|
def randomize(self, image, color, hue_range, sat_range, val_range):
|
|
h, s, v, _ = Color(color).hsv()
|
|
h += hue_range * (random.random() - 0.5)
|
|
s += sat_range * (random.random() - 0.5)
|
|
v += val_range * (random.random() - 0.5)
|
|
|
|
color_rgba = Color()
|
|
color_rgba.RGBA = hsv_to_rgb(h, s, v) + (1.0,)
|
|
print(f"Random Color:\t{color_rgba.hex()}")
|
|
|
|
print(type(image))
|
|
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
|
|
result = torch.from_numpy(_image).unsqueeze(0)
|
|
print(type(result))
|
|
return (result,)
|
|
|
|
|
|
class ColorFill:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"color": ("STRING", {"default": "#7f7f7fff"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "fill"
|
|
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}")
|
|
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_rgba = _color
|
|
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 = np.array(_image.convert(color_mode)).astype(np.float32) / 255.0
|
|
result = torch.from_numpy(_image).unsqueeze(0)
|
|
|
|
return (result,)
|
|
|
|
|
|
import re
|
|
import os
|
|
|
|
RETYPE = type(re.compile(''))
|
|
|
|
class Pool:
|
|
def __init__(self, content):
|
|
self.regex = re.compile('', re.I)
|
|
self.content = content
|
|
self.filtered = []
|
|
|
|
def sieve(self, pattern):
|
|
if self.regex.pattern!=pattern or not self.filtered:
|
|
try: self.regex = re.compile(pattern, re.I)
|
|
except Exception: return None
|
|
self.filtered = [name for name in self.content if self.regex.search(name)]
|
|
return True
|
|
|
|
class LoraPool(Pool):
|
|
name = 'Loras'
|
|
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 = "Models"
|
|
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
|
|
self.loop = 0
|
|
self.every = 1
|
|
self.error = ''
|
|
self.pause = False
|
|
|
|
def out(self):
|
|
if not self.pause: self.loop += 1
|
|
warn = ': '.join(e for e in (
|
|
f"{self.error}" if self.error else '',
|
|
f"No {self.pool.name} Found! Using the last available.\n" if not self.pool.filtered else ''
|
|
) if e)
|
|
return '\n'.join((
|
|
f"{warn}Loop-<{self.loop}/{self.every}>: {self.path}",
|
|
"\nHistory (has generated with):",
|
|
"\n".join(f"{v:>4d}: {os.path.basename(k)}" for k, v in sorted(self.pool.history.items(), key=lambda e: e[-1], reverse=True) if k in self.pool.filtered),
|
|
"\nPool (hasn't generated with):",
|
|
"\n".join(os.path.basename(e) for e in self.pool.filtered if e not in self.pool.history),
|
|
))
|
|
|
|
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, force=False):
|
|
self.loop %= self.every
|
|
# Execute the randomization only after every # generations
|
|
if not force and (self.pause or self.loop != 0): return self.path
|
|
if self.pool.sieve(pattern) is None:
|
|
self.error = "Invalid Pattern"
|
|
else:
|
|
if self.path not in self.pool.history: self.pool.history[self.path] = 0
|
|
self.pool.history[self.path] += 1
|
|
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."}),
|
|
"skip": ("BOOLEAN", {"default": False, "tooltip": "Skip curent model."}),
|
|
"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, skip, 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 skip and not self.randomizer.yet():
|
|
return self.try_get(model, clip, strength_clip, strength_model)
|
|
else:
|
|
self.randomizer.run(pattern, skip)
|
|
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."}),
|
|
"skip": ("BOOLEAN", {"default": False, "tooltip": "Skip curent model."}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
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, skip, **kwargs):
|
|
self.randomizer.every = every
|
|
self.randomizer.pause = pause
|
|
if not skip and not self.randomizer.yet():
|
|
return self.try_get()
|
|
else:
|
|
self.randomizer.run(pattern, skip)
|
|
return self.try_get()
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"TC_CropTo": CropTo,
|
|
"TC_KorniaGamma": KorniaGamma,
|
|
"TC_EqualizeCLAHE": EqualizeCLAHE,
|
|
"TC_SizeApproximation": SizeApproximation,
|
|
"TC_ImageResize": ImageResize,
|
|
"TC_ImageScale": ImageScale,
|
|
"TC_ColorFill": ColorFill,
|
|
"TC_RandomColorFill": RandomColorFill,
|
|
"TC_RandomLora": RandomLora,
|
|
"TC_RandomModel": RandomModel,
|
|
}
|