317 lines
9.6 KiB
Python
317 lines
9.6 KiB
Python
|
|
import json
|
|
import os
|
|
from server import PromptServer
|
|
from aiohttp import web,ClientSession
|
|
import asyncio
|
|
import zipfile
|
|
import tempfile
|
|
|
|
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
root_dir = os.path.join(CURRENT_DIR, '..')
|
|
|
|
cache_dir = os.path.join(root_dir, '.cache')
|
|
docs_dir = os.path.join(root_dir, 'docs')
|
|
setting_path = os.path.join(cache_dir, '.setting.json')
|
|
device_id_path = os.path.join(cache_dir, '.cache_d_id')
|
|
|
|
# 根据设备名称和设备信息 ip地址，创建一个唯一的设备ID
|
|
def create_device_id():
|
|
import hashlib
|
|
import socket
|
|
import uuid
|
|
|
|
# 获取设备名称
|
|
device_name = socket.gethostname()
|
|
# 获取设备IP地址
|
|
device_ip = socket.gethostbyname(device_name)
|
|
# 获取设备唯一标识
|
|
device_uuid = str(uuid.getnode())
|
|
# 创建设备ID
|
|
device_id = hashlib.md5((device_name + device_ip + device_uuid ).encode('utf-8')).hexdigest()
|
|
return device_id
|
|
|
|
if not os.path.exists(cache_dir):
|
|
os.makedirs(cache_dir, 755)
|
|
|
|
# 获取节点文档内容
|
|
def get_node_doc_file_content(node_name):
|
|
file_path = os.path.join(docs_dir, node_name + '.md')
|
|
if os.path.exists(file_path):
|
|
with open(file_path, 'r', encoding='utf-8') as file:
|
|
return file.read()
|
|
else:
|
|
return ""
|
|
|
|
# 获取节点缓存文档
|
|
def get_node_cache_file_content(node_name):
|
|
file_path = os.path.join(cache_dir, node_name + '.md')
|
|
if os.path.exists(file_path):
|
|
with open(file_path, 'r', encoding='utf-8') as file:
|
|
return file.read()
|
|
else:
|
|
return ""
|
|
|
|
# 写入缓存文件
|
|
def write_cache_file(node_name, content):
|
|
file_path = os.path.join(cache_dir, node_name + '.md')
|
|
print(file_path)
|
|
with open(file_path, 'w', encoding='utf-8') as file:
|
|
file.write(content)
|
|
|
|
def write_device_id():
|
|
# 创建设备ID
|
|
device_id = create_device_id()
|
|
# 写入缓存文件
|
|
with open(device_id_path, 'w', encoding='utf-8') as file:
|
|
file.write(device_id)
|
|
return device_id
|
|
|
|
write_device_id()
|
|
|
|
def get_device_id():
|
|
if os.path.exists(device_id_path):
|
|
with open(device_id_path, 'r', encoding='utf-8') as file:
|
|
return file.read()
|
|
else:
|
|
divce_id = write_device_id()
|
|
return divce_id
|
|
|
|
# win: <>:"/\|?*
|
|
paths_ignore_rules = [
|
|
['/', '%2F'],
|
|
['\\', '%5C'],
|
|
[':', '%3A'],
|
|
['*', '%2A'],
|
|
['|', '%7C'],
|
|
['<', '%3C'],
|
|
['>', '%3E'],
|
|
['"', '%22'],
|
|
['?', '%3F']]
|
|
|
|
def replace_ignore_rules(path):
|
|
for rule in paths_ignore_rules:
|
|
path = path.replace(rule[0], rule[1])
|
|
return path
|
|
|
|
# Add route to fetch node info
|
|
@PromptServer.instance.routes.get("/customnode/getNodeInfo")
|
|
async def fetch_customnode_node_info(request):
|
|
try:
|
|
node_name = request.rel_url.query["nodeName"]
|
|
node_name = replace_ignore_rules(node_name)
|
|
|
|
print(node_name)
|
|
if not node_name:
|
|
return web.json_response({"content": ""})
|
|
|
|
cache_content = get_node_cache_file_content(node_name)
|
|
|
|
if len(cache_content) > 0 :
|
|
return web.json_response({"content": cache_content})
|
|
|
|
content = get_node_doc_file_content(node_name)
|
|
return web.json_response({"content": content})
|
|
except Exception as e:
|
|
return web.json_response({"content": ""})
|
|
|
|
# Add route to cache node info
|
|
@PromptServer.instance.routes.get("/customnode/cacheNodeInfo")
|
|
async def cache_customnode_node_info(request):
|
|
try:
|
|
node_name = request.rel_url.query["nodeName"]
|
|
node_name = replace_ignore_rules(node_name)
|
|
|
|
if not node_name:
|
|
return web.json_response({"success": False, "content": ""})
|
|
|
|
print('cache start')
|
|
if not os.path.exists(os.path.join(cache_dir, node_name + '.md')):
|
|
print('cache file not exists')
|
|
content = get_node_doc_file_content(node_name)
|
|
write_cache_file(node_name, content)
|
|
print('cache success')
|
|
return web.json_response({"success": True, "content": content})
|
|
return web.json_response({"success": True, "content": ''})
|
|
except Exception as e:
|
|
print(e)
|
|
return web.json_response({"success": False, "content": ''})
|
|
|
|
|
|
# Add route to update node info
|
|
@PromptServer.instance.routes.post("/customnode/updateNodeInfo")
|
|
async def update_customnode_node_info(request):
|
|
try:
|
|
json_data = await request.json()
|
|
node_name = json_data["nodeName"]
|
|
node_name = replace_ignore_rules(node_name)
|
|
# node_name = request.rel_url.query["nodeName"]
|
|
content = json_data["content"]
|
|
|
|
if not node_name:
|
|
return web.json_response({"success": False})
|
|
|
|
node_name = replace_ignore_rules(node_name)
|
|
write_cache_file(node_name, content)
|
|
|
|
contribute = get_setting_item('contribute')
|
|
if contribute == True:
|
|
print('send doc to cloud')
|
|
asyncio.create_task(send_doc_to_cloud(node_name, content))
|
|
|
|
return web.json_response({"success": True})
|
|
except Exception as e:
|
|
return web.json_response({"success": False})
|
|
|
|
# ================================== 以下是导出节点文档的代码 ==================================
|
|
def get_all_files(directory):
|
|
""" 获取目录下所有文件的路径 """
|
|
file_paths = []
|
|
for root, _, files in os.walk(directory):
|
|
for file in files:
|
|
file_paths.append(os.path.join(root, file))
|
|
return file_paths
|
|
|
|
def collect_unique_files(dir1, dir2):
|
|
""" 收集两个目录下的所有文件，按文件名去重 """
|
|
unique_files = {}
|
|
for file_path in get_all_files(dir1) + get_all_files(dir2):
|
|
file_name = os.path.basename(file_path)
|
|
if file_name not in unique_files:
|
|
unique_files[file_name] = file_path
|
|
return unique_files.values()
|
|
|
|
def zip_files(file_paths, output_zip):
|
|
""" 将文件打包成ZIP，并平铺存储 """
|
|
with zipfile.ZipFile(output_zip, 'w') as zipf:
|
|
for file_path in file_paths:
|
|
arcname = os.path.basename(file_path) # 只使用文件名，不保留路径
|
|
zipf.write(file_path, arcname=arcname)
|
|
|
|
# Add route to export node info
|
|
@PromptServer.instance.routes.get("/customnode/exportNodeInfo")
|
|
async def export_customnode_node_info(request):
|
|
# 把所有的缓存文件和docs文件夹下的文件打包成zip文件
|
|
""" 生成ZIP文件并返回文件流响应 """
|
|
|
|
unique_files = collect_unique_files(cache_dir, docs_dir)
|
|
|
|
# 使用临时文件存储ZIP
|
|
temp_zip = tempfile.NamedTemporaryFile(delete=False)
|
|
try:
|
|
with zipfile.ZipFile(temp_zip, 'w') as zipf:
|
|
for file_path in unique_files:
|
|
arcname = os.path.basename(file_path)
|
|
zipf.write(file_path, arcname=arcname)
|
|
|
|
temp_zip.seek(0)
|
|
|
|
# 生成StreamResponse响应
|
|
response = web.StreamResponse()
|
|
response.headers['Content-Type'] = 'application/zip'
|
|
response.headers['Content-Disposition'] = 'attachment; filename="output.zip"'
|
|
|
|
await response.prepare(request)
|
|
|
|
with open(temp_zip.name, 'rb') as f:
|
|
while chunk := f.read(8192):
|
|
await response.write(chunk)
|
|
|
|
await response.write_eof()
|
|
print(response)
|
|
return response
|
|
except Exception as e:
|
|
print(e)
|
|
finally:
|
|
os.remove(temp_zip.name) # 删除临时文件
|
|
|
|
# Add route to import node info
|
|
@PromptServer.instance.routes.post("/customnode/importNodeInfo")
|
|
async def import_customnode_node_info(request):
|
|
# 接收上传的ZIP文件并解压到指定目录
|
|
""" 接收上传的ZIP文件并解压到指定目录 """
|
|
try:
|
|
data = await request.post()
|
|
zip_file = data.get('file')
|
|
|
|
# 保存上传的ZIP文件
|
|
zip_path = os.path.join(CURRENT_DIR, 'upload.zip')
|
|
with open(zip_path, 'wb') as f:
|
|
f.write(zip_file.file.read())
|
|
|
|
# 解压ZIP文件到.cache目录，并且覆盖原有文件
|
|
with zipfile.ZipFile(zip_path, 'r') as zipf:
|
|
zipf.extractall(cache_dir)
|
|
|
|
os.remove(zip_path) # 删除上传的ZIP文件
|
|
return web.json_response({"success": True})
|
|
except Exception as e:
|
|
return web.json_response({"success": False})
|
|
|
|
|
|
# 获取设置
|
|
def get_setting():
|
|
if os.path.exists(setting_path):
|
|
with open(setting_path, 'r', encoding='utf-8') as file:
|
|
return file.read()
|
|
else:
|
|
return "{}"
|
|
|
|
# 获取设置的每一项
|
|
def get_setting_item(key):
|
|
setting = get_setting()
|
|
setting_json = json.loads(setting)
|
|
return setting_json[key]
|
|
|
|
# 保存设置到文件
|
|
def save_setting(setting):
|
|
with open(setting_path, 'w', encoding='utf-8') as file:
|
|
file.write(setting)
|
|
|
|
if not os.path.exists(setting_path):
|
|
save_setting('{"contribute": true}')
|
|
|
|
# 更新设置到本地.setting.json
|
|
@PromptServer.instance.routes.post("/customnode/updateSetting")
|
|
async def update_setting(request):
|
|
try:
|
|
json_data = await request.json()
|
|
setting = json_data
|
|
# 跟缓存文件的设置项每一项对比，如果有不同则更新
|
|
cache_setting = get_setting()
|
|
cache_setting_json = json.loads(cache_setting)
|
|
for key in setting:
|
|
cache_setting_json[key] = setting[key]
|
|
|
|
# 保存设置到文件，缩近2空格
|
|
save_setting(json.dumps(cache_setting_json, indent=2))
|
|
|
|
return web.json_response({"success": True})
|
|
except Exception as e:
|
|
return web.json_response({"success": False})
|
|
|
|
# 发送当前文档到云
|
|
async def send_doc_to_cloud(node_type, content):
|
|
url = 'http://comfy.zukmb.cn/api/saveNodesDocs'
|
|
async with ClientSession() as session:
|
|
device_id = get_device_id()
|
|
data = {
|
|
'device_id': device_id,
|
|
'node_type': node_type,
|
|
'content': content
|
|
}
|
|
print('send doc to cloud')
|
|
async with session.post(url, data=data) as response:
|
|
print(await response.text())
|
|
# try:
|
|
|
|
# # url = 'http://localhost:8080/api/saveNodesDocs'
|
|
# # data = {
|
|
# # 'device_id': device_id,
|
|
# # 'node_type': node_type,
|
|
# # 'content': content
|
|
# # }
|
|
# # requests.post(url, data=data)
|
|
# except Exception as e:
|
|
# print(e) |