xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-12 22:17:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update execution.py to allow validation function injection

Browse Source
This commit is contained in:
shawnington 2024-08-03 10:33:11 -04:00 committed by GitHub
parent 593d69deef
commit 6b362724ab
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 66 additions and 42 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

98
execution.py
View File

@ -16,61 +16,82 @@ import comfy.model_management
from comfy.cli_args import args from comfy.cli_args import args
def validate_image_shape(image): import sys
transforms = { import copy
"HW": lambda t: t.unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3), import logging
"BHW": lambda t: t.unsqueeze(-1).expand(-1, -1, -1, 3), import threading
"HWC": lambda t: t.unsqueeze(0), import heapq
} import time
import traceback
import inspect
from typing import List, Literal, NamedTuple, Optional
import torch
import nodes
import comfy.model_management
from comfy.cli_args import args
def validate_image_shape(image, node_type, output_name):
image_len = len(image.shape) image_len = len(image.shape)
if image_len == 2: def type_check():
#HW -> add Batch and Channel dimensions print(f"Checking Image Shape: {image.shape}")
image = transforms["HW"](image) if image_len != 4 or image.shape[-1] > 4 or image.shape[-1] < 3:
return image if "error" == args.type_conformance:
raise ValueError(f"Image Shape Error: Node: {node_type}, Output: {output_name}, [{image.shape}] does not match the expected RGB format: torch.Size[B, H, W, 3] or the RGBA format: torch.Size[B, H, W, 4]")
else:
print(f"Image Shape Error: Node: {node_type}, Output: {output_name}, [{image.shape}] does not match the expected RGB format: torch.Size[B, H, W, 3] or the RGBA format: torch.Size[B, H, W, 4]")
return False
return True
if image_len == 3: if type_check() != True and args.type_conformance in ["all", "images"]:
if image.shape[-1] > 4 or image.shape[-1] < 3: transforms = {
#BHW -> add Channel dimension "HW": lambda t: t.unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3),
image = transforms["BHW"](image) "BHW": lambda t: t.unsqueeze(-1).expand(-1, -1, -1, 3),
"HWC": lambda t: t.unsqueeze(0),
}
if image_len == 2:
#HW -> add Batch and Channel dimensions
image = transforms["HW"](image)
return image return image
if image.shape[-1] == 3 or image.shape[-1] == 4: if image_len == 3:
#HW3 or HW4 -> add Batch dimension --- can misbehave in edge cases where image is BHW but W is 3 or 4 px if image.shape[-1] > 4 or image.shape[-1] < 3:
image = transforms["HWC"](image) #BHW -> add Channel dimension
return image image = transforms["BHW"](image)
return image
if image.shape[-1] == 3 or image.shape[-1] == 4:
#HW3 or HW4 -> add Batch dimension --- can misbehave in edge cases where image is BHW but W is 3 or 4 px
image = transforms["HWC"](image)
return image
return image return image
def generate_type_validator(valid_types, expected_type=None): def generate_type_validator(valid_types, validator=None):
def validate_type(node_type, output_name, value): def validate_type(node_type, output_name, value):
print(type(validator))
if not isinstance(value, valid_types) and expected_type is not "IMAGE": if not isinstance(value, valid_types) and validator is None:
print(f"The value is not of the valid types. Expected one of {[t.__name__ for t in valid_types]}, got {[type(value)]}") print(f"TypeError in Node: {node_type}. Expected {output_name} to be one of type(s): {[t.__name__ for t in valid_types]}, got {[type(value).__name__]} instead")
if "all" == args.type_conformance: if "all" == args.type_conformance:
return valid_types[0](value) return valid_types[0](value)
if "error" == args.type_conformance: if "error" == args.type_conformance:
raise TypeError(f"The value is not of the valid types. Expected output: {output_name}, to be of type(s): {[vtype.__name__ for vtype in valid_types]}, got {[type(value).__name__]} instead") raise TypeError(f"Expected: [{output_name}], to be one of type(s): {[vtype.__name__ for vtype in valid_types]}, got {[type(value).__name__]} instead")
return value return value
if isinstance(value, valid_types) and expected_type is "IMAGE": if isinstance(value, valid_types) and validator is not None:
if args.type_conformance in ["all", "images"]: print(f"Validating {output_name} in Node: {node_type}")
return validate_image_shape(value) value = validator(value, node_type, output_name)
if len(value.shape) != 4 or value.shape[-1] > 4 or value.shape[-1] < 3:
print(f"Image Shape Error: Node: {node_type}, Output: {output_name}, [{value.shape}] does not match the expected RGB format: torch.Size[B, H, W, 3] or the RGBA format: torch.Size[B, H, W, 4]")
if "error" == args.type_conformance:
raise ValueError(f"Image Shape Error: Node: {node_type}, Output: {output_name}, [{value.shape}] does not match the expected RGB format: torch.Size[B, H, W, 3] or the RGBA format: torch.Size[B, H, W, 4]")
return value
return value return value
return value return value
return validate_type return validate_type
@ -78,7 +99,7 @@ def generate_type_validator(valid_types, expected_type=None):
def input_validation(input_data_all, obj): def input_validation(input_data_all, obj):
validation_funcs = { validation_funcs = {
"IMAGE": generate_type_validator((torch.Tensor,), "IMAGE"), "IMAGE": generate_type_validator((torch.Tensor,), validator=validate_image_shape),
"INT": generate_type_validator((int,)), "INT": generate_type_validator((int,)),
"FLOAT": generate_type_validator((float,)), "FLOAT": generate_type_validator((float,)),
"STRING": generate_type_validator((str,)), "STRING": generate_type_validator((str,)),
@ -87,6 +108,9 @@ def input_validation(input_data_all, obj):
for _, v in input_types.items(): for _, v in input_types.items():
if isinstance(v, dict): if isinstance(v, dict):
for k2, v2 in v.items(): for k2, v2 in v.items():
#print(f"input_data_all[{k2}]: {input_data_all[k2]}")
#if v2[0] == "IMAGE":
# input_data_all[k2] = [validate_image_shape(x) for x in input_data_all[k2]]
if tuple(v2[0]) in validation_funcs.keys(): if tuple(v2[0]) in validation_funcs.keys():
input_data_all[k2] = [validation_funcs[v2[0]](obj.__class__.__name__, k2, x) for x in input_data_all[k2]] input_data_all[k2] = [validation_funcs[v2[0]](obj.__class__.__name__, k2, x) for x in input_data_all[k2]]
@ -95,7 +119,7 @@ def input_validation(input_data_all, obj):
def output_validation(results, obj): def output_validation(results, obj):
validation_funcs = { validation_funcs = {
"IMAGE": generate_type_validator((torch.Tensor,), "IMAGE"), "IMAGE": generate_type_validator((torch.Tensor,), validator=validate_image_shape),
"INT": generate_type_validator((int,)), "INT": generate_type_validator((int,)),
"FLOAT": generate_type_validator((float,)), "FLOAT": generate_type_validator((float,)),
"STRING": generate_type_validator((str,)), "STRING": generate_type_validator((str,)),