brunvelop/ComfyUI
2
1import os2import sys3import copy4import json5import logging6import threading7import heapq8import traceback9import gc10 11import torch12import nodes13 14import comfy.model_management15 16def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}):17 valid_inputs = class_def.INPUT_TYPES()18 input_data_all = {}19 for x in inputs:20 input_data = inputs[x]21 if isinstance(input_data, list):22 input_unique_id = input_data[0]23 output_index = input_data[1]24 if input_unique_id not in outputs:25 input_data_all[x] = (None,)26 continue27 obj = outputs[input_unique_id][output_index]28 input_data_all[x] = obj29 else:30 if ("required" in valid_inputs and x in valid_inputs["required"]) or ("optional" in valid_inputs and x in valid_inputs["optional"]):31 input_data_all[x] = [input_data]32 33 if "hidden" in valid_inputs:34 h = valid_inputs["hidden"]35 for x in h:36 if h[x] == "PROMPT":37 input_data_all[x] = [prompt]38 if h[x] == "EXTRA_PNGINFO":39 if "extra_pnginfo" in extra_data:40 input_data_all[x] = [extra_data['extra_pnginfo']]41 if h[x] == "UNIQUE_ID":42 input_data_all[x] = [unique_id]43 return input_data_all44 45def map_node_over_list(obj, input_data_all, func, allow_interrupt=False):46 # check if node wants the lists47 input_is_list = False48 if hasattr(obj, "INPUT_IS_LIST"):49 input_is_list = obj.INPUT_IS_LIST50 51 if len(input_data_all) == 0:52 max_len_input = 053 else:54 max_len_input = max([len(x) for x in input_data_all.values()])55 56 # get a slice of inputs, repeat last input when list isn't long enough57 def slice_dict(d, i):58 d_new = dict()59 for k,v in d.items():60 d_new[k] = v[i if len(v) > i else -1]61 return d_new62 63 results = []64 if input_is_list:65 if allow_interrupt:66 nodes.before_node_execution()67 results.append(getattr(obj, func)(**input_data_all))68 elif max_len_input == 0:69 if allow_interrupt:70 nodes.before_node_execution()71 results.append(getattr(obj, func)())72 else:73 for i in range(max_len_input):74 if allow_interrupt:75 nodes.before_node_execution()76 results.append(getattr(obj, func)(**slice_dict(input_data_all, i)))77 return results78 79def get_output_data(obj, input_data_all):80 81 results = []82 uis = []83 return_values = map_node_over_list(obj, input_data_all, obj.FUNCTION, allow_interrupt=True)84 85 for r in return_values:86 if isinstance(r, dict):87 if 'ui' in r:88 uis.append(r['ui'])89 if 'result' in r:90 results.append(r['result'])91 else:92 results.append(r)93 94 output = []95 if len(results) > 0:96 # check which outputs need concatenating97 output_is_list = [False] * len(results[0])98 if hasattr(obj, "OUTPUT_IS_LIST"):99 output_is_list = obj.OUTPUT_IS_LIST100 101 # merge node execution results102 for i, is_list in zip(range(len(results[0])), output_is_list):103 if is_list:104 output.append([x for o in results for x in o[i]])105 else:106 output.append([o[i] for o in results])107 108 ui = dict() 109 if len(uis) > 0:110 ui = {k: [y for x in uis for y in x[k]] for k in uis[0].keys()}111 return output, ui112 113def format_value(x):114 if x is None:115 return None116 elif isinstance(x, (int, float, bool, str)):117 return x118 else:119 return str(x)120 121def recursive_execute(server, prompt, outputs, current_item, extra_data, executed, prompt_id, outputs_ui, object_storage):122 unique_id = current_item123 inputs = prompt[unique_id]['inputs']124 class_type = prompt[unique_id]['class_type']125 class_def = nodes.NODE_CLASS_MAPPINGS[class_type]126 if unique_id in outputs:127 return (True, None, None)128 129 for x in inputs:130 input_data = inputs[x]131 132 if isinstance(input_data, list):133 input_unique_id = input_data[0]134 output_index = input_data[1]135 if input_unique_id not in outputs:136 result = recursive_execute(server, prompt, outputs, input_unique_id, extra_data, executed, prompt_id, outputs_ui, object_storage)137 if result[0] is not True:138 # Another node failed further upstream139 return result140 141 input_data_all = None142 try:143 input_data_all = get_input_data(inputs, class_def, unique_id, outputs, prompt, extra_data)144 if server.client_id is not None:145 server.last_node_id = unique_id146 server.send_sync("executing", { "node": unique_id, "prompt_id": prompt_id }, server.client_id)147 148 obj = object_storage.get((unique_id, class_type), None)149 if obj is None:150 obj = class_def()151 object_storage[(unique_id, class_type)] = obj152 153 output_data, output_ui = get_output_data(obj, input_data_all)154 outputs[unique_id] = output_data155 if len(output_ui) > 0:156 outputs_ui[unique_id] = output_ui157 if server.client_id is not None:158 server.send_sync("executed", { "node": unique_id, "output": output_ui, "prompt_id": prompt_id }, server.client_id)159 except comfy.model_management.InterruptProcessingException as iex:160 logging.info("Processing interrupted")161 162 # skip formatting inputs/outputs163 error_details = {164 "node_id": unique_id,165 }166 167 return (False, error_details, iex)168 except Exception as ex:169 typ, _, tb = sys.exc_info()170 exception_type = full_type_name(typ)171 input_data_formatted = {}172 if input_data_all is not None:173 input_data_formatted = {}174 for name, inputs in input_data_all.items():175 input_data_formatted[name] = [format_value(x) for x in inputs]176 177 output_data_formatted = {}178 for node_id, node_outputs in outputs.items():179 output_data_formatted[node_id] = [[format_value(x) for x in l] for l in node_outputs]180 181 logging.error("!!! Exception during processing !!!")182 logging.error(traceback.format_exc())183 184 error_details = {185 "node_id": unique_id,186 "exception_message": str(ex),187 "exception_type": exception_type,188 "traceback": traceback.format_tb(tb),189 "current_inputs": input_data_formatted,190 "current_outputs": output_data_formatted191 }192 return (False, error_details, ex)193 194 executed.add(unique_id)195 196 return (True, None, None)197 198def recursive_will_execute(prompt, outputs, current_item):199 unique_id = current_item200 inputs = prompt[unique_id]['inputs']201 will_execute = []202 if unique_id in outputs:203 return []204 205 for x in inputs:206 input_data = inputs[x]207 if isinstance(input_data, list):208 input_unique_id = input_data[0]209 output_index = input_data[1]210 if input_unique_id not in outputs:211 will_execute += recursive_will_execute(prompt, outputs, input_unique_id)212 213 return will_execute + [unique_id]214 215def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item):216 unique_id = current_item217 inputs = prompt[unique_id]['inputs']218 class_type = prompt[unique_id]['class_type']219 class_def = nodes.NODE_CLASS_MAPPINGS[class_type]220 221 is_changed_old = ''222 is_changed = ''223 to_delete = False224 if hasattr(class_def, 'IS_CHANGED'):225 if unique_id in old_prompt and 'is_changed' in old_prompt[unique_id]:226 is_changed_old = old_prompt[unique_id]['is_changed']227 if 'is_changed' not in prompt[unique_id]:228 input_data_all = get_input_data(inputs, class_def, unique_id, outputs)229 if input_data_all is not None:230 try:231 #is_changed = class_def.IS_CHANGED(**input_data_all)232 is_changed = map_node_over_list(class_def, input_data_all, "IS_CHANGED")233 prompt[unique_id]['is_changed'] = is_changed234 except:235 to_delete = True236 else:237 is_changed = prompt[unique_id]['is_changed']238 239 if unique_id not in outputs:240 return True241 242 if not to_delete:243 if is_changed != is_changed_old:244 to_delete = True245 elif unique_id not in old_prompt:246 to_delete = True247 elif inputs == old_prompt[unique_id]['inputs']:248 for x in inputs:249 input_data = inputs[x]250 251 if isinstance(input_data, list):252 input_unique_id = input_data[0]253 output_index = input_data[1]254 if input_unique_id in outputs:255 to_delete = recursive_output_delete_if_changed(prompt, old_prompt, outputs, input_unique_id)256 else:257 to_delete = True258 if to_delete:259 break260 else:261 to_delete = True262 263 if to_delete:264 d = outputs.pop(unique_id)265 del d266 return to_delete267 268class PromptExecutor:269 def __init__(self, server):270 self.outputs = {}271 self.object_storage = {}272 self.outputs_ui = {}273 self.old_prompt = {}274 self.server = server275 276 def handle_execution_error(self, prompt_id, prompt, current_outputs, executed, error, ex):277 node_id = error["node_id"]278 class_type = prompt[node_id]["class_type"]279 280 # First, send back the status to the frontend depending281 # on the exception type282 if isinstance(ex, comfy.model_management.InterruptProcessingException):283 mes = {284 "prompt_id": prompt_id,285 "node_id": node_id,286 "node_type": class_type,287 "executed": list(executed),288 }289 self.server.send_sync("execution_interrupted", mes, self.server.client_id)290 else:291 if self.server.client_id is not None:292 mes = {293 "prompt_id": prompt_id,294 "node_id": node_id,295 "node_type": class_type,296 "executed": list(executed),297 298 "exception_message": error["exception_message"],299 "exception_type": error["exception_type"],300 "traceback": error["traceback"],301 "current_inputs": error["current_inputs"],302 "current_outputs": error["current_outputs"],303 }304 self.server.send_sync("execution_error", mes, self.server.client_id)305 306 # Next, remove the subsequent outputs since they will not be executed307 to_delete = []308 for o in self.outputs:309 if (o not in current_outputs) and (o not in executed):310 to_delete += [o]311 if o in self.old_prompt:312 d = self.old_prompt.pop(o)313 del d314 for o in to_delete:315 d = self.outputs.pop(o)316 del d317 318 def execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):319 nodes.interrupt_processing(False)320 321 if "client_id" in extra_data:322 self.server.client_id = extra_data["client_id"]323 else:324 self.server.client_id = None325 326 if self.server.client_id is not None:327 self.server.send_sync("execution_start", { "prompt_id": prompt_id}, self.server.client_id)328 329 with torch.inference_mode():330 #delete cached outputs if nodes don't exist for them331 to_delete = []332 for o in self.outputs:333 if o not in prompt:334 to_delete += [o]335 for o in to_delete:336 d = self.outputs.pop(o)337 del d338 to_delete = []339 for o in self.object_storage:340 if o[0] not in prompt:341 to_delete += [o]342 else:343 p = prompt[o[0]]344 if o[1] != p['class_type']:345 to_delete += [o]346 for o in to_delete:347 d = self.object_storage.pop(o)348 del d349 350 for x in prompt:351 recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x)352 353 current_outputs = set(self.outputs.keys())354 for x in list(self.outputs_ui.keys()):355 if x not in current_outputs:356 d = self.outputs_ui.pop(x)357 del d358 359 comfy.model_management.cleanup_models()360 if self.server.client_id is not None:361 self.server.send_sync("execution_cached", { "nodes": list(current_outputs) , "prompt_id": prompt_id}, self.server.client_id)362 executed = set()363 output_node_id = None364 to_execute = []365 366 for node_id in list(execute_outputs):367 to_execute += [(0, node_id)]368 369 while len(to_execute) > 0:370 #always execute the output that depends on the least amount of unexecuted nodes first371 to_execute = sorted(list(map(lambda a: (len(recursive_will_execute(prompt, self.outputs, a[-1])), a[-1]), to_execute)))372 output_node_id = to_execute.pop(0)[-1]373 374 # This call shouldn't raise anything if there's an error deep in375 # the actual SD code, instead it will report the node where the376 # error was raised377 success, error, ex = recursive_execute(self.server, prompt, self.outputs, output_node_id, extra_data, executed, prompt_id, self.outputs_ui, self.object_storage)378 if success is not True:379 self.handle_execution_error(prompt_id, prompt, current_outputs, executed, error, ex)380 break381 382 for x in executed:383 self.old_prompt[x] = copy.deepcopy(prompt[x])384 self.server.last_node_id = None385 386 387 388def validate_inputs(prompt, item, validated):389 unique_id = item390 if unique_id in validated:391 return validated[unique_id]392 393 inputs = prompt[unique_id]['inputs']394 class_type = prompt[unique_id]['class_type']395 obj_class = nodes.NODE_CLASS_MAPPINGS[class_type]396 397 class_inputs = obj_class.INPUT_TYPES()398 required_inputs = class_inputs['required']399 400 errors = []401 valid = True402 403 for x in required_inputs:404 if x not in inputs:405 error = {406 "type": "required_input_missing",407 "message": "Required input is missing",408 "details": f"{x}",409 "extra_info": {410 "input_name": x411 }412 }413 errors.append(error)414 continue415 416 val = inputs[x]417 info = required_inputs[x]418 type_input = info[0]419 if isinstance(val, list):420 if len(val) != 2:421 error = {422 "type": "bad_linked_input",423 "message": "Bad linked input, must be a length-2 list of [node_id, slot_index]",424 "details": f"{x}",425 "extra_info": {426 "input_name": x,427 "input_config": info,428 "received_value": val429 }430 }431 errors.append(error)432 continue433 434 o_id = val[0]435 o_class_type = prompt[o_id]['class_type']436 r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES437 if r[val[1]] != type_input:438 received_type = r[val[1]]439 details = f"{x}, {received_type} != {type_input}"440 error = {441 "type": "return_type_mismatch",442 "message": "Return type mismatch between linked nodes",443 "details": details,444 "extra_info": {445 "input_name": x,446 "input_config": info,447 "received_type": received_type,448 "linked_node": val449 }450 }451 errors.append(error)452 continue453 try:454 r = validate_inputs(prompt, o_id, validated)455 if r[0] is False:456 # `r` will be set in `validated[o_id]` already457 valid = False458 continue459 except Exception as ex:460 typ, _, tb = sys.exc_info()461 valid = False462 exception_type = full_type_name(typ)463 reasons = [{464 "type": "exception_during_inner_validation",465 "message": "Exception when validating inner node",466 "details": str(ex),467 "extra_info": {468 "input_name": x,469 "input_config": info,470 "exception_message": str(ex),471 "exception_type": exception_type,472 "traceback": traceback.format_tb(tb),473 "linked_node": val474 }475 }]476 validated[o_id] = (False, reasons, o_id)477 continue478 else:479 try:480 if type_input == "INT":481 val = int(val)482 inputs[x] = val483 if type_input == "FLOAT":484 val = float(val)485 inputs[x] = val486 if type_input == "STRING":487 val = str(val)488 inputs[x] = val489 except Exception as ex:490 error = {491 "type": "invalid_input_type",492 "message": f"Failed to convert an input value to a {type_input} value",493 "details": f"{x}, {val}, {ex}",494 "extra_info": {495 "input_name": x,496 "input_config": info,497 "received_value": val,498 "exception_message": str(ex)499 }500 }501 errors.append(error)502 continue503 504 if len(info) > 1:505 if "min" in info[1] and val < info[1]["min"]:506 error = {507 "type": "value_smaller_than_min",508 "message": "Value {} smaller than min of {}".format(val, info[1]["min"]),509 "details": f"{x}",510 "extra_info": {511 "input_name": x,512 "input_config": info,513 "received_value": val,514 }515 }516 errors.append(error)517 continue518 if "max" in info[1] and val > info[1]["max"]:519 error = {520 "type": "value_bigger_than_max",521 "message": "Value {} bigger than max of {}".format(val, info[1]["max"]),522 "details": f"{x}",523 "extra_info": {524 "input_name": x,525 "input_config": info,526 "received_value": val,527 }528 }529 errors.append(error)530 continue531 532 if hasattr(obj_class, "VALIDATE_INPUTS"):533 input_data_all = get_input_data(inputs, obj_class, unique_id)534 #ret = obj_class.VALIDATE_INPUTS(**input_data_all)535 ret = map_node_over_list(obj_class, input_data_all, "VALIDATE_INPUTS")536 for i, r in enumerate(ret):537 if r is not True:538 details = f"{x}"539 if r is not False:540 details += f" - {str(r)}"541 542 error = {543 "type": "custom_validation_failed",544 "message": "Custom validation failed for node",545 "details": details,546 "extra_info": {547 "input_name": x,548 "input_config": info,549 "received_value": val,550 }551 }552 errors.append(error)553 continue554 else:555 if isinstance(type_input, list):556 if val not in type_input:557 input_config = info558 list_info = ""559 560 # Don't send back gigantic lists like if they're lots of561 # scanned model filepaths562 if len(type_input) > 20:563 list_info = f"(list of length {len(type_input)})"564 input_config = None565 else:566 list_info = str(type_input)567 568 error = {569 "type": "value_not_in_list",570 "message": "Value not in list",571 "details": f"{x}: '{val}' not in {list_info}",572 "extra_info": {573 "input_name": x,574 "input_config": input_config,575 "received_value": val,576 }577 }578 errors.append(error)579 continue580 581 if len(errors) > 0 or valid is not True:582 ret = (False, errors, unique_id)583 else:584 ret = (True, [], unique_id)585 586 validated[unique_id] = ret587 return ret588 589def full_type_name(klass):590 module = klass.__module__591 if module == 'builtins':592 return klass.__qualname__593 return module + '.' + klass.__qualname__594 595def validate_prompt(prompt):596 outputs = set()597 for x in prompt:598 class_ = nodes.NODE_CLASS_MAPPINGS[prompt[x]['class_type']]599 if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE == True:600 outputs.add(x)601 602 if len(outputs) == 0:603 error = {604 "type": "prompt_no_outputs",605 "message": "Prompt has no outputs",606 "details": "",607 "extra_info": {}608 }609 return (False, error, [], [])610 611 good_outputs = set()612 errors = []613 node_errors = {}614 validated = {}615 for o in outputs:616 valid = False617 reasons = []618 try:619 m = validate_inputs(prompt, o, validated)620 valid = m[0]621 reasons = m[1]622 except Exception as ex:623 typ, _, tb = sys.exc_info()624 valid = False625 exception_type = full_type_name(typ)626 reasons = [{627 "type": "exception_during_validation",628 "message": "Exception when validating node",629 "details": str(ex),630 "extra_info": {631 "exception_type": exception_type,632 "traceback": traceback.format_tb(tb)633 }634 }]635 validated[o] = (False, reasons, o)636 637 if valid is True:638 good_outputs.add(o)639 else:640 logging.error(f"Failed to validate prompt for output {o}:")641 if len(reasons) > 0:642 logging.error("* (prompt):")643 for reason in reasons:644 logging.error(f" - {reason['message']}: {reason['details']}")645 errors += [(o, reasons)]646 for node_id, result in validated.items():647 valid = result[0]648 reasons = result[1]649 # If a node upstream has errors, the nodes downstream will also650 # be reported as invalid, but there will be no errors attached.651 # So don't return those nodes as having errors in the response.652 if valid is not True and len(reasons) > 0:653 if node_id not in node_errors:654 class_type = prompt[node_id]['class_type']655 node_errors[node_id] = {656 "errors": reasons,657 "dependent_outputs": [],658 "class_type": class_type659 }660 logging.error(f"* {class_type} {node_id}:")661 for reason in reasons:662 logging.error(f" - {reason['message']}: {reason['details']}")663 node_errors[node_id]["dependent_outputs"].append(o)664 logging.error("Output will be ignored")665 666 if len(good_outputs) == 0:667 errors_list = []668 for o, errors in errors:669 for error in errors:670 errors_list.append(f"{error['message']}: {error['details']}")671 errors_list = "\n".join(errors_list)672 673 error = {674 "type": "prompt_outputs_failed_validation",675 "message": "Prompt outputs failed validation",676 "details": errors_list,677 "extra_info": {}678 }679 680 return (False, error, list(good_outputs), node_errors)681 682 return (True, None, list(good_outputs), node_errors)683 684MAXIMUM_HISTORY_SIZE = 10000685 686class PromptQueue:687 def __init__(self, server):688 self.server = server689 self.mutex = threading.RLock()690 self.not_empty = threading.Condition(self.mutex)691 self.task_counter = 0692 self.queue = []693 self.currently_running = {}694 self.history = {}695 server.prompt_queue = self696 697 def put(self, item):698 with self.mutex:699 heapq.heappush(self.queue, item)700 self.server.queue_updated()701 self.not_empty.notify()702 703 def get(self):704 with self.not_empty:705 while len(self.queue) == 0:706 self.not_empty.wait()707 item = heapq.heappop(self.queue)708 i = self.task_counter709 self.currently_running[i] = copy.deepcopy(item)710 self.task_counter += 1711 self.server.queue_updated()712 return (item, i)713 714 def task_done(self, item_id, outputs):715 with self.mutex:716 prompt = self.currently_running.pop(item_id)717 if len(self.history) > MAXIMUM_HISTORY_SIZE:718 self.history.pop(next(iter(self.history)))719 self.history[prompt[1]] = { "prompt": prompt, "outputs": {} }720 for o in outputs:721 self.history[prompt[1]]["outputs"][o] = outputs[o]722 self.server.queue_updated()723 724 def get_current_queue(self):725 with self.mutex:726 out = []727 for x in self.currently_running.values():728 out += [x]729 return (out, copy.deepcopy(self.queue))730 731 def get_tasks_remaining(self):732 with self.mutex:733 return len(self.queue) + len(self.currently_running)734 735 def wipe_queue(self):736 with self.mutex:737 self.queue = []738 self.server.queue_updated()739 740 def delete_queue_item(self, function):741 with self.mutex:742 for x in range(len(self.queue)):743 if function(self.queue[x]):744 if len(self.queue) == 1:745 self.wipe_queue()746 else:747 self.queue.pop(x)748 heapq.heapify(self.queue)749 self.server.queue_updated()750 return True751 return False752 753 def get_history(self, prompt_id=None, max_items=None, offset=-1):754 with self.mutex:755 if prompt_id is None:756 out = {}757 i = 0758 if offset < 0 and max_items is not None:759 offset = len(self.history) - max_items760 for k in self.history:761 if i >= offset:762 out[k] = self.history[k]763 if max_items is not None and len(out) >= max_items:764 break765 i += 1766 return out767 elif prompt_id in self.history:768 return {prompt_id: copy.deepcopy(self.history[prompt_id])}769 else:770 return {}771 772 def wipe_history(self):773 with self.mutex:774 self.history = {}775 776 def delete_history_item(self, id_to_delete):777 with self.mutex:778 self.history.pop(id_to_delete, None)779