| | import torch
|
| | from .tools import VariantSupport
|
| | from comfy_execution.graph_utils import GraphBuilder
|
| |
|
| | class TestLazyMixImages:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "image1": ("IMAGE",{"lazy": True}),
|
| | "image2": ("IMAGE",{"lazy": True}),
|
| | "mask": ("MASK",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "mix"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def check_lazy_status(self, mask, image1, image2):
|
| | mask_min = mask.min()
|
| | mask_max = mask.max()
|
| | needed = []
|
| | if image1 is None and (mask_min != 1.0 or mask_max != 1.0):
|
| | needed.append("image1")
|
| | if image2 is None and (mask_min != 0.0 or mask_max != 0.0):
|
| | needed.append("image2")
|
| | return needed
|
| |
|
| |
|
| | def mix(self, mask, image1, image2):
|
| | mask_min = mask.min()
|
| | mask_max = mask.max()
|
| | if mask_min == 0.0 and mask_max == 0.0:
|
| | return (image1,)
|
| | elif mask_min == 1.0 and mask_max == 1.0:
|
| | return (image2,)
|
| |
|
| | if len(mask.shape) == 2:
|
| | mask = mask.unsqueeze(0)
|
| | if len(mask.shape) == 3:
|
| | mask = mask.unsqueeze(3)
|
| | if mask.shape[3] < image1.shape[3]:
|
| | mask = mask.repeat(1, 1, 1, image1.shape[3])
|
| |
|
| | result = image1 * (1. - mask) + image2 * mask,
|
| | return (result[0],)
|
| |
|
| | class TestVariadicAverage:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("IMAGE",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "variadic_average"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def variadic_average(self, input1, **kwargs):
|
| | inputs = [input1]
|
| | while 'input' + str(len(inputs) + 1) in kwargs:
|
| | inputs.append(kwargs['input' + str(len(inputs) + 1)])
|
| | return (torch.stack(inputs).mean(dim=0),)
|
| |
|
| |
|
| | class TestCustomIsChanged:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "image": ("IMAGE",),
|
| | },
|
| | "optional": {
|
| | "should_change": ("BOOL", {"default": False}),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "custom_is_changed"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def custom_is_changed(self, image, should_change=False):
|
| | return (image,)
|
| |
|
| | @classmethod
|
| | def IS_CHANGED(cls, should_change=False, *args, **kwargs):
|
| | if should_change:
|
| | return float("NaN")
|
| | else:
|
| | return False
|
| |
|
| | class TestIsChangedWithConstants:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "image": ("IMAGE",),
|
| | "value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0}),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "custom_is_changed"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def custom_is_changed(self, image, value):
|
| | return (image * value,)
|
| |
|
| | @classmethod
|
| | def IS_CHANGED(cls, image, value):
|
| | if image is None:
|
| | return value
|
| | else:
|
| | return image.mean().item() * value
|
| |
|
| | class TestCustomValidation1:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("IMAGE,FLOAT",),
|
| | "input2": ("IMAGE,FLOAT",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "custom_validation1"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def custom_validation1(self, input1, input2):
|
| | if isinstance(input1, float) and isinstance(input2, float):
|
| | result = torch.ones([1, 512, 512, 3]) * input1 * input2
|
| | else:
|
| | result = input1 * input2
|
| | return (result,)
|
| |
|
| | @classmethod
|
| | def VALIDATE_INPUTS(cls, input1=None, input2=None):
|
| | if input1 is not None:
|
| | if not isinstance(input1, (torch.Tensor, float)):
|
| | return f"Invalid type of input1: {type(input1)}"
|
| | if input2 is not None:
|
| | if not isinstance(input2, (torch.Tensor, float)):
|
| | return f"Invalid type of input2: {type(input2)}"
|
| |
|
| | return True
|
| |
|
| | class TestCustomValidation2:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("IMAGE,FLOAT",),
|
| | "input2": ("IMAGE,FLOAT",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "custom_validation2"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def custom_validation2(self, input1, input2):
|
| | if isinstance(input1, float) and isinstance(input2, float):
|
| | result = torch.ones([1, 512, 512, 3]) * input1 * input2
|
| | else:
|
| | result = input1 * input2
|
| | return (result,)
|
| |
|
| | @classmethod
|
| | def VALIDATE_INPUTS(cls, input_types, input1=None, input2=None):
|
| | if input1 is not None:
|
| | if not isinstance(input1, (torch.Tensor, float)):
|
| | return f"Invalid type of input1: {type(input1)}"
|
| | if input2 is not None:
|
| | if not isinstance(input2, (torch.Tensor, float)):
|
| | return f"Invalid type of input2: {type(input2)}"
|
| |
|
| | if 'input1' in input_types:
|
| | if input_types['input1'] not in ["IMAGE", "FLOAT"]:
|
| | return f"Invalid type of input1: {input_types['input1']}"
|
| | if 'input2' in input_types:
|
| | if input_types['input2'] not in ["IMAGE", "FLOAT"]:
|
| | return f"Invalid type of input2: {input_types['input2']}"
|
| |
|
| | return True
|
| |
|
| | @VariantSupport()
|
| | class TestCustomValidation3:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("IMAGE,FLOAT",),
|
| | "input2": ("IMAGE,FLOAT",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "custom_validation3"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def custom_validation3(self, input1, input2):
|
| | if isinstance(input1, float) and isinstance(input2, float):
|
| | result = torch.ones([1, 512, 512, 3]) * input1 * input2
|
| | else:
|
| | result = input1 * input2
|
| | return (result,)
|
| |
|
| | class TestCustomValidation4:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("FLOAT",),
|
| | "input2": ("FLOAT",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "custom_validation4"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def custom_validation4(self, input1, input2):
|
| | result = torch.ones([1, 512, 512, 3]) * input1 * input2
|
| | return (result,)
|
| |
|
| | @classmethod
|
| | def VALIDATE_INPUTS(cls, input1, input2):
|
| | if input1 is not None:
|
| | if not isinstance(input1, float):
|
| | return f"Invalid type of input1: {type(input1)}"
|
| | if input2 is not None:
|
| | if not isinstance(input2, float):
|
| | return f"Invalid type of input2: {type(input2)}"
|
| |
|
| | return True
|
| |
|
| | class TestCustomValidation5:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("FLOAT", {"min": 0.0, "max": 1.0}),
|
| | "input2": ("FLOAT", {"min": 0.0, "max": 1.0}),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "custom_validation5"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def custom_validation5(self, input1, input2):
|
| | value = input1 * input2
|
| | return (torch.ones([1, 512, 512, 3]) * value,)
|
| |
|
| | @classmethod
|
| | def VALIDATE_INPUTS(cls, **kwargs):
|
| | if kwargs['input2'] == 7.0:
|
| | return "7s are not allowed. I've never liked 7s."
|
| | return True
|
| |
|
| | class TestDynamicDependencyCycle:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("IMAGE",),
|
| | "input2": ("IMAGE",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE",)
|
| | FUNCTION = "dynamic_dependency_cycle"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def dynamic_dependency_cycle(self, input1, input2):
|
| | g = GraphBuilder()
|
| | mask = g.node("StubMask", value=0.5, height=512, width=512, batch_size=1)
|
| | mix1 = g.node("TestLazyMixImages", image1=input1, mask=mask.out(0))
|
| | mix2 = g.node("TestLazyMixImages", image1=mix1.out(0), image2=input2, mask=mask.out(0))
|
| |
|
| |
|
| | mix1.set_input("image2", mix2.out(0))
|
| |
|
| | return {
|
| | "result": (mix2.out(0),),
|
| | "expand": g.finalize(),
|
| | }
|
| |
|
| | class TestMixedExpansionReturns:
|
| | @classmethod
|
| | def INPUT_TYPES(cls):
|
| | return {
|
| | "required": {
|
| | "input1": ("FLOAT",),
|
| | },
|
| | }
|
| |
|
| | RETURN_TYPES = ("IMAGE","IMAGE")
|
| | FUNCTION = "mixed_expansion_returns"
|
| |
|
| | CATEGORY = "Testing/Nodes"
|
| |
|
| | def mixed_expansion_returns(self, input1):
|
| | white_image = torch.ones([1, 512, 512, 3])
|
| | if input1 <= 0.1:
|
| | return (torch.ones([1, 512, 512, 3]) * 0.1, white_image)
|
| | elif input1 <= 0.2:
|
| | return {
|
| | "result": (torch.ones([1, 512, 512, 3]) * 0.2, white_image),
|
| | }
|
| | else:
|
| | g = GraphBuilder()
|
| | mask = g.node("StubMask", value=0.3, height=512, width=512, batch_size=1)
|
| | black = g.node("StubImage", content="BLACK", height=512, width=512, batch_size=1)
|
| | white = g.node("StubImage", content="WHITE", height=512, width=512, batch_size=1)
|
| | mix = g.node("TestLazyMixImages", image1=black.out(0), image2=white.out(0), mask=mask.out(0))
|
| | return {
|
| | "result": (mix.out(0), white_image),
|
| | "expand": g.finalize(),
|
| | }
|
| |
|
| | TEST_NODE_CLASS_MAPPINGS = {
|
| | "TestLazyMixImages": TestLazyMixImages,
|
| | "TestVariadicAverage": TestVariadicAverage,
|
| | "TestCustomIsChanged": TestCustomIsChanged,
|
| | "TestIsChangedWithConstants": TestIsChangedWithConstants,
|
| | "TestCustomValidation1": TestCustomValidation1,
|
| | "TestCustomValidation2": TestCustomValidation2,
|
| | "TestCustomValidation3": TestCustomValidation3,
|
| | "TestCustomValidation4": TestCustomValidation4,
|
| | "TestCustomValidation5": TestCustomValidation5,
|
| | "TestDynamicDependencyCycle": TestDynamicDependencyCycle,
|
| | "TestMixedExpansionReturns": TestMixedExpansionReturns,
|
| | }
|
| |
|
| | TEST_NODE_DISPLAY_NAME_MAPPINGS = {
|
| | "TestLazyMixImages": "Lazy Mix Images",
|
| | "TestVariadicAverage": "Variadic Average",
|
| | "TestCustomIsChanged": "Custom IsChanged",
|
| | "TestIsChangedWithConstants": "IsChanged With Constants",
|
| | "TestCustomValidation1": "Custom Validation 1",
|
| | "TestCustomValidation2": "Custom Validation 2",
|
| | "TestCustomValidation3": "Custom Validation 3",
|
| | "TestCustomValidation4": "Custom Validation 4",
|
| | "TestCustomValidation5": "Custom Validation 5",
|
| | "TestDynamicDependencyCycle": "Dynamic Dependency Cycle",
|
| | "TestMixedExpansionReturns": "Mixed Expansion Returns",
|
| | }
|
| |
|
