518 lines
15 KiB
Python
518 lines
15 KiB
Python
import json
|
|
import os, stat, errno
|
|
from os import path
|
|
from aiohttp import request, web
|
|
from typing import TypedDict, List
|
|
import shutil
|
|
import time
|
|
import subprocess
|
|
import re
|
|
|
|
import server
|
|
import folder_paths
|
|
|
|
class FileInfoDict(TypedDict):
|
|
type: str
|
|
name: str
|
|
bytes: int
|
|
created_at: float
|
|
folder_path: str
|
|
notes: str
|
|
|
|
routes = server.PromptServer.instance.routes
|
|
|
|
browser_path = path.dirname(__file__)
|
|
collections_path = path.join(browser_path, 'collections')
|
|
config_path = path.join(browser_path, 'config.json')
|
|
sources_path = path.join(browser_path, 'sources')
|
|
|
|
git_remote_name = 'origin'
|
|
|
|
image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.webp']
|
|
video_extensions = ['.mp4', '.mov', '.avi', '.webm', '.mkv']
|
|
|
|
info_file_suffix = '.info'
|
|
|
|
# folder_type = 'outputs', 'collections', 'sources'
|
|
def get_parent_path(folder_type: str):
|
|
if folder_type == 'collections':
|
|
return collections_path
|
|
if folder_type == 'sources':
|
|
return sources_path
|
|
|
|
# outputs
|
|
return folder_paths.output_directory
|
|
|
|
# folder_type = 'outputs', 'collections', 'sources'
|
|
def get_target_folder_files(folder_path: str, folder_type: str = 'outputs'):
|
|
if '..' in folder_path:
|
|
return None
|
|
|
|
parent_path = get_parent_path(folder_type)
|
|
files: List[FileInfoDict] = []
|
|
target_path = path.join(parent_path, folder_path)
|
|
|
|
if not path.exists(target_path):
|
|
return None
|
|
|
|
folder_listing = os.scandir(target_path)
|
|
folder_listing = sorted(folder_listing, key=lambda f: (f.is_file(), -f.stat().st_ctime))
|
|
for item in folder_listing:
|
|
if not path.exists(item.path):
|
|
continue
|
|
|
|
name = path.basename(item.path)
|
|
ext = path.splitext(name)[1].lower()
|
|
if name == '' or name[0] == '.':
|
|
continue
|
|
if item.is_file():
|
|
if not (ext in (image_extensions + video_extensions + ['.json'])):
|
|
continue
|
|
|
|
created_at = item.stat().st_ctime
|
|
info_file_path = get_info_filename(item.path)
|
|
info_data = {}
|
|
if path.exists(info_file_path):
|
|
with open(info_file_path, 'r') as f:
|
|
info_data = json.load(f)
|
|
if item.is_file():
|
|
bytes = item.stat().st_size
|
|
files.append({
|
|
"type": "file",
|
|
"name": name,
|
|
"bytes": bytes,
|
|
"created_at": created_at,
|
|
"folder_path": folder_path,
|
|
"notes": info_data.get("notes", "")
|
|
})
|
|
elif item.is_dir():
|
|
files.append({
|
|
"type": "dir",
|
|
"name": name,
|
|
"bytes": 0,
|
|
"created_at": created_at,
|
|
"folder_path": folder_path,
|
|
"notes": info_data.get("notes", "")
|
|
})
|
|
|
|
return files
|
|
|
|
|
|
def get_info_filename(filename):
|
|
return path.splitext(filename)[0] + info_file_suffix
|
|
|
|
def run_cmd(cmd, run_path, log_cmd=True, log_code=True, log_message=True):
|
|
if log_cmd:
|
|
log(f'running: {cmd}')
|
|
|
|
ret = subprocess.run(
|
|
f'cd {run_path} && {cmd}',
|
|
shell=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
encoding="UTF-8"
|
|
)
|
|
if log_code:
|
|
if ret.returncode == 0:
|
|
log('successed')
|
|
else:
|
|
log('failed')
|
|
if log_message:
|
|
if (len(ret.stdout) > 0 or len(ret.stderr) > 0):
|
|
log(ret.stdout + ret.stderr)
|
|
|
|
return ret
|
|
|
|
|
|
def get_config():
|
|
if not path.exists(config_path):
|
|
return {}
|
|
|
|
with open(config_path, 'r') as f:
|
|
return json.load(f)
|
|
|
|
def set_config(config):
|
|
with open(config_path, 'w') as f:
|
|
json.dump(config, f)
|
|
|
|
def git_set_remote_url(remote_url, run_path = collections_path):
|
|
ret = run_cmd('git remote', run_path)
|
|
|
|
if git_remote_name in ret.stdout.split('\n'):
|
|
return run_cmd(f'git remote set-url {git_remote_name} {remote_url}', run_path)
|
|
else:
|
|
return run_cmd(f'git remote add {git_remote_name} {remote_url}', run_path)
|
|
|
|
def git_init(run_path = collections_path):
|
|
if not path.exists(path.join(run_path, '.git')):
|
|
run_cmd('git init', collections_path)
|
|
|
|
def log(message):
|
|
print('[comfyui-browser] ' + message)
|
|
|
|
def add_uuid_to_filename(filename):
|
|
name, ext = path.splitext(filename)
|
|
return f'{name}_{int(time.time())}{ext}'
|
|
|
|
def handle_remove_readonly(func, path, exc):
|
|
excvalue = exc[1]
|
|
if excvalue.errno == errno.EACCES:
|
|
os.chmod(path, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO) # 0777
|
|
func(path)
|
|
else:
|
|
raise
|
|
|
|
# folder_path, folder_type
|
|
@routes.get("/browser/files")
|
|
async def api_get_files(request):
|
|
folder_path = request.query.get('folder_path', '')
|
|
folder_type = request.query.get('folder_type', 'outputs')
|
|
files = get_target_folder_files(folder_path, folder_type=folder_type)
|
|
|
|
if files == None:
|
|
return web.Response(status=404)
|
|
|
|
return web.json_response({
|
|
'files': files
|
|
})
|
|
|
|
|
|
# filename, folder_path, folder_type
|
|
@routes.delete("/browser/files")
|
|
async def api_delete_file(request):
|
|
json_data = await request.json()
|
|
filename = json_data['filename']
|
|
folder_path = json_data.get('folder_path', '')
|
|
folder_type = json_data.get('folder_type', 'outputs')
|
|
|
|
parent_path = get_parent_path(folder_type)
|
|
target_path = path.join(parent_path, folder_path, filename)
|
|
if not path.exists(target_path):
|
|
return web.json_response(status=404)
|
|
|
|
if path.isdir(target_path):
|
|
shutil.rmtree(target_path)
|
|
else:
|
|
os.remove(target_path)
|
|
info_file_path = get_info_filename(target_path)
|
|
if path.exists(info_file_path):
|
|
os.remove(info_file_path)
|
|
|
|
return web.Response(status=201)
|
|
|
|
|
|
# filename, folder_path, folder_type, new_data: {}
|
|
@routes.put("/browser/files")
|
|
async def api_update_file(request):
|
|
json_data = await request.json()
|
|
filename = json_data['filename']
|
|
folder_path = json_data.get('folder_path', '')
|
|
folder_type = json_data.get('folder_type', 'outputs')
|
|
parent_path = get_parent_path(folder_type)
|
|
|
|
new_data = json_data.get('new_data', None)
|
|
if not new_data:
|
|
return web.Response(status=400)
|
|
|
|
new_filename = new_data['filename']
|
|
notes = new_data['notes']
|
|
|
|
old_file_path = path.join(parent_path, folder_path, filename)
|
|
new_file_path = path.join(parent_path, folder_path, new_filename)
|
|
|
|
if not path.exists(old_file_path):
|
|
return web.Response(status=404)
|
|
|
|
if new_filename and filename != new_filename:
|
|
shutil.move(
|
|
old_file_path,
|
|
new_file_path
|
|
)
|
|
old_info_file_path = get_info_filename(old_file_path)
|
|
if path.exists(old_info_file_path):
|
|
new_info_file_path = get_info_filename(new_file_path)
|
|
shutil.move(
|
|
old_info_file_path,
|
|
new_info_file_path
|
|
)
|
|
|
|
if notes:
|
|
extra = {
|
|
"notes": notes
|
|
}
|
|
info_file_path = get_info_filename(new_file_path)
|
|
with open(info_file_path, "w") as outfile:
|
|
json.dump(extra, outfile)
|
|
|
|
return web.Response(status=201)
|
|
|
|
# filename, folder_path, folder_type
|
|
@routes.get("/browser/files/view")
|
|
async def api_view_file(request):
|
|
folder_type = request.query.get("folder_type", "outputs")
|
|
folder_path = request.query.get("folder_path", "")
|
|
filename = request.query.get("filename", None)
|
|
if not filename:
|
|
return web.Response(status=404)
|
|
|
|
parent_path = get_parent_path(folder_type)
|
|
file_path = path.join(parent_path, folder_path, filename)
|
|
|
|
if not path.exists(file_path):
|
|
return web.Response(status=404)
|
|
|
|
with open(file_path, 'rb') as f:
|
|
media_file = f.read()
|
|
|
|
content_type = 'application/json'
|
|
file_extension = path.splitext(filename)[1].lower()
|
|
if file_extension in image_extensions:
|
|
content_type = f'image/{file_extension[1:]}'
|
|
if file_extension in video_extensions:
|
|
content_type = f'video/{file_extension[1:]}'
|
|
|
|
return web.Response(
|
|
body=media_file,
|
|
content_type=content_type,
|
|
headers={"Content-Disposition": f"filename=\"{filename}\""}
|
|
)
|
|
|
|
# filename, folder_path, folder_type = 'outputs' | 'sources'
|
|
@routes.post("/browser/collections")
|
|
async def api_add_to_collections(request):
|
|
json_data = await request.json()
|
|
filename = json_data.get('filename')
|
|
if not filename:
|
|
return web.Response(status=404)
|
|
|
|
folder_path = json_data.get('folder_path', '')
|
|
|
|
folder_type = json_data.get("folder_type", "outputs")
|
|
parent_path = get_parent_path(folder_type)
|
|
|
|
if not path.exists(collections_path):
|
|
os.mkdir(collections_path)
|
|
|
|
source_file_path = path.join(parent_path, folder_path, filename)
|
|
if not path.exists(source_file_path):
|
|
return web.Response(status=404)
|
|
|
|
new_filepath = path.join(
|
|
collections_path,
|
|
add_uuid_to_filename(filename)
|
|
)
|
|
|
|
if path.isdir(source_file_path):
|
|
shutil.copytree(source_file_path, new_filepath)
|
|
else:
|
|
shutil.copy(source_file_path, new_filepath)
|
|
|
|
return web.Response(status=201)
|
|
|
|
# filename, content
|
|
@routes.post("/browser/collections/workflows")
|
|
async def api_create_new_workflow(request):
|
|
json_data = await request.json()
|
|
filename = json_data.get('filename')
|
|
content = json_data.get('content')
|
|
|
|
if not (filename and content):
|
|
return web.Response(status=404)
|
|
|
|
new_filepath = path.join(
|
|
collections_path,
|
|
add_uuid_to_filename(filename)
|
|
)
|
|
with open(new_filepath, 'w') as f:
|
|
f.write(content)
|
|
|
|
return web.Response(status=201)
|
|
|
|
@routes.get("/browser/config")
|
|
async def api_get_browser_config(_):
|
|
config = get_config()
|
|
|
|
return web.json_response(config)
|
|
|
|
# git_repo
|
|
@routes.put("/browser/config")
|
|
async def api_update_browser_config(request):
|
|
json_data = await request.json()
|
|
config = get_config()
|
|
git_repo = json_data.get('git_repo', config.get('git_repo'))
|
|
|
|
git_init()
|
|
|
|
if git_repo == '':
|
|
ret = run_cmd(f'git remote remove {git_remote_name}', collections_path)
|
|
if not ret.returncode == 0:
|
|
return web.json_response(
|
|
{ 'message': ret.stderr },
|
|
status=500,
|
|
)
|
|
|
|
set_config({ 'git_repo': git_repo })
|
|
return web.Response(status=200)
|
|
|
|
ret = git_set_remote_url(git_repo)
|
|
if not ret.returncode == 0:
|
|
return web.json_response(
|
|
{ 'message': ret.stderr },
|
|
status=500,
|
|
)
|
|
|
|
set_config({ 'git_repo': git_repo })
|
|
return web.Response(status=200)
|
|
|
|
@routes.post("/browser/collections/sync")
|
|
async def api_sync_my_collections(_):
|
|
if not path.exists(config_path):
|
|
return web.Response(status=404)
|
|
|
|
config = get_config()
|
|
git_repo = config.get('git_repo')
|
|
if not git_repo:
|
|
return web.Response(status=404)
|
|
|
|
git_init()
|
|
|
|
cmd = 'git status -s'
|
|
ret = run_cmd(cmd, collections_path)
|
|
if len(ret.stdout) > 0:
|
|
cmd = f'git add . && git commit -m "sync by comfyui-browser at {int(time.time())}"'
|
|
ret = run_cmd(cmd, collections_path)
|
|
if not ret.returncode == 0:
|
|
return web.json_response(
|
|
{ 'message': "\n".join([ret.stdout, ret.stderr]) },
|
|
status=500,
|
|
)
|
|
|
|
cmd = f'git fetch {git_remote_name} -v'
|
|
ret = run_cmd(cmd, collections_path)
|
|
if not ret.returncode == 0:
|
|
return web.json_response(
|
|
{ 'message': "\n".join([ret.stdout, ret.stderr]) },
|
|
status=500,
|
|
)
|
|
|
|
cmd = 'git branch --show-current'
|
|
ret = run_cmd(cmd, collections_path)
|
|
branch = ret.stdout.replace('\n', '')
|
|
|
|
cmd = f'git merge {git_remote_name}/{branch}'
|
|
ret = run_cmd(cmd, collections_path, log_code=False)
|
|
|
|
cmd = f'git push {git_remote_name} {branch}'
|
|
ret = run_cmd(cmd, collections_path)
|
|
if not ret.returncode == 0:
|
|
return web.json_response(
|
|
{ 'message': "\n".join([ret.stdout, ret.stderr]) },
|
|
status=500,
|
|
)
|
|
|
|
return web.Response(status=200)
|
|
|
|
@routes.get("/browser/sources")
|
|
async def api_get_sources(_):
|
|
if not path.exists(sources_path):
|
|
return web.json_response({
|
|
'sources': []
|
|
})
|
|
|
|
sources = []
|
|
source_list = os.scandir(sources_path)
|
|
source_list = sorted(source_list, key=lambda f: (-f.stat().st_ctime))
|
|
for item in source_list:
|
|
if not path.exists(item.path):
|
|
continue
|
|
|
|
if item.is_file():
|
|
continue
|
|
|
|
cmd = f'git remote get-url {git_remote_name}'
|
|
ret = run_cmd(cmd, path.join(item.path), log_cmd=False, log_code=False, log_message=False)
|
|
if not (ret.returncode == 0 and len(ret.stdout) > 0):
|
|
continue
|
|
|
|
url = ret.stdout.split('\n')[0]
|
|
name = path.basename(item.path)
|
|
created_at = item.stat().st_ctime
|
|
sources.append({
|
|
"name": name,
|
|
"created_at": created_at,
|
|
"url": url
|
|
})
|
|
|
|
return web.json_response({
|
|
'sources': sources
|
|
})
|
|
|
|
# repo_url
|
|
@routes.post("/browser/sources")
|
|
async def api_create_source(request):
|
|
json_data = await request.json()
|
|
repo_url = json_data['repo_url']
|
|
|
|
if not repo_url:
|
|
return web.Response(status=400)
|
|
|
|
if not path.exists(sources_path):
|
|
os.mkdir(sources_path)
|
|
|
|
pattern = r'[\:\/]([a-zA-Z0-9-_]+)\/([a-zA-Z0-9-_]+)(\.git)?'
|
|
ret = re.search(pattern, repo_url)
|
|
author = ret.group(1)
|
|
name = ret.group(2)
|
|
if not (author and name):
|
|
return web.Response(status=400, text='wrong url')
|
|
|
|
cmd = f'git clone --depth 1 {repo_url} {author}-{name}'
|
|
ret = run_cmd(cmd, sources_path)
|
|
if ret.returncode != 0:
|
|
return web.Response(status=400, text=ret.stdout + ret.stderr)
|
|
|
|
return web.Response(status=201)
|
|
|
|
|
|
# name
|
|
@routes.delete("/browser/sources/{name}")
|
|
async def api_delete_source(request):
|
|
name = request.match_info.get('name', None)
|
|
|
|
if not name:
|
|
return web.Response(status=401)
|
|
|
|
target_path = path.join(sources_path, name)
|
|
if not path.exists(target_path):
|
|
return web.Response(status=404)
|
|
|
|
shutil.rmtree(target_path, onerror=handle_remove_readonly)
|
|
return web.Response(status=200)
|
|
|
|
# name
|
|
@routes.post("/browser/sources/sync/{name}")
|
|
async def api_sync_source(request):
|
|
name = request.match_info.get('name', None)
|
|
|
|
if not name:
|
|
return web.Response(status=401)
|
|
if not path.exists(path.join(sources_path, name)):
|
|
return web.Response(status=404)
|
|
|
|
cmd = f'git pull'
|
|
ret = run_cmd(cmd, path.join(sources_path, name))
|
|
|
|
if ret.returncode == 0:
|
|
return web.Response(status=200)
|
|
else:
|
|
return web.Response(status=400, text=ret.stdout + ret.stderr)
|
|
|
|
|
|
routes.static(
|
|
'/browser/web',
|
|
path.join(browser_path, 'web/build')
|
|
)
|
|
|
|
WEB_DIRECTORY = "web"
|
|
NODE_CLASS_MAPPINGS = {}
|