mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-11 23:00:51 +08:00
6d09dd70f8
Additionally, if `VALIDATE_INPUTS` takes an argument named `input_types`, that variable will be a dictionary of the socket type of all incoming connections. If that argument exists, normal socket type validation will not occur. This removes the last hurdle for enabling variant types entirely from custom nodes, so I've removed that command-line option. I've added appropriate unit tests for these changes.
106 lines
2.8 KiB
Python
106 lines
2.8 KiB
Python
import torch
|
|
|
|
class StubImage:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"content": (['WHITE', 'BLACK', 'NOISE'],),
|
|
"height": ("INT", {"default": 512, "min": 1, "max": 1024 ** 3, "step": 1}),
|
|
"width": ("INT", {"default": 512, "min": 1, "max": 4096 ** 3, "step": 1}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1024 ** 3, "step": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "stub_image"
|
|
|
|
CATEGORY = "Testing/Stub Nodes"
|
|
|
|
def stub_image(self, content, height, width, batch_size):
|
|
if content == "WHITE":
|
|
return (torch.ones(batch_size, height, width, 3),)
|
|
elif content == "BLACK":
|
|
return (torch.zeros(batch_size, height, width, 3),)
|
|
elif content == "NOISE":
|
|
return (torch.rand(batch_size, height, width, 3),)
|
|
|
|
class StubMask:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
"height": ("INT", {"default": 512, "min": 1, "max": 1024 ** 3, "step": 1}),
|
|
"width": ("INT", {"default": 512, "min": 1, "max": 4096 ** 3, "step": 1}),
|
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1024 ** 3, "step": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
FUNCTION = "stub_mask"
|
|
|
|
CATEGORY = "Testing/Stub Nodes"
|
|
|
|
def stub_mask(self, value, height, width, batch_size):
|
|
return (torch.ones(batch_size, height, width) * value,)
|
|
|
|
class StubInt:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("INT", {"default": 0, "min": -0xffffffff, "max": 0xffffffff, "step": 1}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
FUNCTION = "stub_int"
|
|
|
|
CATEGORY = "Testing/Stub Nodes"
|
|
|
|
def stub_int(self, value):
|
|
return (value,)
|
|
|
|
class StubFloat:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"value": ("FLOAT", {"default": 0.0, "min": -1.0e38, "max": 1.0e38, "step": 0.01}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
FUNCTION = "stub_float"
|
|
|
|
CATEGORY = "Testing/Stub Nodes"
|
|
|
|
def stub_float(self, value):
|
|
return (value,)
|
|
|
|
TEST_STUB_NODE_CLASS_MAPPINGS = {
|
|
"StubImage": StubImage,
|
|
"StubMask": StubMask,
|
|
"StubInt": StubInt,
|
|
"StubFloat": StubFloat,
|
|
}
|
|
TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"StubImage": "Stub Image",
|
|
"StubMask": "Stub Mask",
|
|
"StubInt": "Stub Int",
|
|
"StubFloat": "Stub Float",
|
|
}
|