fred-dev/comfy_ui_ali
0
1import os2import sys3import asyncio4import traceback5 6import nodes7import folder_paths8import execution9import uuid10import urllib11import json12import glob13import struct14import ssl15import socket16import ipaddress17from PIL import Image, ImageOps18from PIL.PngImagePlugin import PngInfo19from io import BytesIO20 21import aiohttp22from aiohttp import web23import logging24 25import mimetypes26from comfy.cli_args import args27import comfy.utils28import comfy.model_management29import node_helpers30from comfyui_version import __version__31from app.frontend_management import FrontendManager32from app.user_manager import UserManager33from app.model_manager import ModelFileManager34from app.custom_node_manager import CustomNodeManager35from typing import Optional36from api_server.routes.internal.internal_routes import InternalRoutes37 38class BinaryEventTypes:39 PREVIEW_IMAGE = 140 UNENCODED_PREVIEW_IMAGE = 241 42async def send_socket_catch_exception(function, message):43 try:44 await function(message)45 except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError, BrokenPipeError, ConnectionError) as err:46 logging.warning("send error: {}".format(err))47 48@web.middleware49async def cache_control(request: web.Request, handler):50 response: web.Response = await handler(request)51 if request.path.endswith('.js') or request.path.endswith('.css'):52 response.headers.setdefault('Cache-Control', 'no-cache')53 return response54 55 56@web.middleware57async def compress_body(request: web.Request, handler):58 accept_encoding = request.headers.get("Accept-Encoding", "")59 response: web.Response = await handler(request)60 if not isinstance(response, web.Response):61 return response62 if response.content_type not in ["application/json", "text/plain"]:63 return response64 if response.body and "gzip" in accept_encoding:65 response.enable_compression()66 return response67 68 69def create_cors_middleware(allowed_origin: str):70 @web.middleware71 async def cors_middleware(request: web.Request, handler):72 if request.method == "OPTIONS":73 # Pre-flight request. Reply successfully:74 response = web.Response()75 else:76 response = await handler(request)77 78 response.headers['Access-Control-Allow-Origin'] = allowed_origin79 response.headers['Access-Control-Allow-Methods'] = 'POST, GET, DELETE, PUT, OPTIONS'80 response.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization'81 response.headers['Access-Control-Allow-Credentials'] = 'true'82 return response83 84 return cors_middleware85 86def is_loopback(host):87 if host is None:88 return False89 try:90 if ipaddress.ip_address(host).is_loopback:91 return True92 else:93 return False94 except:95 pass96 97 loopback = False98 for family in (socket.AF_INET, socket.AF_INET6):99 try:100 r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM)101 for family, _, _, _, sockaddr in r:102 if not ipaddress.ip_address(sockaddr[0]).is_loopback:103 return loopback104 else:105 loopback = True106 except socket.gaierror:107 pass108 109 return loopback110 111 112def create_origin_only_middleware():113 @web.middleware114 async def origin_only_middleware(request: web.Request, handler):115 #this code is used to prevent the case where a random website can queue comfy workflows by making a POST to 127.0.0.1 which browsers don't prevent for some dumb reason.116 #in that case the Host and Origin hostnames won't match117 #I know the proper fix would be to add a cookie but this should take care of the problem in the meantime118 if 'Host' in request.headers and 'Origin' in request.headers:119 host = request.headers['Host']120 origin = request.headers['Origin']121 host_domain = host.lower()122 parsed = urllib.parse.urlparse(origin)123 origin_domain = parsed.netloc.lower()124 host_domain_parsed = urllib.parse.urlsplit('//' + host_domain)125 126 #limit the check to when the host domain is localhost, this makes it slightly less safe but should still prevent the exploit127 loopback = is_loopback(host_domain_parsed.hostname)128 129 if parsed.port is None: #if origin doesn't have a port strip it from the host to handle weird browsers, same for host130 host_domain = host_domain_parsed.hostname131 if host_domain_parsed.port is None:132 origin_domain = parsed.hostname133 134 if loopback and host_domain is not None and origin_domain is not None and len(host_domain) > 0 and len(origin_domain) > 0:135 if host_domain != origin_domain:136 logging.warning("WARNING: request with non matching host and origin {} != {}, returning 403".format(host_domain, origin_domain))137 return web.Response(status=403)138 139 if request.method == "OPTIONS":140 response = web.Response()141 else:142 response = await handler(request)143 144 return response145 146 return origin_only_middleware147 148class PromptServer():149 def __init__(self, loop):150 PromptServer.instance = self151 152 mimetypes.init()153 mimetypes.add_type('application/javascript; charset=utf-8', '.js')154 mimetypes.add_type('image/webp', '.webp')155 156 self.user_manager = UserManager()157 self.model_file_manager = ModelFileManager()158 self.custom_node_manager = CustomNodeManager()159 self.internal_routes = InternalRoutes(self)160 self.supports = ["custom_nodes_from_web"]161 self.prompt_queue = None162 self.loop = loop163 self.messages = asyncio.Queue()164 self.client_session:Optional[aiohttp.ClientSession] = None165 self.number = 0166 167 middlewares = [cache_control]168 if args.enable_compress_response_body:169 middlewares.append(compress_body)170 171 if args.enable_cors_header:172 middlewares.append(create_cors_middleware(args.enable_cors_header))173 else:174 middlewares.append(create_origin_only_middleware())175 176 max_upload_size = round(args.max_upload_size * 1024 * 1024)177 self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)178 self.sockets = dict()179 self.web_root = (180 FrontendManager.init_frontend(args.front_end_version)181 if args.front_end_root is None182 else args.front_end_root183 )184 logging.info(f"[Prompt Server] web root: {self.web_root}")185 routes = web.RouteTableDef()186 self.routes = routes187 self.last_node_id = None188 self.client_id = None189 190 self.on_prompt_handlers = []191 192 @routes.get('/ws')193 async def websocket_handler(request):194 ws = web.WebSocketResponse()195 await ws.prepare(request)196 sid = request.rel_url.query.get('clientId', '')197 if sid:198 # Reusing existing session, remove old199 self.sockets.pop(sid, None)200 else:201 sid = uuid.uuid4().hex202 203 self.sockets[sid] = ws204 205 try:206 # Send initial state to the new client207 await self.send("status", { "status": self.get_queue_info(), 'sid': sid }, sid)208 # On reconnect if we are the currently executing client send the current node209 if self.client_id == sid and self.last_node_id is not None:210 await self.send("executing", { "node": self.last_node_id }, sid)211 212 async for msg in ws:213 if msg.type == aiohttp.WSMsgType.ERROR:214 logging.warning('ws connection closed with exception %s' % ws.exception())215 finally:216 self.sockets.pop(sid, None)217 return ws218 219 @routes.get("/")220 async def get_root(request):221 response = web.FileResponse(os.path.join(self.web_root, "index.html"))222 response.headers['Cache-Control'] = 'no-cache'223 response.headers["Pragma"] = "no-cache"224 response.headers["Expires"] = "0"225 return response226 227 @routes.get("/embeddings")228 def get_embeddings(self):229 embeddings = folder_paths.get_filename_list("embeddings")230 return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings)))231 232 @routes.get("/models")233 def list_model_types(request):234 model_types = list(folder_paths.folder_names_and_paths.keys())235 236 return web.json_response(model_types)237 238 @routes.get("/models/{folder}")239 async def get_models(request):240 folder = request.match_info.get("folder", None)241 if not folder in folder_paths.folder_names_and_paths:242 return web.Response(status=404)243 files = folder_paths.get_filename_list(folder)244 return web.json_response(files)245 246 @routes.get("/extensions")247 async def get_extensions(request):248 files = glob.glob(os.path.join(249 glob.escape(self.web_root), 'extensions/**/*.js'), recursive=True)250 251 extensions = list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files))252 253 for name, dir in nodes.EXTENSION_WEB_DIRS.items():254 files = glob.glob(os.path.join(glob.escape(dir), '**/*.js'), recursive=True)255 extensions.extend(list(map(lambda f: "/extensions/" + urllib.parse.quote(256 name) + "/" + os.path.relpath(f, dir).replace("\\", "/"), files)))257 258 return web.json_response(extensions)259 260 def get_dir_by_type(dir_type):261 if dir_type is None:262 dir_type = "input"263 264 if dir_type == "input":265 type_dir = folder_paths.get_input_directory()266 elif dir_type == "temp":267 type_dir = folder_paths.get_temp_directory()268 elif dir_type == "output":269 type_dir = folder_paths.get_output_directory()270 271 return type_dir, dir_type272 273 def compare_image_hash(filepath, image):274 hasher = node_helpers.hasher()275 276 # function to compare hashes of two images to see if it already exists, fix to #3465277 if os.path.exists(filepath):278 a = hasher()279 b = hasher()280 with open(filepath, "rb") as f:281 a.update(f.read())282 b.update(image.file.read())283 image.file.seek(0)284 f.close()285 return a.hexdigest() == b.hexdigest()286 return False287 288 def image_upload(post, image_save_function=None):289 image = post.get("image")290 overwrite = post.get("overwrite")291 image_is_duplicate = False292 293 image_upload_type = post.get("type")294 upload_dir, image_upload_type = get_dir_by_type(image_upload_type)295 296 if image and image.file:297 filename = image.filename298 if not filename:299 return web.Response(status=400)300 301 subfolder = post.get("subfolder", "")302 full_output_folder = os.path.join(upload_dir, os.path.normpath(subfolder))303 filepath = os.path.abspath(os.path.join(full_output_folder, filename))304 305 if os.path.commonpath((upload_dir, filepath)) != upload_dir:306 return web.Response(status=400)307 308 if not os.path.exists(full_output_folder):309 os.makedirs(full_output_folder)310 311 split = os.path.splitext(filename)312 313 if overwrite is not None and (overwrite == "true" or overwrite == "1"):314 pass315 else:316 i = 1317 while os.path.exists(filepath):318 if compare_image_hash(filepath, image): #compare hash to prevent saving of duplicates with same name, fix for #3465319 image_is_duplicate = True320 break321 filename = f"{split[0]} ({i}){split[1]}"322 filepath = os.path.join(full_output_folder, filename)323 i += 1324 325 if not image_is_duplicate:326 if image_save_function is not None:327 image_save_function(image, post, filepath)328 else:329 with open(filepath, "wb") as f:330 f.write(image.file.read())331 332 return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type})333 else:334 return web.Response(status=400)335 336 @routes.post("/upload/image")337 async def upload_image(request):338 post = await request.post()339 return image_upload(post)340 341 342 @routes.post("/upload/mask")343 async def upload_mask(request):344 post = await request.post()345 346 def image_save_function(image, post, filepath):347 original_ref = json.loads(post.get("original_ref"))348 filename, output_dir = folder_paths.annotated_filepath(original_ref['filename'])349 350 if not filename:351 return web.Response(status=400)352 353 # validation for security: prevent accessing arbitrary path354 if filename[0] == '/' or '..' in filename:355 return web.Response(status=400)356 357 if output_dir is None:358 type = original_ref.get("type", "output")359 output_dir = folder_paths.get_directory_by_type(type)360 361 if output_dir is None:362 return web.Response(status=400)363 364 if original_ref.get("subfolder", "") != "":365 full_output_dir = os.path.join(output_dir, original_ref["subfolder"])366 if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:367 return web.Response(status=403)368 output_dir = full_output_dir369 370 file = os.path.join(output_dir, filename)371 372 if os.path.isfile(file):373 with Image.open(file) as original_pil:374 metadata = PngInfo()375 if hasattr(original_pil,'text'):376 for key in original_pil.text:377 metadata.add_text(key, original_pil.text[key])378 original_pil = original_pil.convert('RGBA')379 mask_pil = Image.open(image.file).convert('RGBA')380 381 # alpha copy382 new_alpha = mask_pil.getchannel('A')383 original_pil.putalpha(new_alpha)384 original_pil.save(filepath, compress_level=4, pnginfo=metadata)385 386 return image_upload(post, image_save_function)387 388 @routes.get("/view")389 async def view_image(request):390 if "filename" in request.rel_url.query:391 filename = request.rel_url.query["filename"]392 filename,output_dir = folder_paths.annotated_filepath(filename)393 394 if not filename:395 return web.Response(status=400)396 397 # validation for security: prevent accessing arbitrary path398 if filename[0] == '/' or '..' in filename:399 return web.Response(status=400)400 401 if output_dir is None:402 type = request.rel_url.query.get("type", "output")403 output_dir = folder_paths.get_directory_by_type(type)404 405 if output_dir is None:406 return web.Response(status=400)407 408 if "subfolder" in request.rel_url.query:409 full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"])410 if os.path.commonpath((os.path.abspath(full_output_dir), output_dir)) != output_dir:411 return web.Response(status=403)412 output_dir = full_output_dir413 414 filename = os.path.basename(filename)415 file = os.path.join(output_dir, filename)416 417 if os.path.isfile(file):418 if 'preview' in request.rel_url.query:419 with Image.open(file) as img:420 preview_info = request.rel_url.query['preview'].split(';')421 image_format = preview_info[0]422 if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''):423 image_format = 'webp'424 425 quality = 90426 if preview_info[-1].isdigit():427 quality = int(preview_info[-1])428 429 buffer = BytesIO()430 if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb':431 img = img.convert("RGB")432 img.save(buffer, format=image_format, quality=quality)433 buffer.seek(0)434 435 return web.Response(body=buffer.read(), content_type=f'image/{image_format}',436 headers={"Content-Disposition": f"filename=\"{filename}\""})437 438 if 'channel' not in request.rel_url.query:439 channel = 'rgba'440 else:441 channel = request.rel_url.query["channel"]442 443 if channel == 'rgb':444 with Image.open(file) as img:445 if img.mode == "RGBA":446 r, g, b, a = img.split()447 new_img = Image.merge('RGB', (r, g, b))448 else:449 new_img = img.convert("RGB")450 451 buffer = BytesIO()452 new_img.save(buffer, format='PNG')453 buffer.seek(0)454 455 return web.Response(body=buffer.read(), content_type='image/png',456 headers={"Content-Disposition": f"filename=\"{filename}\""})457 458 elif channel == 'a':459 with Image.open(file) as img:460 if img.mode == "RGBA":461 _, _, _, a = img.split()462 else:463 a = Image.new('L', img.size, 255)464 465 # alpha img466 alpha_img = Image.new('RGBA', img.size)467 alpha_img.putalpha(a)468 alpha_buffer = BytesIO()469 alpha_img.save(alpha_buffer, format='PNG')470 alpha_buffer.seek(0)471 472 return web.Response(body=alpha_buffer.read(), content_type='image/png',473 headers={"Content-Disposition": f"filename=\"{filename}\""})474 else:475 # Get content type from mimetype, defaulting to 'application/octet-stream'476 content_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream'477 478 # For security, force certain extensions to download instead of display479 file_extension = os.path.splitext(filename)[1].lower()480 if file_extension in {'.html', '.htm', '.js', '.css'}:481 content_type = 'application/octet-stream' # Forces download482 483 return web.FileResponse(484 file,485 headers={486 "Content-Disposition": f"filename=\"{filename}\"",487 "Content-Type": content_type488 }489 )490 491 return web.Response(status=404)492 493 @routes.get("/view_metadata/{folder_name}")494 async def view_metadata(request):495 folder_name = request.match_info.get("folder_name", None)496 if folder_name is None:497 return web.Response(status=404)498 if not "filename" in request.rel_url.query:499 return web.Response(status=404)500 501 filename = request.rel_url.query["filename"]502 if not filename.endswith(".safetensors"):503 return web.Response(status=404)504 505 safetensors_path = folder_paths.get_full_path(folder_name, filename)506 if safetensors_path is None:507 return web.Response(status=404)508 out = comfy.utils.safetensors_header(safetensors_path, max_size=1024*1024)509 if out is None:510 return web.Response(status=404)511 dt = json.loads(out)512 if not "__metadata__" in dt:513 return web.Response(status=404)514 return web.json_response(dt["__metadata__"])515 516 @routes.get("/system_stats")517 async def system_stats(request):518 device = comfy.model_management.get_torch_device()519 device_name = comfy.model_management.get_torch_device_name(device)520 cpu_device = comfy.model_management.torch.device("cpu")521 ram_total = comfy.model_management.get_total_memory(cpu_device)522 ram_free = comfy.model_management.get_free_memory(cpu_device)523 vram_total, torch_vram_total = comfy.model_management.get_total_memory(device, torch_total_too=True)524 vram_free, torch_vram_free = comfy.model_management.get_free_memory(device, torch_free_too=True)525 526 system_stats = {527 "system": {528 "os": os.name,529 "ram_total": ram_total,530 "ram_free": ram_free,531 "comfyui_version": __version__,532 "python_version": sys.version,533 "pytorch_version": comfy.model_management.torch_version,534 "embedded_python": os.path.split(os.path.split(sys.executable)[0])[1] == "python_embeded",535 "argv": sys.argv536 },537 "devices": [538 {539 "name": device_name,540 "type": device.type,541 "index": device.index,542 "vram_total": vram_total,543 "vram_free": vram_free,544 "torch_vram_total": torch_vram_total,545 "torch_vram_free": torch_vram_free,546 }547 ]548 }549 return web.json_response(system_stats)550 551 @routes.get("/prompt")552 async def get_prompt(request):553 return web.json_response(self.get_queue_info())554 555 def node_info(node_class):556 obj_class = nodes.NODE_CLASS_MAPPINGS[node_class]557 info = {}558 info['input'] = obj_class.INPUT_TYPES()559 info['input_order'] = {key: list(value.keys()) for (key, value) in obj_class.INPUT_TYPES().items()}560 info['output'] = obj_class.RETURN_TYPES561 info['output_is_list'] = obj_class.OUTPUT_IS_LIST if hasattr(obj_class, 'OUTPUT_IS_LIST') else [False] * len(obj_class.RETURN_TYPES)562 info['output_name'] = obj_class.RETURN_NAMES if hasattr(obj_class, 'RETURN_NAMES') else info['output']563 info['name'] = node_class564 info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class565 info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else ''566 info['python_module'] = getattr(obj_class, "RELATIVE_PYTHON_MODULE", "nodes")567 info['category'] = 'sd'568 if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True:569 info['output_node'] = True570 else:571 info['output_node'] = False572 573 if hasattr(obj_class, 'CATEGORY'):574 info['category'] = obj_class.CATEGORY575 576 if hasattr(obj_class, 'OUTPUT_TOOLTIPS'):577 info['output_tooltips'] = obj_class.OUTPUT_TOOLTIPS578 579 if getattr(obj_class, "DEPRECATED", False):580 info['deprecated'] = True581 if getattr(obj_class, "EXPERIMENTAL", False):582 info['experimental'] = True583 return info584 585 @routes.get("/object_info")586 async def get_object_info(request):587 with folder_paths.cache_helper:588 out = {}589 for x in nodes.NODE_CLASS_MAPPINGS:590 try:591 out[x] = node_info(x)592 except Exception:593 logging.error(f"[ERROR] An error occurred while retrieving information for the '{x}' node.")594 logging.error(traceback.format_exc())595 return web.json_response(out)596 597 @routes.get("/object_info/{node_class}")598 async def get_object_info_node(request):599 node_class = request.match_info.get("node_class", None)600 out = {}601 if (node_class is not None) and (node_class in nodes.NODE_CLASS_MAPPINGS):602 out[node_class] = node_info(node_class)603 return web.json_response(out)604 605 @routes.get("/history")606 async def get_history(request):607 max_items = request.rel_url.query.get("max_items", None)608 if max_items is not None:609 max_items = int(max_items)610 return web.json_response(self.prompt_queue.get_history(max_items=max_items))611 612 @routes.get("/history/{prompt_id}")613 async def get_history_prompt_id(request):614 prompt_id = request.match_info.get("prompt_id", None)615 return web.json_response(self.prompt_queue.get_history(prompt_id=prompt_id))616 617 @routes.get("/queue")618 async def get_queue(request):619 queue_info = {}620 current_queue = self.prompt_queue.get_current_queue()621 queue_info['queue_running'] = current_queue[0]622 queue_info['queue_pending'] = current_queue[1]623 return web.json_response(queue_info)624 625 @routes.post("/prompt")626 async def post_prompt(request):627 logging.info("got prompt")628 json_data = await request.json()629 json_data = self.trigger_on_prompt(json_data)630 631 if "number" in json_data:632 number = float(json_data['number'])633 else:634 number = self.number635 if "front" in json_data:636 if json_data['front']:637 number = -number638 639 self.number += 1640 641 if "prompt" in json_data:642 prompt = json_data["prompt"]643 valid = execution.validate_prompt(prompt)644 extra_data = {}645 if "extra_data" in json_data:646 extra_data = json_data["extra_data"]647 648 if "client_id" in json_data:649 extra_data["client_id"] = json_data["client_id"]650 if valid[0]:651 prompt_id = str(uuid.uuid4())652 outputs_to_execute = valid[2]653 self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))654 response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}655 return web.json_response(response)656 else:657 logging.warning("invalid prompt: {}".format(valid[1]))658 return web.json_response({"error": valid[1], "node_errors": valid[3]}, status=400)659 else:660 return web.json_response({"error": "no prompt", "node_errors": []}, status=400)661 662 @routes.post("/queue")663 async def post_queue(request):664 json_data = await request.json()665 if "clear" in json_data:666 if json_data["clear"]:667 self.prompt_queue.wipe_queue()668 if "delete" in json_data:669 to_delete = json_data['delete']670 for id_to_delete in to_delete:671 delete_func = lambda a: a[1] == id_to_delete672 self.prompt_queue.delete_queue_item(delete_func)673 674 return web.Response(status=200)675 676 @routes.post("/interrupt")677 async def post_interrupt(request):678 nodes.interrupt_processing()679 return web.Response(status=200)680 681 @routes.post("/free")682 async def post_free(request):683 json_data = await request.json()684 unload_models = json_data.get("unload_models", False)685 free_memory = json_data.get("free_memory", False)686 if unload_models:687 self.prompt_queue.set_flag("unload_models", unload_models)688 if free_memory:689 self.prompt_queue.set_flag("free_memory", free_memory)690 return web.Response(status=200)691 692 @routes.post("/history")693 async def post_history(request):694 json_data = await request.json()695 if "clear" in json_data:696 if json_data["clear"]:697 self.prompt_queue.wipe_history()698 if "delete" in json_data:699 to_delete = json_data['delete']700 for id_to_delete in to_delete:701 self.prompt_queue.delete_history_item(id_to_delete)702 703 return web.Response(status=200)704 705 async def setup(self):706 timeout = aiohttp.ClientTimeout(total=None) # no timeout707 self.client_session = aiohttp.ClientSession(timeout=timeout)708 709 def add_routes(self):710 self.user_manager.add_routes(self.routes)711 self.model_file_manager.add_routes(self.routes)712 self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items())713 self.app.add_subapp('/internal', self.internal_routes.get_app())714 715 # Prefix every route with /api for easier matching for delegation.716 # This is very useful for frontend dev server, which need to forward717 # everything except serving of static files.718 # Currently both the old endpoints without prefix and new endpoints with719 # prefix are supported.720 api_routes = web.RouteTableDef()721 for route in self.routes:722 # Custom nodes might add extra static routes. Only process non-static723 # routes to add /api prefix.724 if isinstance(route, web.RouteDef):725 api_routes.route(route.method, "/api" + route.path)(route.handler, **route.kwargs)726 self.app.add_routes(api_routes)727 self.app.add_routes(self.routes)728 729 # Add routes from web extensions.730 for name, dir in nodes.EXTENSION_WEB_DIRS.items():731 self.app.add_routes([web.static('/extensions/' + name, dir)])732 733 self.app.add_routes([734 web.static('/', self.web_root),735 ])736 737 def get_queue_info(self):738 prompt_info = {}739 exec_info = {}740 exec_info['queue_remaining'] = self.prompt_queue.get_tasks_remaining()741 prompt_info['exec_info'] = exec_info742 return prompt_info743 744 async def send(self, event, data, sid=None):745 if event == BinaryEventTypes.UNENCODED_PREVIEW_IMAGE:746 await self.send_image(data, sid=sid)747 elif isinstance(data, (bytes, bytearray)):748 await self.send_bytes(event, data, sid)749 else:750 await self.send_json(event, data, sid)751 752 def encode_bytes(self, event, data):753 if not isinstance(event, int):754 raise RuntimeError(f"Binary event types must be integers, got {event}")755 756 packed = struct.pack(">I", event)757 message = bytearray(packed)758 message.extend(data)759 return message760 761 async def send_image(self, image_data, sid=None):762 image_type = image_data[0]763 image = image_data[1]764 max_size = image_data[2]765 if max_size is not None:766 if hasattr(Image, 'Resampling'):767 resampling = Image.Resampling.BILINEAR768 else:769 resampling = Image.ANTIALIAS770 771 image = ImageOps.contain(image, (max_size, max_size), resampling)772 type_num = 1773 if image_type == "JPEG":774 type_num = 1775 elif image_type == "PNG":776 type_num = 2777 778 bytesIO = BytesIO()779 header = struct.pack(">I", type_num)780 bytesIO.write(header)781 image.save(bytesIO, format=image_type, quality=95, compress_level=1)782 preview_bytes = bytesIO.getvalue()783 await self.send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid)784 785 async def send_bytes(self, event, data, sid=None):786 message = self.encode_bytes(event, data)787 788 if sid is None:789 sockets = list(self.sockets.values())790 for ws in sockets:791 await send_socket_catch_exception(ws.send_bytes, message)792 elif sid in self.sockets:793 await send_socket_catch_exception(self.sockets[sid].send_bytes, message)794 795 async def send_json(self, event, data, sid=None):796 message = {"type": event, "data": data}797 798 if sid is None:799 sockets = list(self.sockets.values())800 for ws in sockets:801 await send_socket_catch_exception(ws.send_json, message)802 elif sid in self.sockets:803 await send_socket_catch_exception(self.sockets[sid].send_json, message)804 805 def send_sync(self, event, data, sid=None):806 self.loop.call_soon_threadsafe(807 self.messages.put_nowait, (event, data, sid))808 809 def queue_updated(self):810 self.send_sync("status", { "status": self.get_queue_info() })811 812 async def publish_loop(self):813 while True:814 msg = await self.messages.get()815 await self.send(*msg)816 817 async def start(self, address, port, verbose=True, call_on_start=None):818 await self.start_multi_address([(address, port)], call_on_start=call_on_start)819 820 async def start_multi_address(self, addresses, call_on_start=None, verbose=True):821 runner = web.AppRunner(self.app, access_log=None)822 await runner.setup()823 ssl_ctx = None824 scheme = "http"825 if args.tls_keyfile and args.tls_certfile:826 ssl_ctx = ssl.SSLContext(protocol=ssl.PROTOCOL_TLS_SERVER, verify_mode=ssl.CERT_NONE)827 ssl_ctx.load_cert_chain(certfile=args.tls_certfile,828 keyfile=args.tls_keyfile)829 scheme = "https"830 831 if verbose:832 logging.info("Starting server\n")833 for addr in addresses:834 address = addr[0]835 port = addr[1]836 site = web.TCPSite(runner, address, port, ssl_context=ssl_ctx)837 await site.start()838 839 if not hasattr(self, 'address'):840 self.address = address #TODO: remove this841 self.port = port842 843 if ':' in address:844 address_print = "[{}]".format(address)845 else:846 address_print = address847 848 if verbose:849 logging.info("To see the GUI go to: {}://{}:{}".format(scheme, address_print, port))850 851 if call_on_start is not None:852 call_on_start(scheme, self.address, self.port)853 854 def add_on_prompt_handler(self, handler):855 self.on_prompt_handlers.append(handler)856 857 def trigger_on_prompt(self, json_data):858 for handler in self.on_prompt_handlers:859 try:860 json_data = handler(json_data)861 except Exception:862 logging.warning("[ERROR] An error occurred during the on_prompt_handler processing")863 logging.warning(traceback.format_exc())864 865 return json_data866