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

pref: Optimize list loading time

Browse Source
This commit is contained in:
hayden
2024-10-27 17:54:59 +08:00
parent 6b031a50bc
commit 92c55e04fd
4 changed files with 99 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
__init__.py
View File

@@ -68,12 +68,7 @@ async def read_models(request):
Scan all models and read their information.
"""
try:
result = []
model_base_paths = config.model_base_paths
for model_type in model_base_paths:
result.extend(services.scan_models_by_model_type(model_type))
result = [{"id": i, **x} for i, x in enumerate(result)]
result = services.scan_models()
return web.json_response({"success": True, "data": result})
except Exception as e:
error_msg = f"Read models failed: {str(e)}"
@@ -82,6 +77,26 @@ async def read_models(request):
return web.json_response({"success": False, "error": error_msg})
@routes.get("/model-manager/model/{type}/{index}/{filename:.*}")
async def read_model_info(request):
"""
Get the information of the specified model.
"""
model_type = request.match_info.get("type", None)
index = int(request.match_info.get("index", None))
filename = request.match_info.get("filename", None)
try:
model_path = utils.get_valid_full_path(model_type, index, filename)
result = services.get_model_info(model_path)
return web.json_response({"success": True, "data": result})
except Exception as e:
error_msg = f"Read model info failed: {str(e)}"
logging.error(error_msg)
logging.debug(traceback.format_exc())
return web.json_response({"success": False, "error": error_msg})
@routes.put("/model-manager/model/{type}/{index}/{filename:.*}")
async def update_model(request):
"""

103
py/services.py
View File

