| import os |
| import sys |
| import copy |
| import json |
| import threading |
| import heapq |
| import traceback |
| import gc |
|
|
| import torch |
| import nodes |
|
|
| def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}): |
| valid_inputs = class_def.INPUT_TYPES() |
| input_data_all = {} |
| for x in inputs: |
| input_data = inputs[x] |
| if isinstance(input_data, list): |
| input_unique_id = input_data[0] |
| output_index = input_data[1] |
| if input_unique_id not in outputs: |
| return None |
| obj = outputs[input_unique_id][output_index] |
| input_data_all[x] = obj |
| else: |
| if ("required" in valid_inputs and x in valid_inputs["required"]) or ("optional" in valid_inputs and x in valid_inputs["optional"]): |
| input_data_all[x] = input_data |
|
|
| if "hidden" in valid_inputs: |
| h = valid_inputs["hidden"] |
| for x in h: |
| if h[x] == "PROMPT": |
| input_data_all[x] = prompt |
| if h[x] == "EXTRA_PNGINFO": |
| if "extra_pnginfo" in extra_data: |
| input_data_all[x] = extra_data['extra_pnginfo'] |
| if h[x] == "UNIQUE_ID": |
| input_data_all[x] = unique_id |
| return input_data_all |
|
|
| def recursive_execute(server, prompt, outputs, current_item, extra_data={}): |
| unique_id = current_item |
| inputs = prompt[unique_id]['inputs'] |
| class_type = prompt[unique_id]['class_type'] |
| class_def = nodes.NODE_CLASS_MAPPINGS[class_type] |
| if unique_id in outputs: |
| return [] |
|
|
| executed = [] |
|
|
| for x in inputs: |
| input_data = inputs[x] |
|
|
| if isinstance(input_data, list): |
| input_unique_id = input_data[0] |
| output_index = input_data[1] |
| if input_unique_id not in outputs: |
| executed += recursive_execute(server, prompt, outputs, input_unique_id, extra_data) |
|
|
| input_data_all = get_input_data(inputs, class_def, unique_id, outputs, prompt, extra_data) |
| if server.client_id is not None: |
| server.last_node_id = unique_id |
| server.send_sync("executing", { "node": unique_id }, server.client_id) |
| obj = class_def() |
|
|
| nodes.before_node_execution() |
| outputs[unique_id] = getattr(obj, obj.FUNCTION)(**input_data_all) |
| if "ui" in outputs[unique_id]: |
| if server.client_id is not None: |
| server.send_sync("executed", { "node": unique_id, "output": outputs[unique_id]["ui"] }, server.client_id) |
| if "result" in outputs[unique_id]: |
| outputs[unique_id] = outputs[unique_id]["result"] |
| return executed + [unique_id] |
|
|
| def recursive_will_execute(prompt, outputs, current_item): |
| unique_id = current_item |
| inputs = prompt[unique_id]['inputs'] |
| will_execute = [] |
| if unique_id in outputs: |
| return [] |
|
|
| for x in inputs: |
| input_data = inputs[x] |
| if isinstance(input_data, list): |
| input_unique_id = input_data[0] |
| output_index = input_data[1] |
| if input_unique_id not in outputs: |
| will_execute += recursive_will_execute(prompt, outputs, input_unique_id) |
|
|
| return will_execute + [unique_id] |
|
|
| def recursive_output_delete_if_changed(prompt, old_prompt, outputs, current_item): |
| unique_id = current_item |
| inputs = prompt[unique_id]['inputs'] |
| class_type = prompt[unique_id]['class_type'] |
| class_def = nodes.NODE_CLASS_MAPPINGS[class_type] |
|
|
| is_changed_old = '' |
| is_changed = '' |
| if hasattr(class_def, 'IS_CHANGED'): |
| if unique_id in old_prompt and 'is_changed' in old_prompt[unique_id]: |
| is_changed_old = old_prompt[unique_id]['is_changed'] |
| if 'is_changed' not in prompt[unique_id]: |
| input_data_all = get_input_data(inputs, class_def, unique_id, outputs) |
| if input_data_all is not None: |
| is_changed = class_def.IS_CHANGED(**input_data_all) |
| prompt[unique_id]['is_changed'] = is_changed |
| else: |
| is_changed = prompt[unique_id]['is_changed'] |
|
|
| if unique_id not in outputs: |
| return True |
|
|
| to_delete = False |
| if is_changed != is_changed_old: |
| to_delete = True |
| elif unique_id not in old_prompt: |
| to_delete = True |
| elif inputs == old_prompt[unique_id]['inputs']: |
| for x in inputs: |
| input_data = inputs[x] |
|
|
| if isinstance(input_data, list): |
| input_unique_id = input_data[0] |
| output_index = input_data[1] |
| if input_unique_id in outputs: |
| to_delete = recursive_output_delete_if_changed(prompt, old_prompt, outputs, input_unique_id) |
| else: |
| to_delete = True |
| if to_delete: |
| break |
| else: |
| to_delete = True |
|
|
| if to_delete: |
| d = outputs.pop(unique_id) |
| del d |
| return to_delete |
|
|
| class PromptExecutor: |
| def __init__(self, server): |
| self.outputs = {} |
| self.old_prompt = {} |
| self.server = server |
|
|
| def execute(self, prompt, extra_data={}): |
| nodes.interrupt_processing(False) |
|
|
| if "client_id" in extra_data: |
| self.server.client_id = extra_data["client_id"] |
| else: |
| self.server.client_id = None |
|
|
| with torch.inference_mode(): |
| for x in prompt: |
| recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x) |
|
|
| current_outputs = set(self.outputs.keys()) |
| executed = [] |
| try: |
| to_execute = [] |
| for x in prompt: |
| class_ = nodes.NODE_CLASS_MAPPINGS[prompt[x]['class_type']] |
| if hasattr(class_, 'OUTPUT_NODE'): |
| to_execute += [(0, x)] |
|
|
| while len(to_execute) > 0: |
| |
| to_execute = sorted(list(map(lambda a: (len(recursive_will_execute(prompt, self.outputs, a[-1])), a[-1]), to_execute))) |
| x = to_execute.pop(0)[-1] |
|
|
| class_ = nodes.NODE_CLASS_MAPPINGS[prompt[x]['class_type']] |
| if hasattr(class_, 'OUTPUT_NODE'): |
| if class_.OUTPUT_NODE == True: |
| valid = False |
| try: |
| m = validate_inputs(prompt, x) |
| valid = m[0] |
| except: |
| valid = False |
| if valid: |
| executed += recursive_execute(self.server, prompt, self.outputs, x, extra_data) |
| except Exception as e: |
| print(traceback.format_exc()) |
| to_delete = [] |
| for o in self.outputs: |
| if o not in current_outputs: |
| to_delete += [o] |
| if o in self.old_prompt: |
| d = self.old_prompt.pop(o) |
| del d |
| for o in to_delete: |
| d = self.outputs.pop(o) |
| del d |
| else: |
| executed = set(executed) |
| for x in executed: |
| self.old_prompt[x] = copy.deepcopy(prompt[x]) |
| finally: |
| self.server.last_node_id = None |
| if self.server.client_id is not None: |
| self.server.send_sync("executing", { "node": None }, self.server.client_id) |
|
|
| gc.collect() |
| if torch.cuda.is_available(): |
| if torch.version.cuda: |
| torch.cuda.empty_cache() |
| torch.cuda.ipc_collect() |
|
|
|
|
| def validate_inputs(prompt, item): |
| unique_id = item |
| inputs = prompt[unique_id]['inputs'] |
| class_type = prompt[unique_id]['class_type'] |
| obj_class = nodes.NODE_CLASS_MAPPINGS[class_type] |
|
|
| class_inputs = obj_class.INPUT_TYPES() |
| required_inputs = class_inputs['required'] |
| for x in required_inputs: |
| if x not in inputs: |
| return (False, "Required input is missing. {}, {}".format(class_type, x)) |
| val = inputs[x] |
| info = required_inputs[x] |
| type_input = info[0] |
| if isinstance(val, list): |
| if len(val) != 2: |
| return (False, "Bad Input. {}, {}".format(class_type, x)) |
| o_id = val[0] |
| o_class_type = prompt[o_id]['class_type'] |
| r = nodes.NODE_CLASS_MAPPINGS[o_class_type].RETURN_TYPES |
| if r[val[1]] != type_input: |
| return (False, "Return type mismatch. {}, {}, {} != {}".format(class_type, x, r[val[1]], type_input)) |
| r = validate_inputs(prompt, o_id) |
| if r[0] == False: |
| return r |
| else: |
| if type_input == "INT": |
| val = int(val) |
| inputs[x] = val |
| if type_input == "FLOAT": |
| val = float(val) |
| inputs[x] = val |
| if type_input == "STRING": |
| val = str(val) |
| inputs[x] = val |
|
|
| if len(info) > 1: |
| if "min" in info[1] and val < info[1]["min"]: |
| return (False, "Value smaller than min. {}, {}".format(class_type, x)) |
| if "max" in info[1] and val > info[1]["max"]: |
| return (False, "Value bigger than max. {}, {}".format(class_type, x)) |
|
|
| if isinstance(type_input, list): |
| if val not in type_input: |
| return (False, "Value not in list. {}, {}: {} not in {}".format(class_type, x, val, type_input)) |
| return (True, "") |
|
|
| def validate_prompt(prompt): |
| outputs = set() |
| for x in prompt: |
| class_ = nodes.NODE_CLASS_MAPPINGS[prompt[x]['class_type']] |
| if hasattr(class_, 'OUTPUT_NODE') and class_.OUTPUT_NODE == True: |
| outputs.add(x) |
|
|
| if len(outputs) == 0: |
| return (False, "Prompt has no outputs") |
|
|
| good_outputs = set() |
| errors = [] |
| for o in outputs: |
| valid = False |
| reason = "" |
| try: |
| m = validate_inputs(prompt, o) |
| valid = m[0] |
| reason = m[1] |
| except: |
| valid = False |
| reason = "Parsing error" |
|
|
| if valid == True: |
| good_outputs.add(x) |
| else: |
| print("Failed to validate prompt for output {} {}".format(o, reason)) |
| print("output will be ignored") |
| errors += [(o, reason)] |
|
|
| if len(good_outputs) == 0: |
| errors_list = "\n".join(set(map(lambda a: "{}".format(a[1]), errors))) |
| return (False, "Prompt has no properly connected outputs\n {}".format(errors_list)) |
|
|
| return (True, "") |
|
|
|
|
| class PromptQueue: |
| def __init__(self, server): |
| self.server = server |
| self.mutex = threading.RLock() |
| self.not_empty = threading.Condition(self.mutex) |
| self.task_counter = 0 |
| self.queue = [] |
| self.currently_running = {} |
| self.history = {} |
| server.prompt_queue = self |
|
|
| def put(self, item): |
| with self.mutex: |
| heapq.heappush(self.queue, item) |
| self.server.queue_updated() |
| self.not_empty.notify() |
|
|
| def get(self): |
| with self.not_empty: |
| while len(self.queue) == 0: |
| self.not_empty.wait() |
| item = heapq.heappop(self.queue) |
| i = self.task_counter |
| self.currently_running[i] = copy.deepcopy(item) |
| self.task_counter += 1 |
| self.server.queue_updated() |
| return (item, i) |
|
|
| def task_done(self, item_id, outputs): |
| with self.mutex: |
| prompt = self.currently_running.pop(item_id) |
| self.history[prompt[1]] = { "prompt": prompt, "outputs": {} } |
| for o in outputs: |
| if "ui" in outputs[o]: |
| self.history[prompt[1]]["outputs"][o] = outputs[o]["ui"] |
| self.server.queue_updated() |
|
|
| def get_current_queue(self): |
| with self.mutex: |
| out = [] |
| for x in self.currently_running.values(): |
| out += [x] |
| return (out, copy.deepcopy(self.queue)) |
|
|
| def get_tasks_remaining(self): |
| with self.mutex: |
| return len(self.queue) + len(self.currently_running) |
|
|
| def wipe_queue(self): |
| with self.mutex: |
| self.queue = [] |
| self.server.queue_updated() |
|
|
| def delete_queue_item(self, function): |
| with self.mutex: |
| for x in range(len(self.queue)): |
| if function(self.queue[x]): |
| if len(self.queue) == 1: |
| self.wipe_queue() |
| else: |
| self.queue.pop(x) |
| heapq.heapify(self.queue) |
| self.server.queue_updated() |
| return True |
| return False |
|
|
| def get_history(self): |
| with self.mutex: |
| return copy.deepcopy(self.history) |
|
|
| def wipe_history(self): |
| with self.mutex: |
| self.history = {} |
|
|
| def delete_history_item(self, id_to_delete): |
| with self.mutex: |
| self.history.pop(id_to_delete, None) |
|
|
