315 lines
11 KiB
Python
315 lines
11 KiB
Python
|
import os
|
||
|
|
import sys
|
||
|
|
import copy
|
||
|
|
import json
|
||
|
|
import threading
|
||
|
|
import queue
|
||
|
|
import traceback
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
import nodes
|
||
|
|
|
||
|
|
|
||
|
|
def recursive_execute(prompt, outputs, current_item, extra_data={}):
|
||
|
|
unique_id = current_item
|
||
|
|
inputs = prompt[unique_id]['inputs']
|
||
|
|
class_type = prompt[unique_id]['class_type']
|
||
|
|
c_obj = nodes.NODE_CLASS_MAPPINGS[class_type]
|
||
|
|
valid_inputs = c_obj.INPUT_TYPES()
|
||
|
|
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(prompt, outputs, input_unique_id, extra_data)
|
||
|
|
|
||
|
|
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]
|
||
|
|
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
|
||
|
|
|
||
|
|
|
||
|
|
obj = c_obj()
|
||
|
|
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']
|
||
|
|
|
||
|
|
outputs[unique_id] = getattr(obj, obj.FUNCTION)(**input_data_all)
|
||
|
|
return executed + [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']
|
||
|
|
|
||
|
|
if unique_id not in outputs:
|
||
|
|
return True
|
||
|
|
|
||
|
|
to_delete = False
|
||
|
|
if 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:
|
||
|
|
print("deleted", unique_id)
|
||
|
|
d = outputs.pop(unique_id)
|
||
|
|
del d
|
||
|
|
return to_delete
|
||
|
|
|
||
|
|
class PromptExecutor:
|
||
|
|
def __init__(self):
|
||
|
|
self.outputs = {}
|
||
|
|
self.old_prompt = {}
|
||
|
|
|
||
|
|
def execute(self, prompt, extra_data={}):
|
||
|
|
with torch.no_grad():
|
||
|
|
for x in prompt:
|
||
|
|
recursive_output_delete_if_changed(prompt, self.old_prompt, self.outputs, x)
|
||
|
|
|
||
|
|
current_outputs = set(self.outputs.keys())
|
||
|
|
executed = []
|
||
|
|
try:
|
||
|
|
for x in prompt:
|
||
|
|
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(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])
|
||
|
|
|
||
|
|
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 = 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. {}, {}".format(class_type, x))
|
||
|
|
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()
|
||
|
|
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")
|
||
|
|
|
||
|
|
if len(good_outputs) == 0:
|
||
|
|
return (False, "Prompt has no properly connected outputs")
|
||
|
|
|
||
|
|
return (True, "")
|
||
|
|
|
||
|
|
def prompt_worker(q):
|
||
|
|
e = PromptExecutor()
|
||
|
|
while True:
|
||
|
|
item = q.get()
|
||
|
|
e.execute(item[-2], item[-1])
|
||
|
|
q.task_done()
|
||
|
|
|
||
|
|
|
||
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||
|
|
|
||
|
|
class PromptServer(BaseHTTPRequestHandler):
|
||
|
|
def _set_headers(self, code=200, ct='text/html'):
|
||
|
|
self.send_response(code)
|
||
|
|
self.send_header('Content-type', ct)
|
||
|
|
self.end_headers()
|
||
|
|
def log_message(self, format, *args):
|
||
|
|
pass
|
||
|
|
def do_GET(self):
|
||
|
|
if self.path == "/prompt":
|
||
|
|
self._set_headers(ct='application/json')
|
||
|
|
prompt_info = {}
|
||
|
|
exec_info = {}
|
||
|
|
exec_info['queue_remaining'] = self.server.prompt_queue.unfinished_tasks
|
||
|
|
prompt_info['exec_info'] = exec_info
|
||
|
|
self.wfile.write(json.dumps(prompt_info).encode('utf-8'))
|
||
|
|
elif self.path == "/object_info":
|
||
|
|
self._set_headers(ct='application/json')
|
||
|
|
out = {}
|
||
|
|
for x in nodes.NODE_CLASS_MAPPINGS:
|
||
|
|
obj_class = nodes.NODE_CLASS_MAPPINGS[x]
|
||
|
|
info = {}
|
||
|
|
info['input'] = obj_class.INPUT_TYPES()
|
||
|
|
info['output'] = obj_class.RETURN_TYPES
|
||
|
|
info['name'] = x #TODO
|
||
|
|
info['description'] = ''
|
||
|
|
out[x] = info
|
||
|
|
self.wfile.write(json.dumps(out).encode('utf-8'))
|
||
|
|
elif self.path[1:] in os.listdir(self.server.server_dir):
|
||
|
|
self._set_headers()
|
||
|
|
with open(os.path.join(self.server.server_dir, self.path[1:]), "rb") as f:
|
||
|
|
self.wfile.write(f.read())
|
||
|
|
else:
|
||
|
|
self._set_headers()
|
||
|
|
with open(os.path.join(self.server.server_dir, "index.html"), "rb") as f:
|
||
|
|
self.wfile.write(f.read())
|
||
|
|
|
||
|
|
def do_HEAD(self):
|
||
|
|
self._set_headers()
|
||
|
|
|
||
|
|
def do_POST(self):
|
||
|
|
resp_code = 200
|
||
|
|
out_string = ""
|
||
|
|
if self.path == "/prompt":
|
||
|
|
print("got prompt")
|
||
|
|
self.data_string = self.rfile.read(int(self.headers['Content-Length']))
|
||
|
|
json_data = json.loads(self.data_string)
|
||
|
|
if "number" in json_data:
|
||
|
|
number = float(json_data['number'])
|
||
|
|
else:
|
||
|
|
number = self.server.number
|
||
|
|
self.server.number += 1
|
||
|
|
if "prompt" in json_data:
|
||
|
|
prompt = json_data["prompt"]
|
||
|
|
valid = validate_prompt(prompt)
|
||
|
|
extra_data = {}
|
||
|
|
if "extra_data" in json_data:
|
||
|
|
extra_data = json_data["extra_data"]
|
||
|
|
if valid[0]:
|
||
|
|
self.server.prompt_queue.put((number, id(prompt), prompt, extra_data))
|
||
|
|
else:
|
||
|
|
resp_code = 400
|
||
|
|
out_string = valid[1]
|
||
|
|
print("invalid prompt:", valid[1])
|
||
|
|
self._set_headers(code=resp_code)
|
||
|
|
self.end_headers()
|
||
|
|
self.wfile.write(out_string.encode('utf8'))
|
||
|
|
return
|
||
|
|
|
||
|
|
|
||
|
|
def run(prompt_queue, address='', port=8188):
|
||
|
|
server_address = (address, port)
|
||
|
|
httpd = HTTPServer(server_address, PromptServer)
|
||
|
|
httpd.server_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "webshit")
|
||
|
|
httpd.prompt_queue = prompt_queue
|
||
|
|
httpd.number = 0
|
||
|
|
if server_address[0] == '':
|
||
|
|
addr = '0.0.0.0'
|
||
|
|
else:
|
||
|
|
addr = server_address[0]
|
||
|
|
print("Starting server\n")
|
||
|
|
print("To see the GUI go to: http://{}:{}".format(addr, server_address[1]))
|
||
|
|
httpd.serve_forever()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
q = queue.PriorityQueue()
|
||
|
|
threading.Thread(target=prompt_worker, daemon=True, args=(q,)).start()
|
||
|
|
run(q, address='127.0.0.1', port=8188)
|
||
|
|
|
||
|
|
|