diff --git a/py/information.py b/py/information.py index 6593e7e..6b6747a 100644 --- a/py/information.py +++ b/py/information.py @@ -1,10 +1,12 @@ import os import re +import uuid import math import yaml import requests import markdownify + import folder_paths @@ -17,6 +19,7 @@ from io import BytesIO from . import utils from . import config +from . import thread class ModelSearcher(ABC): @@ -307,16 +310,38 @@ class Information: utils.print_error(error_msg) return web.json_response({"success": False, "error": error_msg}) + @routes.get("/model-manager/model-info/scan") + async def get_model_info_download_task(request): + """ + Get model information download task list. + """ + try: + result = self.get_scan_model_info_task_list() + if result is not None: + await self.download_model_info(request) + return web.json_response({"success": True, "data": result}) + except Exception as e: + error_msg = f"Get model info download task list failed: {str(e)}" + utils.print_error(error_msg) + return web.json_response({"success": False, "error": error_msg}) + @routes.post("/model-manager/model-info/scan") - async def download_model_info(request): + async def create_model_info_download_task(request): """ Create a task to download model information. + + - scanMode: The alternatives are diff and full. + - mode: The alternatives are diff and full. + - path: Scanning root path. """ post = await utils.get_request_body(request) try: + # TODO scanMode is deprecated, use mode instead. scan_mode = post.get("scanMode", "diff") - await self.download_model_info(scan_mode, request) - return web.json_response({"success": True}) + scan_mode = post.get("mode", scan_mode) + scan_path = post.get("path", None) + result = await self.create_scan_model_info_task(scan_mode, scan_path, request) + return web.json_response({"success": True, "data": result}) except Exception as e: error_msg = f"Download model info failed: {str(e)}" utils.print_error(error_msg) @@ -440,42 +465,76 @@ class Information: result = model_searcher.search_by_url(model_page) return result - async def download_model_info(self, scan_mode: str, request): - utils.print_info(f"Download model info for {scan_mode}") - model_base_paths = utils.resolve_model_base_paths() - for model_type in model_base_paths: + def get_scan_information_task_filepath(self): + download_dir = utils.get_download_path() + return utils.join_path(download_dir, "scan_information.task") - folders, *others = folder_paths.folder_names_and_paths[model_type] - for path_index, base_path in enumerate(folders): - files = utils.recursive_search_files(base_path, request) + def get_scan_model_info_task_list(self): + scan_info_task_file = self.get_scan_information_task_filepath() + if os.path.isfile(scan_info_task_file): + return utils.load_dict_pickle_file(scan_info_task_file) + return None - models = folder_paths.filter_files_extensions(files, folder_paths.supported_pt_extensions) + async def create_scan_model_info_task(self, scan_mode: str, scan_path: str | None, request): + scan_info_task_file = self.get_scan_information_task_filepath() + scan_info_task_content = {"mode": scan_mode} + scan_models: dict[str, bool] = {} - for fullname in models: - fullname = utils.normalize_path(fullname) - basename = os.path.splitext(fullname)[0] + scan_paths: list[str] = [] + if scan_path is None: + model_base_paths = utils.resolve_model_base_paths() + for model_type in model_base_paths: + folders, *others = folder_paths.folder_names_and_paths[model_type] + for path_index, base_path in enumerate(folders): + scan_paths.append(base_path) + else: + scan_paths = [scan_path] - abs_model_path = utils.join_path(base_path, fullname) + for base_path in scan_paths: + files = utils.recursive_search_files(base_path, request) + models = folder_paths.filter_files_extensions(files, folder_paths.supported_pt_extensions) + for fullname in models: + fullname = utils.normalize_path(fullname) + abs_model_path = utils.join_path(base_path, fullname) + utils.print_debug(f"Found model: {abs_model_path}") + scan_models[abs_model_path] = False - image_name = utils.get_model_preview_name(abs_model_path) - abs_image_path = utils.join_path(base_path, image_name) + scan_info_task_content["models"] = scan_models + utils.save_dict_pickle_file(scan_info_task_file, scan_info_task_content) + await self.download_model_info(request) + return scan_info_task_content - has_preview = os.path.isfile(abs_image_path) + download_thread_pool = thread.DownloadThreadPool() - description_name = utils.get_model_description_name(abs_model_path) - abs_description_path = utils.join_path(base_path, description_name) if description_name else None - has_description = os.path.isfile(abs_description_path) if abs_description_path else False + async def download_model_info(self, request): + async def download_information_task(task_id: str): + scan_info_task_file = self.get_scan_information_task_filepath() + scan_info_task_content = utils.load_dict_pickle_file(scan_info_task_file) + scan_mode = scan_info_task_content.get("mode", "diff") + scan_models: dict[str, bool] = scan_info_task_content.get("models", {}) + for key, value in scan_models.items(): + if value is True: + continue - try: + abs_model_path = key + base_path = os.path.dirname(abs_model_path) - utils.print_info(f"Checking model {abs_model_path}") - utils.print_debug(f"Scan mode: {scan_mode}") - utils.print_debug(f"Has preview: {has_preview}") - utils.print_debug(f"Has description: {has_description}") + image_name = utils.get_model_preview_name(abs_model_path) + abs_image_path = utils.join_path(base_path, image_name) - if scan_mode != "full" and (has_preview and has_description): - continue + has_preview = os.path.isfile(abs_image_path) + description_name = utils.get_model_description_name(abs_model_path) + abs_description_path = utils.join_path(base_path, description_name) if description_name else None + has_description = os.path.isfile(abs_description_path) if abs_description_path else False + + try: + utils.print_info(f"Checking model {abs_model_path}") + utils.print_debug(f"Scan mode: {scan_mode}") + utils.print_debug(f"Has preview: {has_preview}") + utils.print_debug(f"Has description: {has_description}") + + if scan_mode == "full" or not has_preview or not has_description: utils.print_debug(f"Calculate sha256 for {abs_model_path}") hash_value = utils.calculate_sha256(abs_model_path) utils.print_info(f"Searching model info by hash {hash_value}") @@ -490,10 +549,23 @@ class Information: description = model_info.get("description", None) if description: utils.save_model_description(abs_model_path, description) - except Exception as e: - utils.print_error(f"Failed to download model info for {abs_model_path}: {e}") - utils.print_debug("Completed scan model information.") + scan_models[abs_model_path] = True + scan_info_task_content["models"] = scan_models + utils.save_dict_pickle_file(scan_info_task_file, scan_info_task_content) + utils.print_debug(f"Send update scan information task to frontend.") + await utils.send_json("update_scan_information_task", scan_info_task_content) + except Exception as e: + utils.print_error(f"Failed to download model info for {abs_model_path}: {e}") + + os.remove(scan_info_task_file) + utils.print_info("Completed scan model information.") + + try: + task_id = uuid.uuid4().hex + self.download_thread_pool.submit(download_information_task, task_id) + except Exception as e: + utils.print_debug(str(e)) def get_model_searcher_by_url(self, url: str) -> ModelSearcher: parsed_url = urlparse(url) diff --git a/src/App.vue b/src/App.vue index 2954a21..45e7ae3 100644 --- a/src/App.vue +++ b/src/App.vue @@ -9,6 +9,7 @@ import DialogDownload from 'components/DialogDownload.vue' import DialogExplorer from 'components/DialogExplorer.vue' import DialogManager from 'components/DialogManager.vue' +import DialogScanning from 'components/DialogScanning.vue' import GlobalDialogStack from 'components/GlobalDialogStack.vue' import GlobalLoading from 'components/GlobalLoading.vue' import GlobalToast from 'components/GlobalToast.vue' @@ -35,6 +36,19 @@ onMounted(() => { }) } + const openModelScanning = () => { + dialog.open({ + key: 'model-information-scanning', + title: t('batchScanModelInformation'), + content: DialogScanning, + modal: true, + defaultSize: { + width: 680, + height: 490, + }, + }) + } + const openDownloadDialog = () => { dialog.open({ key: 'model-manager-download-list', @@ -64,6 +78,11 @@ onMounted(() => { content: flat.value ? DialogManager : DialogExplorer, keepAlive: true, headerButtons: [ + { + key: 'scanning', + icon: 'mdi mdi-folder-search-outline text-lg', + command: openModelScanning, + }, { key: 'refresh', icon: 'pi pi-refresh', diff --git a/src/components/DialogExplorer.vue b/src/components/DialogExplorer.vue index 49a67be..def83d2 100644 --- a/src/components/DialogExplorer.vue +++ b/src/components/DialogExplorer.vue @@ -133,7 +133,7 @@ import Button from 'primevue/button' import ConfirmDialog from 'primevue/confirmdialog' import ContextMenu from 'primevue/contextmenu' import InputText from 'primevue/inputtext' -import { MenuItem } from 'primevue/menuitem' +import type { MenuItem } from 'primevue/menuitem' import { genModelKey } from 'utils/model' import { computed, ref } from 'vue' import { useI18n } from 'vue-i18n' diff --git a/src/components/DialogScanning.vue b/src/components/DialogScanning.vue new file mode 100644 index 0000000..087158e --- /dev/null +++ b/src/components/DialogScanning.vue @@ -0,0 +1,271 @@ + + + + + + + + + + + + + {{ $t('selectModelType') }} + {{ $t('selectSubdirectory') }} + {{ $t('scanModelInformation') }} + + + + + + + + + + + + + + + + + + + + + + + + + + + + {{ $t('selectedAllPaths') }} + + + + {{ $t('selectedSpecialPath') }} + + + {{ selectedModelFolder }} + + + + + + + + + + + + + + + + + + {{ scanCompleteCount }} / {{ scanTotalCount }} + + + + + + {{ $t('noModelsInCurrentPath') }} + + + + + + + + diff --git a/src/hooks/config.ts b/src/hooks/config.ts index 2760329..fadb9dc 100644 --- a/src/hooks/config.ts +++ b/src/hooks/config.ts @@ -1,10 +1,8 @@ import SettingCardSize from 'components/SettingCardSize.vue' -import { request } from 'hooks/request' import { defineStore } from 'hooks/store' -import { $el, app, ComfyDialog } from 'scripts/comfyAPI' +import { app } from 'scripts/comfyAPI' import { computed, onMounted, onUnmounted, readonly, ref, watch } from 'vue' import { useI18n } from 'vue-i18n' -import { useToast } from './toast' export const useConfig = defineStore('config', (store) => { const { t } = useI18n() @@ -98,41 +96,8 @@ export const configSetting = { } function useAddConfigSettings(store: import('hooks/store').StoreProvider) { - const { toast } = useToast() const { t } = useI18n() - const confirm = (opts: { - message?: string - accept?: () => void - reject?: () => void - }) => { - const dialog = new ComfyDialog('div', []) - - dialog.show( - $el('div', [ - $el('p', { textContent: opts.message }), - $el('div.flex.gap-4', [ - $el('button.flex-1', { - textContent: 'Cancel', - onclick: () => { - opts.reject?.() - dialog.close() - document.body.removeChild(dialog.element) - }, - }), - $el('button.flex-1', { - textContent: 'Continue', - onclick: () => { - opts.accept?.() - dialog.close() - document.body.removeChild(dialog.element) - }, - }), - ]), - ]), - ) - } - onMounted(() => { // API keys app.ui?.settings.addSetting({ @@ -187,101 +152,6 @@ function useAddConfigSettings(store: import('hooks/store').StoreProvider) { }, }) - // Scan information - app.ui?.settings.addSetting({ - id: 'ModelManager.ScanFiles.Full', - category: [t('modelManager'), t('setting.scan'), 'Full'], - name: t('setting.scanAll'), - defaultValue: '', - type: () => { - return $el('button.p-button.p-component.p-button-secondary', { - textContent: 'Full Scan', - onclick: () => { - confirm({ - message: [ - 'This operation will override current files.', - 'This may take a while and generate MANY server requests!', - 'USE AT YOUR OWN RISK! Continue?', - ].join('\n'), - accept: () => { - store.loading.loading.value = true - request('/model-info/scan', { - method: 'POST', - body: JSON.stringify({ scanMode: 'full' }), - }) - .then(() => { - toast.add({ - severity: 'success', - summary: 'Complete download information', - life: 2000, - }) - store.models.refresh() - }) - .catch((err) => { - toast.add({ - severity: 'error', - summary: 'Error', - detail: err.message ?? 'Failed to download information', - life: 15000, - }) - }) - .finally(() => { - store.loading.loading.value = false - }) - }, - }) - }, - }) - }, - }) - - app.ui?.settings.addSetting({ - id: 'ModelManager.ScanFiles.Incremental', - category: [t('modelManager'), t('setting.scan'), 'Incremental'], - name: t('setting.scanMissing'), - defaultValue: '', - type: () => { - return $el('button.p-button.p-component.p-button-secondary', { - textContent: 'Diff Scan', - onclick: () => { - confirm({ - message: [ - 'Download missing information or preview.', - 'This may take a while and generate MANY server requests!', - 'USE AT YOUR OWN RISK! Continue?', - ].join('\n'), - accept: () => { - store.loading.loading.value = true - request('/model-info/scan', { - method: 'POST', - body: JSON.stringify({ scanMode: 'diff' }), - }) - .then(() => { - toast.add({ - severity: 'success', - summary: 'Complete download information', - life: 2000, - }) - store.models.refresh() - }) - .catch((err) => { - toast.add({ - severity: 'error', - summary: 'Error', - detail: err.message ?? 'Failed to download information', - life: 15000, - }) - }) - .finally(() => { - store.loading.loading.value = false - }) - }, - }) - }, - }) - }, - }) - app.ui?.settings.addSetting({ id: configSetting.excludeScanTypes, category: [t('modelManager'), t('setting.scan'), 'ExcludeScanTypes'], diff --git a/src/hooks/model.ts b/src/hooks/model.ts index 8f483e0..df3fcdb 100644 --- a/src/hooks/model.ts +++ b/src/hooks/model.ts @@ -443,7 +443,7 @@ export const useModelBaseInfo = () => { export const useModelFolder = ( option: { - type?: MaybeRefOrGetter + type?: MaybeRefOrGetter } = {}, ) => { const { data: models, folders: modelFolders } = useModels() diff --git a/src/i18n.ts b/src/i18n.ts index ba808d5..509fa4b 100644 --- a/src/i18n.ts +++ b/src/i18n.ts @@ -29,6 +29,18 @@ const messages = { width: 'Width', height: 'Height', reset: 'Reset', + back: 'Back', + next: 'Next', + batchScanModelInformation: 'Batch scan model information', + modelInformationScanning: 'Scanning model information', + selectModelType: 'Select model type', + selectSubdirectory: 'Select subdirectory', + scanModelInformation: 'Scan model information', + selectedAllPaths: 'Selected all model paths', + selectedSpecialPath: 'Selected special path', + scanMissInformation: 'Download missing information', + scanFullInformation: 'Override full information', + noModelsInCurrentPath: 'There are no models available in the current path', sort: { name: 'Name', size: 'Largest', @@ -92,6 +104,18 @@ const messages = { width: '宽度', height: '高度', reset: '重置', + back: '返回', + next: '下一步', + batchScanModelInformation: '批量扫描模型信息', + modelInformationScanning: '扫描模型信息', + selectModelType: '选择模型类型', + selectSubdirectory: '选择子目录', + scanModelInformation: '扫描模型信息', + selectedAllPaths: '已选所有模型路径', + selectedSpecialPath: '已选指定路径', + scanMissInformation: '下载缺失信息', + scanFullInformation: '覆盖所有信息', + noModelsInCurrentPath: '当前路径中没有可用的模型', sort: { name: '名称', size: '最大',