428 lines
21 KiB
Python
428 lines
21 KiB
Python
from server import PromptServer
|
||
from aiohttp import web
|
||
import os
|
||
import json
|
||
from .src.wechat.WechatAuthUtils import *
|
||
from .src.wechat.DataBaseUtil import DataBaseUtil
|
||
import xml.etree.ElementTree as ET
|
||
import time
|
||
import logging
|
||
import queue
|
||
import random
|
||
import types
|
||
import execution
|
||
import uuid
|
||
import folder_paths
|
||
def send_sync(self, event, data, sid=None): #继承父类的send_sync方法
|
||
if sid and hasattr(self, "user_command") and sid in getattr(self,'user_command'):
|
||
if event == "executing" and data['node'] is None and data['prompt_id'] == self.user_command[sid]['prompt_id']:
|
||
history=self.prompt_queue.get_history(prompt_id=data['prompt_id'])[data['prompt_id']]
|
||
#print("历史记录======：",history)
|
||
if Config().wechat['isEnterprise'] and ('isWeb' not in self.user_command[sid] or self.user_command[sid]['isWeb']==False):
|
||
imagePaths=[]
|
||
videoPaths=[]
|
||
for node_id in history['outputs']:
|
||
node_output = history['outputs'][node_id]
|
||
if 'images' in node_output:
|
||
for image in node_output['images']:
|
||
basePath=''
|
||
if 'output' == image['type']:
|
||
basePath=folder_paths.get_output_directory()
|
||
elif 'temp' == image['type']:
|
||
basePath=folder_paths.get_temp_directory()
|
||
elif 'input' == image['type']:
|
||
basePath=folder_paths.get_input_directory()
|
||
|
||
if image['subfolder']:
|
||
basePath=os.path.join(basePath,image['subfolder'])
|
||
|
||
imagePaths.append(os.path.join(basePath,image['filename']))
|
||
if 'gifs' in node_output:
|
||
for image in node_output['gifs']:
|
||
basePath=''
|
||
if 'output' in image['type']:
|
||
basePath=folder_paths.get_output_directory()
|
||
elif 'temp' in image['type']:
|
||
basePath=folder_paths.get_temp_directory()
|
||
elif 'input' in image['type']:
|
||
basePath=folder_paths.get_input_directory()
|
||
|
||
if image['subfolder']:
|
||
basePath=os.path.join(basePath,image['subfolder'])
|
||
|
||
if 'format' in image:
|
||
if image['format'].startswith('video'):
|
||
videoPaths.append(os.path.join(basePath,image['filename']))
|
||
continue
|
||
|
||
imagePaths.append(os.path.join(basePath,image['filename']))
|
||
|
||
logging.info("imagePaths: " + str(imagePaths)+" videoPaths: " + str(videoPaths))
|
||
for imagePath in imagePaths:
|
||
mediaId=getMediaId(imagePath,'image')
|
||
sendServiceImageMessge(mediaId,sid)
|
||
for videoPath in videoPaths:
|
||
mediaId=getMediaId(videoPath,'video')
|
||
sendServiceVideoMessge(mediaId,'AI生成','AI生成视频',sid)
|
||
self.user_command[sid]['status']='wcomplete'
|
||
now = time.localtime()
|
||
end_time = time.strftime("%Y-%m-%d %H:%M:%S", now)
|
||
db=DataBaseUtil()
|
||
db.update_data('wcomplete', end_time, json.dumps(history['outputs']),data['prompt_id'])
|
||
db.close_con()
|
||
|
||
self.loop.call_soon_threadsafe(
|
||
self.messages.put_nowait, (event, data, sid))
|
||
def update_dict(dictionary, keys, value):
|
||
if len(keys) == 1:
|
||
dictionary[keys[0]] = value
|
||
else:
|
||
key = keys.pop(0)
|
||
if isinstance(dictionary[key], dict):
|
||
update_dict(dictionary[key], keys, value)
|
||
elif isinstance(dictionary[key], list):
|
||
for item in dictionary[key]:
|
||
update_dict(item, keys[:], value)
|
||
else:
|
||
raise ValueError("Invalid data structure")
|
||
def setPost(self,FromUserName):
|
||
params=self.user_command[FromUserName]
|
||
logging.info(str(params))
|
||
basePath=folder_paths.folder_names_and_paths['custom_nodes'][0][0]
|
||
comand=Config().wechat['commands'][params['command']]
|
||
filePath = os.path.join(basePath,'ComfyUI_Lam','config','workflow',comand['filename'])
|
||
if os.path.exists(filePath)==False:
|
||
logging.warning("文件不存在："+filePath)
|
||
return None
|
||
f = open(filePath,'r', encoding='utf-8')
|
||
json_data = json.load(f)
|
||
f.close()
|
||
for key in comand['params']:
|
||
keys=comand['params'][key]['keys'][:]
|
||
if key in params:
|
||
update_dict(json_data,keys,params[key])
|
||
|
||
json_data={"prompt":json_data,"client_id":params['openId']}
|
||
logging.info(json_data)
|
||
json_data = self.trigger_on_prompt(json_data)
|
||
if "number" in json_data:
|
||
number = float(json_data['number'])
|
||
else:
|
||
number = self.number
|
||
if "front" in json_data:
|
||
if json_data['front']:
|
||
number = -number
|
||
|
||
self.number += 1
|
||
|
||
if "prompt" in json_data:
|
||
prompt = json_data["prompt"]
|
||
valid = execution.validate_prompt(prompt)
|
||
extra_data = {}
|
||
if "extra_data" in json_data:
|
||
extra_data = json_data["extra_data"]
|
||
|
||
if "client_id" in json_data:
|
||
extra_data["client_id"] = json_data["client_id"]
|
||
if valid[0]:
|
||
prompt_id = str(uuid.uuid4())
|
||
outputs_to_execute = valid[2]
|
||
self.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
|
||
now = time.localtime()
|
||
start_time = time.strftime("%Y-%m-%d %H:%M:%S", now)
|
||
self.user_command[FromUserName]['status']='waiting' #prepare:准备 waiting:待执行 wcomplete完成
|
||
self.user_command[FromUserName]['prompt_id']=prompt_id
|
||
db=DataBaseUtil()
|
||
db.insert_data( params['openId'], json.dumps(params), prompt_id,'waiting',start_time, '', '')
|
||
db.close_con()
|
||
return prompt_id
|
||
else:
|
||
logging.warning("invalid prompt: {}".format(valid[1]))
|
||
return None
|
||
|
||
def getTaskRanking(self,FromUserName):
|
||
if 'prompt_id' not in self.user_command[FromUserName]:
|
||
return '您还没有排队，请先发送您的指令'
|
||
prompt_id=self.user_command[FromUserName]['prompt_id']
|
||
current_queue = self.prompt_queue.get_current_queue()
|
||
print(current_queue)
|
||
for i in range(len(current_queue[0])):
|
||
if current_queue[0][i][1] == prompt_id:
|
||
return '您的任务正在执行。'
|
||
for i in range(len(current_queue[1])):
|
||
if current_queue[1][i][1] == prompt_id:
|
||
return '当前排队人数：'+str(len(current_queue[1]))+'人,您当前的位置在'+current_queue[1][i][0]+'号。'
|
||
return '您还没有排队，请先发送您的指令'
|
||
|
||
@PromptServer.instance.routes.get("/wechatauth/getHistorys")
|
||
async def getHistorys(request):
|
||
if "openId" in request.rel_url.query:
|
||
openId=request.rel_url.query['openId']
|
||
db=DataBaseUtil()
|
||
datas=db.get_many_data(openId)
|
||
db.close_con()
|
||
return web.Response(text=json.dumps(datas), content_type='application/json')
|
||
else:
|
||
return web.Response(status=404)
|
||
|
||
@PromptServer.instance.routes.post("/wechatauth/addTask")
|
||
async def addTask(request):
|
||
post = await request.post()
|
||
print(post.keys())
|
||
command = post.get("command")
|
||
if command == None:
|
||
msg = '指令不能为空！'
|
||
data={'msg':msg,'success':False}
|
||
return web.Response(text=json.dumps(data), content_type='application/json')
|
||
|
||
openId = post.get("openId")
|
||
if openId == None:
|
||
msg = '用户编码不能为空！'
|
||
data={'msg':msg,'success':False}
|
||
return web.Response(text=json.dumps(data), content_type='application/json')
|
||
|
||
if hasattr(PromptServer.instance,'user_command') and openId in PromptServer.instance.user_command and PromptServer.instance.user_command[openId]['status']=='waiting':
|
||
msg = '您已经在队列中，请勿重复提交！'
|
||
data={'msg':msg,'success':False}
|
||
return web.Response(text=json.dumps(data), content_type='application/json')
|
||
|
||
adminNo=base64_decode(Config().wechat['adminNo'])
|
||
if adminNo!=openId:
|
||
db=DataBaseUtil()
|
||
data=db.get_user_frequency(openId)
|
||
if data[0]==None:
|
||
db.user_recharge(openId,Config().wechat['freeSize'])
|
||
data=db.get_user_frequency(openId)
|
||
countd=db.get_user_task_count(openId)
|
||
db.close_con()
|
||
if countd[0] >=data[0]:
|
||
msg='非常抱歉，您的免费使用次数已用完，如需继续使用，请扫描右上角二维码，联系管理员。'
|
||
data={'msg':msg,'success':False}
|
||
return web.Response(text=json.dumps(data), content_type='application/json')
|
||
|
||
params=Config().wechat['commands'][command]['params']
|
||
paramName=''
|
||
userData={'openId':openId,'command':command,'status':'prepare','isWeb':True}
|
||
for param in params:
|
||
val=post.get(param)
|
||
if params[param]['isRequired'] and val==None:
|
||
paramName= params[param]['zhName']
|
||
break
|
||
if val:
|
||
userData[param]=val
|
||
|
||
if paramName:
|
||
msg = '参数"'+paramName+'"不能为空！'
|
||
data={'msg':msg,'success':False}
|
||
return web.Response(text=json.dumps(data), content_type='application/json')
|
||
|
||
if 'seed' in userData:
|
||
if userData['seed'].isdigit()==False or int(userData['seed'])==-1:
|
||
userData['seed']=''.join(random.sample('123456789012345678901234567890',14))
|
||
|
||
if hasattr(PromptServer.instance,"user_command")==False:
|
||
setattr(PromptServer.instance,"user_command",{})
|
||
|
||
PromptServer.instance.user_command[openId]=userData
|
||
resp=setPost(PromptServer.instance,openId)
|
||
if resp!=None:
|
||
data={'msg':'任务已加入队列请等待...','prompt_id':resp,'success':True}
|
||
return web.Response(text=json.dumps(data), content_type='application/json')
|
||
|
||
@PromptServer.instance.routes.get("/wechatauth/getCommands")
|
||
async def getCommands(request):
|
||
commands=Config().wechat['commands']
|
||
return web.Response(text=json.dumps(commands), content_type='application/json')
|
||
|
||
@PromptServer.instance.routes.get("/wechatauth/app")
|
||
async def app(request):
|
||
if "openId" in request.rel_url.query:
|
||
openId=request.rel_url.query['openId']
|
||
openId=base64_decode(openId)
|
||
basePath = folder_paths.folder_names_and_paths['custom_nodes'][0][0]
|
||
htmlPtah = os.path.join(basePath, 'ComfyUI_Lam', 'pages','app.html')
|
||
# 打开文件
|
||
with open(htmlPtah, 'r', encoding='utf-8') as file:
|
||
# 读取文件内容
|
||
html_content = file.read()
|
||
|
||
html_content = html_content.replace('{{openId}}', openId)
|
||
return web.Response(text=html_content, content_type='text/html')
|
||
#return web.FileResponse(htmlPtah)
|
||
else:
|
||
return web.Response(status=404)
|
||
|
||
@PromptServer.instance.routes.get("/wechatauth/handleMessage")
|
||
async def handleMessageGet(request):
|
||
if "echostr" in request.rel_url.query:
|
||
echostr=request.rel_url.query['echostr']
|
||
return web.Response(text=echostr)
|
||
else:
|
||
return web.Response(text="121212")
|
||
|
||
@PromptServer.instance.routes.post("/wechatauth/handleMessage")
|
||
async def handleMessagePost(request):
|
||
# 读取请求体数据
|
||
body = await request.read()
|
||
# 将字节转换为字符串
|
||
data_str = body.decode('utf-8')
|
||
data = ET.fromstring(data_str)
|
||
ToUserName = data.find('ToUserName').text
|
||
FromUserName = data.find('FromUserName').text
|
||
MsgType = data.find('MsgType').text
|
||
logging.info(data_str)
|
||
CreateTime = int(time.time())
|
||
if hasattr(PromptServer.instance,"user_command")==False:
|
||
setattr(PromptServer.instance,"user_command",{})
|
||
|
||
#判断是否为准备状态
|
||
isPrepare=FromUserName in PromptServer.instance.user_command and PromptServer.instance.user_command[FromUserName]['status']=='prepare'
|
||
isWaiting=FromUserName in PromptServer.instance.user_command and PromptServer.instance.user_command[FromUserName]['status']=='waiting'
|
||
# 如果发送的是消息请求，
|
||
if MsgType == 'text': #语音屏蔽or MsgType == 'voice'
|
||
try:
|
||
MsgId = data.find("MsgId").text
|
||
if MsgType == 'text':
|
||
Content = data.find('Content').text # 文本消息内容
|
||
elif MsgType == 'voice':
|
||
Content = data.find('Recognition').text # 语音识别结果，UTF8编码
|
||
# 调用回复函数判断接受的信息，然后返回对应的内容
|
||
reply_content,otherName=receive_msg(Content,isPrepare,isWaiting)
|
||
logging.info("reply_content:"+reply_content+',otherName:'+str(otherName))
|
||
if reply_content=='menu' and otherName:
|
||
if Config().wechat['isEnterprise']:
|
||
sendServiceMenuMessage('您可以执行以下操作：','欢迎再次光临',otherName, FromUserName)
|
||
else:
|
||
msgStr='您可以执行以下操作指令：\n '
|
||
for i in range(len(otherName)):
|
||
msgStr+=otherName[i]['content']+'\n'
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, msg)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
|
||
elif reply_content=='command' and otherName:
|
||
adminNo=base64_decode(Config().wechat['adminNo'])
|
||
msg=''
|
||
isAdopt=True
|
||
if adminNo!=FromUserName:
|
||
db=DataBaseUtil()
|
||
data=db.get_user_frequency(FromUserName)
|
||
if data[0]==None:
|
||
db.user_recharge(FromUserName,Config().wechat['freeSize'])
|
||
data=db.get_user_frequency(FromUserName)
|
||
countd=db.get_user_task_count(FromUserName)
|
||
db.close_con()
|
||
if countd[0] >=data[0]:
|
||
isAdopt=False
|
||
msg='非常抱歉，您的免费使用次数已用完,如需继续使用，请联系管理员。微信号：\n'+Config().wechat['adminWeChat']
|
||
|
||
if isAdopt:
|
||
#添加指令的代码
|
||
PromptServer.instance.user_command[FromUserName]={'openId':FromUserName,'status':'prepare','command':otherName,'seed':''.join(random.sample('123456789012345678901234567890',14))}
|
||
msg,comlist=getCommandMsg(Config().wechat['commands'][otherName],Config().wechat['isEnterprise'])
|
||
if comlist:
|
||
sendServiceMenuMessage(msg,Config().wechat['commands'][otherName]['replyText'],comlist, FromUserName)
|
||
return web.Response(status=200)
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, msg)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
elif reply_content=='query' and otherName:
|
||
if otherName=='查询排队情况':
|
||
msg=''
|
||
if isWaiting:
|
||
msg=getTaskRanking(PromptServer.instance,FromUserName)
|
||
else:
|
||
msg='您还没有排队，请先发送您的指令'
|
||
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, msg)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
elif otherName == '我的编号':
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, base64_encode(FromUserName))
|
||
return web.Response(text=out, content_type='application/xml')
|
||
elif reply_content=='ok':
|
||
if otherName and len(otherName)>0:
|
||
PromptServer.instance.user_command[FromUserName]['prompt']=PromptServer.instance.user_command[FromUserName]['prompt']+otherName if 'prompt' in PromptServer.instance.user_command[FromUserName] else otherName
|
||
params=Config().wechat['commands'][PromptServer.instance.user_command[FromUserName]['command']]['params']
|
||
paramName=''
|
||
for param in params:
|
||
if params[param]['isRequired'] and param not in PromptServer.instance.user_command[FromUserName]:
|
||
paramName= params[param]['zhName']
|
||
break
|
||
if paramName!='':
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, '任务缺乏参数“'+paramName+'”请补充后再提交')
|
||
return web.Response(text=out, content_type='application/xml')
|
||
resp=setPost(PromptServer.instance,FromUserName)
|
||
if resp!=None:
|
||
serverAddress=Config().wechat['serverAddress']+'wechatauth/app?openId='+base64_encode(FromUserName)
|
||
imgUrl='http://mmbiz.qpic.cn/sz_mmbiz_jpg/EpicUicwgk0IGpbYZicuGSXUBaIFWbuDZmBbmDO9PleONOz3FJ3eANEmRicD0eR7mF6PCkxSPxVicbDclQldFKDlHRA/0'
|
||
out = reply_news(FromUserName, ToUserName, CreateTime, '任务加入队列提醒','任务已加入队列请等待。或可点击访问详情，使用更多功能！',imgUrl,serverAddress)
|
||
#out = reply_text(FromUserName, ToUserName, CreateTime, '任务已加入队列请等待,详情请访问：\n'+serverAddress)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
elif reply_content=='recharge' and otherName:
|
||
adminNo=base64_decode(Config().wechat['adminNo'])
|
||
msg=''
|
||
if adminNo==FromUserName:
|
||
data=otherName.split('加')
|
||
openId=base64_decode(data[0])
|
||
tount=int(data[1])
|
||
if openId and tount>0:
|
||
db=DataBaseUtil()
|
||
fdata=db.get_user_frequency(openId)
|
||
if fdata[0] and fdata[0]>0:
|
||
db.user_recharge(openId,tount)
|
||
msg='次数加入成功'
|
||
else:
|
||
msg='用户不存在'
|
||
|
||
db.close_con()
|
||
else:
|
||
msg='用户编码错误'
|
||
else:
|
||
msg='无效的指令'
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, msg)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
elif isPrepare:
|
||
msg=''
|
||
if reply_content and otherName!=None:
|
||
commandName=PromptServer.instance.user_command[FromUserName]['command']
|
||
if reply_content in Config().wechat['commands'][commandName]['params']:
|
||
if 'options' in Config().wechat['commands'][commandName]['params'][reply_content]:
|
||
PromptServer.instance.user_command[FromUserName][reply_content]=Config().wechat['commands'][commandName]['params'][reply_content]['options'][otherName]
|
||
else:
|
||
PromptServer.instance.user_command[FromUserName][reply_content]=otherName
|
||
msg='已填选'+Config().wechat['commands'][commandName]['params'][reply_content]['zhName']+':'+otherName
|
||
else:
|
||
PromptServer.instance.user_command[FromUserName]['prompt']=PromptServer.instance.user_command[FromUserName]['prompt']+reply_content if 'prompt' in PromptServer.instance.user_command[FromUserName] else reply_content
|
||
msg='已填入提示词：'+reply_content
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, msg)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
else:
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, reply_content)
|
||
logging.info(out)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
except Exception as e:
|
||
logging.error('消息处理异常:'+str(e))
|
||
elif MsgType == 'image': #图片
|
||
if isPrepare:
|
||
PicUrl=data.find('PicUrl').text
|
||
PromptServer.instance.user_command[FromUserName]['image']=PicUrl
|
||
else:
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, '请先选择指令')
|
||
return web.Response(text=out, content_type='application/xml')
|
||
elif MsgType == 'event': #事件
|
||
try:
|
||
Event = data.find('Event').text
|
||
Event_key = data.find('EventKey').text
|
||
CreateTime = int(time.time())
|
||
# 判断事件，并返回内容
|
||
reply_content = receive_event(Event,Event_key)
|
||
if reply_content:
|
||
out = reply_text(FromUserName, ToUserName, CreateTime, reply_content)
|
||
return web.Response(text=out, content_type='application/xml')
|
||
except Exception as e:
|
||
logging.error('事件处理异常:'+str(e))
|
||
return web.Response(status=200)
|
||
|
||
PromptServer.instance.send_sync=types.MethodType(send_sync,PromptServer.instance)
|
||
|
||
NODE_CLASS_MAPPINGS = {}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {}
|