mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-13 00:47:03 +08:00
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.
This commit is contained in:
parent
82cae45d44
commit
5749d34917
52
execution.py
52
execution.py
@ -13,6 +13,30 @@ import nodes
|
|||||||
|
|
||||||
import comfy.model_management
|
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={}):
|
def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}):
|
||||||
valid_inputs = class_def.INPUT_TYPES()
|
valid_inputs = class_def.INPUT_TYPES()
|
||||||
input_data_all = {}
|
input_data_all = {}
|
||||||
@ -42,6 +66,14 @@ def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_da
|
|||||||
return input_data_all
|
return input_data_all
|
||||||
|
|
||||||
def map_node_over_list(obj, input_data_all, func, allow_interrupt=False):
|
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
|
# check if node wants the lists
|
||||||
input_is_list = False
|
input_is_list = False
|
||||||
if hasattr(obj, "INPUT_IS_LIST"):
|
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:
|
if allow_interrupt:
|
||||||
nodes.before_node_execution()
|
nodes.before_node_execution()
|
||||||
results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))
|
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
|
return results
|
||||||
|
|
||||||
def get_output_data(obj, input_data_all):
|
def get_output_data(obj, input_data_all):
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user