Spaces:
Running on Zero
Running on Zero
| import torch | |
| from collections import defaultdict, deque | |
| from typing import Dict, Any, List | |
| from comfy_integration.nodes import NODE_CLASS_MAPPINGS | |
| from utils.app_utils import get_value_at_index | |
| def _normalize_dynamic_inputs(node_class, class_type, kwargs): | |
| """Rebuild the structures ComfyUI's executor would have built for this node. | |
| Two V3 dynamic shapes need post-processing when nodes are called directly: | |
| * ``COMFY_DYNAMICCOMBO_V3`` — the widget arrives as a bare string but the | |
| node signature expects ``{id: <option>, **sub_keys}`` (e.g. | |
| ``sampling_mode={"sampling_mode": "off"}``). | |
| * ``COMFY_AUTOGROW_V3`` — slots arrive flattened as ``image1``, ``image2``, | |
| ... but the node signature expects ``images={"image1": <value>, ...}``. | |
| ComfyUI's ``comfy_api.latest._io.build_nested_inputs`` does both; this | |
| Space executes nodes directly and therefore has to redo the work. | |
| """ | |
| schema = getattr(node_class, 'SCHEMA', None) | |
| inputs = getattr(schema, 'inputs', None) if schema is not None else None | |
| if not inputs: | |
| return kwargs | |
| # --- COMFY_AUTOGROW_V3 first: consume image1/video1/... slots --- | |
| for spec in inputs: | |
| sub_id = getattr(spec, 'id', None) | |
| template = getattr(spec, 'template', None) | |
| if not sub_id or template is None: | |
| continue | |
| names = getattr(template, 'names', None) or [] | |
| prefix = getattr(template, 'prefix', None) | |
| candidates = list(names) | |
| if prefix: | |
| max_n = getattr(template, 'max', 0) or 0 | |
| if not candidates: | |
| candidates = [f"{prefix}{i}" for i in range(max_n)] | |
| collected = {} | |
| for name in candidates: | |
| if name in kwargs: | |
| collected[name] = kwargs.pop(name) | |
| if collected and not isinstance(kwargs.get(sub_id), dict): | |
| kwargs[sub_id] = collected | |
| # --- COMFY_DYNAMICCOMBO_V3 --- | |
| for spec in inputs: | |
| sub_id = getattr(spec, 'id', None) | |
| options = getattr(spec, 'options', None) | |
| if not sub_id or not options or sub_id not in kwargs: | |
| continue | |
| value = kwargs[sub_id] | |
| if isinstance(value, dict): | |
| continue # already structured | |
| option_keys = [] | |
| for opt in options: | |
| if isinstance(opt, dict): | |
| option_keys.append(opt.get('key')) | |
| else: | |
| option_keys.append(getattr(opt, 'key', None)) | |
| if value not in option_keys: | |
| print(f"Warning: '{sub_id}' value {value!r} is not one of {option_keys} for {class_type}.") | |
| continue | |
| # Pull sibling sub-keys (e.g. "sampling_mode.temperature") into the dict. | |
| nested = {sub_id: str(value)} | |
| prefix = f"{sub_id}." | |
| for key in list(kwargs.keys()): | |
| if key.startswith(prefix): | |
| nested[key[len(prefix):]] = kwargs.pop(key) | |
| kwargs[sub_id] = nested | |
| print(f" -> Normalized dynamic input '{sub_id}' for {class_type}: {nested}") | |
| return kwargs | |
| class WorkflowExecutor: | |
| def topological_sort(workflow: Dict[str, Any]) -> List[str]: | |
| graph = defaultdict(list) | |
| in_degree = {node_id: 0 for node_id in workflow} | |
| for node_id, node_info in workflow.items(): | |
| for input_value in node_info.get('inputs', {}).values(): | |
| if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str): | |
| source_node_id = input_value[0] | |
| if source_node_id in workflow: | |
| graph[source_node_id].append(node_id) | |
| in_degree[node_id] += 1 | |
| queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0]) | |
| sorted_nodes = [] | |
| while queue: | |
| current_node_id = queue.popleft() | |
| sorted_nodes.append(current_node_id) | |
| for neighbor_node_id in graph[current_node_id]: | |
| in_degree[neighbor_node_id] -= 1 | |
| if in_degree[neighbor_node_id] == 0: | |
| queue.append(neighbor_node_id) | |
| if len(sorted_nodes) != len(workflow): | |
| raise RuntimeError("Workflow contains a cycle and cannot be executed.") | |
| return sorted_nodes | |
| def execute_workflow(workflow: Dict[str, Any], initial_objects: Dict[str, Any]): | |
| with torch.no_grad(): | |
| computed_outputs = initial_objects | |
| try: | |
| sorted_node_ids = WorkflowExecutor.topological_sort(workflow) | |
| final_node_id = None | |
| for node_id in reversed(sorted_node_ids): | |
| if workflow[node_id].get('class_type') == 'SaveImage': | |
| final_node_id = node_id | |
| break | |
| if final_node_id: | |
| required_nodes = set() | |
| nodes_to_visit = [final_node_id] | |
| while nodes_to_visit: | |
| curr_id = nodes_to_visit.pop() | |
| if curr_id in required_nodes: | |
| continue | |
| required_nodes.add(curr_id) | |
| curr_info = workflow.get(curr_id, {}) | |
| for input_val in curr_info.get('inputs', {}).values(): | |
| if isinstance(input_val, list) and len(input_val) == 2 and isinstance(input_val[0], str): | |
| src_id = input_val[0] | |
| if src_id in workflow and src_id not in required_nodes: | |
| nodes_to_visit.append(src_id) | |
| sorted_node_ids = [nid for nid in sorted_node_ids if nid in required_nodes] | |
| print(f"--- [Workflow Executor] Execution order: {sorted_node_ids}") | |
| except RuntimeError as e: | |
| print("--- [Workflow Executor] ERROR: Failed to sort workflow. Dumping graph details. ---") | |
| for node_id, node_info in workflow.items(): | |
| print(f" Node {node_id} ({node_info['class_type']}):") | |
| for input_name, input_value in node_info['inputs'].items(): | |
| if isinstance(input_value, list) and len(input_value) == 2 and isinstance(input_value[0], str): | |
| print(f" - {input_name} <- [{input_value[0]}, {input_value[1]}]") | |
| raise e | |
| for node_id in sorted_node_ids: | |
| if node_id in computed_outputs: | |
| continue | |
| node_info = workflow[node_id] | |
| class_type = node_info['class_type'] | |
| is_loader_with_filename = 'Loader' in class_type and any(key.endswith('_name') for key in node_info['inputs']) | |
| if node_id in initial_objects and is_loader_with_filename: | |
| continue | |
| node_class = NODE_CLASS_MAPPINGS.get(class_type) | |
| if node_class is None: | |
| raise RuntimeError(f"Could not find node class '{class_type}'. Is it imported in comfy_integration/nodes.py?") | |
| node_instance = node_class() | |
| kwargs = {} | |
| for param_name, param_value in node_info['inputs'].items(): | |
| if isinstance(param_value, list) and len(param_value) == 2 and isinstance(param_value[0], str): | |
| source_node_id, output_index = param_value | |
| if source_node_id not in computed_outputs: | |
| raise RuntimeError(f"Workflow integrity error: Output of node {source_node_id} needed for {node_id} but not yet computed.") | |
| source_output_tuple = computed_outputs[source_node_id] | |
| actual_value = get_value_at_index(source_output_tuple, output_index) | |
| else: | |
| actual_value = param_value | |
| if '.' in param_name: | |
| parent_key, child_key = param_name.split('.', 1) | |
| if parent_key not in kwargs or not isinstance(kwargs[parent_key], dict): | |
| kwargs[parent_key] = {} | |
| kwargs[parent_key][child_key] = actual_value | |
| else: | |
| kwargs[param_name] = actual_value | |
| # Safety guard for nodes expecting dynamic autogrow dictionary (e.g. TextEncodeMingImageEdit) | |
| if class_type == 'TextEncodeMingImageEdit' and ('images' not in kwargs or kwargs.get('images') is None): | |
| kwargs['images'] = {} | |
| # ComfyUI's own executor restructures V3 dynamic inputs into the | |
| # nested dict the node signature expects (see | |
| # comfy_api.latest._io.build_nested_inputs). This direct-call | |
| # path bypasses that step, so a COMFY_DYNAMICCOMBO_V3 widget | |
| # arrives as a bare string and the node fails on | |
| # `sampling_mode.get(...)`. Rebuild it here. | |
| kwargs = _normalize_dynamic_inputs(node_class, class_type, kwargs) | |
| function_name = getattr(node_class, 'FUNCTION') | |
| execution_method = getattr(node_instance, function_name) | |
| result = execution_method(**kwargs) | |
| computed_outputs[node_id] = result | |
| final_node_id = None | |
| for node_id in reversed(sorted_node_ids): | |
| if workflow[node_id]['class_type'] == 'SaveImage': | |
| final_node_id = node_id | |
| break | |
| if not final_node_id: | |
| raise RuntimeError("Workflow does not contain a 'SaveImage' node as the output.") | |
| save_image_inputs = workflow[final_node_id]['inputs'] | |
| image_source_node_id, image_source_index = save_image_inputs['images'] | |
| return get_value_at_index(computed_outputs[image_source_node_id], image_source_index) | |