330 lines
12 KiB
Python
330 lines
12 KiB
Python
import os
|
|
from aiohttp import web
|
|
import server
|
|
import hashlib
|
|
import folder_paths
|
|
import uuid
|
|
import time
|
|
import json
|
|
from .libs.prompt_parser import get_learned_conditioning_prompt_schedules
|
|
from .sdfx_path import get_gallery_path, get_workflows_path, get_templates_path
|
|
|
|
gallery_path = get_gallery_path()
|
|
workflows_path = get_workflows_path()
|
|
templates_path = get_templates_path()
|
|
|
|
def get_media_url_from_params(url_path, filename, type, gallery):
|
|
return url_path + '/sdfx/media?filename=' + filename + '&type=' + type + '&gallery=' + gallery
|
|
|
|
def get_dir_by_model_id_and_type(model_id, dir_type):
|
|
# === Todo handle other media types ===
|
|
if dir_type is None or dir_type not in ["image"]:
|
|
dir_type = "image"
|
|
|
|
if dir_type == "image":
|
|
type_dir = os.path.join(gallery_path, model_id, "image")
|
|
|
|
return type_dir, dir_type
|
|
|
|
def media_upload(post, url_path):
|
|
image = post.get("file")
|
|
overwrite = post.get("overwrite")
|
|
model_id = post.get("modelId")
|
|
|
|
image_upload_type = post.get("type")
|
|
upload_dir, image_upload_type = get_dir_by_model_id_and_type(model_id, image_upload_type)
|
|
|
|
if image and image.file:
|
|
filename = image.filename
|
|
if not filename:
|
|
return web.Response(status=400)
|
|
|
|
filepath = os.path.abspath(os.path.join(upload_dir, filename))
|
|
|
|
if os.path.commonpath((upload_dir, filepath)) != upload_dir:
|
|
return web.Response(status=400)
|
|
|
|
if not os.path.exists(upload_dir):
|
|
os.makedirs(upload_dir)
|
|
|
|
split = os.path.splitext(filename)
|
|
|
|
if overwrite is not None and (overwrite == "true" or overwrite == "1"):
|
|
pass
|
|
else:
|
|
i = 1
|
|
while os.path.exists(filepath):
|
|
filename = f"{split[0]} ({i}){split[1]}"
|
|
filepath = os.path.join(upload_dir, filename)
|
|
i += 1
|
|
|
|
with open(filepath, "wb") as f:
|
|
f.write(image.file.read())
|
|
return web.json_response({"name" : filename, "type": image_upload_type, "path": get_media_url_from_params(url_path, filename, image_upload_type, model_id)})
|
|
else:
|
|
return web.Response(status=400)
|
|
|
|
def get_domain_from_request(request):
|
|
return request.scheme + '://' + request.host
|
|
|
|
def get_medias_from_model_id(media_id, url_path):
|
|
media_type = 'image'
|
|
# === Todo handle other media types ===
|
|
media_folder = os.path.join(gallery_path, media_id, media_type)
|
|
if os.path.exists(media_folder) and os.path.isdir(media_folder):
|
|
return [{"name": image, "type": media_type, "path": get_media_url_from_params(url_path, image, media_type, media_id)} for image in os.listdir(media_folder)]
|
|
|
|
return []
|
|
|
|
def add_sdfx_routes():
|
|
# === Comfy determines the path of a checkpoint by its name ===
|
|
# === If two different paths contain two checkpoints with the same name, there is an arbitration done to choose the appropriate path ===
|
|
@server.PromptServer.instance.routes.get("/sdfx/model/list")
|
|
def get_models(request):
|
|
url_path = get_domain_from_request(request)
|
|
models = []
|
|
cns = folder_paths.get_filename_list('checkpoints')
|
|
print(cns)
|
|
for cn in cns:
|
|
fp = folder_paths.get_full_path("checkpoints", cn)
|
|
media_id = hashlib.md5(fp.encode()).hexdigest()
|
|
models.append({
|
|
"modelId": media_id,
|
|
"name": cn,
|
|
"type": "checkpoint",
|
|
"gallery": get_medias_from_model_id(media_id, url_path)
|
|
})
|
|
|
|
return web.json_response(models)
|
|
|
|
# === Single API to upload all kind of assets ===
|
|
@server.PromptServer.instance.routes.post("/sdfx/media")
|
|
async def upload_media(request):
|
|
post = await request.post()
|
|
url_path = get_domain_from_request(request)
|
|
return media_upload(post, url_path)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/sdfx/media")
|
|
async def view_media(request):
|
|
if "filename" not in request.rel_url.query:
|
|
return web.Response(status=400)
|
|
if "type" not in request.rel_url.query:
|
|
return web.Response(status=400)
|
|
if "gallery" not in request.rel_url.query:
|
|
return web.Response(status=400)
|
|
|
|
filename = request.rel_url.query["filename"]
|
|
type = request.rel_url.query["type"]
|
|
gallery = request.rel_url.query["gallery"]
|
|
|
|
if filename[0] == '/' or '..' in filename:
|
|
return web.Response(status=400)
|
|
if type[0] == '/' or '..' in type:
|
|
return web.Response(status=400)
|
|
if gallery[0] == '/' or '..' in gallery:
|
|
return web.Response(status=400)
|
|
|
|
file = os.path.join(gallery_path, gallery, type, filename)
|
|
|
|
if os.path.isfile(file):
|
|
if 'preview' in request.rel_url.query:
|
|
with Image.open(file) as img:
|
|
preview_info = request.rel_url.query['preview'].split(';')
|
|
image_format = preview_info[0]
|
|
if image_format not in ['webp', 'jpeg'] or 'a' in request.rel_url.query.get('channel', ''):
|
|
image_format = 'webp'
|
|
|
|
quality = 90
|
|
if preview_info[-1].isdigit():
|
|
quality = int(preview_info[-1])
|
|
|
|
buffer = BytesIO()
|
|
if image_format in ['jpeg'] or request.rel_url.query.get('channel', '') == 'rgb':
|
|
img = img.convert("RGB")
|
|
img.save(buffer, format=image_format, quality=quality)
|
|
buffer.seek(0)
|
|
|
|
return web.Response(body=buffer.read(), content_type=f'image/{image_format}',
|
|
headers={"Content-Disposition": f"filename=\"{filename}\""})
|
|
|
|
if 'channel' not in request.rel_url.query:
|
|
channel = 'rgba'
|
|
else:
|
|
channel = request.rel_url.query["channel"]
|
|
|
|
if channel == 'rgb':
|
|
with Image.open(file) as img:
|
|
if img.mode == "RGBA":
|
|
r, g, b, a = img.split()
|
|
new_img = Image.merge('RGB', (r, g, b))
|
|
else:
|
|
new_img = img.convert("RGB")
|
|
|
|
buffer = BytesIO()
|
|
new_img.save(buffer, format='PNG')
|
|
buffer.seek(0)
|
|
|
|
return web.Response(body=buffer.read(), content_type='image/png', headers={"Content-Disposition": f"filename=\"{filename}\""})
|
|
|
|
elif channel == 'a':
|
|
with Image.open(file) as img:
|
|
if img.mode == "RGBA":
|
|
_, _, _, a = img.split()
|
|
else:
|
|
a = Image.new('L', img.size, 255)
|
|
|
|
# alpha img
|
|
alpha_img = Image.new('RGBA', img.size)
|
|
alpha_img.putalpha(a)
|
|
alpha_buffer = BytesIO()
|
|
alpha_img.save(alpha_buffer, format='PNG')
|
|
alpha_buffer.seek(0)
|
|
|
|
return web.Response(body=alpha_buffer.read(), content_type='image/png', headers={"Content-Disposition": f"filename=\"{filename}\""})
|
|
else:
|
|
return web.FileResponse(file, headers={"Content-Disposition": f"filename=\"{filename}\""})
|
|
|
|
return web.Response(status=400)
|
|
|
|
@server.PromptServer.instance.routes.get("/sdfx/prompt_parser")
|
|
async def prompt_parser(request):
|
|
if "prompt" in request.rel_url.query:
|
|
prompt = request.rel_url.query["prompt"]
|
|
env = request.rel_url.query['env'] if 'env' in request.rel_url.query else 'prod'
|
|
g = lambda p: get_learned_conditioning_prompt_schedules([p], env)[0]
|
|
return web.json_response({'prompt' : prompt, 'data' : g(prompt)})
|
|
return web.Response(status=404)
|
|
|
|
@server.PromptServer.instance.routes.get("/sdfx/workflow/list")
|
|
@server.PromptServer.instance.routes.get("/sdfx/template/list")
|
|
async def workflows_templates(request):
|
|
list = []
|
|
path = request.url.path
|
|
|
|
if path == "/sdfx/workflow/list":
|
|
working_dir = workflows_path
|
|
elif path == "/sdfx/template/list":
|
|
working_dir = templates_path
|
|
|
|
for filename in os.listdir(working_dir):
|
|
if filename.endswith(".json"):
|
|
file_path = os.path.join(working_dir, filename)
|
|
with open(file_path, 'r') as file:
|
|
try:
|
|
data = json.load(file)
|
|
uid = data.get('uid')
|
|
name = data.get('name')
|
|
dateCreated = data.get('dateCreated')
|
|
meta = data.get('meta')
|
|
list.append({ 'uid': uid, 'name': name, 'meta': meta, 'dateCreated': dateCreated})
|
|
|
|
except json.JSONDecodeError as e:
|
|
return web.json_response({'error': 'invalid Json'}, status=400)
|
|
return web.json_response(list)
|
|
|
|
@server.PromptServer.instance.routes.get("/sdfx/workflow")
|
|
@server.PromptServer.instance.routes.post("/sdfx/workflow")
|
|
@server.PromptServer.instance.routes.put("/sdfx/workflow")
|
|
@server.PromptServer.instance.routes.delete("/sdfx/workflow")
|
|
@server.PromptServer.instance.routes.get("/sdfx/template")
|
|
@server.PromptServer.instance.routes.post("/sdfx/template")
|
|
@server.PromptServer.instance.routes.put("/sdfx/template")
|
|
@server.PromptServer.instance.routes.delete("/sdfx/template")
|
|
async def workflow_template(request):
|
|
method = request.method
|
|
path = request.url.path
|
|
if method == 'GET':
|
|
uid = request.query.get('uid')
|
|
if not uid:
|
|
return web.json_response({'error': 'uid missing'}, status=400)
|
|
if path == "/sdfx/workflow":
|
|
filename = os.path.join(workflows_path, f"{uid}.json")
|
|
elif path == "/sdfx/template":
|
|
filename = os.path.join(templates_path, f"{uid}.json")
|
|
if not os.path.isfile(filename):
|
|
return web.json_response({'error': 'file missing'}, status=404)
|
|
with open(filename, 'r') as file:
|
|
file_content = json.load(file)
|
|
return web.json_response(file_content)
|
|
elif method == 'POST' or method == 'PUT':
|
|
try:
|
|
data = await request.json()
|
|
data['dateCreated'] = int(time.time())
|
|
if(method == 'POST'):
|
|
hash = str(uuid.uuid4())
|
|
else:
|
|
hash = request.query.get('uid')
|
|
if not hash:
|
|
return web.json_response({'error': 'uid missing'}, status=400)
|
|
data['uid'] = hash
|
|
if path == "/sdfx/workflow":
|
|
filename = os.path.join(workflows_path, f"{hash}.json")
|
|
elif path == "/sdfx/template":
|
|
filename = os.path.join(templates_path, f"{hash}.json")
|
|
with open(filename, 'w') as file:
|
|
json.dump(data, file)
|
|
response = {'uid': data.get('uid'), 'name': data.get('name'), 'meta': data.get('meta'), 'dateCreated': data.get('dateCreated')}
|
|
return web.json_response(response)
|
|
except json.JSONDecodeError:
|
|
return web.json_response({'error': 'invalid Json'}, status=400)
|
|
elif method == 'DELETE':
|
|
uid = request.query.get('uid')
|
|
if not uid:
|
|
return web.json_response({'error': "Fatal: Parameter 'uid' missing"}, status=400)
|
|
if path == "/sdfx/workflow":
|
|
filename = os.path.join(workflows_path, f"{uid}.json")
|
|
elif path == "/sdfx/template":
|
|
filename = os.path.join(templates_path, f"{uid}.json")
|
|
if os.path.isfile(filename):
|
|
os.remove(filename)
|
|
return web.json_response({'uid': uid, 'status':'deleted'})
|
|
else:
|
|
return web.Response(status=404)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/sdfx/input")
|
|
@server.PromptServer.instance.routes.delete("/sdfx/input")
|
|
@server.PromptServer.instance.routes.get("/sdfx/output")
|
|
@server.PromptServer.instance.routes.delete("/sdfx/output")
|
|
async def input_output(request):
|
|
method = request.method
|
|
path = request.url.path
|
|
if method == 'GET':
|
|
name = request.query.get('name')
|
|
if not name:
|
|
return web.json_response({'error': 'name missing'}, status=400)
|
|
if path == '/sdfx/input':
|
|
filename = os.path.join(folder_paths.get_input_directory(), name)
|
|
elif path == '/sdfx/output':
|
|
filename = os.path.join(folder_paths.get_output_directory(), name)
|
|
if not os.path.isfile(filename):
|
|
return web.json_response({'error': 'file missing'}, status=404)
|
|
return web.json_response({'name': name})
|
|
if method == 'DELETE':
|
|
name = request.query.get('name')
|
|
if not name:
|
|
return web.json_response({'error': 'name missing'}, status=400)
|
|
if path == '/sdfx/input':
|
|
filename = os.path.join(folder_paths.get_input_directory(), name)
|
|
elif path == '/sdfx/output':
|
|
filename = os.path.join(folder_paths.get_output_directory(), name)
|
|
if os.path.isfile(filename):
|
|
os.remove(filename)
|
|
return web.json_response({'name': name, 'status':'deleted'})
|
|
|
|
@server.PromptServer.instance.routes.get("/sdfx/input/list")
|
|
@server.PromptServer.instance.routes.get("/sdfx/output/list")
|
|
async def inputs_outputs(request):
|
|
list = []
|
|
path = request.url.path
|
|
|
|
if path == '/sdfx/input/list':
|
|
working_dir = folder_paths.get_input_directory()
|
|
elif path == '/sdfx/output/list':
|
|
working_dir = folder_paths.get_output_directory()
|
|
|
|
for filename in os.listdir(working_dir):
|
|
if '.' in filename:
|
|
list.append({'name': filename})
|
|
return web.json_response(list) |