140 lines
4.5 KiB
Python
140 lines
4.5 KiB
Python
import datetime
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from multiprocessing import Process
|
|
from threading import Lock
|
|
|
|
import requests
|
|
import uvicorn
|
|
from fastapi import APIRouter, FastAPI, Request, HTTPException
|
|
|
|
TIMEOUT_KEEP_ALIVE = 30
|
|
COMFY_PORT = 8188
|
|
SAGEMAKER_PORT = 8080
|
|
LOCALHOST = '0.0.0.0'
|
|
PHY_LOCALHOST = '127.0.0.1'
|
|
|
|
SLEEP_TIME = 30
|
|
|
|
app = FastAPI()
|
|
|
|
logger = logging.getLogger(__name__)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
async def invocations(request: Request):
|
|
req = await request.json()
|
|
logger.info(f"invocations start req:{req} url:{PHY_LOCALHOST}:{COMFY_PORT}/invocations")
|
|
response = requests.post(f"http://{PHY_LOCALHOST}:{COMFY_PORT}/invocations", json=req)
|
|
if response.status_code != 200:
|
|
raise HTTPException(status_code=response.status_code,
|
|
detail=f"COMFY service returned an error: {response.text}")
|
|
return response.json()
|
|
|
|
|
|
def ping():
|
|
init_already = os.environ.get('ALREADY_INIT')
|
|
if init_already and init_already.lower() == 'false':
|
|
raise HTTPException(status_code=500)
|
|
else:
|
|
return {'status': 'Healthy'}
|
|
|
|
|
|
class Api:
|
|
def add_api_route(self, path: str, endpoint, **kwargs):
|
|
return self.app.add_api_route(path, endpoint, **kwargs)
|
|
|
|
def __init__(self, app: FastAPI, queue_lock: Lock):
|
|
self.router = APIRouter()
|
|
self.app = app
|
|
self.queue_lock = queue_lock
|
|
self.add_api_route("/invocations", invocations, methods=["POST"])
|
|
self.add_api_route("/ping", ping, methods=["GET"], response_model={})
|
|
|
|
def launch(self, server_name, port):
|
|
self.app.include_router(self.router)
|
|
uvicorn.run(self.app, host=server_name, port=port, timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
|
|
|
|
|
|
class ComfyApp:
|
|
def __init__(self, host=LOCALHOST, port=COMFY_PORT):
|
|
self.host = host
|
|
self.port = port
|
|
self.process = None
|
|
|
|
def start(self):
|
|
# cmd = "python main.py --listen {} --port {}".format(self.host, self.port)
|
|
# self.process = Process(target=os.system, args=(cmd,))
|
|
# self.process.start()
|
|
cmd = ["python", "main.py", "--listen", self.host, "--port", str(self.port)]
|
|
self.process = subprocess.Popen(cmd)
|
|
os.environ['ALREADY_INIT'] = 'true'
|
|
|
|
def restart(self):
|
|
# if self.process and self.process.is_alive():
|
|
# self.process.terminate()
|
|
# self.start()
|
|
# else:
|
|
# logger.info("Comfy app process is not running.")
|
|
logger.info("Comfy app process is going to restart")
|
|
if self.process and self.process.poll() is None:
|
|
os.environ['ALREADY_INIT'] = 'false'
|
|
self.process.terminate()
|
|
self.process.wait()
|
|
self.start()
|
|
|
|
|
|
def check_sync():
|
|
logger.info("start check_sync!")
|
|
while True:
|
|
try:
|
|
logger.info("start check_sync! checking function-------")
|
|
response = requests.post(f"http://{PHY_LOCALHOST}:{COMFY_PORT}/sync_instance")
|
|
logger.info(f"sync response:{response.json()} time : {datetime.datetime.now()}")
|
|
|
|
logger.info("start check_reboot! checking function-------")
|
|
response2 = requests.post(f"http://{PHY_LOCALHOST}:{COMFY_PORT}/reboot")
|
|
logger.info(f"reboot response:{response.json()} time : {datetime.datetime.now()}")
|
|
time.sleep(SLEEP_TIME)
|
|
except Exception as e:
|
|
logger.info(f"check_and_reboot error:{e}")
|
|
time.sleep(SLEEP_TIME)
|
|
|
|
|
|
def check_reboot():
|
|
logger.info("start check_reboot!")
|
|
while True:
|
|
try:
|
|
logger.info("start check_reboot! checking function-------")
|
|
response = requests.post(f"http://{PHY_LOCALHOST}:{COMFY_PORT}/reboot")
|
|
logger.info(f"reboot response:{response.json()} time : {datetime.datetime.now()}")
|
|
time.sleep(SLEEP_TIME)
|
|
except Exception as e:
|
|
logger.info(f"check_and_reboot error:{e}")
|
|
time.sleep(SLEEP_TIME)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
queue_lock = threading.Lock()
|
|
api = Api(app, queue_lock)
|
|
|
|
comfy_app = ComfyApp()
|
|
api_process = Process(target=api.launch, args=(LOCALHOST, SAGEMAKER_PORT))
|
|
|
|
check_sync_thread = threading.Thread(target=check_sync)
|
|
# check_reboot_thread = threading.Thread(target=check_reboot)
|
|
|
|
comfy_app.start()
|
|
api_process.start()
|
|
check_sync_thread.start()
|
|
# check_reboot_thread.start()
|
|
|
|
api_process.join()
|
|
# check_sync_thread.join()
|
|
# check_reboot_thread.join()
|
|
|
|
|