| import os |
| import sys |
| import asyncio |
| import nodes |
| import folder_paths |
| import execution |
| import uuid |
| import json |
| import glob |
| 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() |
|
|
| import mimetypes |
|
|
|
|
| @web.middleware |
| async def cache_control(request: web.Request, handler): |
| response: web.Response = await handler(request) |
| if request.path.endswith('.js') or request.path.endswith('.css'): |
| response.headers.setdefault('Cache-Control', 'no-cache') |
| return response |
|
|
| class PromptServer(): |
| def __init__(self, loop): |
| PromptServer.instance = self |
|
|
| mimetypes.init(); |
| mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8' |
| self.prompt_queue = None |
| self.loop = loop |
| self.messages = asyncio.Queue() |
| self.number = 0 |
| self.app = web.Application(client_max_size=20971520, middlewares=[cache_control]) |
| self.sockets = dict() |
| self.web_root = os.path.join(os.path.dirname( |
| os.path.realpath(__file__)), "web") |
| routes = web.RouteTableDef() |
| self.routes = routes |
| self.last_node_id = None |
| self.client_id = None |
|
|
| @routes.get('/ws') |
| async def websocket_handler(request): |
| ws = web.WebSocketResponse() |
| await ws.prepare(request) |
| sid = request.rel_url.query.get('clientId', '') |
| if sid: |
| |
| self.sockets.pop(sid, None) |
| else: |
| sid = uuid.uuid4().hex |
|
|
| self.sockets[sid] = ws |
|
|
| try: |
| |
| await self.send("status", { "status": self.get_queue_info(), 'sid': sid }, sid) |
| |
| if self.client_id == sid and self.last_node_id is not None: |
| await self.send("executing", { "node": self.last_node_id }, sid) |
| |
| async for msg in ws: |
| if msg.type == aiohttp.WSMsgType.ERROR: |
| print('ws connection closed with exception %s' % ws.exception()) |
| finally: |
| self.sockets.pop(sid, None) |
| return ws |
|
|
| @routes.get("/") |
| async def get_root(request): |
| return web.FileResponse(os.path.join(self.web_root, "index.html")) |
|
|
| @routes.get("/embeddings") |
| def get_embeddings(self): |
| embeddings = folder_paths.get_filename_list("embeddings") |
| return web.json_response(list(map(lambda a: os.path.splitext(a)[0].lower(), embeddings))) |
|
|
| @routes.get("/extensions") |
| async def get_extensions(request): |
| files = glob.glob(os.path.join(self.web_root, 'extensions/**/*.js'), recursive=True) |
| return web.json_response(list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files))) |
|
|
| @routes.post("/upload/image") |
| async def upload_image(request): |
| upload_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") |
|
|
| if not os.path.exists(upload_dir): |
| os.makedirs(upload_dir) |
| |
| post = await request.post() |
| image = post.get("image") |
|
|
| if image and image.file: |
| filename = image.filename |
| if not filename: |
| return web.Response(status=400) |
|
|
| split = os.path.splitext(filename) |
| i = 1 |
| while os.path.exists(os.path.join(upload_dir, filename)): |
| filename = f"{split[0]} ({i}){split[1]}" |
| i += 1 |
|
|
| filepath = os.path.join(upload_dir, filename) |
|
|
| with open(filepath, "wb") as f: |
| f.write(image.file.read()) |
| |
| return web.json_response({"name" : filename}) |
| else: |
| return web.Response(status=400) |
|
|
|
|
| @routes.get("/view") |
| async def view_image(request): |
| if "filename" in request.rel_url.query: |
| type = request.rel_url.query.get("type", "output") |
| if type not in ["output", "input", "temp"]: |
| return web.Response(status=400) |
|
|
| output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type) |
| if "subfolder" in request.rel_url.query: |
| full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"]) |
| if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir: |
| return web.Response(status=403) |
| output_dir = full_output_dir |
|
|
| filename = request.rel_url.query["filename"] |
| filename = os.path.basename(filename) |
| file = os.path.join(output_dir, filename) |
|
|
| if os.path.isfile(file): |
| return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""}) |
| |
| return web.Response(status=404) |
|
|
| @routes.get("/prompt") |
| async def get_prompt(request): |
| return web.json_response(self.get_queue_info()) |
|
|
| @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['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output'] |
| info['name'] = x |
| 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("/history") |
| async def get_history(request): |
| return web.json_response(self.prompt_queue.get_history()) |
|
|
| @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 = execution.validate_prompt(prompt) |
| extra_data = {} |
| if "extra_data" in json_data: |
| extra_data = json_data["extra_data"] |
|
|
| if "client_id" in json_data: |
| extra_data["client_id"] = json_data["client_id"] |
| 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) |
|
|
| @routes.post("/interrupt") |
| async def post_interrupt(request): |
| nodes.interrupt_processing() |
| return web.Response(status=200) |
|
|
| @routes.post("/history") |
| async def post_history(request): |
| json_data = await request.json() |
| if "clear" in json_data: |
| if json_data["clear"]: |
| self.prompt_queue.wipe_history() |
| if "delete" in json_data: |
| to_delete = json_data['delete'] |
| for id_to_delete in to_delete: |
| self.prompt_queue.delete_history_item(id_to_delete) |
|
|
| return web.Response(status=200) |
| |
| def add_routes(self): |
| self.app.add_routes(self.routes) |
| self.app.add_routes([ |
| web.static('/', self.web_root), |
| ]) |
|
|
| def get_queue_info(self): |
| prompt_info = {} |
| exec_info = {} |
| exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining() |
| prompt_info['exec_info'] = exec_info |
| return prompt_info |
|
|
| async def send(self, event, data, sid=None): |
| message = {"type": event, "data": data} |
| |
| if isinstance(message, str) == False: |
| message = json.dumps(message) |
|
|
| if sid is None: |
| for ws in self.sockets.values(): |
| await ws.send_str(message) |
| elif sid in self.sockets: |
| await self.sockets[sid].send_str(message) |
|
|
| def send_sync(self, event, data, sid=None): |
| self.loop.call_soon_threadsafe( |
| self.messages.put_nowait, (event, data, sid)) |
|
|
| def queue_updated(self): |
| self.send_sync("status", { "status": self.get_queue_info() }) |
|
|
| async def publish_loop(self): |
| while True: |
| msg = await self.messages.get() |
| await self.send(*msg) |
|
|
| async def start(self, address, port, verbose=True, call_on_start=None): |
| runner = web.AppRunner(self.app) |
| await runner.setup() |
| site = web.TCPSite(runner, address, port) |
| await site.start() |
|
|
| if address == '': |
| address = '0.0.0.0' |
| if verbose: |
| print("Starting server\n") |
| print("To see the GUI go to: http://{}:{}".format(address, port)) |
| if call_on_start is not None: |
| call_on_start(address, port) |
|
|
|
|
