Files
JaredTherriault-ComfyUI-JNodes/py/server_backend.py
T

852 lines
32 KiB
Python

import os
import re
from .logger import *
from .utils import *
from .server_backend_get_subdirectory_images import GetSubdirectoryImages
from app.user_manager import UserManager
import folder_paths
import aiohttp
from aiohttp import web
import server
from PIL import Image
import json
import shutil
from safetensors import safe_open
from pathlib import Path
CANCELLATION_REQUESTED = False
def should_cancel_task():
return CANCELLATION_REQUESTED
def request_task_cancellation():
global CANCELLATION_REQUESTED
CANCELLATION_REQUESTED = True
def reset_cancellation():
global CANCELLATION_REQUESTED
CANCELLATION_REQUESTED = False
async def read_web_request_content(reader):
data = await reader.read()
return data.decode('utf-8')
def get_model_items(request):
reset_cancellation()
type = "loras"
if "type" in request.rel_url.query:
type = request.rel_url.query["type"]
file_list = folder_paths.get_filename_list(type)
subdirectory = request.rel_url.query.get("subdirectory", "")
if subdirectory:
file_list = [f for f in file_list if f.startswith(subdirectory + "/") or f == subdirectory]
image_extension_filter = {'.png', '.jpg', '.jpeg', '.gif', '.webp', '.mp4', '.webm'}
info_extension_filter = {'.info', '.txt', '.json'}
familiar_dictionaries = create_familiar_dictionaries(file_list, type, image_extension_filter, info_extension_filter)
#logger.info(familiar_dictionaries)
return web.json_response(familiar_dictionaries)
def create_familiar_dictionaries(names, type, image_extension_filter, info_extension_filter):
"""
Makes a dictionary (key: item_name) containing a nested dictionary: a list of familiar images and a list of familiar info files.
Intended for loras but can be used for other items with a similar setup.
Works faster when all familiar items are together in one folder for the item itself rather than when all items,
familiar and non-familiar, are in one big folder. For example "loras/MyLora/ contains all "MyLora" familiar files,
and they aren't all just in the base "lora" folder.
"""
def load_pt_metadata(file_path):
try:
data = torch.load(file_path, map_location="cpu") # Load on CPU to avoid GPU dependency
except Exception as e:
print(f"load_pt_metadata: Error loading file: {e}")
return None
metadata = None
# Check if the file contains metadata (depends on how it was saved)
if isinstance(data, dict) and "metadata" in data:
metadata = data["metadata"]
return metadata
def load_safetensors_metadata(file_path):
metadata = None
# Open the safetensors file
with safe_open(file_path, framework="pt") as file:
try:
metadata = file.metadata()
except:
pass
return metadata
familiar_dictionaries = {}
for item_name in names:
if should_cancel_task():
break
item_name = item_name.replace("\\", "/")
# logger.info(f"item_name: {item_name}")
try:
file_path = folder_paths.get_full_path(type, item_name)
# logger.info(f"file_path: {file_path}")
if file_path is None:
logger.warning(f"Unable to get path for {type} {item_name}")
continue
metadata = None
try:
if item_name.endswith(".safetensors"):
metadata = load_safetensors_metadata(file_path)
else:
metadata = load_pt_metadata(file_path)
if metadata:
try: # Sorted dictionaries only available in py 3.7+
metadata = {k: metadata[k] for k in sorted(metadata)}
except:
pass
except Exception as e:
logger.error(f"Error loading {item_name} ({type}): {e}")
file_name_no_ext, file_ext = os.path.splitext(item_name)
# logger.info(f"file_name_no_ext, file_ext: {file_name_no_ext, file_ext}")
# Get time of creation since the last epoch, in seconds
file_age = get_creation_time(file_path)
file_path = file_path.replace("\\", "/")
# logger.info(f'file_path: {file_path}')
file_name_no_ext = file_name_no_ext.replace("\\", "/")
# logger.info(f'file_name_no_ext: {file_name_no_ext}')
parent_directory = os.path.dirname(file_path)
# logger.info(f"parent_directory: {parent_directory}")
containing_directory = None
if "/" in file_name_no_ext:
split = file_name_no_ext.split("/")
if len(split) > 1:
file_name_no_ext = split[len(split) - 1]
split.pop()
containing_directory = "/".join(split)
familiar_images = find_items_with_similar_names(parent_directory, containing_directory, file_name_no_ext, image_extension_filter)
familiar_infos = find_items_with_similar_names(parent_directory, containing_directory, file_name_no_ext, info_extension_filter, True)
# logger.info(f"familiar_images: {familiar_images}")
familiar_dictionaries[file_name_no_ext] = {
"containing_directory": containing_directory,
"full_name": item_name,
"file_age": file_age,
"file_ext": file_ext,
"metadata": metadata,
"familiar_images": familiar_images,
"familiar_infos": familiar_infos
}
except Exception as e:
logger.error(f"Error parsing {item_name} ({type}): {e}")
return familiar_dictionaries
def find_items_with_similar_names(folder_path, containing_directory, base_name, extension_filter, load = False):
# logger.info(
# f'folder_path, containing_directory, base_name, extension_filter, load: {folder_path, containing_directory, base_name, extension_filter, load}'
# )
familiars = []
for file_name in os.listdir(folder_path):
if should_cancel_task():
break
# logger.info(f"file_name in directory: {file_name}")
file_name_no_ext, ext = os.path.splitext(file_name)
# logger.info(f"file_name_no_ext, ext: {file_name_no_ext, ext}")
if ext in extension_filter and base_name in file_name_no_ext:
# logger.info(f"file_name matches criteria: {file_name}")
if load:
try:
with open(folder_path + "/" + file_name, 'r', encoding='utf-8') as opened_file:
loaded_text = opened_file.read()
familiars.append({"file_name": file_name, "loaded_text": loaded_text})
continue
except Exception as e:
log_exception("Error loading text, will append familiar file name. Error:", e)
familiars.append({"file_name": containing_directory + "/" + file_name if containing_directory is not None else file_name, "loaded_text": ""})
familiars.sort(key=lambda x: x["file_name"])
return familiars
def list_subdirectories_recursively(root_directory):
results = []
def walk_through_subdirectories(in_directory):
"""
Recursively list subdirectories within the specified directory.
Args:
in_directory (str): The directory to process.
"""
if should_cancel_task():
return
full_directory = os.path.join(root_directory, in_directory)
if not os.path.isdir(full_directory):
return
results.append(in_directory)
items = os.listdir(full_directory)
for item in items:
if os.path.isdir(os.path.join(full_directory, item)):
walk_through_subdirectories(os.path.join(in_directory, item))
walk_through_subdirectories("")
results = sorted(results)
return results
def list_immediate_subdirectories(root_directory):
results = []
if not os.path.isdir(root_directory):
return results
items = os.listdir(root_directory)
for item in items:
combined_path = os.path.join(root_directory, item)
if os.path.isdir(combined_path):
results.append(item)
results = sorted(results)
return results
def list_comfyui_subdirectories_request(request):
reset_cancellation()
try:
root_directory = convert_relative_comfyui_path_to_full_path(request.rel_url.query["root_directory"])
results = list_subdirectories_recursively(root_directory)
return web.json_response({"success": len(results) > 0, "payload": results })
except Exception as e:
log_exception("Error listing subdirectories:", e)
return web.json_response({"success": False, "error": str(e)})
def list_immediate_subdirectories_request(request):
try:
root_directory = resolve_file_path(request.rel_url.query["root_directory"])
results = list_immediate_subdirectories(root_directory)
return web.json_response({"success": len(results) > 0, "payload": results })
except Exception as e:
log_exception("Error listing subdirectories:", e)
return web.json_response({"success": False, "error": str(e)})
def list_model_subdirectories_request(request):
try:
type = request.rel_url.query.get("type", "loras")
root_directory = convert_relative_comfyui_path_to_full_path(type)
results = list_subdirectories_recursively(root_directory)
return web.json_response({"success": len(results) > 0, "payload": results})
except Exception as e:
log_exception("Error listing model subdirectories:", e)
return web.json_response({"success": False, "error": str(e)})
def get_comfyui_subdirectory_images_request(request):
reset_cancellation()
try:
root_directory = request.rel_url.query["root_directory"] or ""
selected_subdirectory = request.rel_url.query["selected_subdirectory"] or ""
recursive = request.rel_url.query["recursive"] == "true"
results = GetSubdirectoryImages(
os.path.join(convert_relative_comfyui_path_to_full_path(root_directory), selected_subdirectory), recursive,
external_cancel_check=should_cancel_task).get_subdirectory_images()
return web.json_response({"success": True, "payload": results })
except Exception as e:
log_exception("Error listing subdirectory images:", e)
return web.json_response({"success": False, "error": str(e)})
async def request_open_file_manager(request):
try:
result = await validate_and_return_file_from_request(request)
if result["success"] == True:
open_file_manager(os.path.normpath(result["payload"]["file"]))
return web.json_response({"success": True})
except Exception as e:
logger.error(e)
return web.json_response({"success": False, "error": str(e)})
async def validate_and_return_file_from_request(request):
type = "loras"
if "type" in request.rel_url.query:
type = request.rel_url.query["type"]
base_dirs = None
try: # Try to infer base_dir
base_dirs = folder_paths.get_folder_paths(type)
# logger.info(f"base_dirs: {base_dirs}")
except: # If we can't, most likely because it's not a built-in type, assume type is a subfolder in the ComfyUI directory
# logger.info("validate_and_return_file_from_request: failed to infer base_dir")
try:
base_dirs = [convert_relative_comfyui_path_to_full_path(type)]
# logger.info(f"base_dirs: {base_dirs}")
except Exception as e:
log_exception(f"Error finding folder {type}. Error:", e)
if base_dirs is None or len(base_dirs) == 0:
logger.warning(f"Unable to get parent directory for {type}")
return { "success": False, "response": 400 }
for base_dir in base_dirs:
subfolder = ''
if "subfolder" in request.rel_url.query:
subfolder = request.rel_url.query["subfolder"]
base_dir = os.path.join(base_dir, subfolder)
if "filename" in request.rel_url.query:
filename = request.rel_url.query["filename"]
# validation for security: prevent accessing arbitrary path
if filename[0] == '/' or filename.startswith('..') or filename.startswith('./'):
logger.warning(f"Attempting to access an arbitrary path, aborting. filename: {filename}")
return { "success": False, "response": 400 }
#filename = os.path.basename(filename)
file = os.path.join(base_dir, filename)
# Hack for linux/mac/unix
if not os.path.isfile(file):
file = f"/{file}"
if os.path.isfile(file):
return { "success": True, "payload": { "base_dir": base_dir, "subfolder": subfolder, "filename": filename, "file": file } }
return { "success": False, "response": 400 }
async def view_image(request):
result = await validate_and_return_file_from_request(request)
if result["success"] == True:
if 'preview' in request.rel_url.query:
with Image.open(result["payload"]["file"]) as img:
preview_info = request.rel_url.query['preview'].split(';')
image_format = preview_info[0]
if image_format not in ['webp', 'jpeg', 'jpg'] 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', 'jpg'] or request.rel_url.query.get('channel', '') == 'rgb':
img = img.convert("RGB")
img.save(buffer, format=image_format, quality=quality)
buffer.seek(0)
filename = result["payload"]["filename"]
filename = filename.replace('"', '\\"') # Escape double quotes
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(result["payload"]["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)
filename = result["payload"]["filename"]
filename = filename.replace('"', '\\"') # Escape double quotes
return web.Response(
body=buffer.read(),
content_type=f'image/{image_format}',
headers={"Content-Disposition": f'filename="{filename}"'}
)
elif channel == 'a':
with Image.open(result["payload"]["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)
filename = result["payload"]["filename"]
filename = filename.replace('"', '\\"') # Escape double quotes
return web.Response(
body=alpha_buffer.read(),
content_type=f'image/{image_format}',
headers={"Content-Disposition": f'filename="{filename}"'}
)
else:
filename = result["payload"]["filename"]
filename = filename.replace('"', '\\"') # Escape double quotes
return web.FileResponse(result["payload"]["file"], headers={"Content-Disposition": f'filename="{filename}"'})
return web.Response(status= result["response"] if result["response"] else 404)
async def save_image_as_model_preview(request):
try:
# For ease of use, pass the model URI in the URL
result = await validate_and_return_file_from_request(request)
if result["success"] == True:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
new_image_url = request_json["url"]
model_file_path = result["payload"]["file"]
model_directory = os.path.dirname(model_file_path)
model_base_name = os.path.splitext(os.path.basename(model_file_path))[0]
new_image_file_name = None
async with aiohttp.ClientSession() as session:
async with session.get(new_image_url) as response:
if response.status == 200:
new_image_ext = response.content_type.split("/")[1]
new_image_file_name = f"{get_next_base_filename(model_directory, model_base_name)}.{new_image_ext}"
new_image_path = os.path.join(model_directory, new_image_file_name)
with open(new_image_path, "wb") as file:
while chunk := await response.content.read(1024):
file.write(chunk)
else:
return web.json_response({"success": False, "error": "Failed to download image"})
return web.json_response({"success": True, "file_name": new_image_file_name})
except Exception as e:
logger.error(e)
return web.json_response({"success": False, "error": str(e)})
async def copy_item(request):
def copy_file(path_from, path_to):
try:
if os.path.isfile(path_from):
try:
path_to_dir = os.path.dirname(path_to)
if not os.path.exists(path_to_dir):
os.makedirs(path_to_dir, exist_ok=True)
shutil.copy(path_from, path_to)
logger.info(f"Successfully copied file from '{path_from}' to '{path_to}'")
return web.json_response({"success": True})
except Exception as e:
logger.error(f"Error occurred while copying file from '{path_from}' to '{path_to}': {e}")
return web.json_response({"success": False, "message": "Failed to copy file."})
else:
message = f"'{path_from}' is not a valid file."
logger.warning(message)
return web.json_response({"success": False, "message": message})
except Exception as e:
log_exception(f"Error occurred while deleting '{path}':", e)
return web.json_response({"success": False, "message": message})
result = await validate_and_return_file_from_request(request)
if result["success"] == True:
path_from = result["payload"]["file"]
if "destination" in request.rel_url.query:
path_to = request.rel_url.query["destination"]
path_to = resolve_file_path(path_to)
if path_from and path_to:
return copy_file(path_from, path_to)
logger.warning("File could not be copied: file not found or filename and destination not defined")
return web.Response(status=404)
async def upload_image(request):
try:
post = await request.post()
image = post.get("image")
overwrite = post.get("overwrite")
image_upload_type = post.get("type")
upload_dir = convert_relative_comfyui_path_to_full_path(image_upload_type)
if image and image.file:
filename = post.get("filename")
if not filename or filename == "":
filename = image.filename
if not filename:
return web.Response(status=400)
subfolder = post.get("subfolder", "")
full_output_folder = os.path.join(upload_dir, subfolder.replace("\\", "/"))
filepath = os.path.abspath(os.path.join(full_output_folder, filename))
if os.path.commonpath((upload_dir, filepath)) != upload_dir:
return web.Response(status=400)
if not os.path.exists(full_output_folder):
os.makedirs(full_output_folder)
if overwrite is not None and (overwrite == "true" or overwrite == "1"):
pass
else:
split = os.path.splitext(filename)
i = 1
while os.path.exists(filepath):
filename = f"{split[0]} ({i}){split[1]}"
filepath = os.path.join(full_output_folder, filename)
i += 1
with open(filepath, "wb") as f:
f.write(image.file.read())
return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type})
else:
return web.Response(status=400)
except Exception as e:
log_exception("Error uploading image:", e)
return web.json_response({"success": False, "error": str(e)})
async def save_model_user_info(request):
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
type = subfolder = item_name = text_to_save = None
if "type" in request_json:
type = request_json["type"]
if "subfolder" in request_json:
subfolder = request_json["subfolder"]
if "item_name" in request_json:
item_name = request_json["item_name"]
if "text" in request_json:
text_to_save = request_json["text"]
if type and item_name and text_to_save:
name = subfolder + '/' + item_name if subfolder is not None else item_name
full_path = folder_paths.get_full_path(type, name)
folder_path = os.path.dirname(full_path)
file_name_no_ext, file_ext = os.path.splitext(os.path.basename(item_name))
file_path = f"{folder_path}/{file_name_no_ext}.user.info"
with open(file_path, 'w', encoding='utf-8') as file:
file.write(text_to_save)
return web.json_response({"success": True})
assert ValueError
except Exception as e:
log_exception("Error saving text:", e)
return web.json_response({"success": False, "error": str(e)})
def load_info(request):
type = "loras"
if "type" in request.rel_url.query:
type = request.rel_url.query["type"]
file_list = folder_paths.get_filename_list(type)
base_dir = None
for item_name in file_list:
if "/" not in item_name.replace("\\", "/"):
file_path = folder_paths.get_full_path(type, item_name)
base_dir = os.path.dirname(file_path)
if base_dir is None:
logger.warning(f"Unable to get parent directory for {type}")
return web.Response(status=400)
if "filename" in request.rel_url.query:
filename = request.rel_url.query["filename"]
# validation for security: prevent accessing arbitrary path
if filename[0] == '/' or '..' in filename:
return web.Response(status=400)
if base_dir is None:
return web.Response(status=400)
#filename = os.path.basename(filename)
file = os.path.join(base_dir, filename)
if os.path.isfile(file):
try:
with open(file, 'r', encoding='utf-8') as opened_file:
loaded_text = opened_file.read()
return web.json_response({"success": True, "payload" : loaded_text})
except Exception as e:
log_exception("Error loading text:", e)
return web.json_response({"success": False, "error": str(e)})
return web.Response(status=404)
async def find_files(request):
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
# Extract base path
path = None
if "path" in request_json:
path = Path(resolve_file_path(request_json["path"]))
if not path or not os.path.isdir(path):
return web.json_response({"success": False, "error": "Invalid or missing directory path"})
# Optional flags
recursive = bool(request_json.get("recursive", False))
absolute = bool(request_json.get("absolute", False))
extension = request_json.get("extension", "")
if extension and not extension.startswith("."):
extension = "." + extension
# Build file list
if recursive:
iterator = path.rglob("*")
else:
iterator = path.iterdir()
files = []
for f in iterator:
if f.is_file() and (f.suffix == extension if extension else True):
if absolute:
files.append(str(f.resolve()))
else:
files.append(str(f.relative_to(path)))
return web.json_response({"success": True, "files": files})
except Exception as e:
log_exception("Error finding files:", e)
return web.json_response({"success": False, "error": str(e)})
async def save_text(request):
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
path = None
text_to_save = None
if "path" in request_json:
path = Path(resolve_file_path(request_json["path"]))
if "text" in request_json:
text_to_save = request_json["text"]
if path and text_to_save:
# 🔹 ensure parent directories exist
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(text_to_save, encoding="utf-8")
return web.json_response({"success": True})
except Exception as e:
log_exception("Error saving text:", e)
return web.json_response({"success": False, "error": str(e)})
async def load_text(request):
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
if "path" in request_json:
path = resolve_file_path(request_json["path"])
if os.path.exists(path):
with open(path, 'r', encoding='utf-8') as file:
loaded_text = file.read()
return web.json_response({"success": True, "payload" : loaded_text})
except Exception as e:
log_exception("Error loading text:", e)
return web.json_response({"success": False, "error": str(e)})
def save_settings(request, settings):
um = UserManager()
file = um.get_request_user_filepath(
request, "jnodes.settings.json")
with open(file, "w") as f:
f.write(json.dumps(settings, indent=4))
async def post_setting(request):
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
setting_id = None
if "id" in request_json:
setting_id = request_json["id"]
if not setting_id:
return web.json_response({"success": False, "error": "No 'id' found in request!"})
setting_value = None
if "value" in request_json:
setting_value = request_json["value"]
if not setting_value:
return web.json_response({"success": False, "error": "No 'value' found in request!"})
if setting_id and setting_value:
settings = get_settings(request)
settings[setting_id] = setting_value
save_settings(request, settings)
return web.json_response({"success": True})
except Exception as e:
log_exception("Error saving text:", e)
return web.json_response({"success": False, "error": str(e)})
async def post_all_settings(request):
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
settings = None
if "settings" in request_json:
settings = request_json["settings"]
if settings:
save_settings(request, settings)
return web.json_response({"success": True})
else:
return web.json_response({"success": False, "error": "No 'settings' found in request!"})
except Exception as e:
log_exception("Error saving text:", e)
return web.json_response({"success": False, "error": str(e)})
def get_settings(request):
um = UserManager()
file = um.get_request_user_filepath(
request, "jnodes.settings.json")
if os.path.isfile(file):
with open(file) as f:
return json.load(f)
else:
return {}
async def get_setting(request):
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
setting_id = None
if "id" in request_json:
setting_id = request_json["id"]
settings = get_settings(request)
return web.json_response({"success": True, "payload" : settings[setting_id]})
except Exception as e:
log_exception("Error loading text:", e)
return web.json_response({"success": False, "error": str(e)})
async def get_all_settings(request):
try:
settings = get_settings(request)
return web.json_response({"success": True, "payload" : json.dumps(settings)})
except Exception as e:
log_exception("Error loading text:", e)
return web.json_response({"success": False, "error": str(e)})
async def delete_item(request):
def delete_file(path):
try:
if os.path.isfile(path):
try:
# Try to import send2trash module
from send2trash import send2trash
# If import successful, use send2trash to send file to trash
send2trash(path)
message = f"File '{path}' sent to trash/recycle bin successfully."
print(message)
return web.json_response({"success": True, "type": "send2trash", "message": message})
except ImportError:
# If send2trash module not available, delete file permanently
os.remove(path)
message = f"File '{path}' deleted permanently. If you want to send to trash/recycle, consider installing send2trash (pip install send2trash)."
print(message)
return web.json_response({"success": True, "type": "permanent", "message": message})
else:
message = f"'{path}' is not a valid file."
print(message)
return web.json_response({"success": False, "message": message})
except Exception as e:
log_exception(f"Error occurred while deleting '{path}':", e)
return web.json_response({"success": False, "message": message})
result = await validate_and_return_file_from_request(request)
if result["success"] == True:
return delete_file(result["payload"]["file"])
# Alt path
try:
request_data = await read_web_request_content(request.content)
request_json = json.loads(request_data)
path = None
if "path" in request_json:
path = Path(resolve_file_path(request_json["path"]))
return delete_file(path)
except:
pass
logger.warning("File could not be deleted: file not found")
return web.Response(status=404)