@@ -5,6 +5,7 @@ import folder_paths
from typing import Any
from multidict import MultiDictProxy
from . import config
from . import utils
from . import socket
from . import download
@@ -31,62 +32,39 @@ async def connect_websocket(request):
return ws
def scan_models_by_model_type(model_type: str):
"""
Scans all models in the given model type and returns a list of models.
"""
out = []
folders, extensions = folder_paths.folder_names_and_paths[model_type]
for path_index, base_path in enumerate(folders):
files = utils.recursive_search_files(base_path)
def scan_models():
result = []
model_base_paths = config.model_base_paths
for model_type in model_base_paths:
models = folder_paths.filter_files_extensions(files, extensions)
folders, extensions = folder_paths.folder_names_and_paths[model_type]
for path_index, base_path in enumerate(folders):
files = utils.recursive_search_files(base_path)
for fullname in models:
"""
fullname is model path relative to base_path
eg.
abs_path is /path/to/models/stable-diffusion/custom_group/model_name.ckpt
base_path is /path/to/models/stable-diffusion
fullname is custom_group/model_name.ckpt
basename is custom_group/model_name
extension is .ckpt
"""
models = folder_paths.filter_files_extensions(files, extensions)
images = folder_paths.filter_files_content_types(files, ["image"])
image_dict = utils.file_list_to_name_dict(images)
fullname = fullname.replace(os.path.sep, "/")
basename = os.path.splitext(fullname)[0]
extension = os.path.splitext(fullname)[1]
prefix_path = fullname.replace(os.path.basename(fullname), "")
for fullname in models:
fullname = fullname.replace(os.path.sep, "/")
basename = os.path.splitext(fullname)[0]
extension = os.path.splitext(fullname)[1]
abs_path = os.path.join(base_path, fullname)
file_stats = os.stat(abs_path)
abs_path = os.path.join(base_path, fullname)
file_stats = os.stat(abs_path)
# Resolve metadata
metadata = utils.get_model_metadata(abs_path)
# Resolve preview
image_name = image_dict.get(basename, "no-preview.png")
abs_image_path = os.path.join(base_path, image_name)
if os.path.isfile(abs_image_path):
image_state = os.stat(abs_image_path)
image_timestamp = round(image_state.st_mtime_ns / 1000000)
image_name = f"{image_name}?ts={image_timestamp}"
model_preview = (
f"/model-manager/preview/{model_type}/{path_index}/{image_name}"
)
# Resolve preview
image_name = utils.get_model_preview_name(abs_path)
image_name = os.path.join(prefix_path, image_name)
abs_image_path = os.path.join(base_path, image_name)
if os.path.isfile(abs_image_path):
image_state = os.stat(abs_image_path)
image_timestamp = round(image_state.st_mtime_ns / 1000000)
image_name = f"{image_name}?ts={image_timestamp}"
model_preview = (
f"/model-manager/preview/{model_type}/{path_index}/{image_name}"
)
# Resolve description
description_file = utils.get_model_description_name(abs_path)
description_file = os.path.join(prefix_path, description_file)
abs_desc_path = os.path.join(base_path, description_file)
description = None
if os.path.isfile(abs_desc_path):
with open(abs_desc_path, "r", encoding="utf-8") as f:
description = f.read()
out.append(
{
model_info = {
"fullname": fullname,
"basename": basename,
"extension": extension,
@@ -94,14 +72,31 @@ def scan_models_by_model_type(model_type: str):
"pathIndex": path_index,
"sizeBytes": file_stats.st_size,
"preview": model_preview,
"description": description,
"createdAt": round(file_stats.st_ctime_ns / 1000000),
"updatedAt": round(file_stats.st_mtime_ns / 1000000),
"metadata": metadata,
}
)
return out
result.append(model_info)
return result
def get_model_info(model_path: str):
directory = os.path.dirname(model_path)
metadata = utils.get_model_metadata(model_path)
description_file = utils.get_model_description_name(model_path)
description_file = os.path.join(directory, description_file)
description = None
if os.path.isfile(description_file):
with open(description_file, "r", encoding="utf-8") as f:
description = f.read()
return {
"metadata": metadata,
"description": description,
}
def update_model(model_path: str, post: MultiDictProxy):

10
py/utils.py
View File

@@ -58,17 +58,23 @@ def recursive_search_files(directory: str):
files, folder_all = folder_paths.recursive_search(
directory, excluded_dir_names=[".git"]
)
files.sort()
return files
def search_files(directory: str):
entries = os.listdir(directory)
files = [f for f in entries if os.path.isfile(os.path.join(directory, f))]
files.sort()
return files
def file_list_to_name_dict(files: list[str]):
file_dict: dict[str, str] = {}
for file in files:
filename = os.path.splitext(file)[0]
file_dict[filename] = file
return file_dict
def get_model_metadata(filename: str):
if not filename.endswith(".safetensors"):
return {}

23
src/components/DialogModelDetail.vue
View File

@@ -14,7 +14,7 @@
<div class="px-8">
<ModelContent
v-model:editable="editable"
:model="model"
:model="modelContent"
@submit="handleSave"
@reset="handleCancel"
>
@@ -61,8 +61,9 @@ import ModelContent from 'components/ModelContent.vue'
import DialogResizer from 'components/DialogResizer.vue'
import ResponseScrollArea from 'components/ResponseScrollArea.vue'
import { useConfig } from 'hooks/config'
import { computed, ref } from 'vue'
import { computed, ref, watchEffect } from 'vue'
import { useModelNodeAction, useModels } from 'hooks/model'
import { useRequest } from 'hooks/request'
const visible = defineModel<boolean>('visible')
interface Props {
@@ -75,6 +76,24 @@ const { remove, update } = useModels()
const editable = ref(false)
const { data: extraInfo, refresh: fetchExtraInfo } = useRequest(
`/model/${props.model.type}/${props.model.pathIndex}/${props.model.fullname}`,
{
method: 'GET',
manual: true,
},
)
const modelContent = computed(() => {
return Object.assign({}, props.model, extraInfo.value)
})
watchEffect(() => {
if (visible.value === true) {
fetchExtraInfo()
}
})
const filename = computed(() => {
const basename = props.model.fullname.split('/').pop()!
return basename.replace(props.model.extension, '')