TBNilles/ComfyUI-Model-Manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Increase supported image extensions

Browse Source 
- Supports lower priority `<name>.preview.<extension>` preview names.
  - Note that this means that model names should not be of the format `<name>.preview`, else behavior may be unexpected.
- Removed use of splitext.
- Added checks for valid file type extensions.

- More conformal path accessing for `extension_uri` and `comfyui_model_uri`.
This commit is contained in:
Christian Bastian
2024-03-24 20:36:16 -04:00
parent 191d98565e
commit df8fb76f26
2 changed files with 137 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

208
__init__.py
View File

@@ -20,21 +20,51 @@ requests.packages.urllib3.disable_warnings()
import folder_paths import folder_paths
config_loader_path = os.path.join(os.path.dirname(__file__), 'config_loader.py') comfyui_model_uri = folder_paths.models_dir
extension_uri = os.path.dirname(__file__)
config_loader_path = os.path.join(extension_uri, 'config_loader.py')
config_loader_spec = importlib.util.spec_from_file_location('config_loader', config_loader_path) config_loader_spec = importlib.util.spec_from_file_location('config_loader', config_loader_path)
config_loader = importlib.util.module_from_spec(config_loader_spec) config_loader = importlib.util.module_from_spec(config_loader_spec)
config_loader_spec.loader.exec_module(config_loader) config_loader_spec.loader.exec_module(config_loader)
comfyui_model_uri = os.path.join(os.getcwd(), "models")
extension_uri = os.path.join(os.getcwd(), "custom_nodes" + os.path.sep + "ComfyUI-Model-Manager")
no_preview_image = os.path.join(extension_uri, "no-preview.png") no_preview_image = os.path.join(extension_uri, "no-preview.png")
ui_settings_uri = os.path.join(extension_uri, "ui_settings.yaml") ui_settings_uri = os.path.join(extension_uri, "ui_settings.yaml")
server_settings_uri = os.path.join(extension_uri, "server_settings.yaml") server_settings_uri = os.path.join(extension_uri, "server_settings.yaml")
fallback_model_extensions = set([".bin", ".ckpt", ".onnx", ".pt", ".pth", ".safetensors"]) # TODO: magic values fallback_model_extensions = set([".bin", ".ckpt", ".onnx", ".pt", ".pth", ".safetensors"]) # TODO: magic values
image_extensions = (".apng", ".gif", ".jpeg", ".jpg", ".png", ".webp") # TODO: JavaScript does not know about this (x2 states) image_extensions = (
".png", # order matters
".webp",
".jpeg",
".jpg",
".gif",
".apng",
)
stable_diffusion_webui_civitai_helper_image_extensions = (
".preview.png", # order matters
".preview.webp",
".preview.jpeg",
".preview.jpg",
".preview.gif",
".preview.apng",
)
preview_extensions = ( # TODO: JavaScript does not know about this (x2 states)
image_extensions + # order matters
stable_diffusion_webui_civitai_helper_image_extensions
)
model_info_extension = ".txt"
#video_extensions = (".avi", ".mp4", ".webm") # TODO: Requires ffmpeg or cv2. Cache preview frame? #video_extensions = (".avi", ".mp4", ".webm") # TODO: Requires ffmpeg or cv2. Cache preview frame?
def split_valid_ext(s, *arg_exts):
sl = s.lower()
for exts in arg_exts:
for ext in exts:
if sl.endswith(ext.lower()):
return (s[:-len(ext)], ext)
return (s, "")
_folder_names_and_paths = None # dict[str, tuple[list[str], list[str]]] _folder_names_and_paths = None # dict[str, tuple[list[str], list[str]]]
def folder_paths_folder_names_and_paths(refresh = False): def folder_paths_folder_names_and_paths(refresh = False):
global _folder_names_and_paths global _folder_names_and_paths
@@ -194,18 +224,18 @@ async def get_model_preview(request):
uri = request.query.get("uri") uri = request.query.get("uri")
image_path = no_preview_image image_path = no_preview_image
image_extension = "png" image_type = "png"
image_data = None image_data = None
if uri != "no-preview": if uri != "no-preview":
sep = os.path.sep sep = os.path.sep
uri = uri.replace("/" if sep == "\\" else "/", sep) uri = uri.replace("/" if sep == "\\" else "/", sep)
path, _ = search_path_to_system_path(uri) path, _ = search_path_to_system_path(uri)
head, extension = os.path.splitext(path) head, extension = split_valid_ext(path, preview_extensions)
if os.path.exists(path): if os.path.exists(path):
image_extension = extension[1:] image_type = extension.rsplit(".", 1)[1]
image_path = path image_path = path
elif os.path.exists(head) and os.path.splitext(head)[1] == ".safetensors": elif os.path.exists(head) and head.endswith(".safetensors"):
image_extension = extension[1:] image_type = extension.rsplit(".", 1)[1]
header = get_safetensor_header(head) header = get_safetensor_header(head)
metadata = header.get("__metadata__", None) metadata = header.get("__metadata__", None)
if metadata is not None: if metadata is not None:
@@ -218,23 +248,32 @@ async def get_model_preview(request):
with open(image_path, "rb") as file: with open(image_path, "rb") as file:
image_data = file.read() image_data = file.read()
return web.Response(body=image_data, content_type="image/" + image_extension) return web.Response(body=image_data, content_type="image/" + image_type)
@server.PromptServer.instance.routes.get("/model-manager/image/extensions")
async def get_image_extensions(request):
return web.json_response(image_extensions)
def download_model_preview(formdata): def download_model_preview(formdata):
path = formdata.get("path", None) path = formdata.get("path", None)
if type(path) is not str: if type(path) is not str:
raise ("Invalid path!") raise ("Invalid path!")
path, _ = search_path_to_system_path(path) path, model_type = search_path_to_system_path(path)
path_without_extension, _ = os.path.splitext(path) model_type_extensions = folder_paths_get_supported_pt_extensions(model_type)
path_without_extension, _ = split_valid_ext(path, model_type_extensions)
overwrite = formdata.get("overwrite", "true").lower() overwrite = formdata.get("overwrite", "true").lower()
overwrite = True if overwrite == "true" else False overwrite = True if overwrite == "true" else False
image = formdata.get("image", None) image = formdata.get("image", None)
if type(image) is str: if type(image) is str:
image_path = download_image(image, path, overwrite) _, image_extension = split_valid_ext(image, image_extensions) # TODO: doesn't work for https://civitai.com/images/...
_, image_extension = os.path.splitext(image_path) if image_extension == "":
raise ValueError("Invalid image type!")
image_path = path_without_extension + image_extension
download_file(image, image_path, overwrite)
else: else:
content_type = image.content_type content_type = image.content_type
if not content_type.startswith("image/"): if not content_type.startswith("image/"):
@@ -251,7 +290,7 @@ def download_model_preview(formdata):
with open(image_path, "wb") as f: with open(image_path, "wb") as f:
f.write(image_data) f.write(image_data)
delete_same_name_files(path_without_extension, image_extensions, image_extension) delete_same_name_files(path_without_extension, preview_extensions, image_extension)
@server.PromptServer.instance.routes.post("/model-manager/preview/set") @server.PromptServer.instance.routes.post("/model-manager/preview/set")
@@ -272,10 +311,11 @@ async def delete_model_preview(request):
return web.json_response({ "success": False }) return web.json_response({ "success": False })
model_path = urllib.parse.unquote(model_path) model_path = urllib.parse.unquote(model_path)
file, _ = search_path_to_system_path(model_path) model_path, model_type = search_path_to_system_path(model_path)
path_and_name, _ = os.path.splitext(file) model_extensions = folder_paths_get_supported_pt_extensions(model_type)
delete_same_name_files(path_and_name, image_extensions) path_and_name, _ = split_valid_ext(model_path, model_extensions)
delete_same_name_files(path_and_name, preview_extensions)
return web.json_response({ "success": True }) return web.json_response({ "success": True })
@@ -297,26 +337,32 @@ async def get_model_list(request):
for base_path_index, model_base_path in enumerate(folder_paths_get_folder_paths(model_type)): for base_path_index, model_base_path in enumerate(folder_paths_get_folder_paths(model_type)):
if not os.path.exists(model_base_path): # TODO: Bug in main code? ("ComfyUI\output\checkpoints", "ComfyUI\output\clip", "ComfyUI\models\t2i_adapter", "ComfyUI\output\vae") if not os.path.exists(model_base_path): # TODO: Bug in main code? ("ComfyUI\output\checkpoints", "ComfyUI\output\clip", "ComfyUI\models\t2i_adapter", "ComfyUI\output\vae")
continue continue
for cwd, _subdirs, files in os.walk(model_base_path): for cwd, subdirs, files in os.walk(model_base_path):
dir_models = [] dir_models = []
dir_images = [] dir_images = []
for file in files: for file in files:
if file.lower().endswith(model_extensions): if file.lower().endswith(model_extensions):
dir_models.append(file) dir_models.append(file)
elif file.lower().endswith(image_extensions): elif file.lower().endswith(preview_extensions):
dir_images.append(file) dir_images.append(file)
for model in dir_models: for model in dir_models:
model_name, model_ext = os.path.splitext(model) model_name, model_ext = split_valid_ext(model, model_extensions)
image = None image = None
image_modified = None image_modified = None
for iImage in range(len(dir_images)-1, -1, -1): for ext in preview_extensions: # order matters
image_name, _ = os.path.splitext(dir_images[iImage]) for iImage in range(len(dir_images)-1, -1, -1):
if model_name == image_name: image_name = dir_images[iImage]
image = end_swap_and_pop(dir_images, iImage) if not image_name.lower().endswith(ext.lower()):
img_abs_path = os.path.join(cwd, image) continue
image_modified = pathlib.Path(img_abs_path).stat().st_mtime_ns image_name = image_name[:-len(ext)]
if model_name == image_name:
image = end_swap_and_pop(dir_images, iImage)
img_abs_path = os.path.join(cwd, image)
image_modified = pathlib.Path(img_abs_path).stat().st_mtime_ns
break
if image is not None:
break break
abs_path = os.path.join(cwd, model) abs_path = os.path.join(cwd, model)
stats = pathlib.Path(abs_path).stat() stats = pathlib.Path(abs_path).stat()
@@ -378,7 +424,7 @@ def linear_directory_hierarchy(refresh = False):
for dir_path_index, dir_path in enumerate(model_dirs): for dir_path_index, dir_path in enumerate(model_dirs):
if not os.path.exists(dir_path) or os.path.isfile(dir_path): if not os.path.exists(dir_path) or os.path.isfile(dir_path):
continue continue
#dir_list.append({ "name": str(dir_path_index), "childIndex": None, "childCount": 0 }) #dir_list.append({ "name": str(dir_path_index), "childIndex": None, "childCount": 0 })
dir_stack = [(dir_path, model_dir_child_index + dir_path_index)] dir_stack = [(dir_path, model_dir_child_index + dir_path_index)]
while len(dir_stack) > 0: # DEPTH-FIRST while len(dir_stack) > 0: # DEPTH-FIRST
@@ -403,8 +449,7 @@ def linear_directory_hierarchy(refresh = False):
dir_child_count += 1 dir_child_count += 1
else: else:
# file # file
_, file_extension = os.path.splitext(item_name) if extension_whitelist is None or split_valid_ext(item_name, extension_whitelist)[1] != "":
if extension_whitelist is None or file_extension in extension_whitelist:
dir_list.append({ "name": item_name }) dir_list.append({ "name": item_name })
dir_child_count += 1 dir_child_count += 1
if dir_child_count > 0: if dir_child_count > 0:
@@ -514,16 +559,6 @@ def download_file(url, filename, overwrite):
os.rename(filename_temp, filename) os.rename(filename_temp, filename)
def download_image(image_uri, model_path, overwrite):
_, extension = os.path.splitext(image_uri) # TODO: doesn't work for https://civitai.com/images/...
if not extension in image_extensions:
raise ValueError("Invalid image type!")
path_without_extension, _ = os.path.splitext(model_path)
file = path_without_extension + extension
download_file(image_uri, file, overwrite)
return file
@server.PromptServer.instance.routes.get("/model-manager/model/info") @server.PromptServer.instance.routes.get("/model-manager/model/info")
async def get_model_info(request): async def get_model_info(request):
model_path = request.query.get("path", None) model_path = request.query.get("path", None)
@@ -531,35 +566,36 @@ async def get_model_info(request):
return web.json_response({ "success": False }) return web.json_response({ "success": False })
model_path = urllib.parse.unquote(model_path) model_path = urllib.parse.unquote(model_path)
file, _ = search_path_to_system_path(model_path) abs_path, model_type = search_path_to_system_path(model_path)
if file is None: if abs_path is None:
return web.json_response({}) return web.json_response({})
info = {} info = {}
path, name = os.path.split(model_path) comfyui_directory, name = os.path.split(model_path)
info["File Name"] = name info["File Name"] = name
info["File Directory"] = path info["File Directory"] = comfyui_directory
info["File Size"] = str(os.path.getsize(file)) + " bytes" info["File Size"] = str(os.path.getsize(abs_path)) + " bytes"
stats = pathlib.Path(file).stat() stats = pathlib.Path(abs_path).stat()
date_format = "%Y-%m-%d %H:%M:%S" date_format = "%Y-%m-%d %H:%M:%S"
date_modified = datetime.fromtimestamp(stats.st_mtime).strftime(date_format) date_modified = datetime.fromtimestamp(stats.st_mtime).strftime(date_format)
info["Date Modified"] = date_modified info["Date Modified"] = date_modified
info["Date Created"] = datetime.fromtimestamp(stats.st_ctime).strftime(date_format) info["Date Created"] = datetime.fromtimestamp(stats.st_ctime).strftime(date_format)
file_name, _ = os.path.splitext(file) model_extensions = folder_paths_get_supported_pt_extensions(model_type)
abs_name , _ = split_valid_ext(abs_path, model_extensions)
for extension in image_extensions: for extension in preview_extensions:
maybe_image = file_name + extension maybe_preview = abs_name + extension
if os.path.isfile(maybe_image): if os.path.isfile(maybe_preview):
image_path, _ = os.path.splitext(model_path) preview_path, _ = split_valid_ext(model_path, model_extensions)
image_modified = pathlib.Path(maybe_image).stat().st_mtime_ns preview_modified = pathlib.Path(maybe_preview).stat().st_mtime_ns
info["Preview"] = { info["Preview"] = {
"path": urllib.parse.quote_plus(image_path + extension), "path": urllib.parse.quote_plus(preview_path + extension),
"dateModified": urllib.parse.quote_plus(str(image_modified)), "dateModified": urllib.parse.quote_plus(str(preview_modified)),
} }
break break
header = get_safetensor_header(file) header = get_safetensor_header(abs_path)
metadata = header.get("__metadata__", None) metadata = header.get("__metadata__", None)
#json.dump(metadata, sys.stdout, indent=4) #json.dump(metadata, sys.stdout, indent=4)
#print() #print()
@@ -622,7 +658,7 @@ async def get_model_info(request):
training_comment if training_comment != "None" else "" training_comment if training_comment != "None" else ""
).strip() ).strip()
txt_file = file_name + ".txt" txt_file = abs_name + model_info_extension
notes = "" notes = ""
if os.path.isfile(txt_file): if os.path.isfile(txt_file):
with open(txt_file, 'r', encoding="utf-8") as f: with open(txt_file, 'r', encoding="utf-8") as f:
@@ -687,8 +723,9 @@ async def download_model(request):
return web.json_response(result) return web.json_response(result)
name = formdata.get("name") name = formdata.get("name")
_, model_extension = os.path.splitext(name) model_extensions = folder_paths_get_supported_pt_extensions(model_type)
if not model_extension in folder_paths_get_supported_pt_extensions(model_type): _, model_extension = split_valid_ext(name, model_extensions)
if model_extension == "":
result["invalid"] = "name" result["invalid"] = "name"
return web.json_response(result) return web.json_response(result)
file_name = os.path.join(directory, name) file_name = os.path.join(directory, name)
@@ -725,8 +762,9 @@ async def move_model(request):
old_file, old_model_type = search_path_to_system_path(old_file) old_file, old_model_type = search_path_to_system_path(old_file)
if not os.path.isfile(old_file): if not os.path.isfile(old_file):
return web.json_response({ "success": False }) return web.json_response({ "success": False })
_, model_extension = os.path.splitext(old_file) old_model_extensions = folder_paths_get_supported_pt_extensions(old_model_type)
if not model_extension in folder_paths_get_supported_pt_extensions(old_model_type): old_file_without_extension, model_extension = split_valid_ext(old_file, old_model_extensions)
if model_extension == "":
# cannot move arbitrary files # cannot move arbitrary files
return web.json_response({ "success": False }) return web.json_response({ "success": False })
@@ -740,7 +778,10 @@ async def move_model(request):
if os.path.isfile(new_file): if os.path.isfile(new_file):
# cannot overwrite existing file # cannot overwrite existing file
return web.json_response({ "success": False }) return web.json_response({ "success": False })
if not model_extension in folder_paths_get_supported_pt_extensions(new_model_type): new_model_extensions = folder_paths_get_supported_pt_extensions(new_model_type)
new_file_without_extension, new_model_extension = split_valid_ext(new_file, new_model_extensions)
if model_extension != new_model_extension:
# cannot change extension
return web.json_response({ "success": False }) return web.json_response({ "success": False })
new_file_dir, _ = os.path.split(new_file) new_file_dir, _ = os.path.split(new_file)
if not os.path.isdir(new_file_dir): if not os.path.isdir(new_file_dir):
@@ -754,11 +795,8 @@ async def move_model(request):
print(e, file=sys.stderr, flush=True) print(e, file=sys.stderr, flush=True)
return web.json_response({ "success": False }) return web.json_response({ "success": False })
old_file_without_extension, _ = os.path.splitext(old_file) # TODO: this could overwrite existing files in destination...
new_file_without_extension, _ = os.path.splitext(new_file) for extension in preview_extensions + (model_info_extension,):
# TODO: this could overwrite existing files...
for extension in image_extensions + (".txt",):
old_file = old_file_without_extension + extension old_file = old_file_without_extension + extension
if os.path.isfile(old_file): if os.path.isfile(old_file):
try: try:
@@ -772,9 +810,9 @@ async def move_model(request):
def delete_same_name_files(path_without_extension, extensions, keep_extension=None): def delete_same_name_files(path_without_extension, extensions, keep_extension=None):
for extension in extensions: for extension in extensions:
if extension == keep_extension: continue if extension == keep_extension: continue
image_file = path_without_extension + extension file = path_without_extension + extension
if os.path.isfile(image_file): if os.path.isfile(file):
os.remove(image_file) os.remove(file)
@server.PromptServer.instance.routes.post("/model-manager/model/delete") @server.PromptServer.instance.routes.post("/model-manager/model/delete")
@@ -785,27 +823,22 @@ async def delete_model(request):
if model_path is None: if model_path is None:
return web.json_response(result) return web.json_response(result)
model_path = urllib.parse.unquote(model_path) model_path = urllib.parse.unquote(model_path)
model_path, model_type = search_path_to_system_path(model_path)
file, model_type = search_path_to_system_path(model_path) if model_path is None:
if file is None:
return web.json_response(result) return web.json_response(result)
_, extension = os.path.splitext(file) model_extensions = folder_paths_get_supported_pt_extensions(model_type)
if not extension in folder_paths_get_supported_pt_extensions(model_type): path_and_name, model_extension = split_valid_ext(model_path, model_extensions)
if model_extension == "":
# cannot delete arbitrary files # cannot delete arbitrary files
return web.json_response(result) return web.json_response(result)
if os.path.isfile(file): if os.path.isfile(model_path):
os.remove(file) os.remove(model_path)
result["success"] = True result["success"] = True
path_and_name, _ = os.path.splitext(file) delete_same_name_files(path_and_name, preview_extensions)
delete_same_name_files(path_and_name, (model_info_extension,))
delete_same_name_files(path_and_name, image_extensions)
txt_file = path_and_name + ".txt"
if os.path.isfile(txt_file):
os.remove(txt_file)
return web.json_response(result) return web.json_response(result)
@@ -821,9 +854,10 @@ async def set_notes(request):
model_path = body.get("path", None) model_path = body.get("path", None)
if type(model_path) is not str: if type(model_path) is not str:
return web.json_response({ "success": False }) return web.json_response({ "success": False })
model_path, _ = search_path_to_system_path(model_path) model_path, model_type = search_path_to_system_path(model_path)
file_path_without_extension, _ = os.path.splitext(model_path) model_extensions = folder_paths_get_supported_pt_extensions(model_type)
filename = os.path.normpath(file_path_without_extension + ".txt") file_path_without_extension, _ = split_valid_ext(model_path, model_extensions)
filename = os.path.normpath(file_path_without_extension + model_info_extension)
if text.isspace() or text == "": if text.isspace() or text == "":
if os.path.exists(filename): if os.path.exists(filename):
os.remove(filename) os.remove(filename)

