From 5749d3491749e2498754b917850612d632dac2ad Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Tue, 30 Jul 2024 21:32:24 -0400 Subject: [PATCH 1/6] implemented BHW3 compliance for IMAGE types for all nodes. This change transparently implements BHW3 compliance for IMAGE type inputs and outputs for all nodes. No changes need to be made to existing nodes, the conformance is applied before the IMAGE is passed into the node, and after it is passed out of the node. This change was discussed on the #backend_developers discord with @comfyanonymous and @mcmonkey4eva More rigorous testing needs to be done to ensure that there are no egregious failure modes and that all formats are covered. Im sure the `force_bhw3` function can also be made better, but this is currently a working implementation. --- execution.py | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/execution.py b/execution.py index 0a2e62e7e..106df36b5 100644 --- a/execution.py +++ b/execution.py @@ -13,6 +13,30 @@ import nodes import comfy.model_management +def force_bhw3(image): + #convert [CHW, BCHW, CWH] to BHW3 + was_list = False + + while isinstance(image, list): + was_list = True + image = image[0] + + if len(image.shape) == 3: + #add batch dimension + image = image.unsqueeze(0) + + if image.shape[1] == 3: + #BCHW color + image = image.permute(0, 2, 3, 1) + return image if not was_list else [image] + + if image.shape[1] == 1: + #BCWH black and white + image = image.permute(0, 3, 2, 1).expand(-1, -1, -1, 3) + return image if not was_list else [image] + + return image if not was_list else [image] + def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}): valid_inputs = class_def.INPUT_TYPES() input_data_all = {} @@ -42,6 +66,14 @@ def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_da return input_data_all def map_node_over_list(obj, input_data_all, func, allow_interrupt=False): + #Ensure image inputs are in BHW3 format + input_types = obj.INPUT_TYPES() + for _, v in input_types.items(): + if isinstance(v, dict): + for k2, v2 in v.items(): + if v2[0] == "IMAGE": + input_data_all[k2] = [force_bhw3(x) for x in input_data_all[k2]] + # check if node wants the lists input_is_list = False if hasattr(obj, "INPUT_IS_LIST"): @@ -73,6 +105,26 @@ def map_node_over_list(obj, input_data_all, func, allow_interrupt=False): if allow_interrupt: nodes.before_node_execution() results.append(getattr(obj, func)(**slice_dict(input_data_all, i))) + + #Ensure IMAGE outputs conform to BHWC + return_indexs = {} + formated_results = [] + + if hasattr(obj, "RETURN_NAMES") and hasattr(obj, "RETURN_TYPES"): + for i, t in enumerate(obj.RETURN_TYPES): + return_indexs[i] = t + + for i, r in enumerate(results[0]): + if return_indexs[i] == "IMAGE": + print(f"Result: {force_bhw3(r).shape}") + formated_results.append(force_bhw3(r)) + else: + formated_results.append(r) + + results = [tuple(formated_results)] + + del formated_results + return results def get_output_data(obj, input_data_all): From 40a86b3dc8edf5639a489e51e6639ae7d30ed2b9 Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Tue, 30 Jul 2024 22:39:24 -0400 Subject: [PATCH 2/6] made formated_results formatted_results for guill --- execution.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/execution.py b/execution.py index 106df36b5..71430a14b 100644 --- a/execution.py +++ b/execution.py @@ -108,7 +108,7 @@ def map_node_over_list(obj, input_data_all, func, allow_interrupt=False): #Ensure IMAGE outputs conform to BHWC return_indexs = {} - formated_results = [] + formatted_results = [] if hasattr(obj, "RETURN_NAMES") and hasattr(obj, "RETURN_TYPES"): for i, t in enumerate(obj.RETURN_TYPES): @@ -117,13 +117,13 @@ def map_node_over_list(obj, input_data_all, func, allow_interrupt=False): for i, r in enumerate(results[0]): if return_indexs[i] == "IMAGE": print(f"Result: {force_bhw3(r).shape}") - formated_results.append(force_bhw3(r)) + formatted_results.append(force_bhw3(r)) else: - formated_results.append(r) + formatted_results.append(r) - results = [tuple(formated_results)] + results = [tuple(formatted_results)] - del formated_results + del formatted_results return results From 5c5a3333f6fa53cb3e753fd6a4f4de04a889de9d Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Fri, 2 Aug 2024 16:24:31 -0400 Subject: [PATCH 3/6] Update cli_args.py added cli args --type-confromance with 3 options `all` - attempts to convert all values to the input/output types specified by the node `images` - only converts images to bwh3 format `error` - raises an error if the type does not match what is specified. `none` - prints error messages to console and does nothing else. --- comfy/cli_args.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 2397de3d6..f1d04059d 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -128,6 +128,9 @@ parser.add_argument("--multi-user", action="store_true", help="Enables per-user parser.add_argument("--verbose", action="store_true", help="Enables more debug prints.") +parser.add_argument("--type-conformance", type=str, choices=['all', 'images', 'error', 'none'], default='none', help="Level of type conformance for input/output values between nodes. 'all/images' will attempt to convert all or only image values to the expected type, 'error' will raise an error if the types do not match.") + + # The default built-in provider hosted under web/ DEFAULT_VERSION_STRING = "comfyanonymous/ComfyUI@latest" From 593d69deef27e117a090bfb6707c965530316d8e Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Fri, 2 Aug 2024 16:36:51 -0400 Subject: [PATCH 4/6] Update execution.py for flexible conformance handling. A validation generator function has been aded along with the ability to specify custom functions for validating individual types. cli args have been added in the form of --type-conformance='arg' with the following args: all : attempts to convert any mistyped inputs/outputs to the type specified by the node images : only enforces BHW3 on images, doing the minimum verifiably accurate transforms such as MASK HW and BHW to BHW3 error: raises an error with information about the type mismatch at the node where it occurs. none: the default value, it outputs errors to console, and does nothing else. This should address most of the areas of discussion. --- execution.py | 153 +++++++++++++++++++++++++++++++++++---------------- 1 file changed, 105 insertions(+), 48 deletions(-) diff --git a/execution.py b/execution.py index 71430a14b..5e3d38b1d 100644 --- a/execution.py +++ b/execution.py @@ -13,29 +13,108 @@ import nodes import comfy.model_management -def force_bhw3(image): - #convert [CHW, BCHW, CWH] to BHW3 - was_list = False +from comfy.cli_args import args + + +def validate_image_shape(image): + transforms = { + "HW": lambda t: t.unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3), + "BHW": lambda t: t.unsqueeze(-1).expand(-1, -1, -1, 3), + "HWC": lambda t: t.unsqueeze(0), + } - while isinstance(image, list): - was_list = True - image = image[0] + image_len = len(image.shape) - if len(image.shape) == 3: - #add batch dimension - image = image.unsqueeze(0) + if image_len == 2: + #HW -> add Batch and Channel dimensions + image = transforms["HW"](image) + return image - if image.shape[1] == 3: - #BCHW color - image = image.permute(0, 2, 3, 1) - return image if not was_list else [image] + if image_len == 3: + if image.shape[-1] > 4 or image.shape[-1] < 3: + #BHW -> add Channel dimension + 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 + + +def generate_type_validator(valid_types, expected_type=None): + def validate_type(node_type, output_name, value): + + if not isinstance(value, valid_types) and expected_type is not "IMAGE": + print(f"The value is not of the valid types. Expected one of {[t.__name__ for t in valid_types]}, got {[type(value)]}") + + if "all" == args.type_conformance: + return valid_types[0](value) + + 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") + + return value + + if isinstance(value, valid_types) and expected_type is "IMAGE": + if args.type_conformance in ["all", "images"]: + return validate_image_shape(value) + + 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 validate_type + + +def input_validation(input_data_all, obj): + validation_funcs = { + "IMAGE": generate_type_validator((torch.Tensor,), "IMAGE"), + "INT": generate_type_validator((int,)), + "FLOAT": generate_type_validator((float,)), + "STRING": generate_type_validator((str,)), + } + input_types = obj.INPUT_TYPES() + for _, v in input_types.items(): + if isinstance(v, dict): + for k2, v2 in v.items(): + 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]] + + return input_data_all + + +def output_validation(results, obj): + validation_funcs = { + "IMAGE": generate_type_validator((torch.Tensor,), "IMAGE"), + "INT": generate_type_validator((int,)), + "FLOAT": generate_type_validator((float,)), + "STRING": generate_type_validator((str,)), + } + if hasattr(obj, "RETURN_NAMES") and hasattr(obj, "RETURN_TYPES"): + return_indexs = {} + formatted_results = [] + + for i, return_type in enumerate(obj.RETURN_TYPES): + return_indexs[i] = return_type + + for i, result in enumerate(results[0]): + return_type = return_indexs[i] + formatted_results.append(validation_funcs[return_type](obj.__class__.__name__, obj.RETURN_NAMES[i], result) if return_type in validation_funcs.keys() else result) + + results = [tuple(formatted_results)] + del formatted_results + return results - if image.shape[1] == 1: - #BCWH black and white - image = image.permute(0, 3, 2, 1).expand(-1, -1, -1, 3) - return image if not was_list else [image] - - return image if not was_list else [image] def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}): valid_inputs = class_def.INPUT_TYPES() @@ -52,7 +131,7 @@ def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_da input_data_all[x] = obj else: if ("required" in valid_inputs and x in valid_inputs["required"]) or ("optional" in valid_inputs and x in valid_inputs["optional"]): - input_data_all[x] = [input_data] + input_data_all[x] = [input_data] if "hidden" in valid_inputs: h = valid_inputs["hidden"] @@ -63,16 +142,12 @@ def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_da input_data_all[x] = [extra_data.get('extra_pnginfo', None)] if h[x] == "UNIQUE_ID": input_data_all[x] = [unique_id] + return input_data_all + def map_node_over_list(obj, input_data_all, func, allow_interrupt=False): - #Ensure image inputs are in BHW3 format - input_types = obj.INPUT_TYPES() - for _, v in input_types.items(): - if isinstance(v, dict): - for k2, v2 in v.items(): - if v2[0] == "IMAGE": - input_data_all[k2] = [force_bhw3(x) for x in input_data_all[k2]] + input_data_all = input_validation(input_data_all, obj) # check if node wants the lists input_is_list = False @@ -106,33 +181,15 @@ def map_node_over_list(obj, input_data_all, func, allow_interrupt=False): nodes.before_node_execution() results.append(getattr(obj, func)(**slice_dict(input_data_all, i))) - #Ensure IMAGE outputs conform to BHWC - return_indexs = {} - formatted_results = [] - - if hasattr(obj, "RETURN_NAMES") and hasattr(obj, "RETURN_TYPES"): - for i, t in enumerate(obj.RETURN_TYPES): - return_indexs[i] = t - - for i, r in enumerate(results[0]): - if return_indexs[i] == "IMAGE": - print(f"Result: {force_bhw3(r).shape}") - formatted_results.append(force_bhw3(r)) - else: - formatted_results.append(r) - - results = [tuple(formatted_results)] - - del formatted_results - + results = output_validation(results, obj) return results + def get_output_data(obj, input_data_all): - results = [] uis = [] return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True) - + for r in return_values: if isinstance(r, dict): if 'ui' in r: From 6b362724ab0d09b18be873d4ee35680c35cc3523 Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Sat, 3 Aug 2024 10:33:11 -0400 Subject: [PATCH 5/6] Update execution.py to allow validation function injection --- execution.py | 108 +++++++++++++++++++++++++++++++-------------------- 1 file changed, 66 insertions(+), 42 deletions(-) diff --git a/execution.py b/execution.py index 5e3d38b1d..fae882e02 100644 --- a/execution.py +++ b/execution.py @@ -16,61 +16,82 @@ import comfy.model_management from comfy.cli_args import args -def validate_image_shape(image): - transforms = { - "HW": lambda t: t.unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3), - "BHW": lambda t: t.unsqueeze(-1).expand(-1, -1, -1, 3), - "HWC": lambda t: t.unsqueeze(0), - } - +import sys +import copy +import logging +import threading +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) - if image_len == 2: - #HW -> add Batch and Channel dimensions - image = transforms["HW"](image) - return image - - if image_len == 3: - if image.shape[-1] > 4 or image.shape[-1] < 3: - #BHW -> add Channel dimension - 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) + def type_check(): + print(f"Checking Image Shape: {image.shape}") + if image_len != 4 or image.shape[-1] > 4 or image.shape[-1] < 3: + 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 type_check() != True and args.type_conformance in ["all", "images"]: + transforms = { + "HW": lambda t: t.unsqueeze(0).unsqueeze(-1).expand(-1, -1, -1, 3), + "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 + if image_len == 3: + if image.shape[-1] > 4 or image.shape[-1] < 3: + #BHW -> add Channel dimension + 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 - - -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): - - if not isinstance(value, valid_types) and expected_type is not "IMAGE": - print(f"The value is not of the valid types. Expected one of {[t.__name__ for t in valid_types]}, got {[type(value)]}") + print(type(validator)) + if not isinstance(value, valid_types) and validator is None: + 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: return valid_types[0](value) 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 - if isinstance(value, valid_types) and expected_type is "IMAGE": - if args.type_conformance in ["all", "images"]: - return validate_image_shape(value) - - 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 - + if isinstance(value, valid_types) and validator is not None: + print(f"Validating {output_name} in Node: {node_type}") + value = validator(value, node_type, output_name) return value + return value return validate_type @@ -78,7 +99,7 @@ def generate_type_validator(valid_types, expected_type=None): def input_validation(input_data_all, obj): validation_funcs = { - "IMAGE": generate_type_validator((torch.Tensor,), "IMAGE"), + "IMAGE": generate_type_validator((torch.Tensor,), validator=validate_image_shape), "INT": generate_type_validator((int,)), "FLOAT": generate_type_validator((float,)), "STRING": generate_type_validator((str,)), @@ -87,6 +108,9 @@ def input_validation(input_data_all, obj): for _, v in input_types.items(): if isinstance(v, dict): 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(): 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): validation_funcs = { - "IMAGE": generate_type_validator((torch.Tensor,), "IMAGE"), + "IMAGE": generate_type_validator((torch.Tensor,), validator=validate_image_shape), "INT": generate_type_validator((int,)), "FLOAT": generate_type_validator((float,)), "STRING": generate_type_validator((str,)), From d4ebb8f750f968d451f0ed3607f374ece3c48afb Mon Sep 17 00:00:00 2001 From: shawnington <88048838+shawnington@users.noreply.github.com> Date: Sat, 3 Aug 2024 11:16:38 -0400 Subject: [PATCH 6/6] removed commented out lines in execution.py --- execution.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/execution.py b/execution.py index fae882e02..b361a57ee 100644 --- a/execution.py +++ b/execution.py @@ -108,9 +108,6 @@ def input_validation(input_data_all, obj): for _, v in input_types.items(): if isinstance(v, dict): 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(): input_data_all[k2] = [validation_funcs[v2[0]](obj.__class__.__name__, k2, x) for x in input_data_all[k2]]