From 28ff3676b67c88460b42045f98994ca43701c51b Mon Sep 17 00:00:00 2001 From: Kijai <40791699+kijai@users.noreply.github.com> Date: Fri, 3 May 2024 11:27:01 +0300 Subject: [PATCH] Continue restructuring --- __init__.py | 295 ++++------- nodes/mask_nodes.py | 1165 +++++++++++++++++++++++++++++++++++++++++++ nodes/nodes.py | 1156 +----------------------------------------- 3 files changed, 1273 insertions(+), 1343 deletions(-) create mode 100644 nodes/mask_nodes.py diff --git a/__init__.py b/__init__.py index dac9113..adcf11e 100644 --- a/__init__.py +++ b/__init__.py @@ -4,207 +4,122 @@ from .nodes.batchcrop_nodes import * from .nodes.audioscheduler_nodes import * from .nodes.image_nodes import * from .nodes.intrinsic_lora_nodes import * -NODE_CLASS_MAPPINGS = { +from .nodes.mask_nodes import * +NODE_CONFIG = { #constants - "INTConstant": INTConstant, - "FloatConstant": FloatConstant, - "StringConstant": StringConstant, - "StringConstantMultiline": StringConstantMultiline, + "INTConstant": {"class": INTConstant, "name": "INT Constant"}, + "FloatConstant": {"class": FloatConstant, "name": "Float Constant"}, + "StringConstant": {"class": StringConstant, "name": "String Constant"}, + "StringConstantMultiline": {"class": StringConstantMultiline, "name": "String Constant Multiline"}, #conditioning - "ConditioningMultiCombine": ConditioningMultiCombine, - "ConditioningSetMaskAndCombine": ConditioningSetMaskAndCombine, - "ConditioningSetMaskAndCombine3": ConditioningSetMaskAndCombine3, - "ConditioningSetMaskAndCombine4": ConditioningSetMaskAndCombine4, - "ConditioningSetMaskAndCombine5": ConditioningSetMaskAndCombine5, - "CondPassThrough": CondPassThrough, + "ConditioningMultiCombine": {"class": ConditioningMultiCombine, "name": "Conditioning Multi Combine"}, + "ConditioningSetMaskAndCombine": {"class": ConditioningSetMaskAndCombine, "name": "ConditioningSetMaskAndCombine"}, + "ConditioningSetMaskAndCombine3": {"class": ConditioningSetMaskAndCombine3, "name": "ConditioningSetMaskAndCombine3"}, + "ConditioningSetMaskAndCombine4": {"class": ConditioningSetMaskAndCombine4, "name": "ConditioningSetMaskAndCombine4"}, + "ConditioningSetMaskAndCombine5": {"class": ConditioningSetMaskAndCombine5, "name": "ConditioningSetMaskAndCombine5"}, + "CondPassThrough": {"class": CondPassThrough}, #masking - "BatchCLIPSeg": BatchCLIPSeg, - "RoundMask": RoundMask, - "ResizeMask": ResizeMask, - "OffsetMask": OffsetMask, - "MaskBatchMulti": MaskBatchMulti, - "GrowMaskWithBlur": GrowMaskWithBlur, - "ColorToMask": ColorToMask, - "CreateGradientMask": CreateGradientMask, - "CreateTextMask": CreateTextMask, - "CreateAudioMask": CreateAudioMask, - "CreateFadeMask": CreateFadeMask, - "CreateFadeMaskAdvanced": CreateFadeMaskAdvanced, - "CreateFluidMask" :CreateFluidMask, - "CreateShapeMask": CreateShapeMask, - "CreateVoronoiMask": CreateVoronoiMask, - "CreateMagicMask": CreateMagicMask, - "RemapMaskRange": RemapMaskRange, - "GetMaskSize": GetMaskSize, + "BatchCLIPSeg": {"class": BatchCLIPSeg, "name": "Batch CLIPSeg"}, + "ColorToMask": {"class": ColorToMask, "name": "Color To Mask"}, + "CreateGradientMask": {"class": CreateGradientMask, "name": "Create Gradient Mask"}, + "CreateTextMask": {"class": CreateTextMask, "name": "Create Text Mask"}, + "CreateAudioMask": {"class": CreateAudioMask, "name": "Create Audio Mask"}, + "CreateFadeMask": {"class": CreateFadeMask, "name": "Create Fade Mask"}, + "CreateFadeMaskAdvanced": {"class": CreateFadeMaskAdvanced, "name": "Create Fade Mask Advanced"}, + "CreateFluidMask": {"class": CreateFluidMask, "name": "Create Fluid Mask"}, + "CreateShapeMask": {"class": CreateShapeMask, "name": "Create Shape Mask"}, + "CreateVoronoiMask": {"class": CreateVoronoiMask, "name": "Create Voronoi Mask"}, + "CreateMagicMask": {"class": CreateMagicMask, "name": "Create Magic Mask"}, + "GetMaskSize": {"class": GetMaskSize, "name": "Get Mask Size"}, + "GrowMaskWithBlur": {"class": GrowMaskWithBlur, "name": "Grow Mask With Blur"}, + "MaskBatchMulti": {"class": MaskBatchMulti, "name": "Mask Batch Multi"}, + "OffsetMask": {"class": OffsetMask, "name": "Offset Mask"}, + "RemapMaskRange": {"class": RemapMaskRange, "name": "Remap Mask Range"}, + "ResizeMask": {"class": ResizeMask, "name": "Resize Mask"}, + "RoundMask": {"class": RoundMask, "name": "Round Mask"}, #images - "ImageBatchMulti": ImageBatchMulti, - "ColorMatch": ColorMatch, - "CrossFadeImages": CrossFadeImages, - "GetImageRangeFromBatch": GetImageRangeFromBatch, - "SaveImageWithAlpha": SaveImageWithAlpha, - "ReverseImageBatch": ReverseImageBatch, - "ImageGridComposite2x2": ImageGridComposite2x2, - "ImageGridComposite3x3": ImageGridComposite3x3, - "ImageConcanate": ImageConcanate, - "ImageBatchTestPattern": ImageBatchTestPattern, - "ReplaceImagesInBatch": ReplaceImagesInBatch, - "ImageGrabPIL": ImageGrabPIL, - "AddLabel": AddLabel, - "ImageUpscaleWithModelBatched": ImageUpscaleWithModelBatched, - "GetImagesFromBatchIndexed": GetImagesFromBatchIndexed, - "InsertImagesToBatchIndexed": InsertImagesToBatchIndexed, - "ImageBatchRepeatInterleaving": ImageBatchRepeatInterleaving, - "ImageNormalize_Neg1_To_1": ImageNormalize_Neg1_To_1, - "RemapImageRange": RemapImageRange, - "ImagePass": ImagePass, - "ImagePadForOutpaintMasked": ImagePadForOutpaintMasked, - "ImageAndMaskPreview": ImageAndMaskPreview, - "SplitImageChannels": SplitImageChannels, - "MergeImageChannels": MergeImageChannels, + "ColorMatch": {"class": ColorMatch, "name": "Color Match"}, + "AddLabel": {"class": AddLabel, "name": "Add Label"}, + "ImageBatchMulti": {"class": ImageBatchMulti, "name": "Image Batch Multi"}, + "ImageBatchTestPattern": {"class": ImageBatchTestPattern, "name": "Image Batch Test Pattern"}, + "ImageConcanate": {"class": ImageConcanate, "name": "Image Concanate"}, + "ImageGrabPIL": {"class": ImageGrabPIL, "name": "Image Grab PIL"}, + "ImageGridComposite2x2": {"class": ImageGridComposite2x2, "name": "Image Grid Composite 2x2"}, + "ImageGridComposite3x3": {"class": ImageGridComposite3x3, "name": "Image Grid Composite 3x3"}, + "ImageNormalize_Neg1_To_1": {"class": ImageNormalize_Neg1_To_1, "name": "Image Normalize -1 to 1"}, + "ImagePass": {"class": ImagePass}, + "ImagePadForOutpaintMasked": {"class": ImagePadForOutpaintMasked, "name": "Image Pad For Outpaint Masked"}, + "ImageUpscaleWithModelBatched": {"class": ImageUpscaleWithModelBatched, "name": "Image Upscale With Model Batched"}, + "InsertImagesToBatchIndexed": {"class": InsertImagesToBatchIndexed, "name": "Insert Images To Batch Indexed"}, + "MergeImageChannels": {"class": MergeImageChannels, "name": "Merge Image Channels"}, + "ReverseImageBatch": {"class": ReverseImageBatch, "name": "Reverse Image Batch"}, + "RemapImageRange": {"class": RemapImageRange, "name": "Remap Image Range"}, + "SaveImageWithAlpha": {"class": SaveImageWithAlpha, "name": "Save Image With Alpha"}, + "SplitImageChannels": {"class": SplitImageChannels, "name": "Split Image Channels"}, + "CrossFadeImages": {"class": CrossFadeImages, "name": "Cross Fade Images"}, + "GetImageRangeFromBatch": {"class": GetImageRangeFromBatch, "name": "Get Image Range From Batch"}, #batch cropping - "BatchCropFromMask": BatchCropFromMask, - "BatchCropFromMaskAdvanced": BatchCropFromMaskAdvanced, - "FilterZeroMasksAndCorrespondingImages": FilterZeroMasksAndCorrespondingImages, - "InsertImageBatchByIndexes": InsertImageBatchByIndexes, - "BatchUncrop": BatchUncrop, - "BatchUncropAdvanced": BatchUncropAdvanced, - "SplitBboxes": SplitBboxes, - "BboxToInt": BboxToInt, - "BboxVisualize": BboxVisualize, + "BatchCropFromMask": {"class": BatchCropFromMask, "name": "Batch Crop From Mask"}, + "BatchCropFromMaskAdvanced": {"class": BatchCropFromMaskAdvanced, "name": "Batch Crop From Mask Advanced"}, + "FilterZeroMasksAndCorrespondingImages": {"class": FilterZeroMasksAndCorrespondingImages}, + "InsertImageBatchByIndexes": {"class": InsertImageBatchByIndexes, "name": "Insert Image Batch By Indexes"}, + "BatchUncrop": {"class": BatchUncrop, "name": "Batch Uncrop"}, + "BatchUncropAdvanced": {"class": BatchUncropAdvanced, "name": "Batch Uncrop Advanced"}, + "SplitBboxes": {"class": SplitBboxes, "name": "Split Bboxes"}, + "BboxToInt": {"class": BboxToInt, "name": "Bbox To Int"}, + "BboxVisualize": {"class": BboxVisualize, "name": "Bbox Visualize"}, #noise - "GenerateNoise": GenerateNoise, - "FlipSigmasAdjusted": FlipSigmasAdjusted, - "InjectNoiseToLatent": InjectNoiseToLatent, - "CustomSigmas": CustomSigmas, + "GenerateNoise": {"class": GenerateNoise, "name": "Generate Noise"}, + "FlipSigmasAdjusted": {"class": FlipSigmasAdjusted, "name": "Flip Sigmas Adjusted"}, + "InjectNoiseToLatent": {"class": InjectNoiseToLatent, "name": "Inject Noise To Latent"}, + "CustomSigmas": {"class": CustomSigmas, "name": "Custom Sigmas"}, #utility - "WidgetToString": WidgetToString, - "DummyLatentOut": DummyLatentOut, - "GetLatentsFromBatchIndexed": GetLatentsFromBatchIndexed, - "ScaleBatchPromptSchedule": ScaleBatchPromptSchedule, - "CameraPoseVisualizer": CameraPoseVisualizer, - "JoinStrings": JoinStrings, - "JoinStringMulti": JoinStringMulti, - "Sleep": Sleep, - "VRAM_Debug" : VRAM_Debug, - "SomethingToString" : SomethingToString, - "EmptyLatentImagePresets": EmptyLatentImagePresets, + "WidgetToString": {"class": WidgetToString, "name": "Widget To String"}, + "DummyLatentOut": {"class": DummyLatentOut, "name": "Dummy Latent Out"}, + "GetLatentsFromBatchIndexed": {"class": GetLatentsFromBatchIndexed, "name": "Get Latents From Batch Indexed"}, + "ScaleBatchPromptSchedule": {"class": ScaleBatchPromptSchedule, "name": "Scale Batch Prompt Schedule"}, + "CameraPoseVisualizer": {"class": CameraPoseVisualizer, "name": "Camera Pose Visualizer"}, + "JoinStrings": {"class": JoinStrings, "name": "Join Strings"}, + "JoinStringMulti": {"class": JoinStringMulti, "name": "Join String Multi"}, + "Sleep": {"class": Sleep, "name": "Sleep"}, + "VRAM_Debug": {"class": VRAM_Debug, "name": "VRAM Debug"}, + "SomethingToString": {"class": SomethingToString, "name": "Something To String"}, + "EmptyLatentImagePresets": {"class": EmptyLatentImagePresets, "name": "Empty Latent Image Presets"}, #audioscheduler stuff - "NormalizedAmplitudeToMask": NormalizedAmplitudeToMask, - "OffsetMaskByNormalizedAmplitude": OffsetMaskByNormalizedAmplitude, - "ImageTransformByNormalizedAmplitude": ImageTransformByNormalizedAmplitude, + "NormalizedAmplitudeToMask": {"class": NormalizedAmplitudeToMask}, + "OffsetMaskByNormalizedAmplitude": {"class": OffsetMaskByNormalizedAmplitude}, + "ImageTransformByNormalizedAmplitude": {"class": ImageTransformByNormalizedAmplitude}, #curve nodes - "SplineEditor": SplineEditor, - "CreateShapeMaskOnPath": CreateShapeMaskOnPath, - "WeightScheduleExtend": WeightScheduleExtend, - "MaskOrImageToWeight": MaskOrImageToWeight, - "WeightScheduleConvert": WeightScheduleConvert, - "FloatToMask": FloatToMask, - "FloatToSigmas": FloatToSigmas, + "SplineEditor": {"class": SplineEditor, "name": "Spline Editor"}, + "CreateShapeMaskOnPath": {"class": CreateShapeMaskOnPath, "name": "Create Shape Mask On Path"}, + "WeightScheduleExtend": {"class": WeightScheduleExtend, "name": "Weight Schedule Extend"}, + "MaskOrImageToWeight": {"class": MaskOrImageToWeight, "name": "Mask Or Image To Weight"}, + "WeightScheduleConvert": {"class": WeightScheduleConvert, "name": "Weight Schedule Convert"}, + "FloatToMask": {"class": FloatToMask, "name": "Float To Mask"}, + "FloatToSigmas": {"class": FloatToSigmas, "name": "Float To Sigmas"}, #experimental - "StabilityAPI_SD3": StabilityAPI_SD3, - "SoundReactive": SoundReactive, - "StableZero123_BatchSchedule": StableZero123_BatchSchedule, - "SV3D_BatchSchedule": SV3D_BatchSchedule, - "LoadResAdapterNormalization": LoadResAdapterNormalization, - "Superprompt": Superprompt, - "GLIGENTextBoxApplyBatchCoords": GLIGENTextBoxApplyBatchCoords, - "Intrinsic_lora_sampling": Intrinsic_lora_sampling, + "StabilityAPI_SD3": {"class": StabilityAPI_SD3, "name": "Stability API SD3"}, + "SoundReactive": {"class": SoundReactive, "name": "Sound Reactive"}, + "StableZero123_BatchSchedule": {"class": StableZero123_BatchSchedule, "name": "Stable Zero123 Batch Schedule"}, + "SV3D_BatchSchedule": {"class": SV3D_BatchSchedule, "name": "SV3D Batch Schedule"}, + "LoadResAdapterNormalization": {"class": LoadResAdapterNormalization}, + "Superprompt": {"class": Superprompt, "name": "Superprompt"}, + "GLIGENTextBoxApplyBatchCoords": {"class": GLIGENTextBoxApplyBatchCoords}, + "Intrinsic_lora_sampling": {"class": Intrinsic_lora_sampling, "name": "Intrinsic Lora Sampling"}, +} + +def generate_node_mappings(node_config): + node_class_mappings = {} + node_display_name_mappings = {} + + for node_name, node_info in node_config.items(): + node_class_mappings[node_name] = node_info["class"] + node_display_name_mappings[node_name] = node_info.get("name", node_info["class"].__name__) + + return node_class_mappings, node_display_name_mappings + +NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = generate_node_mappings(NODE_CONFIG) -} -NODE_DISPLAY_NAME_MAPPINGS = { - "INTConstant": "INT Constant", - "FloatConstant": "Float Constant", - "ImageBatchMulti": "Image Batch Multi", - "MaskBatchMulti": "Mask Batch Multi", - "ConditioningMultiCombine": "Conditioning Multi Combine", - "ConditioningSetMaskAndCombine": "ConditioningSetMaskAndCombine", - "ConditioningSetMaskAndCombine3": "ConditioningSetMaskAndCombine3", - "ConditioningSetMaskAndCombine4": "ConditioningSetMaskAndCombine4", - "ConditioningSetMaskAndCombine5": "ConditioningSetMaskAndCombine5", - "GrowMaskWithBlur": "GrowMaskWithBlur", - "ColorToMask": "ColorToMask", - "CreateGradientMask": "CreateGradientMask", - "CreateTextMask" : "CreateTextMask", - "CreateFadeMask" : "CreateFadeMask (Deprecated)", - "CreateFadeMaskAdvanced" : "CreateFadeMaskAdvanced", - "CreateFluidMask" : "CreateFluidMask", - "CreateAudioMask" : "CreateAudioMask (Deprecated)", - "VRAM_Debug" : "VRAM Debug", - "CrossFadeImages": "CrossFadeImages", - "SomethingToString": "SomethingToString", - "EmptyLatentImagePresets": "EmptyLatentImagePresets", - "ColorMatch": "ColorMatch", - "GetImageRangeFromBatch": "GetImageRangeFromBatch", - "InsertImagesToBatchIndexed": "InsertImagesToBatchIndexed", - "SaveImageWithAlpha": "SaveImageWithAlpha", - "ReverseImageBatch": "ReverseImageBatch", - "ImageGridComposite2x2": "ImageGridComposite2x2", - "ImageGridComposite3x3": "ImageGridComposite3x3", - "ImageConcanate": "ImageConcatenate", - "ImageBatchTestPattern": "ImageBatchTestPattern", - "ReplaceImagesInBatch": "ReplaceImagesInBatch", - "BatchCropFromMask": "BatchCropFromMask", - "BatchCropFromMaskAdvanced": "BatchCropFromMaskAdvanced", - "FilterZeroMasksAndCorrespondingImages": "FilterZeroMasksAndCorrespondingImages", - "InsertImageBatchByIndexes": "InsertImageBatchByIndexes", - "BatchUncrop": "BatchUncrop", - "BatchUncropAdvanced": "BatchUncropAdvanced", - "BatchCLIPSeg": "BatchCLIPSeg", - "RoundMask": "RoundMask", - "ResizeMask": "ResizeMask", - "OffsetMask": "OffsetMask", - "WidgetToString": "WidgetToString", - "CreateShapeMask": "CreateShapeMask", - "CreateVoronoiMask": "CreateVoronoiMask", - "CreateMagicMask": "CreateMagicMask", - "BboxToInt": "BboxToInt", - "SplitBboxes": "SplitBboxes", - "ImageGrabPIL": "ImageGrabPIL", - "DummyLatentOut": "DummyLatentOut", - "FlipSigmasAdjusted": "FlipSigmasAdjusted", - "InjectNoiseToLatent": "InjectNoiseToLatent", - "AddLabel": "AddLabel", - "SoundReactive": "SoundReactive", - "GenerateNoise": "GenerateNoise", - "StableZero123_BatchSchedule": "StableZero123_BatchSchedule", - "SV3D_BatchSchedule": "SV3D_BatchSchedule", - "GetImagesFromBatchIndexed": "GetImagesFromBatchIndexed", - "ImageBatchRepeatInterleaving": "ImageBatchRepeatInterleaving", - "NormalizedAmplitudeToMask": "NormalizedAmplitudeToMask", - "OffsetMaskByNormalizedAmplitude": "OffsetMaskByNormalizedAmplitude", - "ImageTransformByNormalizedAmplitude": "ImageTransformByNormalizedAmplitude", - "GetLatentsFromBatchIndexed": "GetLatentsFromBatchIndexed", - "StringConstant": "StringConstant", - "CondPassThrough": "CondPassThrough", - "ImageUpscaleWithModelBatched": "ImageUpscaleWithModelBatched", - "ScaleBatchPromptSchedule": "ScaleBatchPromptSchedule", - "ImageNormalize_Neg1_To_1": "ImageNormalize_Neg1_To_1", - "Intrinsic_lora_sampling": "Intrinsic_lora_sampling", - "RemapMaskRange": "RemapMaskRange", - "LoadResAdapterNormalization": "LoadResAdapterNormalization", - "Superprompt": "Superprompt", - "RemapImageRange": "RemapImageRange", - "CameraPoseVisualizer": "CameraPoseVisualizer", - "BboxVisualize": "BboxVisualize", - "StringConstantMultiline": "StringConstantMultiline", - "JoinStrings": "JoinStrings", - "Sleep": "🛌 Sleep 🛌", - "ImagePadForOutpaintMasked": "Pad Image For Outpaint Masked", - "ImageAndMaskPreview": "Image & Mask Preview", - "StabilityAPI_SD3": "Stability API SD3", - "MaskOrImageToWeight": "Mask Or Image To Weight", - "WeightScheduleConvert": "Weight Schedule Convert", - "FloatToMask": "Float To Mask", - "FloatToSigmas": "Float To Sigmas", - "CustomSigmas": "Custom Sigmas", - "ImagePass": "ImagePass", - "SplitImageChannels": "Split Image Channels", - "MergeImageChannels": "Merge Image Channels", - #curve nodes - "SplineEditor": "Spline Editor", - "CreateShapeMaskOnPath": "Create Shape Mask On Path", - "WeightScheduleExtend": "Weight Schedule Extend" -} __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] WEB_DIRECTORY = "./web" diff --git a/nodes/mask_nodes.py b/nodes/mask_nodes.py new file mode 100644 index 0000000..9ea1f65 --- /dev/null +++ b/nodes/mask_nodes.py @@ -0,0 +1,1165 @@ +import torch +import torch.nn.functional as F +from torchvision.transforms import functional as TF +from PIL import Image, ImageDraw, ImageFilter, ImageFont +import scipy.ndimage +import numpy as np + +import matplotlib.pyplot as plt +from contextlib import nullcontext +import os + +import model_management +from comfy.utils import ProgressBar +from nodes import MAX_RESOLUTION + +import folder_paths + +from ..utility.utility import tensor2pil, pil2tensor + +script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + +class BatchCLIPSeg: + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + + return {"required": + { + "images": ("IMAGE",), + "text": ("STRING", {"multiline": False}), + "threshold": ("FLOAT", {"default": 0.1,"min": 0.0, "max": 10.0, "step": 0.001}), + "binary_mask": ("BOOLEAN", {"default": True}), + "combine_mask": ("BOOLEAN", {"default": False}), + "use_cuda": ("BOOLEAN", {"default": True}), + }, + } + + CATEGORY = "KJNodes/masking" + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("Mask",) + FUNCTION = "segment_image" + DESCRIPTION = """ +Segments an image or batch of images using CLIPSeg. +""" + + def segment_image(self, images, text, threshold, binary_mask, combine_mask, use_cuda): + from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation + out = [] + height, width, _ = images[0].shape + if use_cuda and torch.cuda.is_available(): + device = torch.device("cuda") + else: + device = torch.device("cpu") + dtype = model_management.unet_dtype() + model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") + model.to(dtype) + model.to(device) + images = images.to(device) + processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") + pbar = ProgressBar(images.shape[0]) + autocast_condition = (dtype != torch.float32) and not model_management.is_device_mps(device) + with torch.autocast(model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): + for image in images: + image = (image* 255).type(torch.uint8) + prompt = text + input_prc = processor(text=prompt, images=image, return_tensors="pt") + # Move the processed input to the device + for key in input_prc: + input_prc[key] = input_prc[key].to(device) + + outputs = model(**input_prc) + + tensor = torch.sigmoid(outputs[0]) + tensor_thresholded = torch.where(tensor > threshold, tensor, torch.tensor(0, dtype=torch.float)) + tensor_normalized = (tensor_thresholded - tensor_thresholded.min()) / (tensor_thresholded.max() - tensor_thresholded.min()) + tensor = tensor_normalized + + # Resize the mask + if len(tensor.shape) == 3: + tensor = tensor.unsqueeze(0) + resized_tensor = F.interpolate(tensor, size=(height, width), mode='nearest') + + # Remove the extra dimensions + resized_tensor = resized_tensor[0, 0, :, :] + pbar.update(1) + out.append(resized_tensor) + + results = torch.stack(out).cpu().float() + + if combine_mask: + combined_results = torch.max(results, dim=0)[0] + results = combined_results.unsqueeze(0).repeat(len(images),1,1) + + if binary_mask: + results = results.round() + + return results, + +class CreateTextMask: + + RETURN_TYPES = ("IMAGE", "MASK",) + FUNCTION = "createtextmask" + CATEGORY = "KJNodes/text" + DESCRIPTION = """ +Creates a text image and mask. +Looks for fonts from this folder: +ComfyUI/custom_nodes/ComfyUI-KJNodes/fonts + +If start_rotation and/or end_rotation are different values, +creates animation between them. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), + "text_x": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "text_y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + "font_size": ("INT", {"default": 32,"min": 8, "max": 4096, "step": 1}), + "font_color": ("STRING", {"default": "white"}), + "text": ("STRING", {"default": "HELLO!", "multiline": True}), + "font": (folder_paths.get_filename_list("kjnodes_fonts"), ), + "width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "start_rotation": ("INT", {"default": 0,"min": 0, "max": 359, "step": 1}), + "end_rotation": ("INT", {"default": 0,"min": -359, "max": 359, "step": 1}), + }, + } + + def createtextmask(self, frames, width, height, invert, text_x, text_y, text, font_size, font_color, font, start_rotation, end_rotation): + # Define the number of images in the batch + batch_size = frames + out = [] + masks = [] + rotation = start_rotation + if start_rotation != end_rotation: + rotation_increment = (end_rotation - start_rotation) / (batch_size - 1) + + font_path = folder_paths.get_full_path("kjnodes_fonts", font) + # Generate the text + for i in range(batch_size): + image = Image.new("RGB", (width, height), "black") + draw = ImageDraw.Draw(image) + font = ImageFont.truetype(font_path, font_size) + + # Split the text into words + words = text.split() + + # Initialize variables for line creation + lines = [] + current_line = [] + current_line_width = 0 + try: #new pillow + # Iterate through words to create lines + for word in words: + word_width = font.getbbox(word)[2] + if current_line_width + word_width <= width - 2 * text_x: + current_line.append(word) + current_line_width += word_width + font.getbbox(" ")[2] # Add space width + else: + lines.append(" ".join(current_line)) + current_line = [word] + current_line_width = word_width + except: #old pillow + for word in words: + word_width = font.getsize(word)[0] + if current_line_width + word_width <= width - 2 * text_x: + current_line.append(word) + current_line_width += word_width + font.getsize(" ")[0] # Add space width + else: + lines.append(" ".join(current_line)) + current_line = [word] + current_line_width = word_width + + # Add the last line if it's not empty + if current_line: + lines.append(" ".join(current_line)) + + # Draw each line of text separately + y_offset = text_y + for line in lines: + text_width = font.getlength(line) + text_height = font_size + text_center_x = text_x + text_width / 2 + text_center_y = y_offset + text_height / 2 + try: + draw.text((text_x, y_offset), line, font=font, fill=font_color, features=['-liga']) + except: + draw.text((text_x, y_offset), line, font=font, fill=font_color) + y_offset += text_height # Move to the next line + + if start_rotation != end_rotation: + image = image.rotate(rotation, center=(text_center_x, text_center_y)) + rotation += rotation_increment + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + mask = image[:, :, :, 0] + masks.append(mask) + out.append(image) + + if invert: + return (1.0 - torch.cat(out, dim=0), 1.0 - torch.cat(masks, dim=0),) + return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) + +class ColorToMask: + + RETURN_TYPES = ("MASK",) + FUNCTION = "clip" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Converts chosen RGB value to a mask. +With batch inputs, the **per_batch** +controls the number of images processed at once. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "invert": ("BOOLEAN", {"default": False}), + "red": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "green": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "blue": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "threshold": ("INT", {"default": 10,"min": 0, "max": 255, "step": 1}), + "per_batch": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}), + }, + } + + def clip(self, images, red, green, blue, threshold, invert, per_batch): + + color = torch.tensor([red, green, blue], dtype=torch.uint8) + black = torch.tensor([0, 0, 0], dtype=torch.uint8) + white = torch.tensor([255, 255, 255], dtype=torch.uint8) + + if invert: + black, white = white, black + + steps = images.shape[0] + pbar = ProgressBar(steps) + tensors_out = [] + + for start_idx in range(0, images.shape[0], per_batch): + + # Calculate color distances + color_distances = torch.norm(images[start_idx:start_idx+per_batch] * 255 - color, dim=-1) + + # Create a mask based on the threshold + mask = color_distances <= threshold + + # Apply the mask to create new images + mask_out = torch.where(mask.unsqueeze(-1), white, black).float() + mask_out = mask_out.mean(dim=-1) + + tensors_out.append(mask_out.cpu()) + batch_count = mask_out.shape[0] + pbar.update(batch_count) + + tensors_out = torch.cat(tensors_out, dim=0) + tensors_out = torch.clamp(tensors_out, min=0.0, max=1.0) + return tensors_out, + +class CreateFluidMask: + + RETURN_TYPES = ("IMAGE", "MASK") + FUNCTION = "createfluidmask" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "inflow_count": ("INT", {"default": 3,"min": 0, "max": 255, "step": 1}), + "inflow_velocity": ("INT", {"default": 1,"min": 0, "max": 255, "step": 1}), + "inflow_radius": ("INT", {"default": 8,"min": 0, "max": 255, "step": 1}), + "inflow_padding": ("INT", {"default": 50,"min": 0, "max": 255, "step": 1}), + "inflow_duration": ("INT", {"default": 60,"min": 0, "max": 255, "step": 1}), + }, + } + #using code from https://github.com/GregTJ/stable-fluids + def createfluidmask(self, frames, width, height, invert, inflow_count, inflow_velocity, inflow_radius, inflow_padding, inflow_duration): + from ..utility.fluid import Fluid + from scipy.spatial import erf + out = [] + masks = [] + RESOLUTION = width, height + DURATION = frames + + INFLOW_PADDING = inflow_padding + INFLOW_DURATION = inflow_duration + INFLOW_RADIUS = inflow_radius + INFLOW_VELOCITY = inflow_velocity + INFLOW_COUNT = inflow_count + + print('Generating fluid solver, this may take some time.') + fluid = Fluid(RESOLUTION, 'dye') + + center = np.floor_divide(RESOLUTION, 2) + r = np.min(center) - INFLOW_PADDING + + points = np.linspace(-np.pi, np.pi, INFLOW_COUNT, endpoint=False) + points = tuple(np.array((np.cos(p), np.sin(p))) for p in points) + normals = tuple(-p for p in points) + points = tuple(r * p + center for p in points) + + inflow_velocity = np.zeros_like(fluid.velocity) + inflow_dye = np.zeros(fluid.shape) + for p, n in zip(points, normals): + mask = np.linalg.norm(fluid.indices - p[:, None, None], axis=0) <= INFLOW_RADIUS + inflow_velocity[:, mask] += n[:, None] * INFLOW_VELOCITY + inflow_dye[mask] = 1 + + + for f in range(DURATION): + print(f'Computing frame {f + 1} of {DURATION}.') + if f <= INFLOW_DURATION: + fluid.velocity += inflow_velocity + fluid.dye += inflow_dye + + curl = fluid.step()[1] + # Using the error function to make the contrast a bit higher. + # Any other sigmoid function e.g. smoothstep would work. + curl = (erf(curl * 2) + 1) / 4 + + color = np.dstack((curl, np.ones(fluid.shape), fluid.dye)) + color = (np.clip(color, 0, 1) * 255).astype('uint8') + image = np.array(color).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + mask = image[:, :, :, 0] + masks.append(mask) + out.append(image) + + if invert: + return (1.0 - torch.cat(out, dim=0),1.0 - torch.cat(masks, dim=0),) + return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) + +class CreateAudioMask: + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "createaudiomask" + CATEGORY = "KJNodes/deprecated" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 16,"min": 1, "max": 255, "step": 1}), + "scale": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 2.0, "step": 0.01}), + "audio_path": ("STRING", {"default": "audio.wav"}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + }, + } + + def createaudiomask(self, frames, width, height, invert, audio_path, scale): + try: + import librosa + except ImportError: + raise Exception("Can not import librosa. Install it with 'pip install librosa'") + batch_size = frames + out = [] + masks = [] + if audio_path == "audio.wav": #I don't know why relative path won't work otherwise... + audio_path = os.path.join(script_directory, audio_path) + audio, sr = librosa.load(audio_path) + spectrogram = np.abs(librosa.stft(audio)) + + for i in range(batch_size): + image = Image.new("RGB", (width, height), "black") + draw = ImageDraw.Draw(image) + frame = spectrogram[:, i] + circle_radius = int(height * np.mean(frame)) + circle_radius *= scale + circle_center = (width // 2, height // 2) # Calculate the center of the image + + draw.ellipse([(circle_center[0] - circle_radius, circle_center[1] - circle_radius), + (circle_center[0] + circle_radius, circle_center[1] + circle_radius)], + fill='white') + + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + mask = image[:, :, :, 0] + masks.append(mask) + out.append(image) + + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) + +class CreateGradientMask: + + RETURN_TYPES = ("MASK",) + FUNCTION = "createmask" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + }, + } + def createmask(self, frames, width, height, invert): + # Define the number of images in the batch + batch_size = frames + out = [] + # Create an empty array to store the image batch + image_batch = np.zeros((batch_size, height, width), dtype=np.float32) + # Generate the black to white gradient for each image + for i in range(batch_size): + gradient = np.linspace(1.0, 0.0, width, dtype=np.float32) + time = i / frames # Calculate the time variable + offset_gradient = gradient - time # Offset the gradient values based on time + image_batch[i] = offset_gradient.reshape(1, -1) + output = torch.from_numpy(image_batch) + mask = output + out.append(mask) + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),) + +class CreateFadeMask: + + RETURN_TYPES = ("MASK",) + FUNCTION = "createfademask" + CATEGORY = "KJNodes/deprecated" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 2,"min": 2, "max": 255, "step": 1}), + "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out"],), + "start_level": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 1.0, "step": 0.01}), + "midpoint_level": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}), + "end_level": ("FLOAT", {"default": 0.0,"min": 0.0, "max": 1.0, "step": 0.01}), + "midpoint_frame": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), + }, + } + + def createfademask(self, frames, width, height, invert, interpolation, start_level, midpoint_level, end_level, midpoint_frame): + def ease_in(t): + return t * t + + def ease_out(t): + return 1 - (1 - t) * (1 - t) + + def ease_in_out(t): + return 3 * t * t - 2 * t * t * t + + batch_size = frames + out = [] + image_batch = np.zeros((batch_size, height, width), dtype=np.float32) + + if midpoint_frame == 0: + midpoint_frame = batch_size // 2 + + for i in range(batch_size): + if i <= midpoint_frame: + t = i / midpoint_frame + if interpolation == "ease_in": + t = ease_in(t) + elif interpolation == "ease_out": + t = ease_out(t) + elif interpolation == "ease_in_out": + t = ease_in_out(t) + color = start_level - t * (start_level - midpoint_level) + else: + t = (i - midpoint_frame) / (batch_size - midpoint_frame) + if interpolation == "ease_in": + t = ease_in(t) + elif interpolation == "ease_out": + t = ease_out(t) + elif interpolation == "ease_in_out": + t = ease_in_out(t) + color = midpoint_level - t * (midpoint_level - end_level) + + color = np.clip(color, 0, 255) + image = np.full((height, width), color, dtype=np.float32) + image_batch[i] = image + + output = torch.from_numpy(image_batch) + mask = output + out.append(mask) + + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),) + +class CreateFadeMaskAdvanced: + + RETURN_TYPES = ("MASK",) + FUNCTION = "createfademask" + CATEGORY = "KJNodes/masking/generate" + DESCRIPTION = """ +Create a batch of masks interpolated between given frames and values. +Uses same syntax as Fizz' BatchValueSchedule. +First value is the frame index (not that this starts from 0, not 1) +and the second value inside the brackets is the float value of the mask in range 0.0 - 1.0 + +For example the default values: +0:(0.0) +7:(1.0) +15:(0.0) + +Would create a mask batch fo 16 frames, starting from black, +interpolating with the chosen curve to fully white at the 8th frame, +and interpolating from that to fully black at the 16th frame. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "points_string": ("STRING", {"default": "0:(0.0),\n7:(1.0),\n15:(0.0)\n", "multiline": True}), + "invert": ("BOOLEAN", {"default": False}), + "frames": ("INT", {"default": 16,"min": 2, "max": 255, "step": 1}), + "width": ("INT", {"default": 512,"min": 1, "max": 4096, "step": 1}), + "height": ("INT", {"default": 512,"min": 1, "max": 4096, "step": 1}), + "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out"],), + }, + } + + def createfademask(self, frames, width, height, invert, points_string, interpolation): + def ease_in(t): + return t * t + + def ease_out(t): + return 1 - (1 - t) * (1 - t) + + def ease_in_out(t): + return 3 * t * t - 2 * t * t * t + + # Parse the input string into a list of tuples + points = [] + points_string = points_string.rstrip(',\n') + for point_str in points_string.split(','): + frame_str, color_str = point_str.split(':') + frame = int(frame_str.strip()) + color = float(color_str.strip()[1:-1]) # Remove parentheses around color + points.append((frame, color)) + + # Check if the last frame is already in the points + if len(points) == 0 or points[-1][0] != frames - 1: + # If not, add it with the color of the last specified frame + points.append((frames - 1, points[-1][1] if points else 0)) + + # Sort the points by frame number + points.sort(key=lambda x: x[0]) + + batch_size = frames + out = [] + image_batch = np.zeros((batch_size, height, width), dtype=np.float32) + + # Index of the next point to interpolate towards + next_point = 1 + + for i in range(batch_size): + while next_point < len(points) and i > points[next_point][0]: + next_point += 1 + + # Interpolate between the previous point and the next point + prev_point = next_point - 1 + t = (i - points[prev_point][0]) / (points[next_point][0] - points[prev_point][0]) + if interpolation == "ease_in": + t = ease_in(t) + elif interpolation == "ease_out": + t = ease_out(t) + elif interpolation == "ease_in_out": + t = ease_in_out(t) + elif interpolation == "linear": + pass # No need to modify `t` for linear interpolation + + color = points[prev_point][1] - t * (points[prev_point][1] - points[next_point][1]) + color = np.clip(color, 0, 255) + image = np.full((height, width), color, dtype=np.float32) + image_batch[i] = image + + output = torch.from_numpy(image_batch) + mask = output + out.append(mask) + + if invert: + return (1.0 - torch.cat(out, dim=0),) + return (torch.cat(out, dim=0),) + +class CreateMagicMask: + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "createmagicmask" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}), + "depth": ("INT", {"default": 12,"min": 1, "max": 500, "step": 1}), + "distortion": ("FLOAT", {"default": 1.5,"min": 0.0, "max": 100.0, "step": 0.01}), + "seed": ("INT", {"default": 123,"min": 0, "max": 99999999, "step": 1}), + "transitions": ("INT", {"default": 1,"min": 1, "max": 20, "step": 1}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + }, + } + + def createmagicmask(self, frames, transitions, depth, distortion, seed, frame_width, frame_height): + from ..utility.magictex import coordinate_grid, random_transform, magic + rng = np.random.default_rng(seed) + out = [] + coords = coordinate_grid((frame_width, frame_height)) + + # Calculate the number of frames for each transition + frames_per_transition = frames // transitions + + # Generate a base set of parameters + base_params = { + "coords": random_transform(coords, rng), + "depth": depth, + "distortion": distortion, + } + for t in range(transitions): + # Generate a second set of parameters that is at most max_diff away from the base parameters + params1 = base_params.copy() + params2 = base_params.copy() + + params1['coords'] = random_transform(coords, rng) + params2['coords'] = random_transform(coords, rng) + + for i in range(frames_per_transition): + # Compute the interpolation factor + alpha = i / frames_per_transition + + # Interpolate between the two sets of parameters + params = params1.copy() + params['coords'] = (1 - alpha) * params1['coords'] + alpha * params2['coords'] + + tex = magic(**params) + + dpi = frame_width / 10 + fig = plt.figure(figsize=(10, 10), dpi=dpi) + + ax = fig.add_subplot(111) + plt.subplots_adjust(left=0, right=1, bottom=0, top=1) + + ax.get_yaxis().set_ticks([]) + ax.get_xaxis().set_ticks([]) + ax.imshow(tex, aspect='auto') + + fig.canvas.draw() + img = np.array(fig.canvas.renderer._renderer) + + plt.close(fig) + + pil_img = Image.fromarray(img).convert("L") + mask = torch.tensor(np.array(pil_img)) / 255.0 + + out.append(mask) + + return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) + +class CreateShapeMask: + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "createshapemask" + CATEGORY = "KJNodes/masking/generate" + DESCRIPTION = """ +Creates a mask or batch of masks with the specified shape. +Locations are center locations. +Grow value is the amount to grow the shape on each frame, creating animated masks. +""" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "shape": ( + [ 'circle', + 'square', + 'triangle', + ], + { + "default": 'circle' + }), + "frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), + "location_x": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), + "location_y": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), + "grow": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "shape_width": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + "shape_height": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), + }, + } + + def createshapemask(self, frames, frame_width, frame_height, location_x, location_y, shape_width, shape_height, grow, shape): + # Define the number of images in the batch + batch_size = frames + out = [] + color = "white" + for i in range(batch_size): + image = Image.new("RGB", (frame_width, frame_height), "black") + draw = ImageDraw.Draw(image) + + # Calculate the size for this frame and ensure it's not less than 0 + current_width = max(0, shape_width + i*grow) + current_height = max(0, shape_height + i*grow) + + if shape == 'circle' or shape == 'square': + # Define the bounding box for the shape + left_up_point = (location_x - current_width // 2, location_y - current_height // 2) + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) + two_points = [left_up_point, right_down_point] + + if shape == 'circle': + draw.ellipse(two_points, fill=color) + elif shape == 'square': + draw.rectangle(two_points, fill=color) + + elif shape == 'triangle': + # Define the points for the triangle + left_up_point = (location_x - current_width // 2, location_y + current_height // 2) # bottom left + right_down_point = (location_x + current_width // 2, location_y + current_height // 2) # bottom right + top_point = (location_x, location_y - current_height // 2) # top point + draw.polygon([top_point, left_up_point, right_down_point], fill=color) + + image = pil2tensor(image) + mask = image[:, :, :, 0] + out.append(mask) + outstack = torch.cat(out, dim=0) + return (outstack, 1.0 - outstack,) + +class CreateVoronoiMask: + + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "createvoronoi" + CATEGORY = "KJNodes/masking/generate" + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}), + "num_points": ("INT", {"default": 15,"min": 1, "max": 4096, "step": 1}), + "line_width": ("INT", {"default": 4,"min": 1, "max": 4096, "step": 1}), + "speed": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}), + "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), + }, + } + + def createvoronoi(self, frames, num_points, line_width, speed, frame_width, frame_height): + from scipy.spatial import Voronoi + # Define the number of images in the batch + batch_size = frames + out = [] + + # Calculate aspect ratio + aspect_ratio = frame_width / frame_height + + # Create start and end points for each point, considering the aspect ratio + start_points = np.random.rand(num_points, 2) + start_points[:, 0] *= aspect_ratio + + end_points = np.random.rand(num_points, 2) + end_points[:, 0] *= aspect_ratio + + for i in range(batch_size): + # Interpolate the points' positions based on the current frame + t = (i * speed) / (batch_size - 1) # normalize to [0, 1] over the frames + t = np.clip(t, 0, 1) # ensure t is in [0, 1] + points = (1 - t) * start_points + t * end_points # lerp + + # Adjust points for aspect ratio + points[:, 0] *= aspect_ratio + + vor = Voronoi(points) + + # Create a blank image with a white background + fig, ax = plt.subplots() + plt.subplots_adjust(left=0, right=1, bottom=0, top=1) + ax.set_xlim([0, aspect_ratio]); ax.set_ylim([0, 1]) # adjust x limits + ax.axis('off') + ax.margins(0, 0) + fig.set_size_inches(aspect_ratio * frame_height/100, frame_height/100) # adjust figure size + ax.fill_between([0, 1], [0, 1], color='white') + + # Plot each Voronoi ridge + for simplex in vor.ridge_vertices: + simplex = np.asarray(simplex) + if np.all(simplex >= 0): + plt.plot(vor.vertices[simplex, 0], vor.vertices[simplex, 1], 'k-', linewidth=line_width) + + fig.canvas.draw() + img = np.array(fig.canvas.renderer._renderer) + + plt.close(fig) + + pil_img = Image.fromarray(img).convert("L") + mask = torch.tensor(np.array(pil_img)) / 255.0 + + out.append(mask) + + return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) + +class GetMaskSize: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "mask": ("MASK",), + }} + + RETURN_TYPES = ("MASK","INT", "INT", ) + RETURN_NAMES = ("mask", "width", "height",) + FUNCTION = "getsize" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Returns the width and height of the mask, +and passes through the mask unchanged. + +""" + + def getsize(self, mask): + width = mask.shape[2] + height = mask.shape[1] + return {"ui": { + "text": [f"{width}x{height}"]}, + "result": (mask, width, height) + } + +class GrowMaskWithBlur: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mask": ("MASK",), + "expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), + "incremental_expandrate": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), + "tapered_corners": ("BOOLEAN", {"default": True}), + "flip_input": ("BOOLEAN", {"default": False}), + "blur_radius": ("FLOAT", { + "default": 0.0, + "min": 0.0, + "max": 100, + "step": 0.1 + }), + "lerp_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "decay_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + "optional": { + "fill_holes": ("BOOLEAN", {"default": False}), + }, + } + + CATEGORY = "KJNodes/masking" + RETURN_TYPES = ("MASK", "MASK",) + RETURN_NAMES = ("mask", "mask_inverted",) + FUNCTION = "expand_mask" + DESCRIPTION = """ +# GrowMaskWithBlur +- mask: Input mask or mask batch +- expand: Expand or contract mask or mask batch by a given amount +- incremental_expandrate: increase expand rate by a given amount per frame +- tapered_corners: use tapered corners +- flip_input: flip input mask +- blur_radius: value higher than 0 will blur the mask +- lerp_alpha: alpha value for interpolation between frames +- decay_factor: decay value for interpolation between frames +- fill_holes: fill holes in the mask (slow)""" + + def expand_mask(self, mask, expand, tapered_corners, flip_input, blur_radius, incremental_expandrate, lerp_alpha, decay_factor, fill_holes=False): + alpha = lerp_alpha + decay = decay_factor + if flip_input: + mask = 1.0 - mask + c = 0 if tapered_corners else 1 + kernel = np.array([[c, 1, c], + [1, 1, 1], + [c, 1, c]]) + growmask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).cpu() + out = [] + previous_output = None + current_expand = expand + for m in growmask: + output = m.numpy() + for _ in range(abs(round(current_expand))): + if current_expand < 0: + output = scipy.ndimage.grey_erosion(output, footprint=kernel) + else: + output = scipy.ndimage.grey_dilation(output, footprint=kernel) + if current_expand < 0: + current_expand -= abs(incremental_expandrate) + else: + current_expand += abs(incremental_expandrate) + if fill_holes: + binary_mask = output > 0 + output = scipy.ndimage.binary_fill_holes(binary_mask) + output = output.astype(np.float32) * 255 + output = torch.from_numpy(output) + if alpha < 1.0 and previous_output is not None: + # Interpolate between the previous and current frame + output = alpha * output + (1 - alpha) * previous_output + if decay < 1.0 and previous_output is not None: + # Add the decayed previous output to the current frame + output += decay * previous_output + output = output / output.max() + previous_output = output + out.append(output) + + if blur_radius != 0: + # Convert the tensor list to PIL images, apply blur, and convert back + for idx, tensor in enumerate(out): + # Convert tensor to PIL image + pil_image = tensor2pil(tensor.cpu().detach())[0] + # Apply Gaussian blur + pil_image = pil_image.filter(ImageFilter.GaussianBlur(blur_radius)) + # Convert back to tensor + out[idx] = pil2tensor(pil_image) + blurred = torch.cat(out, dim=0) + return (blurred, 1.0 - blurred) + else: + return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) + +class MaskBatchMulti: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), + "mask_1": ("MASK", ), + "mask_2": ("MASK", ), + }, + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("masks",) + FUNCTION = "combine" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Creates an image batch from multiple masks. +You can set how many inputs the node has, +with the **inputcount** and clicking update. +""" + + def combine(self, inputcount, **kwargs): + mask = kwargs["mask_1"] + for c in range(1, inputcount): + new_mask = kwargs[f"mask_{c + 1}"] + if mask.shape[1:] != new_mask.shape[1:]: + new_mask = F.interpolate(new_mask.unsqueeze(1), size=(mask.shape[1], mask.shape[2]), mode="bicubic").squeeze(1) + mask = torch.cat((mask, new_mask), dim=0) + return (mask,) + +class OffsetMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + "x": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "y": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), + "angle": ("INT", { "default": 0, "min": -360, "max": 360, "step": 1, "display": "number" }), + "duplication_factor": ("INT", { "default": 1, "min": 1, "max": 1000, "step": 1, "display": "number" }), + "roll": ("BOOLEAN", { "default": False }), + "incremental": ("BOOLEAN", { "default": False }), + "padding_mode": ( + [ + 'empty', + 'border', + 'reflection', + + ], { + "default": 'empty' + }), + } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "offset" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Offsets the mask by the specified amount. + - mask: Input mask or mask batch + - x: Horizontal offset + - y: Vertical offset + - angle: Angle in degrees + - roll: roll edge wrapping + - duplication_factor: Number of times to duplicate the mask to form a batch + - border padding_mode: Padding mode for the mask +""" + + def offset(self, mask, x, y, angle, roll=False, incremental=False, duplication_factor=1, padding_mode="empty"): + # Create duplicates of the mask batch + mask = mask.repeat(duplication_factor, 1, 1).clone() + + batch_size, height, width = mask.shape + + if angle != 0 and incremental: + for i in range(batch_size): + rotation_angle = angle * (i+1) + mask[i] = TF.rotate(mask[i].unsqueeze(0), rotation_angle).squeeze(0) + elif angle > 0: + for i in range(batch_size): + mask[i] = TF.rotate(mask[i].unsqueeze(0), angle).squeeze(0) + + if roll: + if incremental: + for i in range(batch_size): + shift_x = min(x*(i+1), width-1) + shift_y = min(y*(i+1), height-1) + if shift_x != 0: + mask[i] = torch.roll(mask[i], shifts=shift_x, dims=1) + if shift_y != 0: + mask[i] = torch.roll(mask[i], shifts=shift_y, dims=0) + else: + shift_x = min(x, width-1) + shift_y = min(y, height-1) + if shift_x != 0: + mask = torch.roll(mask, shifts=shift_x, dims=2) + if shift_y != 0: + mask = torch.roll(mask, shifts=shift_y, dims=1) + else: + + for i in range(batch_size): + if incremental: + temp_x = min(x * (i+1), width-1) + temp_y = min(y * (i+1), height-1) + else: + temp_x = min(x, width-1) + temp_y = min(y, height-1) + if temp_x > 0: + if padding_mode == 'empty': + mask[i] = torch.cat([torch.zeros((height, temp_x)), mask[i, :, :-temp_x]], dim=1) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, :, :-temp_x], (0, temp_x), mode=padding_mode) + elif temp_x < 0: + if padding_mode == 'empty': + mask[i] = torch.cat([mask[i, :, :temp_x], torch.zeros((height, -temp_x))], dim=1) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, :, -temp_x:], (temp_x, 0), mode=padding_mode) + + if temp_y > 0: + if padding_mode == 'empty': + mask[i] = torch.cat([torch.zeros((temp_y, width)), mask[i, :-temp_y, :]], dim=0) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, :-temp_y, :], (0, temp_y), mode=padding_mode) + elif temp_y < 0: + if padding_mode == 'empty': + mask[i] = torch.cat([mask[i, :temp_y, :], torch.zeros((-temp_y, width))], dim=0) + elif padding_mode in ['replicate', 'reflect']: + mask[i] = F.pad(mask[i, -temp_y:, :], (temp_y, 0), mode=padding_mode) + + return mask, + +class RoundMask: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "mask": ("MASK",), + }} + + RETURN_TYPES = ("MASK",) + FUNCTION = "round" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Rounds the mask or batch of masks to a binary mask. +RoundMask example + +""" + + def round(self, mask): + mask = mask.round() + return (mask,) + +class ResizeMask: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, "display": "number" }), + "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, "display": "number" }), + "keep_proportions": ("BOOLEAN", { "default": False }), + } + } + + RETURN_TYPES = ("MASK", "INT", "INT",) + RETURN_NAMES = ("mask", "width", "height",) + FUNCTION = "resize" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Resizes the mask or batch of masks to the specified width and height. +""" + + def resize(self, mask, width, height, keep_proportions): + if keep_proportions: + _, oh, ow, _ = mask.shape + width = ow if width == 0 else width + height = oh if height == 0 else height + ratio = min(width / ow, height / oh) + width = round(ow*ratio) + height = round(oh*ratio) + + outputs = mask.unsqueeze(0) # Add an extra dimension for batch size + outputs = F.interpolate(outputs, size=(height, width), mode="nearest") + outputs = outputs.squeeze(0) # Remove the extra dimension after interpolation + + return(outputs, outputs.shape[2], outputs.shape[1],) + +class RemapMaskRange: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mask": ("MASK",), + "min": ("FLOAT", {"default": 0.0,"min": -10.0, "max": 1.0, "step": 0.01}), + "max": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 10.0, "step": 0.01}), + } + } + + RETURN_TYPES = ("MASK",) + RETURN_NAMES = ("mask",) + FUNCTION = "remap" + CATEGORY = "KJNodes/masking" + DESCRIPTION = """ +Sets new min and max values for the mask. +""" + + def remap(self, mask, min, max): + + # Find the maximum value in the mask + mask_max = torch.max(mask) + + # If the maximum mask value is zero, avoid division by zero by setting it to 1 + mask_max = mask_max if mask_max > 0 else 1 + + # Scale the mask values to the new range defined by min and max + # The highest pixel value in the mask will be scaled to max + scaled_mask = (mask / mask_max) * (max - min) + min + + # Clamp the values to ensure they are within [0.0, 1.0] + scaled_mask = torch.clamp(scaled_mask, min=0.0, max=1.0) + + return (scaled_mask, ) \ No newline at end of file diff --git a/nodes/nodes.py b/nodes/nodes.py index 154345a..2b60ed8 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -1,24 +1,14 @@ import torch import torch.nn.functional as F -from torchvision.transforms import functional as TF - -import scipy.ndimage import matplotlib.pyplot as plt import numpy as np -from PIL import ImageFilter, Image, ImageDraw, ImageFont -from contextlib import nullcontext +from PIL import Image -import json -import re -import os -import io +import json, re, os, io, time import model_management -from nodes import MAX_RESOLUTION - import folder_paths -from ..utility.utility import tensor2pil, pil2tensor -from comfy.utils import ProgressBar +from nodes import MAX_RESOLUTION script_directory = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) folder_paths.add_model_folder_path("kjnodes_fonts", os.path.join(script_directory, "fonts")) @@ -99,340 +89,7 @@ class StringConstantMultiline: return (new_string, ) -class CreateFluidMask: - - RETURN_TYPES = ("IMAGE", "MASK") - FUNCTION = "createfluidmask" - CATEGORY = "KJNodes/masking/generate" - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "invert": ("BOOLEAN", {"default": False}), - "frames": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), - "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - "inflow_count": ("INT", {"default": 3,"min": 0, "max": 255, "step": 1}), - "inflow_velocity": ("INT", {"default": 1,"min": 0, "max": 255, "step": 1}), - "inflow_radius": ("INT", {"default": 8,"min": 0, "max": 255, "step": 1}), - "inflow_padding": ("INT", {"default": 50,"min": 0, "max": 255, "step": 1}), - "inflow_duration": ("INT", {"default": 60,"min": 0, "max": 255, "step": 1}), - }, - } - #using code from https://github.com/GregTJ/stable-fluids - def createfluidmask(self, frames, width, height, invert, inflow_count, inflow_velocity, inflow_radius, inflow_padding, inflow_duration): - from ..utility.fluid import Fluid - from scipy.spatial import erf - out = [] - masks = [] - RESOLUTION = width, height - DURATION = frames - - INFLOW_PADDING = inflow_padding - INFLOW_DURATION = inflow_duration - INFLOW_RADIUS = inflow_radius - INFLOW_VELOCITY = inflow_velocity - INFLOW_COUNT = inflow_count - - print('Generating fluid solver, this may take some time.') - fluid = Fluid(RESOLUTION, 'dye') - - center = np.floor_divide(RESOLUTION, 2) - r = np.min(center) - INFLOW_PADDING - - points = np.linspace(-np.pi, np.pi, INFLOW_COUNT, endpoint=False) - points = tuple(np.array((np.cos(p), np.sin(p))) for p in points) - normals = tuple(-p for p in points) - points = tuple(r * p + center for p in points) - - inflow_velocity = np.zeros_like(fluid.velocity) - inflow_dye = np.zeros(fluid.shape) - for p, n in zip(points, normals): - mask = np.linalg.norm(fluid.indices - p[:, None, None], axis=0) <= INFLOW_RADIUS - inflow_velocity[:, mask] += n[:, None] * INFLOW_VELOCITY - inflow_dye[mask] = 1 - - - for f in range(DURATION): - print(f'Computing frame {f + 1} of {DURATION}.') - if f <= INFLOW_DURATION: - fluid.velocity += inflow_velocity - fluid.dye += inflow_dye - - curl = fluid.step()[1] - # Using the error function to make the contrast a bit higher. - # Any other sigmoid function e.g. smoothstep would work. - curl = (erf(curl * 2) + 1) / 4 - - color = np.dstack((curl, np.ones(fluid.shape), fluid.dye)) - color = (np.clip(color, 0, 1) * 255).astype('uint8') - image = np.array(color).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] - mask = image[:, :, :, 0] - masks.append(mask) - out.append(image) - - if invert: - return (1.0 - torch.cat(out, dim=0),1.0 - torch.cat(masks, dim=0),) - return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) - -class CreateAudioMask: - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "createaudiomask" - CATEGORY = "KJNodes/deprecated" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "invert": ("BOOLEAN", {"default": False}), - "frames": ("INT", {"default": 16,"min": 1, "max": 255, "step": 1}), - "scale": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 2.0, "step": 0.01}), - "audio_path": ("STRING", {"default": "audio.wav"}), - "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - }, - } - - def createaudiomask(self, frames, width, height, invert, audio_path, scale): - try: - import librosa - except ImportError: - raise Exception("Can not import librosa. Install it with 'pip install librosa'") - batch_size = frames - out = [] - masks = [] - if audio_path == "audio.wav": #I don't know why relative path won't work otherwise... - audio_path = os.path.join(script_directory, audio_path) - audio, sr = librosa.load(audio_path) - spectrogram = np.abs(librosa.stft(audio)) - - for i in range(batch_size): - image = Image.new("RGB", (width, height), "black") - draw = ImageDraw.Draw(image) - frame = spectrogram[:, i] - circle_radius = int(height * np.mean(frame)) - circle_radius *= scale - circle_center = (width // 2, height // 2) # Calculate the center of the image - - draw.ellipse([(circle_center[0] - circle_radius, circle_center[1] - circle_radius), - (circle_center[0] + circle_radius, circle_center[1] + circle_radius)], - fill='white') - - image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] - mask = image[:, :, :, 0] - masks.append(mask) - out.append(image) - - if invert: - return (1.0 - torch.cat(out, dim=0),) - return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) - -class CreateGradientMask: - - RETURN_TYPES = ("MASK",) - FUNCTION = "createmask" - CATEGORY = "KJNodes/masking/generate" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "invert": ("BOOLEAN", {"default": False}), - "frames": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), - "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - }, - } - def createmask(self, frames, width, height, invert): - # Define the number of images in the batch - batch_size = frames - out = [] - # Create an empty array to store the image batch - image_batch = np.zeros((batch_size, height, width), dtype=np.float32) - # Generate the black to white gradient for each image - for i in range(batch_size): - gradient = np.linspace(1.0, 0.0, width, dtype=np.float32) - time = i / frames # Calculate the time variable - offset_gradient = gradient - time # Offset the gradient values based on time - image_batch[i] = offset_gradient.reshape(1, -1) - output = torch.from_numpy(image_batch) - mask = output - out.append(mask) - if invert: - return (1.0 - torch.cat(out, dim=0),) - return (torch.cat(out, dim=0),) - -class CreateFadeMask: - - RETURN_TYPES = ("MASK",) - FUNCTION = "createfademask" - CATEGORY = "KJNodes/deprecated" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "invert": ("BOOLEAN", {"default": False}), - "frames": ("INT", {"default": 2,"min": 2, "max": 255, "step": 1}), - "width": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - "height": ("INT", {"default": 256,"min": 16, "max": 4096, "step": 1}), - "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out"],), - "start_level": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 1.0, "step": 0.01}), - "midpoint_level": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}), - "end_level": ("FLOAT", {"default": 0.0,"min": 0.0, "max": 1.0, "step": 0.01}), - "midpoint_frame": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), - }, - } - - def createfademask(self, frames, width, height, invert, interpolation, start_level, midpoint_level, end_level, midpoint_frame): - def ease_in(t): - return t * t - - def ease_out(t): - return 1 - (1 - t) * (1 - t) - - def ease_in_out(t): - return 3 * t * t - 2 * t * t * t - - batch_size = frames - out = [] - image_batch = np.zeros((batch_size, height, width), dtype=np.float32) - - if midpoint_frame == 0: - midpoint_frame = batch_size // 2 - - for i in range(batch_size): - if i <= midpoint_frame: - t = i / midpoint_frame - if interpolation == "ease_in": - t = ease_in(t) - elif interpolation == "ease_out": - t = ease_out(t) - elif interpolation == "ease_in_out": - t = ease_in_out(t) - color = start_level - t * (start_level - midpoint_level) - else: - t = (i - midpoint_frame) / (batch_size - midpoint_frame) - if interpolation == "ease_in": - t = ease_in(t) - elif interpolation == "ease_out": - t = ease_out(t) - elif interpolation == "ease_in_out": - t = ease_in_out(t) - color = midpoint_level - t * (midpoint_level - end_level) - - color = np.clip(color, 0, 255) - image = np.full((height, width), color, dtype=np.float32) - image_batch[i] = image - - output = torch.from_numpy(image_batch) - mask = output - out.append(mask) - - if invert: - return (1.0 - torch.cat(out, dim=0),) - return (torch.cat(out, dim=0),) - -class CreateFadeMaskAdvanced: - - RETURN_TYPES = ("MASK",) - FUNCTION = "createfademask" - CATEGORY = "KJNodes/masking/generate" - DESCRIPTION = """ -Create a batch of masks interpolated between given frames and values. -Uses same syntax as Fizz' BatchValueSchedule. -First value is the frame index (not that this starts from 0, not 1) -and the second value inside the brackets is the float value of the mask in range 0.0 - 1.0 - -For example the default values: -0:(0.0) -7:(1.0) -15:(0.0) - -Would create a mask batch fo 16 frames, starting from black, -interpolating with the chosen curve to fully white at the 8th frame, -and interpolating from that to fully black at the 16th frame. -""" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "points_string": ("STRING", {"default": "0:(0.0),\n7:(1.0),\n15:(0.0)\n", "multiline": True}), - "invert": ("BOOLEAN", {"default": False}), - "frames": ("INT", {"default": 16,"min": 2, "max": 255, "step": 1}), - "width": ("INT", {"default": 512,"min": 1, "max": 4096, "step": 1}), - "height": ("INT", {"default": 512,"min": 1, "max": 4096, "step": 1}), - "interpolation": (["linear", "ease_in", "ease_out", "ease_in_out"],), - }, - } - - def createfademask(self, frames, width, height, invert, points_string, interpolation): - def ease_in(t): - return t * t - - def ease_out(t): - return 1 - (1 - t) * (1 - t) - - def ease_in_out(t): - return 3 * t * t - 2 * t * t * t - - # Parse the input string into a list of tuples - points = [] - points_string = points_string.rstrip(',\n') - for point_str in points_string.split(','): - frame_str, color_str = point_str.split(':') - frame = int(frame_str.strip()) - color = float(color_str.strip()[1:-1]) # Remove parentheses around color - points.append((frame, color)) - - # Check if the last frame is already in the points - if len(points) == 0 or points[-1][0] != frames - 1: - # If not, add it with the color of the last specified frame - points.append((frames - 1, points[-1][1] if points else 0)) - - # Sort the points by frame number - points.sort(key=lambda x: x[0]) - - batch_size = frames - out = [] - image_batch = np.zeros((batch_size, height, width), dtype=np.float32) - - # Index of the next point to interpolate towards - next_point = 1 - - for i in range(batch_size): - while next_point < len(points) and i > points[next_point][0]: - next_point += 1 - - # Interpolate between the previous point and the next point - prev_point = next_point - 1 - t = (i - points[prev_point][0]) / (points[next_point][0] - points[prev_point][0]) - if interpolation == "ease_in": - t = ease_in(t) - elif interpolation == "ease_out": - t = ease_out(t) - elif interpolation == "ease_in_out": - t = ease_in_out(t) - elif interpolation == "linear": - pass # No need to modify `t` for linear interpolation - - color = points[prev_point][1] - t * (points[prev_point][1] - points[next_point][1]) - color = np.clip(color, 0, 255) - image = np.full((height, width), color, dtype=np.float32) - image_batch[i] = image - - output = torch.from_numpy(image_batch) - mask = output - out.append(mask) - - if invert: - return (1.0 - torch.cat(out, dim=0),) - return (torch.cat(out, dim=0),) class ScaleBatchPromptSchedule: @@ -515,266 +172,7 @@ Selects and returns the latents at the specified indices as an latent batch. samples["samples"] = chosen_latents return (samples,) -class CreateTextMask: - RETURN_TYPES = ("IMAGE", "MASK",) - FUNCTION = "createtextmask" - CATEGORY = "KJNodes/text" - DESCRIPTION = """ -Creates a text image and mask. -Looks for fonts from this folder: -ComfyUI/custom_nodes/ComfyUI-KJNodes/fonts - -If start_rotation and/or end_rotation are different values, -creates animation between them. -""" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "invert": ("BOOLEAN", {"default": False}), - "frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), - "text_x": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), - "text_y": ("INT", {"default": 0,"min": 0, "max": 4096, "step": 1}), - "font_size": ("INT", {"default": 32,"min": 8, "max": 4096, "step": 1}), - "font_color": ("STRING", {"default": "white"}), - "text": ("STRING", {"default": "HELLO!", "multiline": True}), - "font": (folder_paths.get_filename_list("kjnodes_fonts"), ), - "width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - "height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - "start_rotation": ("INT", {"default": 0,"min": 0, "max": 359, "step": 1}), - "end_rotation": ("INT", {"default": 0,"min": -359, "max": 359, "step": 1}), - }, - } - - def createtextmask(self, frames, width, height, invert, text_x, text_y, text, font_size, font_color, font, start_rotation, end_rotation): - # Define the number of images in the batch - batch_size = frames - out = [] - masks = [] - rotation = start_rotation - if start_rotation != end_rotation: - rotation_increment = (end_rotation - start_rotation) / (batch_size - 1) - - font_path = folder_paths.get_full_path("kjnodes_fonts", font) - # Generate the text - for i in range(batch_size): - image = Image.new("RGB", (width, height), "black") - draw = ImageDraw.Draw(image) - font = ImageFont.truetype(font_path, font_size) - - # Split the text into words - words = text.split() - - # Initialize variables for line creation - lines = [] - current_line = [] - current_line_width = 0 - try: #new pillow - # Iterate through words to create lines - for word in words: - word_width = font.getbbox(word)[2] - if current_line_width + word_width <= width - 2 * text_x: - current_line.append(word) - current_line_width += word_width + font.getbbox(" ")[2] # Add space width - else: - lines.append(" ".join(current_line)) - current_line = [word] - current_line_width = word_width - except: #old pillow - for word in words: - word_width = font.getsize(word)[0] - if current_line_width + word_width <= width - 2 * text_x: - current_line.append(word) - current_line_width += word_width + font.getsize(" ")[0] # Add space width - else: - lines.append(" ".join(current_line)) - current_line = [word] - current_line_width = word_width - - # Add the last line if it's not empty - if current_line: - lines.append(" ".join(current_line)) - - # Draw each line of text separately - y_offset = text_y - for line in lines: - text_width = font.getlength(line) - text_height = font_size - text_center_x = text_x + text_width / 2 - text_center_y = y_offset + text_height / 2 - try: - draw.text((text_x, y_offset), line, font=font, fill=font_color, features=['-liga']) - except: - draw.text((text_x, y_offset), line, font=font, fill=font_color) - y_offset += text_height # Move to the next line - - if start_rotation != end_rotation: - image = image.rotate(rotation, center=(text_center_x, text_center_y)) - rotation += rotation_increment - - image = np.array(image).astype(np.float32) / 255.0 - image = torch.from_numpy(image)[None,] - mask = image[:, :, :, 0] - masks.append(mask) - out.append(image) - - if invert: - return (1.0 - torch.cat(out, dim=0), 1.0 - torch.cat(masks, dim=0),) - return (torch.cat(out, dim=0),torch.cat(masks, dim=0),) - -class GrowMaskWithBlur: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "mask": ("MASK",), - "expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}), - "incremental_expandrate": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100.0, "step": 0.1}), - "tapered_corners": ("BOOLEAN", {"default": True}), - "flip_input": ("BOOLEAN", {"default": False}), - "blur_radius": ("FLOAT", { - "default": 0.0, - "min": 0.0, - "max": 100, - "step": 0.1 - }), - "lerp_alpha": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - "decay_factor": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), - }, - "optional": { - "fill_holes": ("BOOLEAN", {"default": False}), - }, - } - - CATEGORY = "KJNodes/masking" - RETURN_TYPES = ("MASK", "MASK",) - RETURN_NAMES = ("mask", "mask_inverted",) - FUNCTION = "expand_mask" - DESCRIPTION = """ -# GrowMaskWithBlur -- mask: Input mask or mask batch -- expand: Expand or contract mask or mask batch by a given amount -- incremental_expandrate: increase expand rate by a given amount per frame -- tapered_corners: use tapered corners -- flip_input: flip input mask -- blur_radius: value higher than 0 will blur the mask -- lerp_alpha: alpha value for interpolation between frames -- decay_factor: decay value for interpolation between frames -- fill_holes: fill holes in the mask (slow)""" - - def expand_mask(self, mask, expand, tapered_corners, flip_input, blur_radius, incremental_expandrate, lerp_alpha, decay_factor, fill_holes=False): - alpha = lerp_alpha - decay = decay_factor - if flip_input: - mask = 1.0 - mask - c = 0 if tapered_corners else 1 - kernel = np.array([[c, 1, c], - [1, 1, 1], - [c, 1, c]]) - growmask = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).cpu() - out = [] - previous_output = None - current_expand = expand - for m in growmask: - output = m.numpy() - for _ in range(abs(round(current_expand))): - if current_expand < 0: - output = scipy.ndimage.grey_erosion(output, footprint=kernel) - else: - output = scipy.ndimage.grey_dilation(output, footprint=kernel) - if current_expand < 0: - current_expand -= abs(incremental_expandrate) - else: - current_expand += abs(incremental_expandrate) - if fill_holes: - binary_mask = output > 0 - output = scipy.ndimage.binary_fill_holes(binary_mask) - output = output.astype(np.float32) * 255 - output = torch.from_numpy(output) - if alpha < 1.0 and previous_output is not None: - # Interpolate between the previous and current frame - output = alpha * output + (1 - alpha) * previous_output - if decay < 1.0 and previous_output is not None: - # Add the decayed previous output to the current frame - output += decay * previous_output - output = output / output.max() - previous_output = output - out.append(output) - - if blur_radius != 0: - # Convert the tensor list to PIL images, apply blur, and convert back - for idx, tensor in enumerate(out): - # Convert tensor to PIL image - pil_image = tensor2pil(tensor.cpu().detach())[0] - # Apply Gaussian blur - pil_image = pil_image.filter(ImageFilter.GaussianBlur(blur_radius)) - # Convert back to tensor - out[idx] = pil2tensor(pil_image) - blurred = torch.cat(out, dim=0) - return (blurred, 1.0 - blurred) - else: - return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) - -class ColorToMask: - - RETURN_TYPES = ("MASK",) - FUNCTION = "clip" - CATEGORY = "KJNodes/masking" - DESCRIPTION = """ -Converts chosen RGB value to a mask. -With batch inputs, the **per_batch** -controls the number of images processed at once. -""" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "images": ("IMAGE",), - "invert": ("BOOLEAN", {"default": False}), - "red": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), - "green": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), - "blue": ("INT", {"default": 0,"min": 0, "max": 255, "step": 1}), - "threshold": ("INT", {"default": 10,"min": 0, "max": 255, "step": 1}), - "per_batch": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}), - }, - } - - def clip(self, images, red, green, blue, threshold, invert, per_batch): - - color = torch.tensor([red, green, blue], dtype=torch.uint8) - black = torch.tensor([0, 0, 0], dtype=torch.uint8) - white = torch.tensor([255, 255, 255], dtype=torch.uint8) - - if invert: - black, white = white, black - - steps = images.shape[0] - pbar = ProgressBar(steps) - tensors_out = [] - - for start_idx in range(0, images.shape[0], per_batch): - - # Calculate color distances - color_distances = torch.norm(images[start_idx:start_idx+per_batch] * 255 - color, dim=-1) - - # Create a mask based on the threshold - mask = color_distances <= threshold - - # Apply the mask to create new images - mask_out = torch.where(mask.unsqueeze(-1), white, black).float() - mask_out = mask_out.mean(dim=-1) - - tensors_out.append(mask_out.cpu()) - batch_count = mask_out.shape[0] - pbar.update(batch_count) - - tensors_out = torch.cat(tensors_out, dim=0) - tensors_out = torch.clamp(tensors_out, min=0.0, max=1.0) - return tensors_out, - class ConditioningMultiCombine: @classmethod def INPUT_TYPES(s): @@ -804,37 +202,6 @@ Combines multiple conditioning nodes into one return (cond, inputcount,) - -class MaskBatchMulti: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "inputcount": ("INT", {"default": 2, "min": 2, "max": 1000, "step": 1}), - "mask_1": ("MASK", ), - "mask_2": ("MASK", ), - }, - } - - RETURN_TYPES = ("MASK",) - RETURN_NAMES = ("masks",) - FUNCTION = "combine" - CATEGORY = "KJNodes/masking" - DESCRIPTION = """ -Creates an image batch from multiple masks. -You can set how many inputs the node has, -with the **inputcount** and clicking update. -""" - - def combine(self, inputcount, **kwargs): - mask = kwargs["mask_1"] - for c in range(1, inputcount): - new_mask = kwargs[f"mask_{c + 1}"] - if mask.shape[1:] != new_mask.shape[1:]: - new_mask = F.interpolate(new_mask.unsqueeze(1), size=(mask.shape[1], mask.shape[2]), mode="bicubic").squeeze(1) - mask = torch.cat((mask, new_mask), dim=0) - return (mask,) - class JoinStrings: @classmethod def INPUT_TYPES(cls): @@ -1311,264 +678,6 @@ class EmptyLatentImagePresets: return (latent, int(width), int(height),) -class BatchCLIPSeg: - - def __init__(self): - pass - - @classmethod - def INPUT_TYPES(s): - - return {"required": - { - "images": ("IMAGE",), - "text": ("STRING", {"multiline": False}), - "threshold": ("FLOAT", {"default": 0.1,"min": 0.0, "max": 10.0, "step": 0.001}), - "binary_mask": ("BOOLEAN", {"default": True}), - "combine_mask": ("BOOLEAN", {"default": False}), - "use_cuda": ("BOOLEAN", {"default": True}), - }, - } - - CATEGORY = "KJNodes/masking" - RETURN_TYPES = ("MASK",) - RETURN_NAMES = ("Mask",) - FUNCTION = "segment_image" - DESCRIPTION = """ -Segments an image or batch of images using CLIPSeg. -""" - - def segment_image(self, images, text, threshold, binary_mask, combine_mask, use_cuda): - from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation - out = [] - height, width, _ = images[0].shape - if use_cuda and torch.cuda.is_available(): - device = torch.device("cuda") - else: - device = torch.device("cpu") - dtype = model_management.unet_dtype() - model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") - model.to(dtype) - model.to(device) - images = images.to(device) - processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") - pbar = ProgressBar(images.shape[0]) - autocast_condition = (dtype != torch.float32) and not model_management.is_device_mps(device) - with torch.autocast(model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): - for image in images: - image = (image* 255).type(torch.uint8) - prompt = text - input_prc = processor(text=prompt, images=image, return_tensors="pt") - # Move the processed input to the device - for key in input_prc: - input_prc[key] = input_prc[key].to(device) - - outputs = model(**input_prc) - - tensor = torch.sigmoid(outputs[0]) - tensor_thresholded = torch.where(tensor > threshold, tensor, torch.tensor(0, dtype=torch.float)) - tensor_normalized = (tensor_thresholded - tensor_thresholded.min()) / (tensor_thresholded.max() - tensor_thresholded.min()) - tensor = tensor_normalized - - # Resize the mask - if len(tensor.shape) == 3: - tensor = tensor.unsqueeze(0) - resized_tensor = F.interpolate(tensor, size=(height, width), mode='nearest') - - # Remove the extra dimensions - resized_tensor = resized_tensor[0, 0, :, :] - pbar.update(1) - out.append(resized_tensor) - - results = torch.stack(out).cpu().float() - - if combine_mask: - combined_results = torch.max(results, dim=0)[0] - results = combined_results.unsqueeze(0).repeat(len(images),1,1) - - if binary_mask: - results = results.round() - - return results, - -class GetMaskSize: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "mask": ("MASK",), - }} - - RETURN_TYPES = ("MASK","INT", "INT", ) - RETURN_NAMES = ("mask", "width", "height",) - FUNCTION = "getsize" - CATEGORY = "KJNodes/masking" - DESCRIPTION = """ -Returns the width and height of the mask, -and passes through the mask unchanged. - -""" - - def getsize(self, mask): - width = mask.shape[2] - height = mask.shape[1] - return (mask, width, height,) - -class RoundMask: - @classmethod - def INPUT_TYPES(s): - return {"required": { - "mask": ("MASK",), - }} - - RETURN_TYPES = ("MASK",) - FUNCTION = "round" - CATEGORY = "KJNodes/masking" - DESCRIPTION = """ -Rounds the mask or batch of masks to a binary mask. -RoundMask example - -""" - - def round(self, mask): - mask = mask.round() - return (mask,) - -class ResizeMask: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "mask": ("MASK",), - "width": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, "display": "number" }), - "height": ("INT", { "default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8, "display": "number" }), - "keep_proportions": ("BOOLEAN", { "default": False }), - } - } - - RETURN_TYPES = ("MASK", "INT", "INT",) - RETURN_NAMES = ("mask", "width", "height",) - FUNCTION = "resize" - CATEGORY = "KJNodes/masking" - DESCRIPTION = """ -Resizes the mask or batch of masks to the specified width and height. -""" - - def resize(self, mask, width, height, keep_proportions): - if keep_proportions: - _, oh, ow, _ = mask.shape - width = ow if width == 0 else width - height = oh if height == 0 else height - ratio = min(width / ow, height / oh) - width = round(ow*ratio) - height = round(oh*ratio) - - outputs = mask.unsqueeze(0) # Add an extra dimension for batch size - outputs = F.interpolate(outputs, size=(height, width), mode="nearest") - outputs = outputs.squeeze(0) # Remove the extra dimension after interpolation - - return(outputs, outputs.shape[2], outputs.shape[1],) - -class OffsetMask: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "mask": ("MASK",), - "x": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), - "y": ("INT", { "default": 0, "min": -4096, "max": MAX_RESOLUTION, "step": 1, "display": "number" }), - "angle": ("INT", { "default": 0, "min": -360, "max": 360, "step": 1, "display": "number" }), - "duplication_factor": ("INT", { "default": 1, "min": 1, "max": 1000, "step": 1, "display": "number" }), - "roll": ("BOOLEAN", { "default": False }), - "incremental": ("BOOLEAN", { "default": False }), - "padding_mode": ( - [ - 'empty', - 'border', - 'reflection', - - ], { - "default": 'empty' - }), - } - } - - RETURN_TYPES = ("MASK",) - RETURN_NAMES = ("mask",) - FUNCTION = "offset" - CATEGORY = "KJNodes/masking" - DESCRIPTION = """ -Offsets the mask by the specified amount. - - mask: Input mask or mask batch - - x: Horizontal offset - - y: Vertical offset - - angle: Angle in degrees - - roll: roll edge wrapping - - duplication_factor: Number of times to duplicate the mask to form a batch - - border padding_mode: Padding mode for the mask -""" - - def offset(self, mask, x, y, angle, roll=False, incremental=False, duplication_factor=1, padding_mode="empty"): - # Create duplicates of the mask batch - mask = mask.repeat(duplication_factor, 1, 1).clone() - - batch_size, height, width = mask.shape - - if angle != 0 and incremental: - for i in range(batch_size): - rotation_angle = angle * (i+1) - mask[i] = TF.rotate(mask[i].unsqueeze(0), rotation_angle).squeeze(0) - elif angle > 0: - for i in range(batch_size): - mask[i] = TF.rotate(mask[i].unsqueeze(0), angle).squeeze(0) - - if roll: - if incremental: - for i in range(batch_size): - shift_x = min(x*(i+1), width-1) - shift_y = min(y*(i+1), height-1) - if shift_x != 0: - mask[i] = torch.roll(mask[i], shifts=shift_x, dims=1) - if shift_y != 0: - mask[i] = torch.roll(mask[i], shifts=shift_y, dims=0) - else: - shift_x = min(x, width-1) - shift_y = min(y, height-1) - if shift_x != 0: - mask = torch.roll(mask, shifts=shift_x, dims=2) - if shift_y != 0: - mask = torch.roll(mask, shifts=shift_y, dims=1) - else: - - for i in range(batch_size): - if incremental: - temp_x = min(x * (i+1), width-1) - temp_y = min(y * (i+1), height-1) - else: - temp_x = min(x, width-1) - temp_y = min(y, height-1) - if temp_x > 0: - if padding_mode == 'empty': - mask[i] = torch.cat([torch.zeros((height, temp_x)), mask[i, :, :-temp_x]], dim=1) - elif padding_mode in ['replicate', 'reflect']: - mask[i] = F.pad(mask[i, :, :-temp_x], (0, temp_x), mode=padding_mode) - elif temp_x < 0: - if padding_mode == 'empty': - mask[i] = torch.cat([mask[i, :, :temp_x], torch.zeros((height, -temp_x))], dim=1) - elif padding_mode in ['replicate', 'reflect']: - mask[i] = F.pad(mask[i, :, -temp_x:], (temp_x, 0), mode=padding_mode) - - if temp_y > 0: - if padding_mode == 'empty': - mask[i] = torch.cat([torch.zeros((temp_y, width)), mask[i, :-temp_y, :]], dim=0) - elif padding_mode in ['replicate', 'reflect']: - mask[i] = F.pad(mask[i, :-temp_y, :], (0, temp_y), mode=padding_mode) - elif temp_y < 0: - if padding_mode == 'empty': - mask[i] = torch.cat([mask[i, :temp_y, :], torch.zeros((-temp_y, width))], dim=0) - elif padding_mode in ['replicate', 'reflect']: - mask[i] = F.pad(mask[i, -temp_y:, :], (temp_y, 0), mode=padding_mode) - - return mask, class WidgetToString: @@ -1621,228 +730,6 @@ To see node id's, enable node id display from Manager badge menu. raise NameError(f"Node not found: {id}") return (', '.join(results).strip(', '), ) -class CreateShapeMask: - - RETURN_TYPES = ("MASK", "MASK",) - RETURN_NAMES = ("mask", "mask_inverted",) - FUNCTION = "createshapemask" - CATEGORY = "KJNodes/masking/generate" - DESCRIPTION = """ -Creates a mask or batch of masks with the specified shape. -Locations are center locations. -Grow value is the amount to grow the shape on each frame, creating animated masks. -""" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "shape": ( - [ 'circle', - 'square', - 'triangle', - ], - { - "default": 'circle' - }), - "frames": ("INT", {"default": 1,"min": 1, "max": 4096, "step": 1}), - "location_x": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), - "location_y": ("INT", {"default": 256,"min": 0, "max": 4096, "step": 1}), - "grow": ("INT", {"default": 0, "min": -512, "max": 512, "step": 1}), - "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - "shape_width": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), - "shape_height": ("INT", {"default": 128,"min": 8, "max": 4096, "step": 1}), - }, - } - - def createshapemask(self, frames, frame_width, frame_height, location_x, location_y, shape_width, shape_height, grow, shape): - # Define the number of images in the batch - batch_size = frames - out = [] - color = "white" - for i in range(batch_size): - image = Image.new("RGB", (frame_width, frame_height), "black") - draw = ImageDraw.Draw(image) - - # Calculate the size for this frame and ensure it's not less than 0 - current_width = max(0, shape_width + i*grow) - current_height = max(0, shape_height + i*grow) - - if shape == 'circle' or shape == 'square': - # Define the bounding box for the shape - left_up_point = (location_x - current_width // 2, location_y - current_height // 2) - right_down_point = (location_x + current_width // 2, location_y + current_height // 2) - two_points = [left_up_point, right_down_point] - - if shape == 'circle': - draw.ellipse(two_points, fill=color) - elif shape == 'square': - draw.rectangle(two_points, fill=color) - - elif shape == 'triangle': - # Define the points for the triangle - left_up_point = (location_x - current_width // 2, location_y + current_height // 2) # bottom left - right_down_point = (location_x + current_width // 2, location_y + current_height // 2) # bottom right - top_point = (location_x, location_y - current_height // 2) # top point - draw.polygon([top_point, left_up_point, right_down_point], fill=color) - - image = pil2tensor(image) - mask = image[:, :, :, 0] - out.append(mask) - outstack = torch.cat(out, dim=0) - return (outstack, 1.0 - outstack,) - -class CreateVoronoiMask: - - RETURN_TYPES = ("MASK", "MASK",) - RETURN_NAMES = ("mask", "mask_inverted",) - FUNCTION = "createvoronoi" - CATEGORY = "KJNodes/masking/generate" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}), - "num_points": ("INT", {"default": 15,"min": 1, "max": 4096, "step": 1}), - "line_width": ("INT", {"default": 4,"min": 1, "max": 4096, "step": 1}), - "speed": ("FLOAT", {"default": 0.5,"min": 0.0, "max": 1.0, "step": 0.01}), - "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - }, - } - - def createvoronoi(self, frames, num_points, line_width, speed, frame_width, frame_height): - from scipy.spatial import Voronoi - # Define the number of images in the batch - batch_size = frames - out = [] - - # Calculate aspect ratio - aspect_ratio = frame_width / frame_height - - # Create start and end points for each point, considering the aspect ratio - start_points = np.random.rand(num_points, 2) - start_points[:, 0] *= aspect_ratio - - end_points = np.random.rand(num_points, 2) - end_points[:, 0] *= aspect_ratio - - for i in range(batch_size): - # Interpolate the points' positions based on the current frame - t = (i * speed) / (batch_size - 1) # normalize to [0, 1] over the frames - t = np.clip(t, 0, 1) # ensure t is in [0, 1] - points = (1 - t) * start_points + t * end_points # lerp - - # Adjust points for aspect ratio - points[:, 0] *= aspect_ratio - - vor = Voronoi(points) - - # Create a blank image with a white background - fig, ax = plt.subplots() - plt.subplots_adjust(left=0, right=1, bottom=0, top=1) - ax.set_xlim([0, aspect_ratio]); ax.set_ylim([0, 1]) # adjust x limits - ax.axis('off') - ax.margins(0, 0) - fig.set_size_inches(aspect_ratio * frame_height/100, frame_height/100) # adjust figure size - ax.fill_between([0, 1], [0, 1], color='white') - - # Plot each Voronoi ridge - for simplex in vor.ridge_vertices: - simplex = np.asarray(simplex) - if np.all(simplex >= 0): - plt.plot(vor.vertices[simplex, 0], vor.vertices[simplex, 1], 'k-', linewidth=line_width) - - fig.canvas.draw() - img = np.array(fig.canvas.renderer._renderer) - - plt.close(fig) - - pil_img = Image.fromarray(img).convert("L") - mask = torch.tensor(np.array(pil_img)) / 255.0 - - out.append(mask) - - return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) - -class CreateMagicMask: - - RETURN_TYPES = ("MASK", "MASK",) - RETURN_NAMES = ("mask", "mask_inverted",) - FUNCTION = "createmagicmask" - CATEGORY = "KJNodes/masking/generate" - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "frames": ("INT", {"default": 16,"min": 2, "max": 4096, "step": 1}), - "depth": ("INT", {"default": 12,"min": 1, "max": 500, "step": 1}), - "distortion": ("FLOAT", {"default": 1.5,"min": 0.0, "max": 100.0, "step": 0.01}), - "seed": ("INT", {"default": 123,"min": 0, "max": 99999999, "step": 1}), - "transitions": ("INT", {"default": 1,"min": 1, "max": 20, "step": 1}), - "frame_width": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - "frame_height": ("INT", {"default": 512,"min": 16, "max": 4096, "step": 1}), - }, - } - - def createmagicmask(self, frames, transitions, depth, distortion, seed, frame_width, frame_height): - from ..utility.magictex import coordinate_grid, random_transform, magic - rng = np.random.default_rng(seed) - out = [] - coords = coordinate_grid((frame_width, frame_height)) - - # Calculate the number of frames for each transition - frames_per_transition = frames // transitions - - # Generate a base set of parameters - base_params = { - "coords": random_transform(coords, rng), - "depth": depth, - "distortion": distortion, - } - for t in range(transitions): - # Generate a second set of parameters that is at most max_diff away from the base parameters - params1 = base_params.copy() - params2 = base_params.copy() - - params1['coords'] = random_transform(coords, rng) - params2['coords'] = random_transform(coords, rng) - - for i in range(frames_per_transition): - # Compute the interpolation factor - alpha = i / frames_per_transition - - # Interpolate between the two sets of parameters - params = params1.copy() - params['coords'] = (1 - alpha) * params1['coords'] + alpha * params2['coords'] - - tex = magic(**params) - - dpi = frame_width / 10 - fig = plt.figure(figsize=(10, 10), dpi=dpi) - - ax = fig.add_subplot(111) - plt.subplots_adjust(left=0, right=1, bottom=0, top=1) - - ax.get_yaxis().set_ticks([]) - ax.get_xaxis().set_ticks([]) - ax.imshow(tex, aspect='auto') - - fig.canvas.draw() - img = np.array(fig.canvas.renderer._renderer) - - plt.close(fig) - - pil_img = Image.fromarray(img).convert("L") - mask = torch.tensor(np.array(pil_img)) / 255.0 - - out.append(mask) - - return (torch.stack(out, dim=0), 1.0 - torch.stack(out, dim=0),) - class BboxToInt: @classmethod @@ -2432,42 +1319,6 @@ https://huggingface.co/stabilityai/sv3d latent = torch.zeros([batch_size, 4, height // 8, width // 8]) return (final_positive, final_negative, {"samples": latent}) -class RemapMaskRange: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "mask": ("MASK",), - "min": ("FLOAT", {"default": 0.0,"min": -10.0, "max": 1.0, "step": 0.01}), - "max": ("FLOAT", {"default": 1.0,"min": 0.0, "max": 10.0, "step": 0.01}), - } - } - - RETURN_TYPES = ("MASK",) - RETURN_NAMES = ("mask",) - FUNCTION = "remap" - CATEGORY = "KJNodes/masking" - DESCRIPTION = """ -Sets new min and max values for the mask. -""" - - def remap(self, mask, min, max): - - # Find the maximum value in the mask - mask_max = torch.max(mask) - - # If the maximum mask value is zero, avoid division by zero by setting it to 1 - mask_max = mask_max if mask_max > 0 else 1 - - # Scale the mask values to the new range defined by min and max - # The highest pixel value in the mask will be scaled to max - scaled_mask = (mask / mask_max) * (max - min) + min - - # Clamp the values to ensure they are within [0.0, 1.0] - scaled_mask = torch.clamp(scaled_mask, min=0.0, max=1.0) - - return (scaled_mask, ) - class LoadResAdapterNormalization: @classmethod def INPUT_TYPES(s): @@ -2553,7 +1404,6 @@ https://huggingface.co/roborovski/superprompt-v1 return (out, ) - class CameraPoseVisualizer: @classmethod