99 lines
3.9 KiB
Python
99 lines
3.9 KiB
Python
import json
|
|
import os
|
|
import server
|
|
import traceback
|
|
#Check for availability of workflow checkpointing
|
|
import aiohttp
|
|
import asyncio
|
|
from .workflowcheckpointing import post_prompt_remote
|
|
|
|
web = server.web
|
|
ps = server.PromptServer.instance
|
|
|
|
finished_startup = False
|
|
original_server_start = ps.start
|
|
async def server_start(address, port, verbose=True, call_on_start=None):
|
|
original_on_start= call_on_start
|
|
def on_start(*args, **kwargs):
|
|
if original_on_start is not None:
|
|
original_on_start(*args, **kwargs)
|
|
global finished_startup
|
|
finished_startup= True
|
|
return await original_server_start(address, port, verbose, on_start)
|
|
ps.start = server_start
|
|
|
|
|
|
@ps.routes.get("/health")
|
|
async def heath(request):
|
|
#while any of the server endpoints could likely be used
|
|
return web.json_response([])
|
|
|
|
@ps.routes.get("/startup")
|
|
async def startup(request):
|
|
if finished_startup:
|
|
return web.json_response([])
|
|
return web.Response(status=503)
|
|
|
|
@ps.routes.get("/ready")
|
|
async def ready(request):
|
|
current_queue = ps.prompt_queue.get_current_queue()
|
|
if len(current_queue[0]) == 0 and len(current_queue[1]) == 0:
|
|
return web.json_response(current_queue)
|
|
return web.json_response(current_queue, status=503)
|
|
|
|
async def websocket_loop():
|
|
async with aiohttp.ClientSession() as session:
|
|
if 'ORCHESTRATION_SERVER' not in os.environ:
|
|
while True:
|
|
await asyncio.sleep(60)
|
|
async with session.ws_connect(os.environ["ORCHESTRATION_SERVER"]) as ws:
|
|
print("connected to server")
|
|
async for msg in ws:
|
|
try:
|
|
print("got command: " + str(msg))
|
|
if msg.type == aiohttp.WSMsgType.TEXT:
|
|
js = msg.json()
|
|
resp = {"message_id": js.get('message_id', 0)}
|
|
match js['command']:
|
|
case 'prompt':
|
|
#wrap as mock request
|
|
class MockRequest:
|
|
async def json(self):
|
|
return js['data']
|
|
out = await post_prompt_remote(MockRequest())
|
|
resp['data'] = json.loads(out.body._value)
|
|
case "queue":
|
|
resp['data'] = ps.prompt_queue.get_current_queue()
|
|
case "files":
|
|
resp['data'] = [f.name for f in os.scandir('fetches') if f.is_file()]
|
|
case "info":
|
|
resp['data'] = {}
|
|
if 'SALAD_MACHINE_ID' in os.environ:
|
|
resp['data']['machine_id'] = os.environ['SALAD_MACHINE_ID']
|
|
else:
|
|
resp['data']['machine_id'] = os.environ.get('HOSTNAME', 'local')
|
|
case "logs":
|
|
with open('comfyui.log', 'r') as f:
|
|
resp['data'] = f.read()
|
|
case _:
|
|
resp = {"error": "Unknown command"}
|
|
print(resp)
|
|
await ws.send_json(resp)
|
|
elif msg.type == aiohttp.WSMsgType.ERROR:
|
|
await ws.send_json({"error": "Received bad message"})
|
|
except Exception as e:
|
|
#NOTE: this will reraise if error was socket closing
|
|
await ws.send_json({"error": str(e)})
|
|
async def try_websocket():
|
|
while True:
|
|
try:
|
|
await websocket_loop()
|
|
except aiohttp.client_exceptions.ClientConnectorError:
|
|
print("disconnected")
|
|
except:
|
|
print(traceback.format_exc())
|
|
await asyncio.sleep(5)
|
|
print("Attempting re connection")
|
|
|
|
process_loop = ps.loop.create_task(try_websocket())
|