Spaces:
Running on Zero
Running on Zero
File size: 9,112 Bytes
6ac5e98 439fb79 6ac5e98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | import yaml
import os
import re
import importlib
from copy import deepcopy
from comfy_integration.nodes import NODE_CLASS_MAPPINGS
from chain_injectors import discover_injectors, get_registered_features
from core.settings import FEATURES_CONFIG
class WorkflowAssembler:
def __init__(self, recipe_path, dynamic_values=None):
self.base_path = os.path.dirname(recipe_path)
self.dynamic_values = dynamic_values or {}
self.node_counter = 0
self.workflow = {}
self.node_map = {}
model_type = self.dynamic_values.get('model_type')
self._load_injector_config(model_type=model_type)
self.recipe = self._load_and_merge_recipe(os.path.basename(recipe_path), self.dynamic_values)
def _load_injector_config(self, model_type=None):
self.global_injectors = discover_injectors()
registered_features = get_registered_features()
order = []
if model_type and model_type in FEATURES_CONFIG:
enabled_features = FEATURES_CONFIG[model_type].get('enabled_chains', [])
for feat in enabled_features:
if feat in registered_features:
chain_key = registered_features[feat]['chain_type']
else:
chain_key = f"dynamic_{feat}_chains"
if chain_key in self.global_injectors and chain_key not in order:
order.append(chain_key)
for chain_key in self.global_injectors.keys():
if chain_key not in order:
order.append(chain_key)
self.injector_order = order
def _get_unique_id(self):
self.node_counter += 1
return str(self.node_counter)
def _get_node_template(self, class_type):
if class_type not in NODE_CLASS_MAPPINGS:
raise ValueError(f"Node class '{class_type}' not found. Ensure it's correctly imported in comfy_integration/nodes.py.")
node_class = NODE_CLASS_MAPPINGS[class_type]
input_types = node_class.INPUT_TYPES()
template = {
"inputs": {},
"class_type": class_type,
"_meta": {"title": node_class.NODE_NAME if hasattr(node_class, 'NODE_NAME') else class_type}
}
all_inputs = {**input_types.get('required', {}), **input_types.get('optional', {})}
for name, details in all_inputs.items():
config = details[1] if len(details) > 1 and isinstance(details[1], dict) else {}
template["inputs"][name] = config.get("default")
return template
def _load_and_merge_recipe(self, recipe_filename, dynamic_values, search_context_dir=None):
search_path = search_context_dir or self.base_path
recipe_path_to_use = os.path.join(search_path, recipe_filename)
if not os.path.exists(recipe_path_to_use):
raise FileNotFoundError(f"Recipe file not found: {recipe_path_to_use}")
with open(recipe_path_to_use, 'r', encoding='utf-8') as f:
content = f.read()
for key, value in dynamic_values.items():
if value is not None:
content = content.replace(f"{{{{ {key} }}}}", str(value))
main_recipe = yaml.safe_load(content)
merged_recipe = {'nodes': {}, 'connections': [], 'ui_map': {}}
for key in self.injector_order:
if key.startswith('dynamic_'):
merged_recipe[key] = {}
parent_recipe_dir = os.path.dirname(recipe_path_to_use)
for import_path_template in main_recipe.get('imports', []):
import_path = import_path_template
for key, value in dynamic_values.items():
if value is not None:
import_path = import_path.replace(f"{{{{ {key} }}}}", str(value))
try:
imported_recipe = self._load_and_merge_recipe(import_path, dynamic_values, search_context_dir=parent_recipe_dir)
merged_recipe['nodes'].update(imported_recipe.get('nodes', {}))
merged_recipe['connections'].extend(imported_recipe.get('connections', []))
merged_recipe['ui_map'].update(imported_recipe.get('ui_map', {}))
for key in self.injector_order:
if key in imported_recipe and key.startswith('dynamic_'):
merged_recipe[key].update(imported_recipe.get(key, {}))
except FileNotFoundError:
print(f"Warning: Optional recipe partial '{import_path}' not found. Skipping.")
merged_recipe['nodes'].update(main_recipe.get('nodes', {}))
merged_recipe['connections'].extend(main_recipe.get('connections', []))
merged_recipe['ui_map'].update(main_recipe.get('ui_map', {}))
for key in self.injector_order:
if key in main_recipe and key.startswith('dynamic_'):
merged_recipe[key].update(main_recipe.get(key, {}))
return merged_recipe
def add_node(self, class_type: str, inputs: dict = None, title: str = None) -> str:
template = self._get_node_template(class_type)
node_data = deepcopy(template)
if title:
node_data['_meta']['title'] = title
if inputs:
for k, v in inputs.items():
node_data['inputs'][k] = v
node_id = self._get_unique_id()
self.workflow[node_id] = node_data
return node_id
def connect(self, from_node: str, from_output_idx: int, to_node: str, to_input_name: str):
from_id = self.node_map.get(from_node, from_node)
to_id = self.node_map.get(to_node, to_node)
if from_id in self.workflow and to_id in self.workflow:
self.workflow[to_id]['inputs'][to_input_name] = [from_id, int(from_output_idx)]
else:
print(f"Warning: Cannot connect '{from_node}' -> '{to_node}'. Node ID not found.")
def assemble(self, ui_values):
self.ui_values = ui_values
for name, details in self.recipe['nodes'].items():
if 'class_type' not in details:
continue
class_type = details['class_type']
match = re.search(r"\{\{\s*(\w+)\s*\}\}", class_type)
if match:
placeholder_key = match.group(1)
if placeholder_key in ui_values and ui_values[placeholder_key] is not None:
class_type = ui_values[placeholder_key]
else:
continue
template = self._get_node_template(class_type)
node_data = deepcopy(template)
unique_id = self._get_unique_id()
self.node_map[name] = unique_id
if 'title' in details:
node_data['_meta']['title'] = details['title']
if 'params' in details:
for param, value in details['params'].items():
if param in node_data['inputs']:
node_data['inputs'][param] = value
self.workflow[unique_id] = node_data
for ui_key, target in self.recipe.get('ui_map', {}).items():
if ui_key in ui_values and ui_values[ui_key] is not None:
target_list = target if isinstance(target, list) else [target]
for t in target_list:
if isinstance(t, str) and ':' in t:
target_name, target_param = t.split(':')
if target_name in self.node_map:
self.workflow[self.node_map[target_name]]['inputs'][target_param] = ui_values[ui_key]
for conn in self.recipe.get('connections', []):
if not isinstance(conn.get('to'), str) or not isinstance(conn.get('from'), str):
continue
from_name, from_output_idx = conn['from'].split(':')
to_name, to_input_name = conn['to'].split(':')
from_id = self.node_map.get(from_name)
to_id = self.node_map.get(to_name)
if from_id and to_id:
self.workflow[to_id]['inputs'][to_input_name] = [from_id, int(from_output_idx)]
print("--- [Assembler] Applying dynamic injectors ---")
recipe_chain_types = {key for key in self.recipe if key.startswith('dynamic_')}
processing_order = [key for key in self.injector_order if key in recipe_chain_types]
for chain_type in processing_order:
injector_func = self.global_injectors.get(chain_type)
if injector_func:
for chain_key, chain_def in self.recipe.get(chain_type, {}).items():
if chain_key in ui_values and ui_values[chain_key]:
print(f" -> Injecting '{chain_type}' for '{chain_key}'...")
chain_items = ui_values[chain_key]
injector_func(self, chain_def, chain_items)
print("--- [Assembler] Finished applying injectors ---")
return self.workflow |