code rewrite, bug fixes, nodes reorganization
This commit is contained in:
cubiq
+32
-1
@@ -1,3 +1,34 @@
|
||||
from .essentials import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
#from .essentials import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .image import IMAGE_CLASS_MAPPINGS, IMAGE_NAME_MAPPINGS
|
||||
from .mask import MASK_CLASS_MAPPINGS, MASK_NAME_MAPPINGS
|
||||
from .sampling import SAMPLING_CLASS_MAPPINGS, SAMPLING_NAME_MAPPINGS
|
||||
from .segmentation import SEG_CLASS_MAPPINGS, SEG_NAME_MAPPINGS
|
||||
from .misc import MISC_CLASS_MAPPINGS, MISC_NAME_MAPPINGS
|
||||
from .conditioning import COND_CLASS_MAPPINGS, COND_NAME_MAPPINGS
|
||||
from .text import TEXT_CLASS_MAPPINGS, TEXT_NAME_MAPPINGS
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(COND_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(COND_NAME_MAPPINGS)
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(IMAGE_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(IMAGE_NAME_MAPPINGS)
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(MASK_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(MASK_NAME_MAPPINGS)
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(SAMPLING_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(SAMPLING_NAME_MAPPINGS)
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(SEG_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(SEG_NAME_MAPPINGS)
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(TEXT_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(TEXT_NAME_MAPPINGS)
|
||||
|
||||
NODE_CLASS_MAPPINGS.update(MISC_CLASS_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(MISC_NAME_MAPPINGS)
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
from nodes import MAX_RESOLUTION
|
||||
|
||||
class CLIPTextEncodeSDXLSimplified:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
|
||||
"height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
|
||||
"size_cond_factor": ("INT", {"default": 4, "min": 1, "max": 16 }),
|
||||
"text": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": ""}),
|
||||
"clip": ("CLIP", ),
|
||||
}}
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/conditioning"
|
||||
|
||||
def execute(self, clip, width, height, size_cond_factor, text):
|
||||
crop_w = 0
|
||||
crop_h = 0
|
||||
width = width*size_cond_factor
|
||||
height = height*size_cond_factor
|
||||
target_width = width
|
||||
target_height = height
|
||||
text_g = text_l = text
|
||||
|
||||
tokens = clip.tokenize(text_g)
|
||||
tokens["l"] = clip.tokenize(text_l)["l"]
|
||||
if len(tokens["l"]) != len(tokens["g"]):
|
||||
empty = clip.tokenize("")
|
||||
while len(tokens["l"]) < len(tokens["g"]):
|
||||
tokens["l"] += empty["l"]
|
||||
while len(tokens["l"]) > len(tokens["g"]):
|
||||
tokens["g"] += empty["g"]
|
||||
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
|
||||
return ([[cond, {"pooled_output": pooled, "width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]], )
|
||||
|
||||
class ConditioningCombineMultiple:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"conditioning_1": ("CONDITIONING",),
|
||||
"conditioning_2": ("CONDITIONING",),
|
||||
}, "optional": {
|
||||
"conditioning_3": ("CONDITIONING",),
|
||||
"conditioning_4": ("CONDITIONING",),
|
||||
"conditioning_5": ("CONDITIONING",),
|
||||
},
|
||||
}
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/conditioning"
|
||||
|
||||
def execute(self, conditioning_1, conditioning_2, conditioning_3=None, conditioning_4=None, conditioning_5=None):
|
||||
c = conditioning_1 + conditioning_2
|
||||
|
||||
if conditioning_3 is not None:
|
||||
c += conditioning_3
|
||||
if conditioning_4 is not None:
|
||||
c += conditioning_4
|
||||
if conditioning_5 is not None:
|
||||
c += conditioning_5
|
||||
|
||||
return (c,)
|
||||
|
||||
COND_CLASS_MAPPINGS = {
|
||||
"CLIPTextEncodeSDXL+": CLIPTextEncodeSDXLSimplified,
|
||||
"ConditioningCombineMultiple+": ConditioningCombineMultiple,
|
||||
}
|
||||
|
||||
COND_NAME_MAPPINGS = {
|
||||
"CLIPTextEncodeSDXL+": "🔧 SDXL CLIPTextEncode",
|
||||
"ConditioningCombineMultiple+": "🔧 Conditionings Combine Multiple",
|
||||
}
|
||||
-2103
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,593 @@
|
||||
from nodes import SaveImage
|
||||
import torch
|
||||
import torchvision.transforms.v2 as T
|
||||
import random
|
||||
import folder_paths
|
||||
import comfy.utils
|
||||
from .image import ImageExpandBatch
|
||||
import numpy as np
|
||||
import scipy
|
||||
from PIL import Image
|
||||
from nodes import MAX_RESOLUTION
|
||||
import math
|
||||
|
||||
class MaskBlur:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"amount": ("INT", { "default": 6, "min": 0, "max": 256, "step": 1, }),
|
||||
"device": (["auto", "cpu", "gpu"],),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, mask, amount, device):
|
||||
if amount == 0:
|
||||
return (mask,)
|
||||
|
||||
if "gpu" == device:
|
||||
mask = mask.to(comfy.model_management.get_torch_device())
|
||||
elif "cpu" == device:
|
||||
mask = mask.to('cpu')
|
||||
|
||||
if amount % 2 == 0:
|
||||
amount+= 1
|
||||
|
||||
if mask.dim() == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
|
||||
mask = T.functional.gaussian_blur(mask.unsqueeze(1), amount).squeeze(1)
|
||||
|
||||
if "gpu" == device or "cpu" == device:
|
||||
mask = mask.to(comfy.model_management.intermediate_device())
|
||||
|
||||
return(mask,)
|
||||
|
||||
class MaskFlip:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"axis": (["x", "y", "xy"],),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, mask, axis):
|
||||
if mask.dim() == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
|
||||
dim = ()
|
||||
if "y" in axis:
|
||||
dim += (1,)
|
||||
if "x" in axis:
|
||||
dim += (2,)
|
||||
mask = torch.flip(mask, dims=dim)
|
||||
|
||||
return(mask,)
|
||||
|
||||
class MaskPreview(SaveImage):
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_temp_directory()
|
||||
self.type = "temp"
|
||||
self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
|
||||
self.compress_level = 4
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {"mask": ("MASK",), },
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, mask, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
|
||||
preview = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
||||
return self.save_images(preview, filename_prefix, prompt, extra_pnginfo)
|
||||
|
||||
class MaskBatch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask1": ("MASK",),
|
||||
"mask2": ("MASK",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask batch"
|
||||
|
||||
def execute(self, mask1, mask2):
|
||||
if mask1.shape[1:] != mask2.shape[1:]:
|
||||
mask2 = comfy.utils.common_upscale(mask2.unsqueeze(1).expand(-1,3,-1,-1), mask1.shape[2], mask1.shape[1], upscale_method='bicubic', crop='center')[:,0,:,:]
|
||||
|
||||
return (torch.cat((mask1, mask2), dim=0),)
|
||||
|
||||
class MaskExpandBatch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"size": ("INT", { "default": 16, "min": 1, "step": 1, }),
|
||||
"method": (["expand", "repeat all", "repeat first", "repeat last"],)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask batch"
|
||||
|
||||
def execute(self, mask, size, method):
|
||||
return (ImageExpandBatch().execute(mask.unsqueeze(1).expand(-1,3,-1,-1), size, method)[0][:,0,:,:],)
|
||||
|
||||
|
||||
class MaskBoundingBox:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"padding": ("INT", { "default": 0, "min": 0, "max": 4096, "step": 1, }),
|
||||
"blur": ("INT", { "default": 0, "min": 0, "max": 256, "step": 1, }),
|
||||
},
|
||||
"optional": {
|
||||
"image_optional": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK", "IMAGE", "INT", "INT", "INT", "INT")
|
||||
RETURN_NAMES = ("MASK", "IMAGE", "x", "y", "width", "height")
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, mask, padding, blur, image_optional=None):
|
||||
if mask.dim() == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
|
||||
if image_optional is None:
|
||||
image_optional = mask.unsqueeze(3).repeat(1, 1, 1, 3)
|
||||
|
||||
# resize the image if it's not the same size as the mask
|
||||
if image_optional.shape[1:] != mask.shape[1:]:
|
||||
image_optional = comfy.utils.common_upscale(image_optional.permute([0,3,1,2]), mask.shape[2], mask.shape[1], upscale_method='bicubic', crop='center').permute([0,2,3,1])
|
||||
|
||||
# match batch size
|
||||
if image_optional.shape[0] < mask.shape[0]:
|
||||
image_optional = torch.cat((image_optional, image_optional[-1].unsqueeze(0).repeat(mask.shape[0]-image_optional.shape[0], 1, 1, 1)), dim=0)
|
||||
elif image_optional.shape[0] > mask.shape[0]:
|
||||
image_optional = image_optional[:mask.shape[0]]
|
||||
|
||||
# blur the mask
|
||||
if blur > 0:
|
||||
if blur % 2 == 0:
|
||||
blur += 1
|
||||
mask = T.functional.gaussian_blur(mask.unsqueeze(1), blur).squeeze(1)
|
||||
|
||||
_, y, x = torch.where(mask)
|
||||
x1 = max(0, x.min().item() - padding)
|
||||
x2 = min(mask.shape[2], x.max().item() + 1 + padding)
|
||||
y1 = max(0, y.min().item() - padding)
|
||||
y2 = min(mask.shape[1], y.max().item() + 1 + padding)
|
||||
|
||||
# crop the mask
|
||||
mask = mask[:, y1:y2, x1:x2]
|
||||
image_optional = image_optional[:, y1:y2, x1:x2, :]
|
||||
|
||||
return (mask, image_optional, x1, y1, x2 - x1, y2 - y1)
|
||||
|
||||
|
||||
class MaskFromColor:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", ),
|
||||
"red": ("INT", { "default": 255, "min": 0, "max": 255, "step": 1, }),
|
||||
"green": ("INT", { "default": 255, "min": 0, "max": 255, "step": 1, }),
|
||||
"blue": ("INT", { "default": 255, "min": 0, "max": 255, "step": 1, }),
|
||||
"threshold": ("INT", { "default": 0, "min": 0, "max": 127, "step": 1, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, image, red, green, blue, threshold):
|
||||
temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int)
|
||||
color = torch.tensor([red, green, blue])
|
||||
lower_bound = (color - threshold).clamp(min=0)
|
||||
upper_bound = (color + threshold).clamp(max=255)
|
||||
lower_bound = lower_bound.view(1, 1, 1, 3)
|
||||
upper_bound = upper_bound.view(1, 1, 1, 3)
|
||||
mask = (temp >= lower_bound) & (temp <= upper_bound)
|
||||
mask = mask.all(dim=-1)
|
||||
mask = mask.float()
|
||||
|
||||
return (mask, )
|
||||
|
||||
|
||||
class MaskFromSegmentation:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", ),
|
||||
"segments": ("INT", { "default": 6, "min": 1, "max": 16, "step": 1, }),
|
||||
"remove_isolated_pixels": ("INT", { "default": 0, "min": 0, "max": 32, "step": 1, }),
|
||||
"remove_small_masks": ("FLOAT", { "default": 0.0, "min": 0., "max": 1., "step": 0.01, }),
|
||||
"fill_holes": ("BOOLEAN", { "default": False }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, image, segments, remove_isolated_pixels, fill_holes, remove_small_masks):
|
||||
im = image[0] # we only work on the first image in the batch
|
||||
im = Image.fromarray((im * 255).to(torch.uint8).cpu().numpy(), mode="RGB")
|
||||
im = im.quantize(palette=im.quantize(colors=segments), dither=Image.Dither.NONE)
|
||||
im = torch.tensor(np.array(im.convert("RGB"))).float() / 255.0
|
||||
|
||||
colors = im.reshape(-1, im.shape[-1])
|
||||
colors = torch.unique(colors, dim=0)
|
||||
|
||||
masks = []
|
||||
for color in colors:
|
||||
mask = (im == color).all(dim=-1).float()
|
||||
# remove isolated pixels
|
||||
if remove_isolated_pixels > 0:
|
||||
mask = torch.from_numpy(scipy.ndimage.binary_opening(mask.cpu().numpy(), structure=np.ones((remove_isolated_pixels, remove_isolated_pixels))))
|
||||
|
||||
# fill holes
|
||||
if fill_holes:
|
||||
mask = torch.from_numpy(scipy.ndimage.binary_fill_holes(mask.cpu().numpy()))
|
||||
|
||||
# if the mask is too small, it's probably noise
|
||||
if mask.sum() / (mask.shape[0]*mask.shape[1]) > remove_small_masks:
|
||||
masks.append(mask)
|
||||
|
||||
if masks == []:
|
||||
masks.append(torch.zeros_like(im)[:,:,0]) # return an empty mask if no masks were found, prevents errors
|
||||
|
||||
mask = torch.stack(masks, dim=0).float()
|
||||
|
||||
return (mask, )
|
||||
|
||||
|
||||
class MaskFix:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"erode_dilate": ("INT", { "default": 0, "min": -256, "max": 256, "step": 1, }),
|
||||
"fill_holes": ("INT", { "default": 0, "min": 0, "max": 128, "step": 1, }),
|
||||
"remove_isolated_pixels": ("INT", { "default": 0, "min": 0, "max": 32, "step": 1, }),
|
||||
"smooth": ("INT", { "default": 0, "min": 0, "max": 256, "step": 1, }),
|
||||
"blur": ("INT", { "default": 0, "min": 0, "max": 256, "step": 1, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, mask, erode_dilate, smooth, remove_isolated_pixels, blur, fill_holes):
|
||||
masks = []
|
||||
for m in mask:
|
||||
# erode and dilate
|
||||
if erode_dilate != 0:
|
||||
if erode_dilate < 0:
|
||||
m = torch.from_numpy(scipy.ndimage.grey_erosion(m.cpu().numpy(), size=(-erode_dilate, -erode_dilate)))
|
||||
else:
|
||||
m = torch.from_numpy(scipy.ndimage.grey_dilation(m.cpu().numpy(), size=(erode_dilate, erode_dilate)))
|
||||
|
||||
# fill holes
|
||||
if fill_holes > 0:
|
||||
#m = torch.from_numpy(scipy.ndimage.binary_fill_holes(m.cpu().numpy(), structure=np.ones((fill_holes,fill_holes)))).float()
|
||||
m = torch.from_numpy(scipy.ndimage.grey_closing(m.cpu().numpy(), size=(fill_holes, fill_holes)))
|
||||
|
||||
# remove isolated pixels
|
||||
if remove_isolated_pixels > 0:
|
||||
m = torch.from_numpy(scipy.ndimage.grey_opening(m.cpu().numpy(), size=(remove_isolated_pixels, remove_isolated_pixels)))
|
||||
|
||||
# smooth the mask
|
||||
if smooth > 0:
|
||||
if smooth % 2 == 0:
|
||||
smooth += 1
|
||||
m = T.functional.gaussian_blur((m > 0.5).unsqueeze(0), smooth).squeeze(0)
|
||||
|
||||
# blur the mask
|
||||
if blur > 0:
|
||||
if blur % 2 == 0:
|
||||
blur += 1
|
||||
m = T.functional.gaussian_blur(m.float().unsqueeze(0), blur).squeeze(0)
|
||||
|
||||
masks.append(m.float())
|
||||
|
||||
masks = torch.stack(masks, dim=0).float()
|
||||
|
||||
return (masks, )
|
||||
|
||||
class MaskSmooth:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK",),
|
||||
"amount": ("INT", { "default": 0, "min": 0, "max": 127, "step": 1, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, mask, amount):
|
||||
if amount == 0:
|
||||
return (mask,)
|
||||
|
||||
if amount % 2 == 0:
|
||||
amount += 1
|
||||
|
||||
mask = mask > 0.5
|
||||
mask = T.functional.gaussian_blur(mask.unsqueeze(1), amount).squeeze(1).float()
|
||||
|
||||
return (mask,)
|
||||
|
||||
class MaskFromBatch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask": ("MASK", ),
|
||||
"start": ("INT", { "default": 0, "min": 0, "step": 1, }),
|
||||
"length": ("INT", { "default": 1, "min": 1, "step": 1, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask batch"
|
||||
|
||||
def execute(self, mask, start, length):
|
||||
if length > mask.shape[0]:
|
||||
length = mask.shape[0]
|
||||
|
||||
start = min(start, mask.shape[0]-1)
|
||||
length = min(mask.shape[0]-start, length)
|
||||
return (mask[start:start + length], )
|
||||
|
||||
class MaskFromList:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"width": ("INT", { "default": 32, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
|
||||
"height": ("INT", { "default": 32, "min": 0, "max": MAX_RESOLUTION, "step": 8, }),
|
||||
}, "optional": {
|
||||
"values": ("INT,FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, }),
|
||||
"str_values": ("STRING", { "default": "", "multiline": True, "placeholder": "0.0, 0.5, 1.0",}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, width, height, values=None, str_values=""):
|
||||
out = []
|
||||
|
||||
if values is not None:
|
||||
if not isinstance(values, list):
|
||||
out = [values]
|
||||
else:
|
||||
out.extend(values)
|
||||
|
||||
if str_values != "":
|
||||
str_values = [float(v) for v in str_values.split(",")]
|
||||
out.extend(str_values)
|
||||
|
||||
if out == []:
|
||||
raise ValueError("No values provided")
|
||||
|
||||
out = torch.tensor(out).float().clamp(0.0, 1.0)
|
||||
out = out.view(-1, 1, 1).expand(-1, height, width)
|
||||
|
||||
values = None
|
||||
str_values = ""
|
||||
|
||||
return (out, )
|
||||
|
||||
class MaskFromRGBCMYBW:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", ),
|
||||
"threshold_r": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1, "step": 0.01, }),
|
||||
"threshold_g": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1, "step": 0.01, }),
|
||||
"threshold_b": ("FLOAT", { "default": 0.15, "min": 0.0, "max": 1, "step": 0.01, }),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK","MASK","MASK","MASK","MASK","MASK","MASK","MASK",)
|
||||
RETURN_NAMES = ("red","green","blue","cyan","magenta","yellow","black","white",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def execute(self, image, threshold_r, threshold_g, threshold_b):
|
||||
red = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] < threshold_g) & (image[..., 2] < threshold_b)).float()
|
||||
green = ((image[..., 0] < threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] < threshold_b)).float()
|
||||
blue = ((image[..., 0] < threshold_r) & (image[..., 1] < threshold_g) & (image[..., 2] >= 1-threshold_b)).float()
|
||||
|
||||
cyan = ((image[..., 0] < threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] >= 1-threshold_b)).float()
|
||||
magenta = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] < threshold_g) & (image[..., 2] > 1-threshold_b)).float()
|
||||
yellow = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] < threshold_b)).float()
|
||||
|
||||
black = ((image[..., 0] <= threshold_r) & (image[..., 1] <= threshold_g) & (image[..., 2] <= threshold_b)).float()
|
||||
white = ((image[..., 0] >= 1-threshold_r) & (image[..., 1] >= 1-threshold_g) & (image[..., 2] >= 1-threshold_b)).float()
|
||||
|
||||
return (red, green, blue, cyan, magenta, yellow, black, white,)
|
||||
|
||||
class TransitionMask:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"width": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, }),
|
||||
"height": ("INT", { "default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1, }),
|
||||
"frames": ("INT", { "default": 16, "min": 1, "max": 9999, "step": 1, }),
|
||||
"start_frame": ("INT", { "default": 0, "min": 0, "step": 1, }),
|
||||
"end_frame": ("INT", { "default": 9999, "min": 0, "step": 1, }),
|
||||
"transition_type": (["horizontal slide", "vertical slide", "horizontal bar", "vertical bar", "center box", "horizontal door", "vertical door", "circle", "fade"],),
|
||||
"timing_function": (["linear", "in", "out", "in-out"],)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/mask"
|
||||
|
||||
def linear(self, i, t):
|
||||
return i/t
|
||||
def ease_in(self, i, t):
|
||||
return pow(i/t, 2)
|
||||
def ease_out(self, i, t):
|
||||
return 1 - pow(1 - i/t, 2)
|
||||
def ease_in_out(self, i, t):
|
||||
if i < t/2:
|
||||
return pow(i/(t/2), 2) / 2
|
||||
else:
|
||||
return 1 - pow(1 - (i - t/2)/(t/2), 2) / 2
|
||||
|
||||
def execute(self, width, height, frames, start_frame, end_frame, transition_type, timing_function):
|
||||
if timing_function == 'in':
|
||||
timing_function = self.ease_in
|
||||
elif timing_function == 'out':
|
||||
timing_function = self.ease_out
|
||||
elif timing_function == 'in-out':
|
||||
timing_function = self.ease_in_out
|
||||
else:
|
||||
timing_function = self.linear
|
||||
|
||||
out = []
|
||||
|
||||
end_frame = min(frames, end_frame)
|
||||
transition = end_frame - start_frame
|
||||
|
||||
if start_frame > 0:
|
||||
out = out + [torch.full((height, width), 0.0, dtype=torch.float32, device="cpu")] * start_frame
|
||||
|
||||
for i in range(transition):
|
||||
frame = torch.full((height, width), 0.0, dtype=torch.float32, device="cpu")
|
||||
progress = timing_function(i, transition-1)
|
||||
|
||||
if "horizontal slide" in transition_type:
|
||||
pos = round(width*progress)
|
||||
frame[:, :pos] = 1.0
|
||||
elif "vertical slide" in transition_type:
|
||||
pos = round(height*progress)
|
||||
frame[:pos, :] = 1.0
|
||||
elif "box" in transition_type:
|
||||
box_w = round(width*progress)
|
||||
box_h = round(height*progress)
|
||||
x1 = (width - box_w) // 2
|
||||
y1 = (height - box_h) // 2
|
||||
x2 = x1 + box_w
|
||||
y2 = y1 + box_h
|
||||
frame[y1:y2, x1:x2] = 1.0
|
||||
elif "circle" in transition_type:
|
||||
radius = math.ceil(math.sqrt(pow(width,2)+pow(height,2))*progress/2)
|
||||
c_x = width // 2
|
||||
c_y = height // 2
|
||||
# is this real life? Am I hallucinating?
|
||||
x = torch.arange(0, width, dtype=torch.float32, device="cpu")
|
||||
y = torch.arange(0, height, dtype=torch.float32, device="cpu")
|
||||
y, x = torch.meshgrid((y, x), indexing="ij")
|
||||
circle = ((x - c_x) ** 2 + (y - c_y) ** 2) <= (radius ** 2)
|
||||
frame[circle] = 1.0
|
||||
elif "horizontal bar" in transition_type:
|
||||
bar = round(height*progress)
|
||||
y1 = (height - bar) // 2
|
||||
y2 = y1 + bar
|
||||
frame[y1:y2, :] = 1.0
|
||||
elif "vertical bar" in transition_type:
|
||||
bar = round(width*progress)
|
||||
x1 = (width - bar) // 2
|
||||
x2 = x1 + bar
|
||||
frame[:, x1:x2] = 1.0
|
||||
elif "horizontal door" in transition_type:
|
||||
bar = math.ceil(height*progress/2)
|
||||
if bar > 0:
|
||||
frame[:bar, :] = 1.0
|
||||
frame[-bar:, :] = 1.0
|
||||
elif "vertical door" in transition_type:
|
||||
bar = math.ceil(width*progress/2)
|
||||
if bar > 0:
|
||||
frame[:, :bar] = 1.0
|
||||
frame[:, -bar:] = 1.0
|
||||
elif "fade" in transition_type:
|
||||
frame[:,:] = progress
|
||||
|
||||
out.append(frame)
|
||||
|
||||
if end_frame < frames:
|
||||
out = out + [torch.full((height, width), 1.0, dtype=torch.float32, device="cpu")] * (frames - end_frame)
|
||||
|
||||
out = torch.stack(out, dim=0)
|
||||
|
||||
return (out, )
|
||||
|
||||
MASK_CLASS_MAPPINGS = {
|
||||
"MaskBlur+": MaskBlur,
|
||||
"MaskBoundingBox+": MaskBoundingBox,
|
||||
"MaskFix+": MaskFix,
|
||||
"MaskFlip+": MaskFlip,
|
||||
"MaskFromColor+": MaskFromColor,
|
||||
"MaskFromList+": MaskFromList,
|
||||
"MaskFromRGBCMYBW+": MaskFromRGBCMYBW,
|
||||
"MaskFromSegmentation+": MaskFromSegmentation,
|
||||
"MaskPreview+": MaskPreview,
|
||||
"MaskSmooth+": MaskSmooth,
|
||||
"TransitionMask+": TransitionMask,
|
||||
|
||||
# Batch
|
||||
"MaskBatch+": MaskBatch,
|
||||
"MaskExpandBatch+": MaskExpandBatch,
|
||||
"MaskFromBatch+": MaskFromBatch,
|
||||
}
|
||||
|
||||
MASK_NAME_MAPPINGS = {
|
||||
"MaskBlur+": "🔧 Mask Blur",
|
||||
"MaskFix+": "🔧 Mask Fix",
|
||||
"MaskFlip+": "🔧 Mask Flip",
|
||||
"MaskFromColor+": "🔧 Mask From Color",
|
||||
"MaskFromList+": "🔧 Mask From List",
|
||||
"MaskFromRGBCMYBW+": "🔧 Mask From RGB/CMY/BW",
|
||||
"MaskFromSegmentation+": "🔧 Mask From Segmentation",
|
||||
"MaskPreview+": "🔧 Mask Preview",
|
||||
"MaskBoundingBox+": "🔧 Mask Bounding Box",
|
||||
"MaskSmooth+": "🔧 Mask Smooth",
|
||||
"TransitionMask+": "🔧 Transition Mask",
|
||||
|
||||
"MaskBatch+": "🔧 Mask Batch",
|
||||
"MaskExpandBatch+": "🔧 Mask Expand Batch",
|
||||
"MaskFromBatch+": "🔧 Mask From Batch",
|
||||
}
|
||||
@@ -0,0 +1,240 @@
|
||||
import math
|
||||
import torch
|
||||
from .utils import AnyType
|
||||
import comfy.model_management
|
||||
|
||||
any = AnyType("*")
|
||||
|
||||
class SimpleMath:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"optional": {
|
||||
"a": ("INT,FLOAT", { "default": 0.0, "step": 0.1 }),
|
||||
"b": ("INT,FLOAT", { "default": 0.0, "step": 0.1 }),
|
||||
},
|
||||
"required": {
|
||||
"value": ("STRING", { "multiline": False, "default": "" }),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT", "FLOAT", )
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/utilities"
|
||||
|
||||
def execute(self, value, a = 0.0, b = 0.0):
|
||||
import ast
|
||||
import operator as op
|
||||
|
||||
operators = {
|
||||
ast.Add: op.add,
|
||||
ast.Sub: op.sub,
|
||||
ast.Mult: op.mul,
|
||||
ast.Div: op.truediv,
|
||||
ast.FloorDiv: op.floordiv,
|
||||
ast.Pow: op.pow,
|
||||
ast.BitXor: op.xor,
|
||||
ast.USub: op.neg,
|
||||
ast.Mod: op.mod,
|
||||
}
|
||||
|
||||
op_functions = {
|
||||
'min': min,
|
||||
'max': max,
|
||||
'round': round,
|
||||
'sum': sum,
|
||||
'len': len,
|
||||
}
|
||||
|
||||
def eval_(node):
|
||||
if isinstance(node, ast.Num): # number
|
||||
return node.n
|
||||
elif isinstance(node, ast.Name): # variable
|
||||
if node.id == "a":
|
||||
return a
|
||||
if node.id == "b":
|
||||
return b
|
||||
elif isinstance(node, ast.BinOp): # <left> <operator> <right>
|
||||
return operators[type(node.op)](eval_(node.left), eval_(node.right))
|
||||
elif isinstance(node, ast.UnaryOp): # <operator> <operand> e.g., -1
|
||||
return operators[type(node.op)](eval_(node.operand))
|
||||
elif isinstance(node, ast.Call): # custom function
|
||||
if node.func.id in op_functions:
|
||||
args =[eval_(arg) for arg in node.args]
|
||||
return op_functions[node.func.id](*args)
|
||||
elif isinstance(node, ast.Subscript): # indexing or slicing
|
||||
value = eval_(node.value)
|
||||
if isinstance(node.slice, ast.Constant):
|
||||
return value[node.slice.value]
|
||||
else:
|
||||
return 0
|
||||
else:
|
||||
return 0
|
||||
|
||||
result = eval_(ast.parse(value, mode='eval').body)
|
||||
|
||||
if math.isnan(result):
|
||||
result = 0.0
|
||||
|
||||
return (round(result), result, )
|
||||
|
||||
class ConsoleDebug:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"value": (any, {}),
|
||||
},
|
||||
"optional": {
|
||||
"prefix": ("STRING", { "multiline": False, "default": "Value:" })
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/utilities"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def execute(self, value, prefix):
|
||||
print(f"\033[96m{prefix} {value}\033[0m")
|
||||
|
||||
return (None,)
|
||||
|
||||
class DebugTensorShape:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"tensor": (any, {}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/utilities"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def execute(self, tensor):
|
||||
shapes = []
|
||||
def tensorShape(tensor):
|
||||
if isinstance(tensor, dict):
|
||||
for k in tensor:
|
||||
tensorShape(tensor[k])
|
||||
elif isinstance(tensor, list):
|
||||
for i in range(len(tensor)):
|
||||
tensorShape(tensor[i])
|
||||
elif hasattr(tensor, 'shape'):
|
||||
shapes.append(list(tensor.shape))
|
||||
|
||||
tensorShape(tensor)
|
||||
|
||||
print(f"\033[96mShapes found: {shapes}\033[0m")
|
||||
|
||||
return (None,)
|
||||
|
||||
class BatchCount:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"batch": (any, {}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("INT",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/utilities"
|
||||
|
||||
def execute(self, batch):
|
||||
count = 0
|
||||
if hasattr(batch, 'shape'):
|
||||
count = batch.shape[0]
|
||||
elif isinstance(batch, dict) and 'samples' in batch:
|
||||
count = batch['samples'].shape[0]
|
||||
elif isinstance(batch, list) or isinstance(batch, dict):
|
||||
count = len(batch)
|
||||
|
||||
return (count, )
|
||||
|
||||
class ModelCompile():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL",),
|
||||
"fullgraph": ("BOOLEAN", { "default": False }),
|
||||
"dynamic": ("BOOLEAN", { "default": False }),
|
||||
"mode": (["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"],),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", )
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/utilities"
|
||||
|
||||
def execute(self, model, fullgraph, dynamic, mode):
|
||||
work_model = model.clone()
|
||||
torch._dynamo.config.suppress_errors = True
|
||||
work_model.model.diffusion_model = torch.compile(work_model.model.diffusion_model, dynamic=dynamic, fullgraph=fullgraph, mode=mode)
|
||||
return (work_model, )
|
||||
|
||||
class RemoveLatentMask:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "samples": ("LATENT",),}}
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "execute"
|
||||
|
||||
CATEGORY = "essentials/utilities"
|
||||
|
||||
def execute(self, samples):
|
||||
s = samples.copy()
|
||||
if "noise_mask" in s:
|
||||
del s["noise_mask"]
|
||||
|
||||
return (s,)
|
||||
|
||||
class SDXLEmptyLatentSizePicker:
|
||||
def __init__(self):
|
||||
self.device = comfy.model_management.intermediate_device()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"resolution": (["704x1408 (0.5)","704x1344 (0.52)","768x1344 (0.57)","768x1280 (0.6)","832x1216 (0.68)","832x1152 (0.72)","896x1152 (0.78)","896x1088 (0.82)","960x1088 (0.88)","960x1024 (0.94)","1024x1024 (1.0)","1024x960 (1.07)","1088x960 (1.13)","1088x896 (1.21)","1152x896 (1.29)","1152x832 (1.38)","1216x832 (1.46)","1280x768 (1.67)","1344x768 (1.75)","1344x704 (1.91)","1408x704 (2.0)","1472x704 (2.09)","1536x640 (2.4)","1600x640 (2.5)","1664x576 (2.89)","1728x576 (3.0)",], {"default": "1024x1024 (1.0)"}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("LATENT","INT","INT",)
|
||||
RETURN_NAMES = ("LATENT","width", "height",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/utilities"
|
||||
|
||||
def execute(self, resolution, batch_size):
|
||||
width, height = resolution.split(" ")[0].split("x")
|
||||
width = int(width)
|
||||
height = int(height)
|
||||
|
||||
latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
|
||||
|
||||
return ({"samples":latent}, width, height,)
|
||||
|
||||
MISC_CLASS_MAPPINGS = {
|
||||
"BatchCount+": BatchCount,
|
||||
"ConsoleDebug+": ConsoleDebug,
|
||||
"DebugTensorShape+": DebugTensorShape,
|
||||
#"ModelCompile+": ModelCompile,
|
||||
"RemoveLatentMask+": RemoveLatentMask,
|
||||
"SDXLEmptyLatentSizePicker+": SDXLEmptyLatentSizePicker,
|
||||
"SimpleMath+": SimpleMath,
|
||||
}
|
||||
|
||||
MISC_NAME_MAPPINGS = {
|
||||
"BatchCount+": "🔧 Batch Count",
|
||||
"ConsoleDebug+": "🔧 Console Debug",
|
||||
"DebugTensorShape+": "🔧 Debug Tensor Shape",
|
||||
#"ModelCompile+": "🔧 Model Compile",
|
||||
"RemoveLatentMask+": "🔧 Remove Latent Mask",
|
||||
"SDXLEmptyLatentSizePicker+": "🔧 SDXL Empty Latent Size Picker",
|
||||
"SimpleMath+": "🔧 Simple Math",
|
||||
}
|
||||
+159
@@ -0,0 +1,159 @@
|
||||
import comfy.samplers
|
||||
import comfy.sample
|
||||
import torch
|
||||
from nodes import common_ksampler
|
||||
from .utils import expand_mask
|
||||
|
||||
class KSamplerVariationsWithNoise:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"model": ("MODEL", ),
|
||||
"latent_image": ("LATENT", ),
|
||||
"main_seed": ("INT:seed", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
|
||||
"cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
|
||||
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
|
||||
"positive": ("CONDITIONING", ),
|
||||
"negative": ("CONDITIONING", ),
|
||||
"variation_strength": ("FLOAT", {"default": 0.17, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}),
|
||||
#"start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
|
||||
#"end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
|
||||
#"return_with_leftover_noise": (["disable", "enable"], ),
|
||||
"variation_seed": ("INT:seed", {"default": 12345, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.01, "round": 0.01}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/sampling"
|
||||
|
||||
# From https://github.com/BlenderNeko/ComfyUI_Noise/
|
||||
def slerp(self, val, low, high):
|
||||
dims = low.shape
|
||||
|
||||
low = low.reshape(dims[0], -1)
|
||||
high = high.reshape(dims[0], -1)
|
||||
|
||||
low_norm = low/torch.norm(low, dim=1, keepdim=True)
|
||||
high_norm = high/torch.norm(high, dim=1, keepdim=True)
|
||||
|
||||
low_norm[low_norm != low_norm] = 0.0
|
||||
high_norm[high_norm != high_norm] = 0.0
|
||||
|
||||
omega = torch.acos((low_norm*high_norm).sum(1))
|
||||
so = torch.sin(omega)
|
||||
res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
|
||||
|
||||
return res.reshape(dims)
|
||||
|
||||
def prepare_mask(self, mask, shape):
|
||||
mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
|
||||
mask = mask.expand((-1,shape[1],-1,-1))
|
||||
if mask.shape[0] < shape[0]:
|
||||
mask = mask.repeat((shape[0] -1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]]
|
||||
return mask
|
||||
|
||||
def execute(self, model, latent_image, main_seed, steps, cfg, sampler_name, scheduler, positive, negative, variation_strength, variation_seed, denoise):
|
||||
if main_seed == variation_seed:
|
||||
variation_seed += 1
|
||||
|
||||
end_at_step = steps #min(steps, end_at_step)
|
||||
start_at_step = round(end_at_step - end_at_step * denoise)
|
||||
|
||||
force_full_denoise = True
|
||||
disable_noise = True
|
||||
|
||||
device = comfy.model_management.get_torch_device()
|
||||
|
||||
# Generate base noise
|
||||
batch_size, _, height, width = latent_image["samples"].shape
|
||||
generator = torch.manual_seed(main_seed)
|
||||
base_noise = torch.randn((1, 4, height, width), dtype=torch.float32, device="cpu", generator=generator).repeat(batch_size, 1, 1, 1).cpu()
|
||||
|
||||
# Generate variation noise
|
||||
generator = torch.manual_seed(variation_seed)
|
||||
variation_noise = torch.randn((batch_size, 4, height, width), dtype=torch.float32, device="cpu", generator=generator).cpu()
|
||||
|
||||
slerp_noise = self.slerp(variation_strength, base_noise, variation_noise)
|
||||
|
||||
# Calculate sigma
|
||||
comfy.model_management.load_model_gpu(model)
|
||||
sampler = comfy.samplers.KSampler(model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler, denoise=1.0, model_options=model.model_options)
|
||||
sigmas = sampler.sigmas
|
||||
sigma = sigmas[start_at_step] - sigmas[end_at_step]
|
||||
sigma /= model.model.latent_format.scale_factor
|
||||
sigma = sigma.detach().cpu().item()
|
||||
|
||||
work_latent = latent_image.copy()
|
||||
work_latent["samples"] = latent_image["samples"].clone() + slerp_noise * sigma
|
||||
|
||||
# if there's a mask we need to expand it to avoid artifacts, 5 pixels should be enough
|
||||
if "noise_mask" in latent_image:
|
||||
noise_mask = self.prepare_mask(latent_image["noise_mask"], latent_image['samples'].shape)
|
||||
work_latent["samples"] = noise_mask * work_latent["samples"] + (1-noise_mask) * latent_image["samples"]
|
||||
work_latent['noise_mask'] = expand_mask(latent_image["noise_mask"].clone(), 5, True)
|
||||
|
||||
return common_ksampler(model, main_seed, steps, cfg, sampler_name, scheduler, positive, negative, work_latent, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
|
||||
|
||||
|
||||
class KSamplerVariationsStochastic:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required":{
|
||||
"model": ("MODEL",),
|
||||
"latent_image": ("LATENT", ),
|
||||
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"steps": ("INT", {"default": 25, "min": 1, "max": 10000}),
|
||||
"cfg": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
|
||||
"sampler": (comfy.samplers.KSampler.SAMPLERS, ),
|
||||
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
|
||||
"positive": ("CONDITIONING", ),
|
||||
"negative": ("CONDITIONING", ),
|
||||
"variation_seed": ("INT:seed", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"variation_strength": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, "step":0.05, "round": 0.01}),
|
||||
#"variation_sampler": (comfy.samplers.KSampler.SAMPLERS, ),
|
||||
"cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.05, "round": 0.01}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("LATENT", )
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/sampling"
|
||||
|
||||
def execute(self, model, latent_image, noise_seed, steps, cfg, sampler, scheduler, positive, negative, variation_seed, variation_strength, cfg_scale, variation_sampler="dpmpp_2m_sde"):
|
||||
# Stage 1: composition sampler
|
||||
force_full_denoise = False # return with leftover noise = "enable"
|
||||
disable_noise = False # add noise = "enable"
|
||||
|
||||
end_at_step = max(int(steps * (1-variation_strength)), 1)
|
||||
start_at_step = 0
|
||||
|
||||
work_latent = latent_image.copy()
|
||||
batch_size = work_latent["samples"].shape[0]
|
||||
work_latent["samples"] = work_latent["samples"][0].unsqueeze(0)
|
||||
|
||||
stage1 = common_ksampler(model, noise_seed, steps, cfg, sampler, scheduler, positive, negative, work_latent, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)[0]
|
||||
|
||||
if batch_size > 1:
|
||||
stage1["samples"] = stage1["samples"].clone().repeat(batch_size, 1, 1, 1)
|
||||
|
||||
# Stage 2: variation sampler
|
||||
force_full_denoise = True
|
||||
disable_noise = True
|
||||
cfg = max(cfg * cfg_scale, 1.0)
|
||||
start_at_step = end_at_step
|
||||
end_at_step = steps
|
||||
|
||||
return common_ksampler(model, variation_seed, steps, cfg, variation_sampler, scheduler, positive, negative, stage1, denoise=1.0, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
|
||||
|
||||
|
||||
SAMPLING_CLASS_MAPPINGS = {
|
||||
"KSamplerVariationsStochastic+": KSamplerVariationsStochastic,
|
||||
"KSamplerVariationsWithNoise+": KSamplerVariationsWithNoise,
|
||||
}
|
||||
|
||||
SAMPLING_NAME_MAPPINGS = {
|
||||
"KSamplerVariationsStochastic+": "🔧 KSampler Stochastic Variations",
|
||||
"KSamplerVariationsWithNoise+": "🔧 KSampler Variations with Noise Injection",
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
import torch
|
||||
import torchvision.transforms.v2 as T
|
||||
import torch.nn.functional as F
|
||||
from .utils import expand_mask
|
||||
|
||||
class LoadCLIPSegModels:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP_SEG",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/segmentation"
|
||||
|
||||
def execute(self):
|
||||
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
|
||||
processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
|
||||
model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
|
||||
|
||||
return ((processor, model),)
|
||||
|
||||
class ApplyCLIPSeg:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"clip_seg": ("CLIP_SEG",),
|
||||
"image": ("IMAGE",),
|
||||
"prompt": ("STRING", { "multiline": False, "default": "" }),
|
||||
"threshold": ("FLOAT", { "default": 0.4, "min": 0.0, "max": 1.0, "step": 0.05 }),
|
||||
"smooth": ("INT", { "default": 9, "min": 0, "max": 32, "step": 1 }),
|
||||
"dilate": ("INT", { "default": 0, "min": -32, "max": 32, "step": 1 }),
|
||||
"blur": ("INT", { "default": 0, "min": 0, "max": 64, "step": 1 }),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/segmentation"
|
||||
|
||||
def execute(self, image, clip_seg, prompt, threshold, smooth, dilate, blur):
|
||||
processor, model = clip_seg
|
||||
|
||||
imagenp = image.mul(255).clamp(0, 255).byte().cpu().numpy()
|
||||
|
||||
outputs = []
|
||||
for i in imagenp:
|
||||
inputs = processor(text=prompt, images=[i], return_tensors="pt")
|
||||
out = model(**inputs)
|
||||
out = out.logits.unsqueeze(1)
|
||||
out = torch.sigmoid(out[0][0])
|
||||
out = (out > threshold)
|
||||
outputs.append(out)
|
||||
|
||||
del imagenp
|
||||
|
||||
outputs = torch.stack(outputs, dim=0)
|
||||
|
||||
if smooth > 0:
|
||||
if smooth % 2 == 0:
|
||||
smooth += 1
|
||||
outputs = T.functional.gaussian_blur(outputs, smooth)
|
||||
|
||||
outputs = outputs.float()
|
||||
|
||||
if dilate != 0:
|
||||
outputs = expand_mask(outputs, dilate, True)
|
||||
|
||||
if blur > 0:
|
||||
if blur % 2 == 0:
|
||||
blur += 1
|
||||
outputs = T.functional.gaussian_blur(outputs, blur)
|
||||
|
||||
# resize to original size
|
||||
outputs = F.interpolate(outputs.unsqueeze(1), size=(image.shape[1], image.shape[2]), mode='bicubic').squeeze(1)
|
||||
|
||||
return (outputs,)
|
||||
|
||||
SEG_CLASS_MAPPINGS = {
|
||||
"ApplyCLIPSeg+": ApplyCLIPSeg,
|
||||
"LoadCLIPSegModels+": LoadCLIPSegModels,
|
||||
}
|
||||
|
||||
SEG_NAME_MAPPINGS = {
|
||||
"ApplyCLIPSeg+": "🔧 Apply CLIPSeg",
|
||||
"LoadCLIPSegModels+": "🔧 Load CLIPSeg Models",
|
||||
}
|
||||
@@ -0,0 +1,110 @@
|
||||
import os
|
||||
import torch
|
||||
from nodes import MAX_RESOLUTION
|
||||
import torchvision.transforms.v2 as T
|
||||
|
||||
FONTS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "fonts")
|
||||
class DrawText:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"text": ("STRING", { "multiline": True, "dynamicPrompts": True, "default": "Hello, World!" }),
|
||||
"font": ([f for f in os.listdir(FONTS_DIR) if f.endswith('.ttf') or f.endswith('.otf')], ),
|
||||
"size": ("INT", { "default": 56, "min": 1, "max": 9999, "step": 1 }),
|
||||
"color": ("STRING", { "multiline": False, "default": "#FFFFFF" }),
|
||||
"background_color": ("STRING", { "multiline": False, "default": "#00000000" }),
|
||||
"shadow_distance": ("INT", { "default": 0, "min": 0, "max": 100, "step": 1 }),
|
||||
"shadow_blur": ("INT", { "default": 0, "min": 0, "max": 100, "step": 1 }),
|
||||
"shadow_color": ("STRING", { "multiline": False, "default": "#000000" }),
|
||||
"horizontal_align": (["left", "center", "right"],),
|
||||
"vertical_align": (["top", "center", "bottom"],),
|
||||
"offset_x": ("INT", { "default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1 }),
|
||||
"offset_y": ("INT", { "default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1 }),
|
||||
},
|
||||
"optional": {
|
||||
"img_composite": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "MASK",)
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "essentials/text"
|
||||
|
||||
def execute(self, text, font, size, color, background_color, shadow_distance, shadow_blur, shadow_color, horizontal_align, vertical_align, offset_x, offset_y, img_composite=None):
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageColor, ImageFilter
|
||||
|
||||
font = ImageFont.truetype(os.path.join(FONTS_DIR, font), size)
|
||||
|
||||
lines = text.split("\n")
|
||||
|
||||
# Calculate the width and height of the text
|
||||
text_width = max(font.getbbox(line)[2] for line in lines)
|
||||
line_height = font.getmask(text).getbbox()[3] + font.getmetrics()[1] # add descent to height
|
||||
text_height = line_height * len(lines)
|
||||
|
||||
if img_composite is not None:
|
||||
img_composite = T.ToPILImage()(img_composite.permute([0,3,1,2])[0]).convert('RGBA')
|
||||
width = img_composite.width
|
||||
height = img_composite.height
|
||||
image = Image.new('RGBA', (width, height), color=background_color)
|
||||
else:
|
||||
width = text_width
|
||||
height = text_height
|
||||
background_color = ImageColor.getrgb(background_color)
|
||||
image = Image.new('RGBA', (width + shadow_distance, height + shadow_distance), color=background_color)
|
||||
|
||||
image_shadow = None
|
||||
if shadow_distance > 0:
|
||||
image_shadow = image.copy()
|
||||
#image_shadow = Image.new('RGBA', (width + shadow_distance, height + shadow_distance), color=background_color)
|
||||
|
||||
for i, line in enumerate(lines):
|
||||
line_width = font.getbbox(line)[2]
|
||||
#text_height =font.getbbox(line)[3]
|
||||
if horizontal_align == "left":
|
||||
x = 0
|
||||
elif horizontal_align == "center":
|
||||
x = (width - line_width) / 2
|
||||
elif horizontal_align == "right":
|
||||
x = width - line_width
|
||||
|
||||
if vertical_align == "top":
|
||||
y = 0
|
||||
elif vertical_align == "center":
|
||||
y = (height - text_height) / 2
|
||||
elif vertical_align == "bottom":
|
||||
y = height - text_height
|
||||
|
||||
x += offset_x
|
||||
y += i * line_height + offset_y
|
||||
|
||||
draw = ImageDraw.Draw(image)
|
||||
draw.text((x, y), line, font=font, fill=color)
|
||||
|
||||
if image_shadow is not None:
|
||||
draw = ImageDraw.Draw(image_shadow)
|
||||
draw.text((x + shadow_distance, y + shadow_distance), line, font=font, fill=shadow_color)
|
||||
|
||||
if image_shadow is not None:
|
||||
image_shadow = image_shadow.filter(ImageFilter.GaussianBlur(shadow_blur))
|
||||
image = Image.alpha_composite(image_shadow, image)
|
||||
|
||||
#image = T.ToTensor()(image).unsqueeze(0).permute([0,2,3,1])
|
||||
mask = T.ToTensor()(image).unsqueeze(0).permute([0,2,3,1])
|
||||
mask = mask[:, :, :, 3] if mask.shape[3] == 4 else torch.ones_like(mask[:, :, :, 0])
|
||||
|
||||
if img_composite is not None:
|
||||
image = Image.alpha_composite(img_composite, image)
|
||||
|
||||
image = T.ToTensor()(image).unsqueeze(0).permute([0,2,3,1])
|
||||
|
||||
return (image[:, :, :, :3], mask,)
|
||||
|
||||
TEXT_CLASS_MAPPINGS = {
|
||||
"DrawText+": DrawText,
|
||||
}
|
||||
|
||||
TEXT_NAME_MAPPINGS = {
|
||||
"DrawText+": "🔧 Draw Text",
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import scipy
|
||||
|
||||
# from https://github.com/pythongosssss/ComfyUI-Custom-Scripts
|
||||
class AnyType(str):
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
def min_(tensor_list):
|
||||
# return the element-wise min of the tensor list.
|
||||
x = torch.stack(tensor_list)
|
||||
mn = x.min(axis=0)[0]
|
||||
return torch.clamp(mn, min=0)
|
||||
|
||||
def max_(tensor_list):
|
||||
# return the element-wise max of the tensor list.
|
||||
x = torch.stack(tensor_list)
|
||||
mx = x.max(axis=0)[0]
|
||||
return torch.clamp(mx, max=1)
|
||||
|
||||
def expand_mask(mask, expand, tapered_corners):
|
||||
c = 0 if tapered_corners else 1
|
||||
kernel = np.array([[c, 1, c],
|
||||
[1, 1, 1],
|
||||
[c, 1, c]])
|
||||
mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
|
||||
out = []
|
||||
for m in mask:
|
||||
output = m.numpy()
|
||||
for _ in range(abs(expand)):
|
||||
if expand < 0:
|
||||
output = scipy.ndimage.grey_erosion(output, footprint=kernel)
|
||||
else:
|
||||
output = scipy.ndimage.grey_dilation(output, footprint=kernel)
|
||||
output = torch.from_numpy(output)
|
||||
out.append(output)
|
||||
|
||||
return torch.stack(out, dim=0)
|
||||
Reference in New Issue
Block a user