18
web/model-manager.js
View File

@@ -35,7 +35,21 @@ const modelNodeType = {
}; };
const MODEL_EXTENSIONS = [".bin", ".ckpt", ".onnx", ".pt", ".pth", ".safetensors"]; // TODO: ask server for? const MODEL_EXTENSIONS = [".bin", ".ckpt", ".onnx", ".pt", ".pth", ".safetensors"]; // TODO: ask server for?
const IMAGE_EXTENSIONS = [".apng", ".gif", ".jpeg", ".jpg", ".png", ".webp"]; // TODO: ask server for? const IMAGE_EXTENSIONS = [
".png",
".webp",
".jpeg",
".jpg",
".gif",
".apng",
".preview.png",
".preview.webp",
".preview.jpeg",
".preview.jpg",
".preview.gif",
".preview.apng",
]; // TODO: /model-manager/image/extensions
class SearchPath { class SearchPath {
/** /**
@@ -1413,7 +1427,7 @@ class ModelInfoView {
$: (el) => (this.elements.setPreviewButton = el), $: (el) => (this.elements.setPreviewButton = el),
textContent: "Set as Preview", textContent: "Set as Preview",
onclick: async(e) => { onclick: async(e) => {
const confirmation = window.confirm("Change preview image PERMANENTLY?"); const confirmation = window.confirm("Change preview image(s) PERMANENTLY?");
let updatedPreview = false; let updatedPreview = false;
if (confirmation) { if (confirmation) {
e.target.disabled = true; e.target.disabled = true;