1674 lines
80 KiB
Python
Executable File
1674 lines
80 KiB
Python
Executable File
import concurrent.futures
|
||
import signal
|
||
import threading
|
||
|
||
import boto3
|
||
import requests
|
||
from aiohttp import web
|
||
|
||
import folder_paths
|
||
import server
|
||
from execution import PromptExecutor
|
||
import execution
|
||
import comfy
|
||
|
||
from watchdog.observers import Observer
|
||
from watchdog.events import FileSystemEventHandler
|
||
import subprocess
|
||
from dotenv import load_dotenv
|
||
|
||
import fcntl
|
||
import hashlib
|
||
|
||
import base64
|
||
import datetime
|
||
import json
|
||
import logging
|
||
import os
|
||
import sys
|
||
import tarfile
|
||
import time
|
||
import uuid
|
||
import gc
|
||
from dataclasses import dataclass
|
||
from typing import Optional
|
||
|
||
|
||
from boto3.dynamodb.conditions import Key
|
||
|
||
|
||
DISABLE_AWS_PROXY = 'DISABLE_AWS_PROXY'
|
||
sync_msg_list = []
|
||
client_release_map = {}
|
||
|
||
logging.basicConfig(level=logging.DEBUG)
|
||
logger = logging.getLogger(__name__)
|
||
logger.setLevel(logging.INFO)
|
||
|
||
is_on_sagemaker = os.getenv('ON_SAGEMAKER') == 'true'
|
||
is_on_ec2 = os.getenv('ON_EC2') == 'true'
|
||
|
||
if is_on_ec2:
|
||
env_path = '/etc/environment'
|
||
|
||
if 'ENV_FILE_PATH' in os.environ and os.environ.get('ENV_FILE_PATH'):
|
||
env_path = os.environ.get('ENV_FILE_PATH')
|
||
|
||
load_dotenv('/etc/environment')
|
||
logger.info(f"env_path{env_path}")
|
||
|
||
env_keys = ['ENV_FILE_PATH', 'COMFY_INPUT_PATH', 'COMFY_MODEL_PATH', 'COMFY_NODE_PATH', 'COMFY_API_URL',
|
||
'COMFY_API_TOKEN', 'COMFY_ENDPOINT', 'COMFY_NEED_SYNC', 'COMFY_NEED_PREPARE', 'COMFY_BUCKET_NAME',
|
||
'MAX_WAIT_TIME', 'MSG_MAX_WAIT_TIME', 'THREAD_MAX_WAIT_TIME', DISABLE_AWS_PROXY, 'DISABLE_AUTO_SYNC']
|
||
|
||
for item in os.environ.keys():
|
||
if item in env_keys:
|
||
logger.info(f'evn key： {item} {os.environ.get(item)}')
|
||
|
||
DIR3 = "input"
|
||
DIR1 = "models"
|
||
DIR2 = "custom_nodes"
|
||
|
||
if 'COMFY_INPUT_PATH' in os.environ and os.environ.get('COMFY_INPUT_PATH'):
|
||
DIR3 = os.environ.get('COMFY_INPUT_PATH')
|
||
if 'COMFY_MODEL_PATH' in os.environ and os.environ.get('COMFY_MODEL_PATH'):
|
||
DIR1 = os.environ.get('COMFY_MODEL_PATH')
|
||
if 'COMFY_NODE_PATH' in os.environ and os.environ.get('COMFY_NODE_PATH'):
|
||
DIR2 = os.environ.get('COMFY_NODE_PATH')
|
||
|
||
|
||
api_url = os.environ.get('COMFY_API_URL')
|
||
api_token = os.environ.get('COMFY_API_TOKEN')
|
||
comfy_endpoint = os.environ.get('COMFY_ENDPOINT')
|
||
comfy_need_sync = os.environ.get('COMFY_NEED_SYNC', True)
|
||
comfy_need_prepare = os.environ.get('COMFY_NEED_PREPARE', False)
|
||
bucket_name = os.environ.get('COMFY_BUCKET_NAME')
|
||
thread_max_wait_time = os.environ.get('THREAD_MAX_WAIT_TIME', 60)
|
||
max_wait_time = os.environ.get('MAX_WAIT_TIME', 86400)
|
||
msg_max_wait_time = os.environ.get('MSG_MAX_WAIT_TIME', 86400)
|
||
is_master_process = os.getenv('MASTER_PROCESS') == 'true'
|
||
program_name = os.getenv('PROGRAM_NAME')
|
||
no_need_sync_files = ['.autosave', '.cache', '.autosave1', '~', '.swp']
|
||
|
||
need_resend_msg_result = []
|
||
PREPARE_ID = 'default'
|
||
# additional
|
||
PREPARE_MODE = 'additional'
|
||
|
||
if not api_url:
|
||
raise ValueError("API_URL environment variables must be set.")
|
||
|
||
if not api_token:
|
||
raise ValueError("API_TOKEN environment variables must be set.")
|
||
|
||
if not comfy_endpoint:
|
||
raise ValueError("COMFY_ENDPOINT environment variables must be set.")
|
||
|
||
headers = {"x-api-key": api_token, "Content-Type": "application/json"}
|
||
|
||
|
||
def save_images_locally(response_json, local_folder):
|
||
try:
|
||
data = response_json.get("data", {})
|
||
prompt_id = data.get("prompt_id")
|
||
image_video_data = data.get("image_video_data", {})
|
||
|
||
if not prompt_id or not image_video_data:
|
||
logger.info("Missing prompt_id or image_video_data in the response.")
|
||
return
|
||
|
||
folder_path = os.path.join(local_folder, prompt_id)
|
||
os.makedirs(folder_path, exist_ok=True)
|
||
|
||
for image_name, image_url in image_video_data.items():
|
||
image_response = requests.get(image_url)
|
||
if image_response.status_code == 200:
|
||
image_path = os.path.join(folder_path, image_name)
|
||
with open(image_path, "wb") as image_file:
|
||
image_file.write(image_response.content)
|
||
logger.info(f"Image '{image_name}' saved to {image_path}")
|
||
else:
|
||
logger.info(
|
||
f"Failed to download image '{image_name}' from {image_url}. Status code: {image_response.status_code}")
|
||
|
||
except Exception as e:
|
||
logger.info(f"Error saving images locally: {e}")
|
||
|
||
|
||
def calculate_file_hash(file_path):
|
||
# 创建一个哈希对象
|
||
hasher = hashlib.sha256()
|
||
# 打开文件并逐块更新哈希对象
|
||
with open(file_path, 'rb') as file:
|
||
buffer = file.read(65536) # 64KB 的缓冲区大小
|
||
while len(buffer) > 0:
|
||
hasher.update(buffer)
|
||
buffer = file.read(65536)
|
||
# 返回哈希值的十六进制表示
|
||
return hasher.hexdigest()
|
||
|
||
|
||
def save_files(prefix, execute, key, target_dir, need_prefix):
|
||
if key in execute['data']:
|
||
temp_files = execute['data'][key]
|
||
for url in temp_files:
|
||
loca_file = get_file_name(url)
|
||
response = requests.get(url)
|
||
# if target_dir not exists, create it
|
||
if not os.path.exists(target_dir):
|
||
os.makedirs(target_dir)
|
||
logger.info(f"Saving file {loca_file} to {target_dir}")
|
||
if loca_file.endswith("output_images_will_be_put_here"):
|
||
continue
|
||
if need_prefix:
|
||
with open(f"./{target_dir}/{prefix}_{loca_file}", 'wb') as f:
|
||
f.write(response.content)
|
||
# current override exist
|
||
with open(f"./{target_dir}/{loca_file}", 'wb') as f:
|
||
f.write(response.content)
|
||
else:
|
||
with open(f"./{target_dir}/{loca_file}", 'wb') as f:
|
||
f.write(response.content)
|
||
|
||
|
||
def get_file_name(url: str):
|
||
file_name = url.split('/')[-1]
|
||
file_name = file_name.split('?')[0]
|
||
return file_name
|
||
|
||
|
||
def send_service_msg(server_use, msg):
|
||
event = msg.get('event')
|
||
data = msg.get('data')
|
||
sid = msg.get('sid') if 'sid' in msg else None
|
||
server_use.send_sync(event, data, sid)
|
||
|
||
|
||
def handle_sync_messages(server_use, msg_array):
|
||
already_synced = False
|
||
global sync_msg_list
|
||
for msg in msg_array:
|
||
for item_msg in msg:
|
||
event = item_msg.get('event')
|
||
data = item_msg.get('data')
|
||
sid = item_msg.get('sid') if 'sid' in item_msg else None
|
||
if data in sync_msg_list:
|
||
continue
|
||
sync_msg_list.append(data)
|
||
if event == 'finish':
|
||
already_synced = True
|
||
elif event == 'executed':
|
||
global need_resend_msg_result
|
||
need_resend_msg_result.append(msg)
|
||
server_use.send_sync(event, data, sid)
|
||
|
||
return already_synced
|
||
|
||
|
||
def execute_proxy(func):
|
||
def wrapper(*args, **kwargs):
|
||
def send_error_msg(executor, prompt_id, msg):
|
||
mes = {
|
||
"prompt_id": prompt_id,
|
||
"node_id": "",
|
||
"node_type": "on cloud",
|
||
"executed": [],
|
||
"exception_message": msg,
|
||
"exception_type": "",
|
||
"traceback": [],
|
||
"current_inputs": "",
|
||
"current_outputs": "",
|
||
}
|
||
executor.add_message("execution_error", mes, broadcast=True)
|
||
|
||
if 'True' == os.environ.get(DISABLE_AWS_PROXY):
|
||
logger.info("disabled aws proxy, use local")
|
||
return func(*args, **kwargs)
|
||
logger.info(f"enable aws proxy, use aws {comfy_endpoint}")
|
||
executor = args[0]
|
||
server_use = executor.server
|
||
prompt = args[1]
|
||
prompt_id = args[2]
|
||
extra_data = args[3]
|
||
|
||
client_id = extra_data['client_id'] if 'client_id' in extra_data else None
|
||
if not client_id:
|
||
send_error_msg(executor, prompt_id, f"Something went wrong when execute,please check your client_id and try again")
|
||
return web.Response()
|
||
global client_release_map
|
||
workflow_name = client_release_map.get(client_id)
|
||
if not workflow_name and not is_master_process:
|
||
send_error_msg(executor, prompt_id, f"Please choose a release env before you execute prompt")
|
||
return web.Response()
|
||
|
||
payload = {
|
||
"number": str(server.PromptServer.instance.number),
|
||
"prompt": prompt,
|
||
"prompt_id": prompt_id,
|
||
"extra_data": extra_data,
|
||
"endpoint_name": comfy_endpoint,
|
||
"need_prepare": comfy_need_prepare,
|
||
"need_sync": comfy_need_sync,
|
||
"multi_async": False,
|
||
"workflow_name": workflow_name,
|
||
}
|
||
|
||
def send_post_request(url, params):
|
||
logger.debug(f"sending post request {url} , params {params}")
|
||
get_response = requests.post(url, json=params, headers=headers)
|
||
return get_response
|
||
|
||
def send_get_request(url):
|
||
get_response = requests.get(url, headers=headers)
|
||
return get_response
|
||
|
||
def check_if_sync_is_already(url):
|
||
get_response = send_get_request(url)
|
||
prepare_response = get_response.json()
|
||
if (prepare_response['statusCode'] == 200 and 'data' in prepare_response and prepare_response['data']
|
||
and prepare_response['data']['prepareSuccess']):
|
||
logger.info(f"sync available")
|
||
return True
|
||
else:
|
||
logger.info(f"no sync available for {url} response {prepare_response}")
|
||
return False
|
||
|
||
|
||
logger.debug(f"payload is: {payload}")
|
||
is_synced = check_if_sync_is_already(f"{api_url}/prepare/{comfy_endpoint}")
|
||
if not is_synced:
|
||
logger.debug(f"is_synced is {is_synced} stop cloud prompt")
|
||
send_error_msg(executor, prompt_id, "Your local environment has not compleated to synchronized on cloud already. Please wait for a moment or click the 'Synchronize' button .")
|
||
return
|
||
|
||
with concurrent.futures.ThreadPoolExecutor() as executorThread:
|
||
execute_future = executorThread.submit(send_post_request, f"{api_url}/executes", payload)
|
||
|
||
save_already = False
|
||
if comfy_need_sync:
|
||
msg_future = executorThread.submit(send_get_request,
|
||
f"{api_url}/sync/{prompt_id}")
|
||
done, _ = concurrent.futures.wait([execute_future, msg_future],
|
||
return_when=concurrent.futures.ALL_COMPLETED)
|
||
already_synced = False
|
||
global sync_msg_list
|
||
sync_msg_list = []
|
||
for future in done:
|
||
if future == msg_future:
|
||
msg_response = future.result()
|
||
logger.info(f"get syc msg: {msg_response.json()}")
|
||
if msg_response.status_code == 200:
|
||
if 'data' not in msg_response.json() or not msg_response.json().get("data"):
|
||
logger.error("there is no response from sync msg by thread ")
|
||
time.sleep(1)
|
||
else:
|
||
logger.debug(msg_response.json())
|
||
already_synced = handle_sync_messages(server_use, msg_response.json().get("data"))
|
||
elif future == execute_future:
|
||
execute_resp = future.result()
|
||
logger.info(f"get execute status: {execute_resp.status_code}")
|
||
if execute_resp.status_code == 200 or execute_resp.status_code == 201 or execute_resp.status_code == 202:
|
||
i = thread_max_wait_time
|
||
while i > 0:
|
||
images_response = send_get_request(f"{api_url}/executes/{prompt_id}")
|
||
response = images_response.json()
|
||
logger.info(f"get execute images: {images_response.status_code}")
|
||
if images_response.status_code == 404:
|
||
logger.info("no images found already ,waiting sagemaker thread result .....")
|
||
time.sleep(3)
|
||
i = i - 2
|
||
elif response['data']['status'] == 'failed':
|
||
logger.error(f"there is no response on sagemaker from execute thread result !!!!!!!! ")
|
||
# send_error_msg(executor, prompt_id,
|
||
# f"There may be some errors when valid and execute the prompt on the cloud. Please check the SageMaker logs. error info: {response['data']['message']}")
|
||
# no need to send msg anymore
|
||
already_synced = True
|
||
break
|
||
elif response['data']['status'] != 'Completed' and response['data']['status'] != 'success':
|
||
logger.info(f"no images found already ,waiting sagemaker thread result, current status is {response['data']['status']}")
|
||
time.sleep(2)
|
||
i = i - 1
|
||
elif 'data' not in response or not response['data'] or 'status' not in response['data'] or not response['data']['status']:
|
||
logger.error(f"there is no response from execute thread result !!!!!!!! {response}")
|
||
# no need to send msg anymore
|
||
already_synced = True
|
||
# send_error_msg(executor, prompt_id,"There may be some errors when executing the prompt on cloud. No images or videos generated.")
|
||
break
|
||
else:
|
||
if ('temp_files' in images_response.json()['data'] and len(
|
||
images_response.json()['data']['temp_files']) > 0) or ((
|
||
'output_files' in images_response.json()['data'] and len(
|
||
images_response.json()['data']['output_files']) > 0)):
|
||
save_files(prompt_id, images_response.json(), 'temp_files', 'temp', False)
|
||
save_files(prompt_id, images_response.json(), 'output_files', 'output', True)
|
||
else:
|
||
send_error_msg(executor, prompt_id,
|
||
"There may be some errors when executing the prompt on the cloud. Please check the SageMaker logs.")
|
||
# no need to send msg anymore
|
||
already_synced = True
|
||
logger.debug(images_response.json())
|
||
save_already = True
|
||
break
|
||
else:
|
||
logger.error(f"get execute error: {execute_resp}")
|
||
# send_error_msg(executor, prompt_id, "Please valid your prompt and try again.")
|
||
# send_error_msg(executor, prompt_id,
|
||
# f"There may be some errors when valid and execute the prompt on the cloud. Please check the SageMaker logs. error info: {response['data']['message']}")
|
||
# no need to send msg anymore
|
||
already_synced = True
|
||
break
|
||
logger.debug(execute_resp.json())
|
||
|
||
m = msg_max_wait_time
|
||
while not already_synced:
|
||
msg_response = send_get_request(f"{api_url}/sync/{prompt_id}")
|
||
# logger.info(msg_response.json())
|
||
if msg_response.status_code == 200:
|
||
if m <= 0:
|
||
logger.error("there is no response from sync msg by timeout")
|
||
already_synced = True
|
||
elif 'data' not in msg_response.json() or not msg_response.json().get("data"):
|
||
logger.error("there is no response from sync msg")
|
||
time.sleep(1)
|
||
m = m - 1
|
||
else:
|
||
logger.debug(msg_response.json())
|
||
already_synced = handle_sync_messages(server_use, msg_response.json().get("data"))
|
||
logger.info(f"already_synced is :{already_synced}")
|
||
time.sleep(1)
|
||
m = m - 1
|
||
logger.info(f"check if images are already synced {save_already}")
|
||
|
||
if not save_already:
|
||
logger.info("check if images are not already synced, please wait")
|
||
execute_resp = execute_future.result()
|
||
logger.debug(f"execute result :{execute_resp.json()}")
|
||
if execute_resp.status_code == 200 or execute_resp.status_code == 201 or execute_resp.status_code == 202:
|
||
i = max_wait_time
|
||
while i > 0:
|
||
images_response = send_get_request(f"{api_url}/executes/{prompt_id}")
|
||
response = images_response.json()
|
||
logger.debug(response)
|
||
if images_response.status_code == 404:
|
||
logger.info(f"{i} no images found already ,waiting sagemaker result .....")
|
||
i = i - 2
|
||
time.sleep(3)
|
||
elif response['data']['status'] == 'failed':
|
||
logger.error(
|
||
f"there is no response on sagemaker from execute result !!!!!!!! ")
|
||
if 'message' in response['data'] and response['data']['message']:
|
||
send_error_msg(executor, prompt_id,
|
||
f"There may be some errors when valid or execute the prompt on the cloud. Please check the SageMaker logs. errors: {response['data']['message']}")
|
||
break
|
||
else:
|
||
logger.error(f"valid error on sagemaker :{response['data']}")
|
||
send_error_msg(executor, prompt_id,
|
||
f"There may be some errors when valid or execute the prompt on the cloud. errors")
|
||
break
|
||
elif response['data']['status'] != 'Completed' and response['data']['status'] != 'success':
|
||
logger.info(f"{i} images not already ,waiting sagemaker result .....{response['data']['status'] }")
|
||
i = i - 1
|
||
time.sleep(3)
|
||
elif 'data' not in response or not response['data'] or 'status' not in response['data'] or not response['data']['status']:
|
||
logger.info(f"{i} there is no response from sync executes {response}")
|
||
send_error_msg(executor, prompt_id, f"There may be some errors when executing the prompt on the cloud. No images or videos generated. {response['message']}")
|
||
break
|
||
elif response['data']['status'] == 'Completed' or response['data']['status'] == 'success':
|
||
if ('temp_files' in images_response.json()['data'] and len(images_response.json()['data']['temp_files']) > 0) or (('output_files' in images_response.json()['data'] and len(images_response.json()['data']['output_files']) > 0)):
|
||
save_files(prompt_id, images_response.json(), 'temp_files', 'temp', False)
|
||
save_files(prompt_id, images_response.json(), 'output_files', 'output', True)
|
||
break
|
||
else:
|
||
send_error_msg(executor, prompt_id,
|
||
"There may be some errors when executing the prompt on the cloud. Please check the SageMaker logs.")
|
||
break
|
||
else:
|
||
# logger.info(
|
||
# f"{i} images not already other,waiting sagemaker result .....{response}")
|
||
# i = i - 1
|
||
# time.sleep(3)
|
||
send_error_msg(executor, prompt_id,
|
||
"You have some errors when execute prompt on cloud . Please check your sagemaker logs.")
|
||
break
|
||
else:
|
||
logger.error(f"get execute error: {execute_resp}")
|
||
send_error_msg(executor, prompt_id, "Please valid your prompt and try again.")
|
||
logger.info("execute finished")
|
||
executorThread.shutdown()
|
||
|
||
return wrapper
|
||
|
||
|
||
PromptExecutor.execute = execute_proxy(PromptExecutor.execute)
|
||
|
||
|
||
def send_sync_proxy(func):
|
||
def wrapper(*args, **kwargs):
|
||
logger.info(f"Sending sync request----- {args}")
|
||
return func(*args, **kwargs)
|
||
return wrapper
|
||
|
||
|
||
server.PromptServer.send_sync = send_sync_proxy(server.PromptServer.send_sync)
|
||
|
||
|
||
def compress_and_upload(folder_path, prepare_version):
|
||
for subdir in next(os.walk(folder_path))[1]:
|
||
subdir_path = os.path.join(folder_path, subdir)
|
||
tar_filename = f"{subdir}.tar.gz"
|
||
logger.info(f"Compressing the {tar_filename}")
|
||
|
||
# 创建 tar 压缩文件
|
||
with tarfile.open(tar_filename, "w:gz") as tar:
|
||
tar.add(subdir_path, arcname=os.path.basename(subdir_path))
|
||
s5cmd_syn_node_command = f's5cmd --log=error cp {tar_filename} "s3://{bucket_name}/comfy/{comfy_endpoint}/{prepare_version}/custom_nodes/"'
|
||
logger.info(s5cmd_syn_node_command)
|
||
os.system(s5cmd_syn_node_command)
|
||
logger.info(f"rm {tar_filename}")
|
||
os.remove(tar_filename)
|
||
|
||
# for root, dirs, files in os.walk(folder_path):
|
||
# for directory in dirs:
|
||
# dir_path = os.path.join(root, directory)
|
||
# logger.info(f"Compressing the {dir_path}")
|
||
# tar_filename = f"{directory}.tar.gz"
|
||
# tar_filepath = os.path.join(root, tar_filename)
|
||
# with tarfile.open(tar_filepath, "w:gz") as tar:
|
||
# tar.add(dir_path, arcname=os.path.basename(dir_path))
|
||
# s5cmd_syn_node_command = f's5cmd --log=error cp {tar_filepath} "s3://{bucket_name}/comfy/{comfy_endpoint}/{timestamp}/custom_nodes/"'
|
||
# logger.info(s5cmd_syn_node_command)
|
||
# os.system(s5cmd_syn_node_command)
|
||
# logger.info(f"rm {tar_filepath}")
|
||
# os.remove(tar_filepath)
|
||
|
||
|
||
def sync_default_files():
|
||
try:
|
||
timestamp = str(int(time.time() * 1000))
|
||
prepare_version = PREPARE_ID if PREPARE_MODE == 'additional' else timestamp
|
||
need_prepare = True
|
||
prepare_type = 'default'
|
||
need_reboot = True
|
||
logger.info(f" sync custom nodes files")
|
||
# s5cmd_syn_node_command = f's5cmd --log=error sync --delete=true --exclude="*comfy_local_proxy.py" {DIR2}/ "s3://{bucket_name}/comfy/{comfy_endpoint}/{timestamp}/custom_nodes/"'
|
||
# logger.info(f"sync custom_nodes files start {s5cmd_syn_node_command}")
|
||
# os.system(s5cmd_syn_node_command)
|
||
compress_and_upload(f"{DIR2}", prepare_version)
|
||
logger.info(f" sync input files")
|
||
s5cmd_syn_input_command = f's5cmd --log=error sync --delete=true {DIR3}/ "s3://{bucket_name}/comfy/{comfy_endpoint}/{prepare_version}/input/"'
|
||
logger.info(f"sync input files start {s5cmd_syn_input_command}")
|
||
os.system(s5cmd_syn_input_command)
|
||
logger.info(f" sync models files")
|
||
s5cmd_syn_model_command = f's5cmd --log=error sync --delete=true {DIR1}/ "s3://{bucket_name}/comfy/{comfy_endpoint}/{prepare_version}/models/"'
|
||
logger.info(f"sync models files start {s5cmd_syn_model_command}")
|
||
os.system(s5cmd_syn_model_command)
|
||
logger.info(f"Files changed in:: {need_prepare} {DIR2} {DIR1} {DIR3}")
|
||
url = api_url + "prepare"
|
||
logger.info(f"URL:{url}")
|
||
data = {"endpoint_name": comfy_endpoint, "need_reboot": need_reboot, "prepare_id": prepare_version,
|
||
"prepare_type": prepare_type}
|
||
logger.info(f"prepare params Data: {json.dumps(data, indent=4)}")
|
||
result = subprocess.run(["curl", "--location", "--request", "POST", url, "--header",
|
||
f"x-api-key: {api_token}", "--data-raw", json.dumps(data)],
|
||
capture_output=True, text=True)
|
||
logger.info(result.stdout)
|
||
return result.stdout
|
||
except Exception as e:
|
||
logger.info(f"sync_files error {e}")
|
||
return None
|
||
|
||
|
||
def sync_files(filepath, is_folder, is_auto):
|
||
try:
|
||
directory = os.path.dirname(filepath)
|
||
logger.info(f"Directory changed in: {directory} {filepath}")
|
||
if not directory:
|
||
logger.info("root path no need to sync files by duplicate opt")
|
||
return None
|
||
logger.info(f"Files changed in: {filepath}")
|
||
timestamp = str(int(time.time() * 1000))
|
||
need_prepare = False
|
||
prepare_type = 'default'
|
||
need_reboot = False
|
||
for ignore_item in no_need_sync_files:
|
||
if filepath.endswith(ignore_item):
|
||
logger.info(f"no need to sync files by ignore files {filepath} ends by {ignore_item}")
|
||
return None
|
||
prepare_version = PREPARE_ID if PREPARE_MODE == 'additional' else timestamp
|
||
if (str(directory).endswith(f"{DIR2}" if DIR2.startswith("/") else f"/{DIR2}")
|
||
or str(filepath) == DIR2 or str(filepath) == f'./{DIR2}' or f"{DIR2}/" in filepath):
|
||
logger.info(f" sync custom nodes files: {filepath}")
|
||
s5cmd_syn_node_command = f's5cmd --log=error sync --delete=true --exclude="*comfy_local_proxy.py" {DIR2}/ "s3://{bucket_name}/comfy/{comfy_endpoint}/{prepare_version}/custom_nodes/"'
|
||
# s5cmd_syn_node_command = f'aws s3 sync {DIR2}/ "s3://{bucket_name}/comfy/{comfy_endpoint}/{timestamp}/custom_nodes/"'
|
||
# s5cmd_syn_node_command = f's5cmd sync {DIR2}/* "s3://{bucket_name}/comfy/{comfy_endpoint}/{timestamp}/custom_nodes/"'
|
||
|
||
# custom_node文件夹有变化 稍后再同步
|
||
if is_auto and not is_folder_unlocked(directory):
|
||
logger.info("sync custom_nodes files is changing ,waiting.... ")
|
||
return None
|
||
logger.info("sync custom_nodes files start")
|
||
logger.info(s5cmd_syn_node_command)
|
||
os.system(s5cmd_syn_node_command)
|
||
need_prepare = True
|
||
need_reboot = True
|
||
prepare_type = 'nodes'
|
||
elif (str(directory).endswith(f"{DIR3}" if DIR3.startswith("/") else f"/{DIR3}")
|
||
or str(filepath) == DIR3 or str(filepath) == f'./{DIR3}' or f"{DIR3}/" in filepath):
|
||
logger.info(f" sync input files: {filepath}")
|
||
s5cmd_syn_input_command = f's5cmd --log=error sync --delete=true {DIR3}/ "s3://{bucket_name}/comfy/{comfy_endpoint}/{prepare_version}/input/"'
|
||
|
||
# 判断文件写完后再同步
|
||
if is_auto:
|
||
if bool(is_folder):
|
||
can_sync = is_folder_unlocked(filepath)
|
||
else:
|
||
can_sync = is_file_unlocked(filepath)
|
||
if not can_sync:
|
||
logger.info("sync input files is changing ,waiting.... ")
|
||
return None
|
||
logger.info("sync input files start")
|
||
logger.info(s5cmd_syn_input_command)
|
||
os.system(s5cmd_syn_input_command)
|
||
need_prepare = True
|
||
prepare_type = 'inputs'
|
||
elif (str(directory).endswith(f"{DIR1}" if DIR1.startswith("/") else f"/{DIR1}")
|
||
or str(filepath) == DIR1 or str(filepath) == f'./{DIR1}' or f"{DIR1}/" in filepath):
|
||
logger.info(f" sync models files: {filepath}")
|
||
s5cmd_syn_model_command = f's5cmd --log=error sync --delete=true {DIR1}/ "s3://{bucket_name}/comfy/{comfy_endpoint}/{prepare_version}/models/"'
|
||
|
||
# 判断文件写完后再同步
|
||
if is_auto:
|
||
if bool(is_folder):
|
||
can_sync = is_folder_unlocked(filepath)
|
||
else:
|
||
can_sync = is_file_unlocked(filepath)
|
||
# logger.info(f'is folder {directory} {is_folder} can_sync {can_sync}')
|
||
if not can_sync:
|
||
logger.info("sync input models is changing ,waiting.... ")
|
||
return None
|
||
|
||
logger.info("sync models files start")
|
||
logger.info(s5cmd_syn_model_command)
|
||
os.system(s5cmd_syn_model_command)
|
||
need_prepare = True
|
||
prepare_type = 'models'
|
||
logger.info(f"Files changed in:: {need_prepare} {str(directory)} {DIR2} {DIR1} {DIR3}")
|
||
if need_prepare:
|
||
url = api_url + "prepare"
|
||
logger.info(f"URL:{url}")
|
||
data = {"endpoint_name": comfy_endpoint, "need_reboot": need_reboot, "prepare_id": prepare_version,
|
||
"prepare_type": prepare_type}
|
||
logger.info(f"prepare params Data: {json.dumps(data, indent=4)}")
|
||
result = subprocess.run(["curl", "--location", "--request", "POST", url, "--header",
|
||
f"x-api-key: {api_token}", "--data-raw", json.dumps(data)],
|
||
capture_output=True, text=True)
|
||
logger.info(result.stdout)
|
||
return result.stdout
|
||
return None
|
||
except Exception as e:
|
||
logger.info(f"sync_files error {e}")
|
||
return None
|
||
|
||
|
||
def is_folder_unlocked(directory):
|
||
# logger.info("check if folder ")
|
||
event_handler = MyHandlerWithCheck()
|
||
observer = Observer()
|
||
observer.schedule(event_handler, directory, recursive=True)
|
||
observer.start()
|
||
time.sleep(1)
|
||
result = False
|
||
try:
|
||
if event_handler.file_changed:
|
||
logger.info(f"folder {directory} is still changing..")
|
||
event_handler.file_changed = False
|
||
time.sleep(1)
|
||
if event_handler.file_changed:
|
||
logger.info(f"folder {directory} is still still changing..")
|
||
else:
|
||
logger.info(f"folder {directory} changing stopped")
|
||
result = True
|
||
else:
|
||
logger.info(f"folder {directory} not stopped")
|
||
result = True
|
||
except (KeyboardInterrupt, Exception) as e:
|
||
logger.info(f"folder {directory} changed exception {e}")
|
||
observer.stop()
|
||
return result
|
||
|
||
|
||
def is_file_unlocked(file_path):
|
||
# logger.info("check if file ")
|
||
try:
|
||
initial_size = os.path.getsize(file_path)
|
||
initial_mtime = os.path.getmtime(file_path)
|
||
time.sleep(1)
|
||
|
||
current_size = os.path.getsize(file_path)
|
||
current_mtime = os.path.getmtime(file_path)
|
||
if current_size != initial_size or current_mtime != initial_mtime:
|
||
logger.info(f"unlock file error {file_path} is changing")
|
||
return False
|
||
|
||
with open(file_path, 'r') as f:
|
||
fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||
return True
|
||
except (IOError, OSError, Exception) as e:
|
||
logger.info(f"unlock file error {file_path} is writing")
|
||
logger.error(e)
|
||
return False
|
||
|
||
|
||
class MyHandlerWithCheck(FileSystemEventHandler):
|
||
def __init__(self):
|
||
self.file_changed = False
|
||
|
||
def on_modified(self, event):
|
||
logger.info(f"custom_node folder is changing {event.src_path}")
|
||
self.file_changed = True
|
||
|
||
def on_deleted(self, event):
|
||
logger.info(f"custom_node folder is changing {event.src_path}")
|
||
self.file_changed = True
|
||
|
||
def on_created(self, event):
|
||
logger.info(f"custom_node folder is changing {event.src_path}")
|
||
self.file_changed = True
|
||
|
||
|
||
class MyHandlerWithSync(FileSystemEventHandler):
|
||
def on_modified(self, event):
|
||
logger.info(f"{datetime.datetime.now()} files modified ，start to sync {event}")
|
||
sync_files(event.src_path, event.is_directory, True)
|
||
|
||
def on_created(self, event):
|
||
logger.info(f"{datetime.datetime.now()} files added ，start to sync {event}")
|
||
sync_files(event.src_path, event.is_directory, True)
|
||
|
||
def on_deleted(self, event):
|
||
logger.info(f"{datetime.datetime.now()} files deleted ，start to sync {event}")
|
||
sync_files(event.src_path, event.is_directory, True)
|
||
|
||
|
||
stop_event = threading.Event()
|
||
|
||
|
||
def check_and_sync():
|
||
logger.info("check_and_sync start")
|
||
event_handler = MyHandlerWithSync()
|
||
observer = Observer()
|
||
try:
|
||
observer.schedule(event_handler, DIR1, recursive=True)
|
||
observer.schedule(event_handler, DIR2, recursive=True)
|
||
observer.schedule(event_handler, DIR3, recursive=True)
|
||
observer.start()
|
||
while True:
|
||
time.sleep(1)
|
||
except KeyboardInterrupt:
|
||
logger.info("sync Shutting down please restart ComfyUI")
|
||
observer.stop()
|
||
observer.join()
|
||
|
||
|
||
def signal_handler(sig, frame):
|
||
logger.info("Received termination signal. Exiting...")
|
||
stop_event.set()
|
||
|
||
|
||
if os.environ.get('DISABLE_AUTO_SYNC') == 'false':
|
||
signal.signal(signal.SIGINT, signal_handler)
|
||
signal.signal(signal.SIGTERM, signal_handler)
|
||
|
||
check_sync_thread = threading.Thread(target=check_and_sync)
|
||
check_sync_thread.start()
|
||
|
||
@server.PromptServer.instance.routes.post('/map_release')
|
||
async def map_release(request):
|
||
logger.info(f"start to map_release {request}")
|
||
json_data = await request.json()
|
||
if (not json_data or 'clientId' not in json_data or not json_data.get('clientId')
|
||
or 'releaseVersion' not in json_data or not json_data.get('releaseVersion')):
|
||
return web.Response(status=500, content_type='application/json', body=json.dumps({"result": False}))
|
||
global client_release_map
|
||
client_release_map[json_data.get('clientId')] = json_data.get('releaseVersion')
|
||
logger.info(f"client_release_map :{client_release_map}")
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps({"result": True}))
|
||
|
||
|
||
@server.PromptServer.instance.routes.get("/reboot")
|
||
async def restart(self):
|
||
if is_action_lock():
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps(
|
||
{"result": False, "message": "reboot is not allowed during sync workflow"}))
|
||
|
||
if not is_master_process:
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": "only master can restart"}))
|
||
logger.info(f"start to reboot {self}")
|
||
try:
|
||
from xmlrpc.client import ServerProxy
|
||
server = ServerProxy('http://localhost:9001/RPC2')
|
||
server.supervisor.restart()
|
||
# server.supervisor.shutdown()
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps({"result": True}))
|
||
except Exception as e:
|
||
logger.info(f"error reboot {e}")
|
||
pass
|
||
return os.execv(sys.executable, [sys.executable] + sys.argv)
|
||
|
||
|
||
@server.PromptServer.instance.routes.get("/check_prepare")
|
||
async def check_prepare(self):
|
||
logger.info(f"start to check_prepare {self}")
|
||
try:
|
||
get_response = requests.get(f"{api_url}/prepare/{comfy_endpoint}", headers=headers)
|
||
response = get_response.json()
|
||
logger.info(f"check sync response is {response}")
|
||
if get_response.status_code == 200 and response['data']['prepareSuccess']:
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps({"result": True}))
|
||
else:
|
||
logger.info(f"check sync response is {response} {response['data']['prepareSuccess']}")
|
||
return web.Response(status=500, content_type='application/json', body=json.dumps({"result": False}))
|
||
except Exception as e:
|
||
logger.info(f"error restart {e}")
|
||
pass
|
||
return os.execv(sys.executable, [sys.executable] + sys.argv)
|
||
|
||
|
||
@server.PromptServer.instance.routes.get("/gc")
|
||
async def gc(self):
|
||
logger.info(f"start to gc {self}")
|
||
try:
|
||
logger.info(f"gc start: {time.time()}")
|
||
server_instance = server.PromptServer.instance
|
||
e = execution.PromptExecutor(server_instance)
|
||
e.reset()
|
||
comfy.model_management.cleanup_models()
|
||
gc.collect()
|
||
comfy.model_management.soft_empty_cache()
|
||
gc_triggered = True
|
||
logger.info(f"gc end: {time.time()}")
|
||
except Exception as e:
|
||
logger.info(f"error restart {e}")
|
||
pass
|
||
return os.execv(sys.executable, [sys.executable] + sys.argv)
|
||
|
||
|
||
def is_action_lock():
|
||
lock_file = f'/container/sync_lock'
|
||
if os.path.exists(lock_file):
|
||
with open(lock_file, 'r') as f:
|
||
content = f.read()
|
||
if content and not check_workflow_exists(content):
|
||
return True
|
||
return False
|
||
|
||
|
||
def action_lock(name: str):
|
||
lock_file = f'/container/sync_lock'
|
||
with open(lock_file, 'w') as f:
|
||
f.write(name)
|
||
|
||
|
||
def action_unlock():
|
||
lock_file = f'/container/sync_lock'
|
||
with open(lock_file, 'w') as f:
|
||
f.write("")
|
||
|
||
|
||
@server.PromptServer.instance.routes.get("/restart")
|
||
async def restart(self):
|
||
if is_action_lock():
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps(
|
||
{"result": False, "message": "restart is not allowed during sync workflow"}))
|
||
|
||
logger.info(f"start to restart {self}")
|
||
try:
|
||
sys.stdout.close_log()
|
||
except Exception as e:
|
||
logger.info(f"error restart {e}")
|
||
pass
|
||
return os.execv(sys.executable, [sys.executable] + sys.argv)
|
||
|
||
|
||
@server.PromptServer.instance.routes.get("/sync_env")
|
||
async def sync_env(request):
|
||
logger.info(f"start to sync_env {request}")
|
||
try:
|
||
result = sync_default_files()
|
||
logger.debug(f"sync result is :{result}")
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps({"result": True}))
|
||
except Exception as e:
|
||
logger.info(f"error sync_env {e}")
|
||
pass
|
||
return web.Response(status=500, content_type='application/json', body=json.dumps({"result": False}))
|
||
|
||
|
||
@server.PromptServer.instance.routes.post("/change_env")
|
||
async def change_env(request):
|
||
logger.info(f"start to change_env {request}")
|
||
json_data = await request.json()
|
||
if DISABLE_AWS_PROXY in json_data and json_data[DISABLE_AWS_PROXY] is not None:
|
||
logger.info(f"origin evn key DISABLE_AWS_PROXY is :{os.environ.get(DISABLE_AWS_PROXY)} {str(json_data[DISABLE_AWS_PROXY])}")
|
||
os.environ[DISABLE_AWS_PROXY] = str(json_data[DISABLE_AWS_PROXY])
|
||
logger.info(f"now evn key DISABLE_AWS_PROXY is :{os.environ.get(DISABLE_AWS_PROXY)}")
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps({"result": True}))
|
||
|
||
|
||
@server.PromptServer.instance.routes.get("/get_env")
|
||
async def get_env(request):
|
||
env = os.environ.get(DISABLE_AWS_PROXY, 'False')
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps({"env": env}))
|
||
|
||
def get_directory_size(directory):
|
||
total_size = 0
|
||
for dirpath, dirnames, filenames in os.walk(directory):
|
||
for filename in filenames:
|
||
filepath = os.path.join(dirpath, filename)
|
||
if not os.path.islink(filepath): # 检查文件是否不是符号链接
|
||
total_size += os.path.getsize(filepath)
|
||
return total_size
|
||
|
||
@server.PromptServer.instance.routes.post("/workflows")
|
||
async def release_workflow(request):
|
||
if is_action_lock():
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps(
|
||
{"result": False, "message": "release is not allowed during sync workflow"}))
|
||
|
||
if not is_master_process:
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": "only master can release workflow"}))
|
||
|
||
logger.info(f"start to release workflow {request}")
|
||
try:
|
||
json_data = await request.json()
|
||
if 'name' not in json_data or not json_data['name']:
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": f"name is required"}))
|
||
|
||
workflow_name = json_data['name']
|
||
if workflow_name == 'default':
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": f"{workflow_name} is not allowed"}))
|
||
|
||
payload_json = ''
|
||
|
||
if 'payload_json' in json_data:
|
||
payload_json = json_data['payload_json']
|
||
|
||
if check_workflow_exists(workflow_name):
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": f"{workflow_name} already exists"}))
|
||
|
||
start_time = time.time()
|
||
|
||
action_lock(workflow_name)
|
||
|
||
base_image = os.getenv('BASE_IMAGE')
|
||
subprocess.check_output(f"echo {workflow_name} > /container/image_target_name", shell=True)
|
||
subprocess.check_output(f"echo {base_image} > /container/image_base", shell=True)
|
||
|
||
cur_workflow_name = os.getenv('WORKFLOW_NAME')
|
||
source_path = f"/container/workflows/{cur_workflow_name}"
|
||
|
||
print(f"source_path is {source_path}")
|
||
total_size_bytes = get_directory_size(source_path)
|
||
source_size = round(total_size_bytes / (1024 ** 3), 2)
|
||
|
||
s5cmd_sync_command = (f's5cmd sync '
|
||
f'--delete=true '
|
||
f'--exclude="*comfy.tar" '
|
||
f'--exclude="*.log" '
|
||
f'--exclude="*__pycache__*" '
|
||
f'--exclude="*/ComfyUI/input/*" '
|
||
f'--exclude="*/ComfyUI/output/*" '
|
||
f'"{source_path}/*" '
|
||
f'"s3://{bucket_name}/comfy/workflows/{workflow_name}/"')
|
||
|
||
s5cmd_lock_command = (f'echo "lock" > lock && '
|
||
f's5cmd sync lock s3://{bucket_name}/comfy/workflows/{workflow_name}/lock')
|
||
|
||
logger.info(f"sync workflows files start {s5cmd_sync_command}")
|
||
|
||
subprocess.check_output(s5cmd_sync_command, shell=True)
|
||
subprocess.check_output(s5cmd_lock_command, shell=True)
|
||
|
||
end_time = time.time()
|
||
cost_time = end_time - start_time
|
||
image_hash = os.getenv('IMAGE_HASH')
|
||
image_uri = f"{image_hash}:{workflow_name}"
|
||
|
||
data = {
|
||
"payload_json": payload_json,
|
||
"image_uri": image_uri,
|
||
"name": workflow_name,
|
||
"size": str(source_size),
|
||
}
|
||
get_response = requests.post(f"{api_url}/workflows", headers=headers, data=json.dumps(data))
|
||
response = get_response.json()
|
||
logger.info(f"release workflow response is {response}")
|
||
action_unlock()
|
||
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": True, "message": "success", "cost_time": cost_time}))
|
||
except Exception as e:
|
||
logger.info(e)
|
||
action_unlock()
|
||
return web.Response(status=500, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": 'Release workflow failed'}))
|
||
|
||
|
||
@server.PromptServer.instance.routes.put("/workflows")
|
||
async def switch_workflow(request):
|
||
if is_action_lock():
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps(
|
||
{"result": False, "message": "switch is not allowed during sync workflow"}))
|
||
|
||
try:
|
||
json_data = await request.json()
|
||
if 'name' not in json_data or not json_data['name']:
|
||
raise ValueError("name is required")
|
||
|
||
workflow_name = json_data['name']
|
||
|
||
if workflow_name == os.getenv('WORKFLOW_NAME'):
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": "workflow is already in use"}))
|
||
|
||
if workflow_name == 'default' and not is_master_process:
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": "slave can not use default workflow"}))
|
||
|
||
if workflow_name != 'default':
|
||
if not check_workflow_exists(workflow_name):
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": f"{workflow_name} not exists"}))
|
||
|
||
name_file = os.getenv('WORKFLOW_NAME_FILE')
|
||
subprocess.check_output(f"echo {workflow_name} > {name_file}", shell=True)
|
||
subprocess.run(["pkill", "-f", "python3"])
|
||
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": True, "message": "Please wait to restart"}))
|
||
except Exception as e:
|
||
logger.info(e)
|
||
return web.Response(status=500, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": 'Switch workflow failed'}))
|
||
|
||
|
||
@server.PromptServer.instance.routes.get("/workflows")
|
||
async def get_workflows(request):
|
||
try:
|
||
workflow_name = os.getenv('WORKFLOW_NAME')
|
||
|
||
response = requests.get(f"{api_url}/workflows", headers=headers, params={"limit": 1000})
|
||
if response.status_code != 200:
|
||
return web.Response(status=500, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": 'Get workflows failed'}))
|
||
|
||
data = response.json()['data']
|
||
workflows = data['workflows']
|
||
list = []
|
||
for workflow in workflows:
|
||
if workflow['status'] != 'Enabled':
|
||
continue
|
||
list.append({
|
||
"name": workflow['name'],
|
||
"size": workflow['size'],
|
||
"payload_json": workflow['payload_json'],
|
||
"in_use": workflow['name'] == workflow_name
|
||
})
|
||
|
||
data['current_workflow'] = workflow_name
|
||
data['workflows'] = list
|
||
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": True, "data": data}))
|
||
except Exception as e:
|
||
logger.info(e)
|
||
return web.Response(status=500, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": 'Switch workflow failed'}))
|
||
|
||
|
||
def check_workflow_exists(name: str):
|
||
get_response = requests.get(f"{api_url}/workflows/{name}", headers=headers)
|
||
return get_response.status_code == 200
|
||
|
||
|
||
def restore_commands():
|
||
subprocess.run(["sleep", "5"])
|
||
subprocess.run(["pkill", "-f", "python3"])
|
||
|
||
@server.PromptServer.instance.routes.post("/restore")
|
||
async def release_rebuild_workflow(request):
|
||
if os.getenv('WORKFLOW_NAME') != 'default':
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": "only default workflow can be restored"}))
|
||
|
||
if is_action_lock():
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps(
|
||
{"result": False, "message": "restore is not allowed during sync workflow"}))
|
||
|
||
if not is_master_process:
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": "only master can restore comfy"}))
|
||
|
||
logger.info(f"start to restore EC2 {request}")
|
||
|
||
try:
|
||
os.system("mv /container/workflows/default/ComfyUI/models/checkpoints/v1-5-pruned-emaonly.ckpt /")
|
||
os.system("mv /container/workflows/default/ComfyUI/models/animatediff_models/mm_sd_v15_v2.ckpt /")
|
||
os.system("rm -rf /container/workflows/default")
|
||
os.system("mkdir -p /container/workflows/default")
|
||
os.system("tar --overwrite -xf /container/default.tar -C /container/workflows/default/")
|
||
os.system("mkdir -p /container/workflows/default/ComfyUI/models/checkpoints")
|
||
os.system("mkdir -p /container/workflows/default/ComfyUI/models/animatediff_models")
|
||
os.system("mv /v1-5-pruned-emaonly.ckpt /container/workflows/default/ComfyUI/models/checkpoints/")
|
||
os.system("mv /mm_sd_v15_v2.ckpt /container/workflows/default/ComfyUI/models/animatediff_models/")
|
||
thread = threading.Thread(target=restore_commands)
|
||
thread.start()
|
||
return web.Response(status=200, content_type='application/json',
|
||
body=json.dumps({"result": True, "message": "comfy will be restored in 5 seconds"}))
|
||
except Exception as e:
|
||
return web.Response(status=500, content_type='application/json',
|
||
body=json.dumps({"result": False, "message": e}))
|
||
|
||
|
||
if is_on_sagemaker:
|
||
|
||
global need_sync
|
||
global prompt_id
|
||
global executing
|
||
executing = False
|
||
|
||
global reboot
|
||
reboot = False
|
||
|
||
global last_call_time
|
||
last_call_time = None
|
||
global gc_triggered
|
||
gc_triggered = False
|
||
|
||
REGION = os.environ.get('AWS_REGION')
|
||
BUCKET = os.environ.get('S3_BUCKET_NAME')
|
||
QUEUE_URL = os.environ.get('COMFY_QUEUE_URL')
|
||
|
||
GEN_INSTANCE_ID = os.environ.get('ENDPOINT_INSTANCE_ID') if 'ENDPOINT_INSTANCE_ID' in os.environ and os.environ.get('ENDPOINT_INSTANCE_ID') else str(uuid.uuid4())
|
||
ENDPOINT_NAME = os.environ.get('ENDPOINT_NAME')
|
||
ENDPOINT_ID = os.environ.get('ENDPOINT_ID')
|
||
|
||
INSTANCE_MONITOR_TABLE_NAME = os.environ.get('COMFY_INSTANCE_MONITOR_TABLE')
|
||
SYNC_TABLE_NAME = os.environ.get('COMFY_SYNC_TABLE')
|
||
|
||
dynamodb = boto3.resource('dynamodb', region_name=REGION)
|
||
sync_table = dynamodb.Table(SYNC_TABLE_NAME)
|
||
instance_monitor_table = dynamodb.Table(INSTANCE_MONITOR_TABLE_NAME)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
logger.setLevel(os.environ.get('LOG_LEVEL') or logging.INFO)
|
||
|
||
ROOT_PATH = '/home/ubuntu/ComfyUI'
|
||
sqs_client = boto3.client('sqs', region_name=REGION)
|
||
|
||
GC_WAIT_TIME = 1800
|
||
|
||
|
||
def print_env():
|
||
for key, value in os.environ.items():
|
||
logger.info(f"{key}: {value}")
|
||
|
||
|
||
@dataclass
|
||
class ComfyResponse:
|
||
statusCode: int
|
||
message: str
|
||
body: Optional[dict]
|
||
|
||
|
||
def ok(body: dict):
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps(body))
|
||
|
||
|
||
def error(body: dict):
|
||
# TODO 500 -》200 because of need resp anyway not exception
|
||
return web.Response(status=200, content_type='application/json', body=json.dumps(body))
|
||
|
||
|
||
def sen_sqs_msg(message_body, prompt_id_key):
|
||
response = sqs_client.send_message(
|
||
QueueUrl=QUEUE_URL,
|
||
MessageBody=json.dumps(message_body),
|
||
MessageGroupId=prompt_id_key
|
||
)
|
||
message_id = response['MessageId']
|
||
return message_id
|
||
|
||
|
||
def sen_finish_sqs_msg(prompt_id_key):
|
||
global need_sync
|
||
# logger.info(f"sen_finish_sqs_msg start... {need_sync},{prompt_id_key}")
|
||
if need_sync and QUEUE_URL and REGION:
|
||
message_body = {'prompt_id': prompt_id_key, 'event': 'finish', 'data': {"node": None, "prompt_id": prompt_id_key},
|
||
'sid': None}
|
||
message_id = sen_sqs_msg(message_body, prompt_id_key)
|
||
logger.info(f"finish message sent {message_id}")
|
||
|
||
|
||
async def prepare_comfy_env(sync_item: dict):
|
||
try:
|
||
request_id = sync_item['request_id']
|
||
logger.info(f"prepare_environment start sync_item:{sync_item}")
|
||
prepare_type = sync_item['prepare_type']
|
||
rlt = True
|
||
if prepare_type in ['default', 'models']:
|
||
sync_models_rlt = sync_s3_files_or_folders_to_local(f'{request_id}/models/*', f'{ROOT_PATH}/models', False)
|
||
if not sync_models_rlt:
|
||
rlt = False
|
||
if prepare_type in ['default', 'inputs']:
|
||
sync_inputs_rlt = sync_s3_files_or_folders_to_local(f'{request_id}/input/*', f'{ROOT_PATH}/input', False)
|
||
if not sync_inputs_rlt:
|
||
rlt = False
|
||
if prepare_type in ['default', 'nodes']:
|
||
sync_nodes_rlt = sync_s3_files_or_folders_to_local(f'{request_id}/custom_nodes/*',
|
||
f'{ROOT_PATH}/custom_nodes', True)
|
||
if not sync_nodes_rlt:
|
||
rlt = False
|
||
if prepare_type == 'custom':
|
||
sync_source_path = sync_item['s3_source_path']
|
||
local_target_path = sync_item['local_target_path']
|
||
if not sync_source_path or not local_target_path:
|
||
logger.info("s3_source_path and local_target_path should not be empty")
|
||
else:
|
||
sync_rlt = sync_s3_files_or_folders_to_local(sync_source_path,
|
||
f'{ROOT_PATH}/{local_target_path}', False)
|
||
if not sync_rlt:
|
||
rlt = False
|
||
elif prepare_type == 'other':
|
||
sync_script = sync_item['sync_script']
|
||
logger.info("sync_script")
|
||
# sync_script.startswith('s5cmd') 不允许
|
||
try:
|
||
if sync_script and (sync_script.startswith("python3 -m pip") or sync_script.startswith("python -m pip")
|
||
or sync_script.startswith("pip install") or sync_script.startswith("apt")
|
||
or sync_script.startswith("os.environ") or sync_script.startswith("ls")
|
||
or sync_script.startswith("env") or sync_script.startswith("source")
|
||
or sync_script.startswith("curl") or sync_script.startswith("wget")
|
||
or sync_script.startswith("print") or sync_script.startswith("cat")
|
||
or sync_script.startswith("sudo chmod") or sync_script.startswith("chmod")
|
||
or sync_script.startswith("/home/ubuntu/ComfyUI/venv/bin/python")):
|
||
os.system(sync_script)
|
||
elif sync_script and (sync_script.startswith("export ") and len(sync_script.split(" ")) > 2):
|
||
sync_script_key = sync_script.split(" ")[1]
|
||
sync_script_value = sync_script.split(" ")[2]
|
||
os.environ[sync_script_key] = sync_script_value
|
||
logger.info(os.environ.get(sync_script_key))
|
||
except Exception as e:
|
||
logger.error(f"Exception while execute sync_scripts : {sync_script}")
|
||
rlt = False
|
||
need_reboot = True if ('need_reboot' in sync_item and sync_item['need_reboot']
|
||
and str(sync_item['need_reboot']).lower() == 'true')else False
|
||
global reboot
|
||
reboot = need_reboot
|
||
if need_reboot:
|
||
os.environ['NEED_REBOOT'] = 'true'
|
||
else:
|
||
os.environ['NEED_REBOOT'] = 'false'
|
||
logger.info("prepare_environment end")
|
||
os.environ['LAST_SYNC_REQUEST_ID'] = sync_item['request_id']
|
||
os.environ['LAST_SYNC_REQUEST_TIME'] = str(sync_item['request_time'])
|
||
return rlt
|
||
except Exception as e:
|
||
return False
|
||
|
||
|
||
def sync_s3_files_or_folders_to_local(s3_path, local_path, need_un_tar):
|
||
logger.info("sync_s3_models_or_inputs_to_local start")
|
||
# s5cmd_command = f'{ROOT_PATH}/tools/s5cmd sync "s3://{bucket_name}/{s3_path}/*" "{local_path}/"'
|
||
if need_un_tar:
|
||
s5cmd_command = f's5cmd sync "s3://{BUCKET}/comfy/{ENDPOINT_NAME}/{s3_path}" "{local_path}/"'
|
||
else:
|
||
s5cmd_command = f's5cmd sync --delete=true "s3://{BUCKET}/comfy/{ENDPOINT_NAME}/{s3_path}" "{local_path}/"'
|
||
# s5cmd_command = f's5cmd sync --delete=true "s3://{BUCKET}/comfy/{ENDPOINT_NAME}/{s3_path}" "{local_path}/"'
|
||
# s5cmd_command = f's5cmd sync "s3://{BUCKET}/comfy/{ENDPOINT_NAME}/{s3_path}" "{local_path}/"'
|
||
try:
|
||
logger.info(s5cmd_command)
|
||
os.system(s5cmd_command)
|
||
logger.info(f'Files copied from "s3://{BUCKET}/comfy/{ENDPOINT_NAME}/{s3_path}" to "{local_path}/"')
|
||
if need_un_tar:
|
||
for filename in os.listdir(local_path):
|
||
if filename.endswith(".tar.gz"):
|
||
tar_filepath = os.path.join(local_path, filename)
|
||
# extract_path = os.path.splitext(os.path.splitext(tar_filepath)[0])[0]
|
||
# os.makedirs(extract_path, exist_ok=True)
|
||
# logger.info(f'Extracting extract_path is {extract_path}')
|
||
|
||
with tarfile.open(tar_filepath, "r:gz") as tar:
|
||
for member in tar.getmembers():
|
||
tar.extract(member, path=local_path)
|
||
os.remove(tar_filepath)
|
||
logger.info(f'File {tar_filepath} extracted and removed')
|
||
return True
|
||
except Exception as e:
|
||
logger.info(f"Error executing s5cmd command: {e}")
|
||
return False
|
||
|
||
|
||
def sync_local_outputs_to_s3(s3_path, local_path):
|
||
logger.info("sync_local_outputs_to_s3 start")
|
||
s5cmd_command = f's5cmd sync "{local_path}/*" "s3://{BUCKET}/comfy/{s3_path}/" '
|
||
try:
|
||
logger.info(s5cmd_command)
|
||
os.system(s5cmd_command)
|
||
logger.info(f'Files copied local to "s3://{BUCKET}/comfy/{s3_path}/" to "{local_path}/"')
|
||
clean_cmd = f'rm -rf {local_path}'
|
||
os.system(clean_cmd)
|
||
logger.info(f'Files removed from local {local_path}')
|
||
except Exception as e:
|
||
logger.info(f"Error executing s5cmd command: {e}")
|
||
|
||
|
||
def sync_local_outputs_to_base64(local_path):
|
||
logger.info("sync_local_outputs_to_base64 start")
|
||
try:
|
||
result = {}
|
||
for root, dirs, files in os.walk(local_path):
|
||
for file in files:
|
||
file_path = os.path.join(root, file)
|
||
with open(file_path, "rb") as f:
|
||
file_content = f.read()
|
||
base64_content = base64.b64encode(file_content).decode('utf-8')
|
||
result[file] = base64_content
|
||
clean_cmd = f'rm -rf {local_path}'
|
||
os.system(clean_cmd)
|
||
logger.info(f'Files removed from local {local_path}')
|
||
return result
|
||
except Exception as e:
|
||
logger.info(f"Error executing s5cmd command: {e}")
|
||
return {}
|
||
|
||
|
||
@server.PromptServer.instance.routes.post("/execute_proxy")
|
||
async def execute_proxy(request):
|
||
json_data = await request.json()
|
||
if 'out_path' in json_data and json_data['out_path'] is not None:
|
||
out_path = json_data['out_path']
|
||
else:
|
||
out_path = None
|
||
logger.info(f"invocations start json_data:{json_data}")
|
||
global need_sync
|
||
need_sync = json_data["need_sync"]
|
||
global prompt_id
|
||
prompt_id = json_data["prompt_id"]
|
||
try:
|
||
global executing
|
||
if executing is True:
|
||
resp = {"prompt_id": prompt_id, "instance_id": GEN_INSTANCE_ID, "status": "fail",
|
||
"message": "the environment is not ready valid[0] is false, need to resync"}
|
||
sen_finish_sqs_msg(prompt_id)
|
||
return error(resp)
|
||
executing = True
|
||
logger.info(
|
||
f'bucket_name: {BUCKET}, region: {REGION}')
|
||
if ('need_prepare' in json_data and json_data['need_prepare']
|
||
and 'prepare_props' in json_data and json_data['prepare_props']):
|
||
sync_already = await prepare_comfy_env(json_data['prepare_props'])
|
||
if not sync_already:
|
||
resp = {"prompt_id": prompt_id, "instance_id": GEN_INSTANCE_ID, "status": "fail",
|
||
"message": "the environment is not ready with sync"}
|
||
executing = False
|
||
sen_finish_sqs_msg(prompt_id)
|
||
return error(resp)
|
||
server_instance = server.PromptServer.instance
|
||
if "number" in json_data:
|
||
number = float(json_data['number'])
|
||
server_instance.number = number
|
||
else:
|
||
number = server_instance.number
|
||
if "front" in json_data:
|
||
if json_data['front']:
|
||
number = -number
|
||
server_instance.number += 1
|
||
valid = execution.validate_prompt(json_data['prompt'])
|
||
logger.info(f"Validating prompt result is {valid}")
|
||
if not valid[0]:
|
||
resp = {"prompt_id": prompt_id, "instance_id": GEN_INSTANCE_ID, "status": "fail",
|
||
"message": "the environment is not ready valid[0] is false, need to resync"}
|
||
executing = False
|
||
response = {"prompt_id": prompt_id, "number": number, "node_errors": valid[3]}
|
||
sen_finish_sqs_msg(prompt_id)
|
||
return error(resp)
|
||
# if len(valid) == 4 and len(valid[3]) > 0:
|
||
# logger.info(f"Validating prompt error there is something error because of :valid: {valid}")
|
||
# resp = {"prompt_id": prompt_id, "instance_id": GEN_INSTANCE_ID, "status": "fail",
|
||
# "message": f"the valid is error, need to resync or check the workflow :{valid}"}
|
||
# executing = False
|
||
# return error(resp)
|
||
extra_data = {}
|
||
client_id = ''
|
||
if "extra_data" in json_data:
|
||
extra_data = json_data["extra_data"]
|
||
if 'client_id' in extra_data and extra_data['client_id']:
|
||
client_id = extra_data['client_id']
|
||
if "client_id" in json_data and json_data["client_id"]:
|
||
extra_data["client_id"] = json_data["client_id"]
|
||
client_id = json_data["client_id"]
|
||
|
||
server_instance.client_id = client_id
|
||
|
||
prompt_id = json_data['prompt_id']
|
||
server_instance.last_prompt_id = prompt_id
|
||
e = execution.PromptExecutor(server_instance)
|
||
outputs_to_execute = valid[2]
|
||
e.execute(json_data['prompt'], prompt_id, extra_data, outputs_to_execute)
|
||
|
||
s3_out_path = f'output/{prompt_id}/{out_path}' if out_path is not None else f'output/{prompt_id}'
|
||
s3_temp_path = f'temp/{prompt_id}/{out_path}' if out_path is not None else f'temp/{prompt_id}'
|
||
local_out_path = f'{ROOT_PATH}/output/{out_path}' if out_path is not None else f'{ROOT_PATH}/output'
|
||
local_temp_path = f'{ROOT_PATH}/temp/{out_path}' if out_path is not None else f'{ROOT_PATH}/temp'
|
||
|
||
logger.info(f"s3_out_path is {s3_out_path} and s3_temp_path is {s3_temp_path} and local_out_path is {local_out_path} and local_temp_path is {local_temp_path}")
|
||
|
||
sync_local_outputs_to_s3(s3_out_path, local_out_path)
|
||
sync_local_outputs_to_s3(s3_temp_path, local_temp_path)
|
||
|
||
response_body = {
|
||
"prompt_id": prompt_id,
|
||
"instance_id": GEN_INSTANCE_ID,
|
||
"status": "success",
|
||
"output_path": f's3://{BUCKET}/comfy/{s3_out_path}',
|
||
"temp_path": f's3://{BUCKET}/comfy/{s3_temp_path}',
|
||
}
|
||
sen_finish_sqs_msg(prompt_id)
|
||
logger.info(f"execute inference response is {response_body}")
|
||
executing = False
|
||
return ok(response_body)
|
||
except Exception as ecp:
|
||
logger.info(f"exception occurred {ecp}")
|
||
resp = {"prompt_id": prompt_id, "instance_id": GEN_INSTANCE_ID, "status": "fail",
|
||
"message": f"exception occurred {ecp}"}
|
||
executing = False
|
||
return error(resp)
|
||
finally:
|
||
logger.info(f"gc check: {time.time()}")
|
||
try:
|
||
global last_call_time, gc_triggered
|
||
gc_triggered = False
|
||
if last_call_time is None:
|
||
logger.info(f"gc check last time is NONE")
|
||
last_call_time = time.time()
|
||
else:
|
||
if time.time() - last_call_time > GC_WAIT_TIME:
|
||
if not gc_triggered:
|
||
logger.info(f"gc start: {time.time()} - {last_call_time}")
|
||
e.reset()
|
||
comfy.model_management.cleanup_models()
|
||
gc.collect()
|
||
comfy.model_management.soft_empty_cache()
|
||
gc_triggered = True
|
||
logger.info(f"gc end: {time.time()} - {last_call_time}")
|
||
last_call_time = time.time()
|
||
else:
|
||
last_call_time = time.time()
|
||
logger.info(f"gc check end: {time.time()}")
|
||
except Exception as e:
|
||
logger.info(f"gc error: {e}")
|
||
|
||
|
||
def get_last_ddb_sync_record():
|
||
sync_response = sync_table.query(
|
||
KeyConditionExpression=Key('endpoint_name').eq(ENDPOINT_NAME),
|
||
Limit=1,
|
||
ScanIndexForward=False
|
||
)
|
||
latest_sync_record = sync_response['Items'][0] if ('Items' in sync_response
|
||
and len(sync_response['Items']) > 0) else None
|
||
if latest_sync_record:
|
||
logger.info(f"latest_sync_record is：{latest_sync_record}")
|
||
return latest_sync_record
|
||
|
||
logger.info("no latest_sync_record found")
|
||
return None
|
||
|
||
|
||
def get_latest_ddb_instance_monitor_record():
|
||
key_condition_expression = ('endpoint_name = :endpoint_name_val '
|
||
'AND gen_instance_id = :gen_instance_id_val')
|
||
expression_attribute_values = {
|
||
':endpoint_name_val': ENDPOINT_NAME,
|
||
':gen_instance_id_val': GEN_INSTANCE_ID
|
||
}
|
||
instance_monitor_response = instance_monitor_table.query(
|
||
KeyConditionExpression=key_condition_expression,
|
||
ExpressionAttributeValues=expression_attribute_values
|
||
)
|
||
instance_monitor_record = instance_monitor_response['Items'][0] \
|
||
if ('Items' in instance_monitor_response and len(instance_monitor_response['Items']) > 0) else None
|
||
|
||
if instance_monitor_record:
|
||
logger.info(f"instance_monitor_record is {instance_monitor_record}")
|
||
return instance_monitor_record
|
||
|
||
logger.info("no instance_monitor_record found")
|
||
return None
|
||
|
||
|
||
def save_sync_instance_monitor(last_sync_request_id: str, sync_status: str):
|
||
item = {
|
||
'endpoint_id': ENDPOINT_ID,
|
||
'endpoint_name': ENDPOINT_NAME,
|
||
'gen_instance_id': GEN_INSTANCE_ID,
|
||
'sync_status': sync_status,
|
||
'last_sync_request_id': last_sync_request_id,
|
||
'last_sync_time': datetime.datetime.now().isoformat(),
|
||
'sync_list': [],
|
||
'create_time': datetime.datetime.now().isoformat(),
|
||
'last_heartbeat_time': datetime.datetime.now().isoformat()
|
||
}
|
||
save_resp = instance_monitor_table.put_item(Item=item)
|
||
logger.info(f"save instance item {save_resp}")
|
||
return save_resp
|
||
|
||
|
||
def update_sync_instance_monitor(instance_monitor_record):
|
||
# 更新记录
|
||
update_expression = ("SET sync_status = :new_sync_status, last_sync_request_id = :sync_request_id, "
|
||
"sync_list = :sync_list, last_sync_time = :sync_time, last_heartbeat_time = :heartbeat_time")
|
||
expression_attribute_values = {
|
||
":new_sync_status": instance_monitor_record['sync_status'],
|
||
":sync_request_id": instance_monitor_record['last_sync_request_id'],
|
||
":sync_list": instance_monitor_record['sync_list'],
|
||
":sync_time": datetime.datetime.now().isoformat(),
|
||
":heartbeat_time": datetime.datetime.now().isoformat(),
|
||
}
|
||
|
||
response = instance_monitor_table.update_item(
|
||
Key={'endpoint_name': ENDPOINT_NAME,
|
||
'gen_instance_id': GEN_INSTANCE_ID},
|
||
UpdateExpression=update_expression,
|
||
ExpressionAttributeValues=expression_attribute_values
|
||
)
|
||
logger.info(f"update_sync_instance_monitor :{response}")
|
||
return response
|
||
|
||
|
||
def sync_instance_monitor_status(need_save: bool):
|
||
try:
|
||
logger.info(f"sync_instance_monitor_status {datetime.datetime.now()}")
|
||
if need_save:
|
||
save_sync_instance_monitor('', 'init')
|
||
else:
|
||
update_expression = ("SET last_heartbeat_time = :heartbeat_time")
|
||
expression_attribute_values = {
|
||
":heartbeat_time": datetime.datetime.now().isoformat(),
|
||
}
|
||
instance_monitor_table.update_item(
|
||
Key={'endpoint_name': ENDPOINT_NAME,
|
||
'gen_instance_id': GEN_INSTANCE_ID},
|
||
UpdateExpression=update_expression,
|
||
ExpressionAttributeValues=expression_attribute_values
|
||
)
|
||
except Exception as e:
|
||
logger.info(f"sync_instance_monitor_status error :{e}")
|
||
|
||
|
||
@server.PromptServer.instance.routes.post("/reboot")
|
||
async def restart(self):
|
||
logger.debug(f"start to reboot!!!!!!!! {self}")
|
||
global executing
|
||
if executing is True:
|
||
logger.info(f"other inference doing cannot reboot!!!!!!!!")
|
||
return ok({"message": "other inference doing cannot reboot"})
|
||
need_reboot = os.environ.get('NEED_REBOOT')
|
||
if need_reboot and need_reboot.lower() != 'true':
|
||
logger.info("no need to reboot by os")
|
||
return ok({"message": "no need to reboot by os"})
|
||
global reboot
|
||
if reboot is False:
|
||
logger.info("no need to reboot by global constant")
|
||
return ok({"message": "no need to reboot by constant"})
|
||
|
||
logger.debug("rebooting !!!!!!!!")
|
||
try:
|
||
sys.stdout.close_log()
|
||
except Exception as e:
|
||
logger.info(f"error reboot!!!!!!!! {e}")
|
||
pass
|
||
return os.execv(sys.executable, [sys.executable] + sys.argv)
|
||
|
||
|
||
# must be sync invoke and use the env to check
|
||
@server.PromptServer.instance.routes.post("/sync_instance")
|
||
async def sync_instance(request):
|
||
if not BUCKET:
|
||
logger.error("No bucket provided ,wait and try again")
|
||
resp = {"status": "success", "message": "syncing"}
|
||
return ok(resp)
|
||
|
||
if 'ALREADY_SYNC' in os.environ and os.environ.get('ALREADY_SYNC').lower() == 'false':
|
||
resp = {"status": "success", "message": "syncing"}
|
||
logger.error("other process doing ,wait and try again")
|
||
return ok(resp)
|
||
|
||
os.environ['ALREADY_SYNC'] = 'false'
|
||
logger.info(f"sync_instance start ！！ {datetime.datetime.now().isoformat()} {request}")
|
||
try:
|
||
last_sync_record = get_last_ddb_sync_record()
|
||
if not last_sync_record:
|
||
logger.info("no last sync record found do not need sync")
|
||
sync_instance_monitor_status(True)
|
||
resp = {"status": "success", "message": "no sync"}
|
||
os.environ['ALREADY_SYNC'] = 'true'
|
||
return ok(resp)
|
||
|
||
if ('request_id' in last_sync_record and last_sync_record['request_id']
|
||
and os.environ.get('LAST_SYNC_REQUEST_ID')
|
||
and os.environ.get('LAST_SYNC_REQUEST_ID') == last_sync_record['request_id']
|
||
and os.environ.get('LAST_SYNC_REQUEST_TIME')
|
||
and os.environ.get('LAST_SYNC_REQUEST_TIME') == str(last_sync_record['request_time'])):
|
||
logger.info("last sync record already sync by os check")
|
||
sync_instance_monitor_status(False)
|
||
resp = {"status": "success", "message": "no sync env"}
|
||
os.environ['ALREADY_SYNC'] = 'true'
|
||
return ok(resp)
|
||
|
||
instance_monitor_record = get_latest_ddb_instance_monitor_record()
|
||
if not instance_monitor_record:
|
||
sync_already = await prepare_comfy_env(last_sync_record)
|
||
if sync_already:
|
||
logger.info("should init prepare instance_monitor_record")
|
||
sync_status = 'success' if sync_already else 'failed'
|
||
save_sync_instance_monitor(last_sync_record['request_id'], sync_status)
|
||
else:
|
||
sync_instance_monitor_status(False)
|
||
else:
|
||
if ('last_sync_request_id' in instance_monitor_record
|
||
and instance_monitor_record['last_sync_request_id']
|
||
and instance_monitor_record['last_sync_request_id'] == last_sync_record['request_id']
|
||
and instance_monitor_record['sync_status']
|
||
and instance_monitor_record['sync_status'] == 'success'
|
||
and os.environ.get('LAST_SYNC_REQUEST_TIME')
|
||
and os.environ.get('LAST_SYNC_REQUEST_TIME') == str(last_sync_record['request_time'])):
|
||
logger.info("last sync record already sync")
|
||
sync_instance_monitor_status(False)
|
||
resp = {"status": "success", "message": "no sync ddb"}
|
||
os.environ['ALREADY_SYNC'] = 'true'
|
||
return ok(resp)
|
||
|
||
sync_already = await prepare_comfy_env(last_sync_record)
|
||
instance_monitor_record['sync_status'] = 'success' if sync_already else 'failed'
|
||
instance_monitor_record['last_sync_request_id'] = last_sync_record['request_id']
|
||
sync_list = instance_monitor_record['sync_list'] if ('sync_list' in instance_monitor_record
|
||
and instance_monitor_record['sync_list']) else []
|
||
sync_list.append(last_sync_record['request_id'])
|
||
|
||
instance_monitor_record['sync_list'] = sync_list
|
||
logger.info("should update prepare instance_monitor_record")
|
||
update_sync_instance_monitor(instance_monitor_record)
|
||
os.environ['ALREADY_SYNC'] = 'true'
|
||
resp = {"status": "success", "message": "sync"}
|
||
return ok(resp)
|
||
except Exception as e:
|
||
logger.info("exception occurred", e)
|
||
os.environ['ALREADY_SYNC'] = 'true'
|
||
resp = {"status": "fail", "message": "sync"}
|
||
return error(resp)
|
||
|
||
|
||
def validate_prompt_proxy(func):
|
||
def wrapper(*args, **kwargs):
|
||
# 在这里添加您的代理逻辑
|
||
logger.info("validate_prompt_proxy start...")
|
||
# 调用原始函数
|
||
result = func(*args, **kwargs)
|
||
# 在这里添加执行后的操作
|
||
logger.info("validate_prompt_proxy end...")
|
||
return result
|
||
|
||
return wrapper
|
||
|
||
|
||
execution.validate_prompt = validate_prompt_proxy(execution.validate_prompt)
|
||
|
||
|
||
def send_sync_proxy(func):
|
||
def wrapper(*args, **kwargs):
|
||
logger.debug(f"Sending sync request!!!!!!! {args}")
|
||
global need_sync
|
||
global prompt_id
|
||
logger.info(f"send_sync_proxy start... {need_sync},{prompt_id} {args}")
|
||
func(*args, **kwargs)
|
||
if need_sync and QUEUE_URL and REGION:
|
||
logger.debug(f"send_sync_proxy params... {QUEUE_URL},{REGION},{need_sync},{prompt_id}")
|
||
event = args[1]
|
||
data = args[2]
|
||
sid = args[3] if len(args) == 4 else None
|
||
message_body = {'prompt_id': prompt_id, 'event': event, 'data': data, 'sid': sid}
|
||
message_id = sen_sqs_msg(message_body, prompt_id)
|
||
logger.info(f'send_sync_proxy message_id :{message_id} message_body: {message_body}')
|
||
logger.debug(f"send_sync_proxy end...")
|
||
|
||
return wrapper
|
||
|
||
|
||
server.PromptServer.send_sync = send_sync_proxy(server.PromptServer.send_sync)
|
||
|
||
|
||
def get_save_imge_path_proxy(func):
|
||
def wrapper(*args, **kwargs):
|
||
logger.info(f"get_save_imge_path_proxy args : {args} kwargs : {kwargs}")
|
||
full_output_folder, filename, counter, subfolder, filename_prefix = func(*args, **kwargs)
|
||
global prompt_id
|
||
filename_prefix_new = filename_prefix + "_" + str(prompt_id)
|
||
logger.info(f"get_save_imge_path_proxy filename_prefix new : {filename_prefix_new}")
|
||
return full_output_folder, filename, counter, subfolder, filename_prefix_new
|
||
|
||
return wrapper
|
||
|
||
|
||
folder_paths.get_save_image_path = get_save_imge_path_proxy(folder_paths.get_save_image_path)
|