diff --git a/__init__.py b/__init__.py index 1525a25..7656106 100644 --- a/__init__.py +++ b/__init__.py @@ -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): """ diff --git a/py/services.py b/py/services.py index 835b095..0edd8ca 100644 --- a/py/services.py +++ b/py/services.py @@ -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): diff --git a/py/utils.py b/py/utils.py index fcdac9f..f560b94 100644 --- a/py/utils.py +++ b/py/utils.py @@ -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 {} diff --git a/src/components/DialogModelDetail.vue b/src/components/DialogModelDetail.vue index ac64eee..3eaad50 100644 --- a/src/components/DialogModelDetail.vue +++ b/src/components/DialogModelDetail.vue @@ -14,7 +14,7 @@
@@ -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('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, '')