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] 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):