235 lines
8.0 KiB
Python
235 lines
8.0 KiB
Python
import execution
|
|
import nodes
|
|
from cozy_comfyui.node import CozyBaseNode, COZY_TYPE_ANY
|
|
import time
|
|
import logging
|
|
import inspect
|
|
|
|
MAX_OUTPUTS = 64
|
|
MAX_INPUTS = 64
|
|
|
|
# 检测ComfyUI版本是否支持新的执行模型
|
|
def is_execution_model_version_supported():
|
|
"""检测是否支持新版ComfyUI的执行模型"""
|
|
try:
|
|
import comfy_execution # noqa: F401
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
# 获取合适的执行阻断器
|
|
def get_execution_blocker():
|
|
"""根据ComfyUI版本获取合适的执行阻断器"""
|
|
if is_execution_model_version_supported():
|
|
try:
|
|
from comfy_execution.graph import ExecutionBlocker
|
|
return ExecutionBlocker(None)
|
|
except ImportError as e:
|
|
logging.warning(f"[WBLESS] Failed to import ExecutionBlocker from comfy_execution.graph: {e}")
|
|
return WBLESSExecutionBlocker()
|
|
else:
|
|
return WBLESSExecutionBlocker()
|
|
|
|
class WBLESSExecutionBlocker:
|
|
"""向后兼容的执行阻断器类"""
|
|
def __init__(self):
|
|
pass
|
|
|
|
def is_execution_blocked(values):
|
|
"""检测值列表中是否包含执行阻断器"""
|
|
if not isinstance(values, list):
|
|
return False
|
|
|
|
# 检查两种类型的阻断器
|
|
for v in values:
|
|
if isinstance(v, WBLESSExecutionBlocker):
|
|
return True
|
|
# 检查新版ComfyUI的ExecutionBlocker
|
|
if is_execution_model_version_supported():
|
|
try:
|
|
from comfy_execution.graph import ExecutionBlocker
|
|
if isinstance(v, ExecutionBlocker):
|
|
return True
|
|
except ImportError:
|
|
pass
|
|
return False
|
|
|
|
_original_get_output_data = execution.get_output_data
|
|
|
|
def _hooked_get_output_data(obj, input_data_all, *args, **kwargs):
|
|
"""钩子函数：拦截ComfyUI的节点执行，实现智能的执行阻断机制"""
|
|
if not isinstance(input_data_all, dict):
|
|
return _original_get_output_data(obj, input_data_all, *args, **kwargs)
|
|
|
|
# Switch节点的特殊处理：只检查活动线路的输入
|
|
if isinstance(obj, Switch):
|
|
path_list = input_data_all.get("Path")
|
|
|
|
if not path_list:
|
|
return _original_get_output_data(obj, input_data_all, *args, **kwargs)
|
|
|
|
path_val = path_list[0]
|
|
selected_input_name = f"Input_{path_val}"
|
|
|
|
# 只检查被选中的输入线路
|
|
selected_input_data = input_data_all.get(selected_input_name)
|
|
if selected_input_data and is_execution_blocked(selected_input_data):
|
|
# 返回合适的阻断器
|
|
blocker = get_execution_blocker()
|
|
return ([[blocker]] * len(obj.RETURN_TYPES), {}, False)
|
|
else:
|
|
return _original_get_output_data(obj, input_data_all, *args, **kwargs)
|
|
|
|
# 普通节点：检查所有输入，任何一个被阻断都阻止执行
|
|
for an_input in input_data_all.values():
|
|
if is_execution_blocked(an_input):
|
|
# 返回合适的阻断器
|
|
blocker = get_execution_blocker()
|
|
return ([[blocker]] * len(obj.RETURN_TYPES), {}, False)
|
|
|
|
return _original_get_output_data(obj, input_data_all, *args, **kwargs)
|
|
|
|
# 应用execution hook来处理执行阻断
|
|
# 在新版ComfyUI中主要用于InversedSwitch的阻断逻辑
|
|
# Switch节点现在使用lazy evaluation
|
|
execution.get_output_data = _hooked_get_output_data
|
|
if is_execution_model_version_supported():
|
|
logging.info("[WBLESS] Applied execution hook for new ComfyUI version (for InversedSwitch blocking)")
|
|
else:
|
|
logging.info("[WBLESS] Applied execution hook for legacy ComfyUI version")
|
|
|
|
|
|
class InversedSwitch(CozyBaseNode):
|
|
NAME = "Inversed Switch"
|
|
FUNCTION = "run"
|
|
OUTPUT_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"Input": (COZY_TYPE_ANY,),
|
|
"Path": ("INT", {"default": 1, "min": 1, "max": MAX_OUTPUTS}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (COZY_TYPE_ANY,) * MAX_OUTPUTS
|
|
RETURN_NAMES = tuple([f"Output_{i+1}" for i in range(MAX_OUTPUTS)])
|
|
|
|
def run(self, Input, Path, **kw):
|
|
"""执行Inversed Switch节点的路由逻辑"""
|
|
if isinstance(Path, list) and len(Path) == 1:
|
|
Path = Path[0]
|
|
|
|
value_to_route = Input
|
|
if isinstance(Input, list) and len(Input) > 0:
|
|
value_to_route = Input[0]
|
|
|
|
# 使用合适的执行阻断器填充结果数组
|
|
blocker = get_execution_blocker()
|
|
results = [blocker] * MAX_OUTPUTS
|
|
|
|
# 只有指定的路径输出真实数据
|
|
if 1 <= Path <= MAX_OUTPUTS:
|
|
results[Path - 1] = value_to_route
|
|
|
|
return tuple(results)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
|
|
class Switch(CozyBaseNode):
|
|
NAME = "Switch"
|
|
FUNCTION = "run"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# 为新版ComfyUI创建动态输入
|
|
dyn_inputs = {"Input_1": (COZY_TYPE_ANY, {"lazy": True, "tooltip": "Any input. When connected, one more input slot is added."})}
|
|
|
|
# 如果是新版ComfyUI，使用特殊的容器类来支持动态输入验证
|
|
if is_execution_model_version_supported():
|
|
stack = inspect.stack()
|
|
if stack[2].function == 'get_input_info':
|
|
# 绕过验证的容器类
|
|
class AllContainer:
|
|
def __contains__(self, item):
|
|
return True
|
|
|
|
def __getitem__(self, key):
|
|
return COZY_TYPE_ANY, {"lazy": True}
|
|
|
|
dyn_inputs = AllContainer()
|
|
|
|
return {
|
|
"required": {
|
|
"Path": ("INT", {"default": 1, "min": 1, "max": MAX_INPUTS, "tooltip": "The input number you want to output among the inputs"}),
|
|
},
|
|
"optional": dyn_inputs,
|
|
"hidden": {"unique_id": "UNIQUE_ID", "extra_pnginfo": "EXTRA_PNGINFO"}
|
|
}
|
|
|
|
RETURN_TYPES = (COZY_TYPE_ANY,)
|
|
RETURN_NAMES = ("output",)
|
|
|
|
def check_lazy_status(self, *args, **kwargs):
|
|
"""告诉ComfyUI只需要加载选中的输入"""
|
|
if 'Path' not in kwargs:
|
|
return []
|
|
|
|
try:
|
|
# 处理Path可能是列表的情况
|
|
path_value = kwargs['Path']
|
|
if isinstance(path_value, list) and len(path_value) > 0:
|
|
selected_index = int(path_value[0])
|
|
else:
|
|
selected_index = int(path_value)
|
|
|
|
input_name = f"Input_{selected_index}"
|
|
|
|
logging.info(f"[WBLESS Switch] check_lazy_status - Path={path_value}, selected_index={selected_index}, input_name={input_name}")
|
|
|
|
# 只返回需要加载的输入
|
|
if input_name in kwargs:
|
|
logging.info(f"[WBLESS Switch] Found input {input_name}, loading it")
|
|
return [input_name]
|
|
else:
|
|
logging.info(f"[WBLESS Switch] Input {input_name} not found in kwargs: {list(kwargs.keys())}")
|
|
return []
|
|
|
|
except Exception as e:
|
|
logging.error(f"[WBLESS Switch] Error in check_lazy_status: {e}, Path={kwargs.get('Path')}, type={type(kwargs.get('Path'))}")
|
|
return []
|
|
|
|
def run(self, Path, unique_id=None, extra_pnginfo=None, **kw):
|
|
"""Switch节点的核心逻辑"""
|
|
if isinstance(Path, list) and len(Path) == 1:
|
|
Path = Path[0]
|
|
|
|
selected_input_name = f"Input_{Path}"
|
|
value = kw.get(selected_input_name)
|
|
|
|
# 处理列表格式的输入
|
|
if isinstance(value, list) and len(value) > 0:
|
|
value = value[0]
|
|
|
|
logging.info(f"[WBLESS Switch] Path={Path}, selected_input={selected_input_name}, value_type={type(value)}")
|
|
|
|
return (value,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
return time.time()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Inversed Switch": InversedSwitch,
|
|
"Switch": Switch
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Inversed Switch": "Inversed Switch",
|
|
"Switch": "Switch"
|
|
}
|