877 lines
30 KiB
Python
877 lines
30 KiB
Python
import asyncio
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import queue
|
|
import random
|
|
import time
|
|
import traceback
|
|
import urllib
|
|
import aiohttp
|
|
import urllib.request
|
|
import urllib.parse
|
|
import collections
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from threading import Lock, Condition
|
|
import websockets
|
|
import threading
|
|
from .public import (
|
|
get_output,
|
|
write_json_to_file,
|
|
get_address,
|
|
get_port,
|
|
get_port_from_cmdline,
|
|
args,
|
|
find_project_root,
|
|
get_workflow,
|
|
get_client_id,
|
|
)
|
|
|
|
os.environ["http_proxy"] = ""
|
|
os.environ["https_proxy"] = ""
|
|
os.environ["no_proxy"] = "*"
|
|
SERVER_1_URI = f"wss://aidep.cn/ws?clientId={str(get_client_id())}"
|
|
ADDRESS = get_address()
|
|
PORT = get_port_from_cmdline()
|
|
HTTP_ADDRESS = "http://{}:{}/".format(ADDRESS, PORT)
|
|
new_client_w_id = f"{str(get_client_id())}:{get_port()}"
|
|
SERVER_2_URI = "ws://{}:{}/ws?clientId={}".format(ADDRESS, PORT, new_client_w_id)
|
|
RECONNECT_DELAY = 1
|
|
MAX_RECONNECT_DELAY = 3
|
|
task_queue_1 = queue.Queue()
|
|
task_queue_2 = queue.Queue()
|
|
task_queue_3 = {}
|
|
websocket_queue = collections.deque()
|
|
websocket_conn1 = None
|
|
websocket_conn2 = None
|
|
websocket_conn3 = None
|
|
history_data = {"queue_running": [], "queue_pending": []}
|
|
history_prompt_ids = []
|
|
|
|
WEBSOCKETS_VERSION = tuple(map(int, websockets.__version__.split('.')))
|
|
|
|
|
|
def is_websocket_connected(websocket_conn):
|
|
if websocket_conn is None:
|
|
return False
|
|
if WEBSOCKETS_VERSION < (14, 0):
|
|
return websocket_conn.open
|
|
else:
|
|
return websocket_conn.state == 1
|
|
|
|
|
|
class MonitoredThreadPoolExecutor(ThreadPoolExecutor):
|
|
def __init__(self, max_workers=None, thread_name_prefix=""):
|
|
super().__init__(max_workers=max_workers, thread_name_prefix=thread_name_prefix)
|
|
self._lock = Lock()
|
|
self._condition = Condition(self._lock)
|
|
self._active_tasks = 0
|
|
self._max_workers = max_workers
|
|
|
|
def submit(self, fn, *args, **kwargs):
|
|
with self._lock:
|
|
while self._active_tasks >= self._max_workers:
|
|
self._condition.wait()
|
|
self._active_tasks += 1
|
|
future = super().submit(self._wrap_task(fn), *args, **kwargs)
|
|
return future
|
|
|
|
def _wrap_task(self, fn):
|
|
def wrapped_fn(*args, **kwargs):
|
|
try:
|
|
return fn(*args, **kwargs)
|
|
finally:
|
|
with self._lock:
|
|
self._active_tasks -= 1
|
|
self._condition.notify_all()
|
|
|
|
return wrapped_fn
|
|
|
|
def active_tasks(self):
|
|
with self._lock:
|
|
return self._active_tasks
|
|
|
|
|
|
executor = MonitoredThreadPoolExecutor(max_workers=20)
|
|
|
|
|
|
def print_exception_in_chinese(e):
|
|
tb = traceback.extract_tb(e.__traceback__)
|
|
if tb:
|
|
filename, line_number, function_name, text = tb[-1]
|
|
traceback.print_exception(type(e), e, e.__traceback__)
|
|
else:
|
|
pass
|
|
|
|
|
|
async def websocket_connect(uri, conn_identifier):
|
|
global websocket_conn1, websocket_conn2, send_time
|
|
reconnect_delay = RECONNECT_DELAY
|
|
while True:
|
|
try:
|
|
async with websockets.connect(uri) as websocket:
|
|
print(f"{conn_identifier} 连接成功")
|
|
if conn_identifier == 1:
|
|
websocket_conn1 = websocket
|
|
else:
|
|
websocket_conn2 = websocket
|
|
for key, val in task_queue_3.items():
|
|
is_set = key in history_prompt_ids
|
|
if is_set:
|
|
pass
|
|
else:
|
|
task_queue_2.put(
|
|
{
|
|
"type": "send",
|
|
"prompt_id": key,
|
|
}
|
|
)
|
|
reconnect_delay = RECONNECT_DELAY
|
|
tasks = [
|
|
asyncio.create_task(receive_messages(websocket, conn_identifier)),
|
|
asyncio.create_task(send_heartbeat()),
|
|
]
|
|
await asyncio.gather(*tasks)
|
|
except websockets.ConnectionClosedError as e:
|
|
print_exception_in_chinese(e)
|
|
await asyncio.sleep(reconnect_delay)
|
|
except websockets.ConnectionClosedOK as e:
|
|
print_exception_in_chinese(e)
|
|
await asyncio.sleep(reconnect_delay)
|
|
except Exception as e:
|
|
await asyncio.sleep(reconnect_delay)
|
|
reconnect_delay = min(reconnect_delay * 2, MAX_RECONNECT_DELAY)
|
|
|
|
|
|
def get_history_prompt(prompt_id):
|
|
try:
|
|
if is_websocket_connected(websocket_conn2):
|
|
with urllib.request.urlopen(
|
|
HTTP_ADDRESS + "history" + "/" + prompt_id
|
|
) as response:
|
|
return json.loads(response.read())
|
|
else:
|
|
return {}
|
|
except Exception as e:
|
|
print(f"\033[91m 服务正在连接中{get_time()} \033[0m")
|
|
return {}
|
|
|
|
|
|
async def getHistoryPrompt(prompt_id, type_a=""):
|
|
result_data = [{"type": "str", "k": "prompt_id", "v": prompt_id}]
|
|
result = get_history_prompt(prompt_id)
|
|
response_status = None
|
|
try:
|
|
if prompt_id in result:
|
|
result = result[prompt_id]
|
|
status = result.get("status", {})
|
|
if status.get("completed", False):
|
|
result_data.append({"type": "str", "k": "ok", "v": "1"})
|
|
for output in result.get("outputs", {}).values():
|
|
for media in ["images", "gifs", "videos"]:
|
|
if media in output:
|
|
for item in output[media]:
|
|
if "filename" in item:
|
|
if item["subfolder"] is not "":
|
|
item["filename"] = (
|
|
item["subfolder"] + "/" + item["filename"]
|
|
)
|
|
result_data.append(
|
|
{
|
|
"type": "images",
|
|
"k": "file",
|
|
"v": (
|
|
args.output_directory
|
|
if args.output_directory
|
|
else find_project_root() + "output"
|
|
)
|
|
+ "/"
|
|
+ item["filename"],
|
|
}
|
|
)
|
|
else:
|
|
result_data.append(
|
|
{"type": "str", "k": "ok", "v": "0", "text": "completed状态不对"}
|
|
)
|
|
else:
|
|
is_set = prompt_id in history_prompt_ids
|
|
if is_set:
|
|
return
|
|
result_data.append(
|
|
{"type": "str", "k": "ok", "v": "0", "text": "prompt_id没有找到"}
|
|
)
|
|
response_status = 200
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
result_data.append({"type": "str", "k": "ok", "v": "0", "text": "异常的信息"})
|
|
response_status = 500
|
|
submit_url = "https://aidep.cn/task/completed/?i=66&t=0&v=1.0&from=wxapp&tech_client=tt&tech_scene=990001&c=entry&a=wxapp&do=ttapp&r=comfyui.resultv2.formSubmitForComfyUi&m=tech_huise"
|
|
connector = aiohttp.TCPConnector()
|
|
async with aiohttp.ClientSession(connector=connector) as session:
|
|
try:
|
|
form_res_data = await send_form_data(
|
|
session, submit_url, result_data, prompt_id
|
|
)
|
|
except json.JSONDecodeError as e:
|
|
print_exception_in_chinese(e)
|
|
result_data.append(
|
|
{"type": "str", "k": "ok", "v": "0", "text": "json异常的信息"}
|
|
)
|
|
response_status = 400
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
result_data.append(
|
|
{
|
|
"type": "str",
|
|
"k": "ok",
|
|
"v": "0",
|
|
"text": "aiohttpException异常的信息",
|
|
}
|
|
)
|
|
response_status = 500
|
|
finally:
|
|
if "session" in locals():
|
|
await session.close()
|
|
return {
|
|
"status": response_status,
|
|
"message": "操作完成." if response_status == 200 else "发生错误.",
|
|
}
|
|
|
|
|
|
async def send_form_data(session, url, data, prompt_id=None):
|
|
global websocket_conn1
|
|
form_data = aiohttp.FormData()
|
|
try:
|
|
for item in data:
|
|
if item["type"] == "str":
|
|
form_data.add_field(item["k"], item["v"])
|
|
if (
|
|
item["type"] == "images"
|
|
or item["type"] == "gifs"
|
|
or item["type"] == "videos"
|
|
or item["type"] == "files"
|
|
):
|
|
if os.path.exists(item["v"]):
|
|
with open(item["v"], "rb") as f:
|
|
file_content = f.read()
|
|
form_data.add_field(
|
|
item["k"],
|
|
file_content,
|
|
filename=os.path.basename(item["v"]),
|
|
content_type="application/octet-stream",
|
|
)
|
|
pass
|
|
else:
|
|
pass
|
|
if item["type"] == "file":
|
|
if os.path.exists(item["v"]):
|
|
with open(item["v"], "rb") as f:
|
|
file_content = f.read()
|
|
form_data.add_field(
|
|
item["k"],
|
|
file_content,
|
|
filename=os.path.basename(item["v"]),
|
|
content_type="application/octet-stream",
|
|
)
|
|
else:
|
|
pass
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
async with session.post(url, data=form_data) as response:
|
|
if response.status == 200:
|
|
resp_text = await response.text()
|
|
if prompt_id and is_websocket_connected(websocket_conn1):
|
|
websocket_queue.append(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.executed_success",
|
|
"data": {"prompt_id": prompt_id},
|
|
},
|
|
}
|
|
)
|
|
return resp_text
|
|
else:
|
|
return None
|
|
pass
|
|
|
|
|
|
async def server1_receive_messages(websocket, message_type, message_json):
|
|
if message_type == "init":
|
|
await websocket.send(
|
|
json.dumps(
|
|
{
|
|
"type": "crystools.bind",
|
|
"data": {
|
|
"client_id": new_client_w_id,
|
|
},
|
|
}
|
|
)
|
|
)
|
|
if message_type == "prompt":
|
|
prompt_data = message_json["data"]
|
|
jilu_id = prompt_data["jilu_id"]
|
|
uniqueid = message_json["uniqueid"]
|
|
output = get_output(uniqueid + ".json")
|
|
print('recieve prompt message', jilu_id, uniqueid, output)
|
|
workflow = get_workflow(uniqueid + ".json")
|
|
if output:
|
|
executor.submit(run_async_task, output, prompt_data, workflow, jilu_id)
|
|
else:
|
|
if is_websocket_connected(websocket):
|
|
websocket_queue.append(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.prompt_error",
|
|
"data": {"jilu_id": jilu_id, "msg": "作品工作流找不到了"},
|
|
},
|
|
}
|
|
)
|
|
pass
|
|
|
|
|
|
def optimized_process_history_data(history_data_1):
|
|
running = []
|
|
pending = []
|
|
if history_data_1:
|
|
queue_running = history_data_1.get("queue_running", [])
|
|
if queue_running:
|
|
running.append(queue_running[0][1])
|
|
queue_pending = history_data_1.get("queue_pending", [])
|
|
if queue_pending:
|
|
pending = sorted(queue_pending, key=lambda x: int(x[0]))
|
|
pending = [item[1] for item in pending]
|
|
return running, pending
|
|
|
|
|
|
async def server2_receive_messages(websocket, message_type, message_json):
|
|
global send_time
|
|
if message_type and message_type != "crystools.monitor":
|
|
if message_type == "status" and message_json["data"]["status"]["exec_info"]:
|
|
websocket_queue.append(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.queue",
|
|
"data": {
|
|
"client_id": new_client_w_id,
|
|
"queue_remaining": message_json["data"]["status"][
|
|
"exec_info"
|
|
]["queue_remaining"],
|
|
},
|
|
},
|
|
}
|
|
)
|
|
await send_heartbeat_to_server2()
|
|
if message_type == "execution_start":
|
|
pass
|
|
if message_type == "executing":
|
|
pass
|
|
if message_type == "execution_error":
|
|
task_queue_2.put(
|
|
{
|
|
"type": "send",
|
|
"prompt_id": message_json["data"]["prompt_id"],
|
|
}
|
|
)
|
|
if message_type == "executed":
|
|
task_queue_2.put(
|
|
{
|
|
"type": "send",
|
|
"prompt_id": message_json["data"]["prompt_id"],
|
|
}
|
|
)
|
|
if message_type == "progress":
|
|
pass
|
|
if message_type == "execution_cached" and "prompt_id" in message_json["data"]:
|
|
task_queue_2.put(
|
|
{
|
|
"type": "send",
|
|
"prompt_id": message_json["data"]["prompt_id"],
|
|
}
|
|
)
|
|
|
|
|
|
async def receive_messages(websocket, conn_identifier):
|
|
if is_websocket_connected(websocket):
|
|
try:
|
|
async for message in websocket:
|
|
if type(message) != bytes:
|
|
message_dict = json.loads(message)
|
|
message_type = message_dict.get("type")
|
|
if conn_identifier == 1:
|
|
await server1_receive_messages(
|
|
websocket, message_type, message_dict
|
|
)
|
|
elif conn_identifier == 2:
|
|
await server2_receive_messages(
|
|
websocket, message_type, message_dict
|
|
)
|
|
except json.JSONDecodeError as e:
|
|
print_exception_in_chinese(e)
|
|
finally:
|
|
await asyncio.sleep(0.5)
|
|
|
|
|
|
async def send_heartbeat():
|
|
while True:
|
|
try:
|
|
if is_websocket_connected(websocket_conn1) and is_websocket_connected(websocket_conn2):
|
|
await send_heartbeat_to_server2()
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
finally:
|
|
await asyncio.sleep(10)
|
|
|
|
|
|
def get_history():
|
|
global last_get_history_time
|
|
try:
|
|
if is_websocket_connected(websocket_conn2):
|
|
last_get_history_time = time.time()
|
|
with urllib.request.urlopen(HTTP_ADDRESS + "queue") as response:
|
|
return json.loads(response.read())
|
|
else:
|
|
return {
|
|
"queue_running": [],
|
|
"queue_pending": [],
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"queue_running": [],
|
|
"queue_pending": [],
|
|
}
|
|
|
|
|
|
def get_filenames(directory):
|
|
if os.path.exists(directory):
|
|
all_entries = os.listdir(directory)
|
|
all_entries = [
|
|
name
|
|
for name in all_entries
|
|
if os.path.isfile(os.path.join(directory, name))
|
|
]
|
|
all_entries = [name.split(".")[0] for name in all_entries]
|
|
return all_entries
|
|
else:
|
|
return []
|
|
|
|
|
|
send_time = "0"
|
|
|
|
|
|
def get_time():
|
|
return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
|
|
|
|
|
async def send_heartbeat_to_server2():
|
|
running, pending = optimized_process_history_data(history_data)
|
|
try:
|
|
file_names = get_filenames(
|
|
find_project_root() + "custom_nodes/ComfyUI_Cash/config/json/api/"
|
|
)
|
|
websocket_queue.append(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.monitor",
|
|
"data": {
|
|
"files": file_names,
|
|
"running": running,
|
|
"pending": pending,
|
|
"client_id": new_client_w_id,
|
|
},
|
|
},
|
|
}
|
|
)
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
pass
|
|
|
|
|
|
def run_task_in_loop(task, *args, **kwargs):
|
|
while True:
|
|
task(*args, **kwargs)
|
|
time.sleep(1)
|
|
|
|
|
|
loop_num = 0
|
|
|
|
|
|
async def run_websocket_task_in_loop():
|
|
global loop_num
|
|
while True:
|
|
try:
|
|
if len(websocket_queue) > 0:
|
|
websocket_info = websocket_queue.popleft()
|
|
if "conn_identifier" in websocket_info:
|
|
if is_websocket_connected(websocket_conn3) and is_websocket_connected(websocket_conn1):
|
|
websocket_info["data"]["zhu_client_id"] = new_client_w_id
|
|
if websocket_info["conn_identifier"] == 1:
|
|
await websocket_conn3.send(
|
|
json.dumps(websocket_info["data"])
|
|
)
|
|
else:
|
|
loop_num = loop_num + 1
|
|
if loop_num > 100:
|
|
loop_num = 0
|
|
await websocket_conn3.send(
|
|
json.dumps(
|
|
{
|
|
"time": get_time(),
|
|
}
|
|
)
|
|
)
|
|
except Exception as e:
|
|
break
|
|
finally:
|
|
await asyncio.sleep(0.02)
|
|
|
|
|
|
def generate_md5_uid_timestamp_filename(original_filename):
|
|
timestamp = str(time.time())
|
|
random_number = str(generate_large_random_number(32))
|
|
combined_string = original_filename + timestamp + random_number
|
|
md5_hash = hashlib.md5(combined_string.encode("utf-8")).hexdigest()
|
|
file_extension = os.path.splitext(original_filename)[1]
|
|
filename = md5_hash + file_extension
|
|
return filename
|
|
|
|
|
|
async def loca_download_image(url, filename):
|
|
dir_name = find_project_root() + "input"
|
|
no_proxy_handler = urllib.request.ProxyHandler({})
|
|
opener = urllib.request.build_opener(no_proxy_handler)
|
|
file_new_name = generate_md5_uid_timestamp_filename(filename)
|
|
try:
|
|
response = opener.open(url)
|
|
if response.getcode() == 200:
|
|
full_path = os.path.join(dir_name, file_new_name)
|
|
if os.path.exists(full_path):
|
|
return {
|
|
"code": True,
|
|
"filename": file_new_name,
|
|
}
|
|
with open(full_path, "wb") as f:
|
|
f.write(response.read())
|
|
return {
|
|
"code": True,
|
|
"filename": file_new_name,
|
|
}
|
|
else:
|
|
return {
|
|
"code": False,
|
|
"filename": file_new_name,
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"code": False,
|
|
"filename": file_new_name,
|
|
}
|
|
|
|
|
|
def generate_large_random_number(num_bits):
|
|
return random.getrandbits(num_bits)
|
|
|
|
|
|
def queue_prompt(prompt, workflow, new_client_id):
|
|
try:
|
|
if is_websocket_connected(websocket_conn2):
|
|
p = {
|
|
"prompt": prompt,
|
|
"client_id": new_client_id,
|
|
"extra_data": workflow["extra_data"],
|
|
}
|
|
data = json.dumps(p).encode("utf-8")
|
|
req = urllib.request.Request(HTTP_ADDRESS + "prompt", data=data)
|
|
return json.loads(urllib.request.urlopen(req).read())
|
|
else:
|
|
return {}
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
return {}
|
|
|
|
|
|
async def process_json_elements(json_data, prompt_data, workflow, jilu_id):
|
|
global websocket_conn1
|
|
try:
|
|
if "cs_imgs" in prompt_data and prompt_data["cs_imgs"]:
|
|
for item in prompt_data["cs_imgs"]:
|
|
filename = os.path.basename(item["upImage"])
|
|
download_info = await loca_download_image(item["upImage"], filename)
|
|
download_status = download_info["code"]
|
|
file_new_name = download_info["filename"]
|
|
if download_status == False:
|
|
raise Exception("图片下载失败")
|
|
if (
|
|
str(item["node"]) in json_data
|
|
and "inputs" in json_data[str(item["node"])]
|
|
and "image" in json_data[str(item["node"])]["inputs"]
|
|
):
|
|
json_data[str(item["node"])]["inputs"]["image"] = file_new_name
|
|
if "cs_videos" in prompt_data and prompt_data["cs_videos"]:
|
|
for item in prompt_data["cs_videos"]:
|
|
filename = os.path.basename(item["upImage"])
|
|
download_info = await loca_download_image(item["upImage"], filename)
|
|
download_status = download_info["code"]
|
|
file_new_name = download_info["filename"]
|
|
if download_status == False:
|
|
raise Exception("视频下载失败")
|
|
if (
|
|
str(item["node"]) in json_data
|
|
and "inputs" in json_data[str(item["node"])]
|
|
and "video" in json_data[str(item["node"])]["inputs"]
|
|
):
|
|
json_data[str(item["node"])]["inputs"]["video"] = file_new_name
|
|
if "cs_texts" in prompt_data and prompt_data["cs_texts"]:
|
|
for item in prompt_data["cs_texts"]:
|
|
json_data[str(item["node"])]["inputs"]["text"] = item["value"]
|
|
except KeyError as e:
|
|
print_exception_in_chinese(e)
|
|
websocket_queue.appendleft(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.prompt_error",
|
|
"data": {"jilu_id": jilu_id, "msg": "发送指令失败1"},
|
|
},
|
|
}
|
|
)
|
|
return {"code": 0, "jilu_id": jilu_id}
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
websocket_queue.appendleft(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.prompt_error",
|
|
"data": {"jilu_id": jilu_id, "msg": "发送指令失败2"},
|
|
},
|
|
}
|
|
)
|
|
return {"code": 0, "jilu_id": jilu_id}
|
|
|
|
async def print_item(key, value):
|
|
try:
|
|
if value["class_type"] == "KSampler" and "inputs" in json_data[key]:
|
|
json_data[key]["inputs"]["seed"] = generate_large_random_number(15)
|
|
if (
|
|
value["class_type"] == "VHS_VideoCombine"
|
|
and "inputs" in json_data[key]
|
|
and "crf" in json_data[key]["inputs"]
|
|
):
|
|
if json_data[key]["inputs"]["crf"] == 0:
|
|
json_data[key]["inputs"]["crf"] = 1
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
websocket_queue.appendleft(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.prompt_error",
|
|
"data": {"jilu_id": jilu_id, "msg": "发送指令失败3"},
|
|
},
|
|
}
|
|
)
|
|
|
|
tasks = [print_item(key, value) for key, value in json_data.items()]
|
|
await asyncio.gather(*tasks)
|
|
try:
|
|
result = queue_prompt(json_data, workflow, new_client_w_id)
|
|
if "node_errors" in result:
|
|
if result["node_errors"]:
|
|
raise Exception("发送指令失败")
|
|
try:
|
|
websocket_queue.appendleft(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.prompt_ok",
|
|
"data": {
|
|
"prompt_id": result["prompt_id"],
|
|
"jilu_id": jilu_id,
|
|
"msg": "发送指令成功",
|
|
},
|
|
},
|
|
}
|
|
)
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
task_queue_3[result["prompt_id"]] = {"jilu_id": jilu_id}
|
|
return {"code": 1, "prompt_id": result["prompt_id"]}
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
websocket_queue.appendleft(
|
|
{
|
|
"conn_identifier": 1,
|
|
"data": {
|
|
"type": "crystools.prompt_error",
|
|
"data": {"jilu_id": jilu_id, "msg": "发送指令失败3"},
|
|
},
|
|
}
|
|
)
|
|
return {"code": 0, "prompt_id": jilu_id}
|
|
|
|
|
|
def run_async_task(json_data, prompt_data, workflow, jilu_id):
|
|
return asyncio.run(process_json_elements(json_data, prompt_data, workflow, jilu_id))
|
|
|
|
|
|
def run_async_task2(prompt_id):
|
|
asyncio.run(getHistoryPrompt(prompt_id))
|
|
|
|
|
|
def task_3():
|
|
while True:
|
|
try:
|
|
task_info = task_queue_1.get()
|
|
output = get_output(task_info["uniqueid"] + ".json")
|
|
if output:
|
|
executor.submit(
|
|
run_async_task,
|
|
output,
|
|
task_info["prompt_data"],
|
|
task_info["jilu_id"],
|
|
)
|
|
task_queue_1.task_done()
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
finally:
|
|
time.sleep(1)
|
|
|
|
|
|
def task_4():
|
|
global history_data
|
|
|
|
while True:
|
|
try:
|
|
task_info = task_queue_2.get()
|
|
if "prompt_id" in task_info:
|
|
history_data = get_history()
|
|
preprocess_history_data(history_data)
|
|
task_queue_3.pop(task_info["prompt_id"], None)
|
|
executor.submit(run_async_task2, task_info["prompt_id"])
|
|
task_queue_2.task_done()
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
finally:
|
|
time.sleep(0.1)
|
|
|
|
|
|
def print_thread_status():
|
|
while True:
|
|
print("\n当前活动线程:")
|
|
for thread in threading.enumerate():
|
|
print(
|
|
f"线程名: {thread.name}, 线程ID: {thread.ident}, 活动状态: {thread.is_alive()}"
|
|
)
|
|
time.sleep(5)
|
|
|
|
|
|
def main_task():
|
|
for i in range(10):
|
|
time.sleep(1)
|
|
|
|
|
|
def websocket_thread(uri, conn_identifier):
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
loop.run_until_complete(websocket_connect(uri, conn_identifier))
|
|
|
|
|
|
def websocket_thread_fu(uri, conn_identifier):
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
loop.run_until_complete(websocket_connect_fu(uri, conn_identifier))
|
|
|
|
|
|
def preprocess_history_data(history_data):
|
|
global history_prompt_ids
|
|
|
|
prompt_ids = set()
|
|
if history_data is None:
|
|
history_prompt_ids = prompt_ids
|
|
return
|
|
for queue in ["queue_running", "queue_pending"]:
|
|
for item in history_data.get(queue, []):
|
|
prompt_ids.add(item[1])
|
|
history_prompt_ids = prompt_ids
|
|
|
|
|
|
last_get_history_time = 0
|
|
|
|
|
|
async def task5():
|
|
global history_data
|
|
while True:
|
|
try:
|
|
history_data = get_history()
|
|
preprocess_history_data(history_data)
|
|
except Exception as e:
|
|
print_exception_in_chinese(e)
|
|
await asyncio.sleep(1)
|
|
|
|
|
|
def task5_thread():
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
loop.run_until_complete(task5())
|
|
|
|
|
|
def start_async_task_in_thread(async_func):
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
loop.run_until_complete(async_func())
|
|
|
|
|
|
async def websocket_connect_fu(uri, conn_identifier):
|
|
global websocket_conn3
|
|
reconnect_delay = RECONNECT_DELAY
|
|
while True:
|
|
try:
|
|
async with websockets.connect(uri) as websocket:
|
|
print(f"{conn_identifier} 连接成功")
|
|
websocket_conn3 = websocket
|
|
await websocket_conn3.send(
|
|
json.dumps(
|
|
{
|
|
"type": "crystools.bind",
|
|
"data": {
|
|
"client_id": new_client_w_id + "_fu",
|
|
},
|
|
}
|
|
)
|
|
)
|
|
reconnect_delay = RECONNECT_DELAY
|
|
tasks = [
|
|
asyncio.create_task(run_websocket_task_in_loop()),
|
|
]
|
|
await asyncio.gather(*tasks)
|
|
except websockets.ConnectionClosedError as e:
|
|
print(f"\033[91m 3 服务正在连接中{get_time()} \033[0m")
|
|
await asyncio.sleep(reconnect_delay)
|
|
except websockets.ConnectionClosedOK as e:
|
|
await asyncio.sleep(reconnect_delay)
|
|
except Exception as e:
|
|
await asyncio.sleep(reconnect_delay)
|
|
reconnect_delay = min(reconnect_delay * 2, MAX_RECONNECT_DELAY)
|
|
|
|
|
|
def thread_run():
|
|
threading.Thread(
|
|
target=websocket_thread, args=(SERVER_1_URI, 1), daemon=True
|
|
).start()
|
|
threading.Thread(
|
|
target=websocket_thread, args=(SERVER_2_URI, 2), daemon=True
|
|
).start()
|
|
threading.Thread(
|
|
target=websocket_thread_fu, args=(SERVER_1_URI + '-3', 3), daemon=True
|
|
).start()
|
|
threading.Thread(target=task5_thread).start()
|
|
executor.submit(run_task_in_loop, task_4)
|
|
|
|
|
|
async def update_worker_flow(uniqueid, data, flow_type="api/"):
|
|
write_json_to_file(data, uniqueid + ".json", "json/" + flow_type, "json")
|