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

@@ -13,29 +13,108 @@ import nodes
 
 import comfy.model_management
 
-def force_bhw3(image):
+from comfy.cli_args import args
-    #convert [CHW, BCHW, CWH] to BHW3
-    was_list = False
 
-    while isinstance(image, list):
-        was_list = True
-        image = image[0]
 
-    if len(image.shape) == 3:
+def validate_image_shape(image):
-        #add batch dimension
+    transforms = {
-        image = image.unsqueeze(0)
+        "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.shape[1] == 3:
+    image_len = len(image.shape)
-        #BCHW color
-        image = image.permute(0, 2, 3, 1)
-        return image if not was_list else [image]
     
-    if image.shape[1] == 1:
+    if image_len == 2:
-        #BCWH black and white
+        #HW -> add Batch and Channel dimensions
-        image = image.permute(0, 3, 2, 1).expand(-1, -1, -1, 3)
+        image = transforms["HW"](image)
-        return image if not was_list else [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 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
     
-    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()
@@ -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_data_all = input_validation(input_data_all, obj)
-    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
@@ -106,29 +181,11 @@ 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
+    results = output_validation(results, obj)           
-    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        
-    
     return results
 
-def get_output_data(obj, input_data_all):
 
+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)