537 lines
19 KiB
Python
537 lines
19 KiB
Python
import os
|
|
import sys
|
|
import copy
|
|
import json
|
|
import threading
|
|
import heapq
|
|
import traceback
|
|
import asyncio
|
|
|
|
try:
|
|
import aiohttp
|
|
from aiohttp import web
|
|
except ImportError:
|
|
print("Module 'aiohttp' not installed. Please install it via:")
|
|
print("pip install aiohttp")
|
|
print("or")
|
|
print("pip install -r requirements.txt")
|
|
sys.exit()
|
|
|
|
if __name__ == "__main__":
|
|
if '--help' in sys.argv:
|
|
print("Valid Command line Arguments:")
|
|
print("\t--listen\t\t\tListen on 0.0.0.0 so the UI can be accessed from other computers.")
|
|
print("\t--port 8188\t\t\tSet the listen port.")
|
|
print("\t--dont-upcast-attention\t\tDisable upcasting of attention \n\t\t\t\t\tcan boost speed but increase the chances of black images.\n")
|
|
print("\t--use-split-cross-attention\tUse the split cross attention optimization instead of the sub-quadratic one.\n\t\t\t\t\tIgnored when xformers is used.")
|
|
print()
|
|
print("\t--normalvram\t\t\tUsed to force normal vram use if lowvram gets automatically enabled.")
|
|
print("\t--lowvram\t\t\tSplit the unet in parts to use less vram.")
|
|
print("\t--novram\t\t\tWhen lowvram isn't enough.")
|
|
print()
|
|
exit()
|
|
|
|
if '--dont-upcast-attention' in sys.argv:
|
|
print("disabling upcasting of attention")
|
|
os.environ['ATTN_PRECISION'] = "fp16"
|
|
|
|
import torch
|
|
import nodes
|
|
|
|
def get_input_data(inputs, class_def, 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]
|
|
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']
|
|
return input_data_all
|
|
|
|
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']
|
|
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(prompt, outputs, input_unique_id, extra_data)
|
|
|
|
input_data_all = get_input_data(inputs, class_def, outputs, prompt, extra_data)
|
|
obj = class_def()
|
|
|
|
outputs[unique_id] = getattr(obj, obj.FUNCTION)(**input_data_all)
|
|
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)
|
|
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:
|
|
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:
|
|
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:
|
|
#always execute the output that depends on the least amount of unexecuted nodes first
|
|
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(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. {}, {}: {} 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(map(lambda a: "{}".format(a[1]), errors))
|
|
return (False, "Prompt has no properly connected outputs\n {}".format(errors_list))
|
|
|
|
return (True, "")
|
|
|
|
def prompt_worker(q):
|
|
e = PromptExecutor()
|
|
while True:
|
|
item, item_id = q.get()
|
|
e.execute(item[-2], item[-1])
|
|
q.task_done(item_id)
|
|
|
|
class PromptQueue:
|
|
def __init__(self, socket_handler):
|
|
self.socket_handler = socket_handler
|
|
self.mutex = threading.RLock()
|
|
self.not_empty = threading.Condition(self.mutex)
|
|
self.task_counter = 0
|
|
self.queue = []
|
|
self.currently_running = {}
|
|
socket_handler.prompt_queue = self
|
|
|
|
def put(self, item):
|
|
with self.mutex:
|
|
heapq.heappush(self.queue, item)
|
|
self.socket_handler.queue_updated(self)
|
|
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.socket_handler.queue_updated(self)
|
|
return (item, i)
|
|
|
|
def task_done(self, item_id):
|
|
with self.mutex:
|
|
self.currently_running.pop(item_id)
|
|
self.socket_handler.queue_updated(self)
|
|
|
|
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.socket_handler.queue_updated(self)
|
|
|
|
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.socket_handler.queue_updated(self)
|
|
return True
|
|
return False
|
|
|
|
def get_queue_info(prompt_queue):
|
|
prompt_info = {}
|
|
exec_info = {}
|
|
exec_info['queue_remaining'] = prompt_queue.get_tasks_remaining()
|
|
prompt_info['exec_info'] = exec_info
|
|
return prompt_info
|
|
|
|
class SocketHandler():
|
|
def __init__(self, loop):
|
|
self.connected = set()
|
|
self.messages = asyncio.Queue()
|
|
self.loop = loop
|
|
|
|
async def publish_loop(self):
|
|
while True:
|
|
msg = await self.messages.get()
|
|
await self.send(msg)
|
|
|
|
def queue_updated(self, queue):
|
|
# This is called by the queue processing thread so we need to make it thread safe
|
|
loop.call_soon_threadsafe(self.messages.put_nowait, { 'type': 'status', 'status': get_queue_info(queue) })
|
|
|
|
async def send(self, message, socket = None):
|
|
if isinstance(message, str) == False:
|
|
message = json.dumps(message)
|
|
|
|
if socket is None:
|
|
for ws in self.connected:
|
|
await ws.send_str(message)
|
|
else:
|
|
await socket.send_str(message)
|
|
|
|
async def process(self, request):
|
|
ws = web.WebSocketResponse()
|
|
await ws.prepare(request)
|
|
self.connected.add(ws)
|
|
try:
|
|
# Send initial state to the new client
|
|
await self.send({ 'type': 'status', 'status': get_queue_info(self.prompt_queue) }, ws)
|
|
async for msg in ws:
|
|
if msg.type == aiohttp.WSMsgType.ERROR:
|
|
print('ws connection closed with exception %s' % ws.exception())
|
|
finally:
|
|
self.connected.remove(ws)
|
|
|
|
return ws
|
|
|
|
class PromptServer():
|
|
def __init__(self, prompt_queue, socket_handler):
|
|
self.prompt_queue = prompt_queue
|
|
self.socket_handler = socket_handler
|
|
self.number = 0
|
|
self.app = web.Application()
|
|
self.web_root = os.path.join(os.path.dirname(os.path.realpath(__file__)), "webshit")
|
|
routes = web.RouteTableDef()
|
|
|
|
@routes.get('/ws')
|
|
async def websocket_handler(request):
|
|
return await self.socket_handler.process(request)
|
|
|
|
@routes.get("/")
|
|
async def get_root(request):
|
|
return web.FileResponse(os.path.join(self.web_root, "index.html"))
|
|
|
|
@routes.get("/prompt")
|
|
async def get_prompt(request):
|
|
return web.json_response(get_queue_info(self.prompt_queue))
|
|
|
|
@routes.get("/object_info")
|
|
async def get_object_info(request):
|
|
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'] = ''
|
|
info['category'] = 'sd'
|
|
if hasattr(obj_class, 'CATEGORY'):
|
|
info['category'] = obj_class.CATEGORY
|
|
out[x] = info
|
|
return web.json_response(out)
|
|
|
|
@routes.get("/queue")
|
|
async def get_queue(request):
|
|
queue_info = {}
|
|
current_queue = self.prompt_queue.get_current_queue()
|
|
queue_info['queue_running'] = current_queue[0]
|
|
queue_info['queue_pending'] = current_queue[1]
|
|
return web.json_response(queue_info)
|
|
|
|
@routes.post("/prompt")
|
|
async def post_prompt(request):
|
|
print("got prompt")
|
|
resp_code = 200
|
|
out_string = ""
|
|
json_data = await request.json()
|
|
|
|
if "number" in json_data:
|
|
number = float(json_data['number'])
|
|
else:
|
|
number = self.number
|
|
if "front" in json_data:
|
|
if json_data['front']:
|
|
number = -number
|
|
|
|
self.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.prompt_queue.put((number, id(prompt), prompt, extra_data))
|
|
else:
|
|
resp_code = 400
|
|
out_string = valid[1]
|
|
print("invalid prompt:", valid[1])
|
|
|
|
return web.Response(body=out_string, status=resp_code)
|
|
|
|
@routes.post("/queue")
|
|
async def post_queue(request):
|
|
json_data = await request.json()
|
|
if "clear" in json_data:
|
|
if json_data["clear"]:
|
|
self.prompt_queue.wipe_queue()
|
|
if "delete" in json_data:
|
|
to_delete = json_data['delete']
|
|
for id_to_delete in to_delete:
|
|
delete_func = lambda a: a[1] == int(id_to_delete)
|
|
self.prompt_queue.delete_queue_item(delete_func)
|
|
|
|
return web.Response(status=200)
|
|
|
|
self.app.add_routes(routes)
|
|
self.app.add_routes([
|
|
web.static('/', self.web_root),
|
|
])
|
|
|
|
async def start_server(server, address, port):
|
|
runner = web.AppRunner(server.app)
|
|
await runner.setup()
|
|
site = web.TCPSite(runner, address, port)
|
|
await site.start()
|
|
|
|
if address == '':
|
|
address = '0.0.0.0'
|
|
print("Starting server\n")
|
|
print("To see the GUI go to: http://{}:{}".format(address, port))
|
|
|
|
async def run(prompt_queue, socket_handler, address='', port=8188):
|
|
server = PromptServer(prompt_queue, socket_handler)
|
|
await asyncio.gather(start_server(server, address, port), socket_handler.publish_loop())
|
|
|
|
if __name__ == "__main__":
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
socket_handler = SocketHandler(loop)
|
|
q = PromptQueue(socket_handler)
|
|
threading.Thread(target=prompt_worker, daemon=True, args=(q,)).start()
|
|
if '--listen' in sys.argv:
|
|
address = '0.0.0.0'
|
|
else:
|
|
address = '127.0.0.1'
|
|
|
|
port = 8188
|
|
try:
|
|
p_index = sys.argv.index('--port')
|
|
port = int(sys.argv[p_index + 1])
|
|
except:
|
|
pass
|
|
|
|
loop.run_until_complete(run(q, socket_handler, address=address, port=port))
|
|
|