From 40a1a7f43ad53d24e583a74c3cd7a7c520859937 Mon Sep 17 00:00:00 2001
From: Hayden <48267247+hayden-fr@users.noreply.github.com>
Date: Mon, 13 Jan 2025 15:15:32 +0800
Subject: [PATCH] feat: add exclude scan model types (#92)

---
 __init__.py         | 2 +-
 py/config.py        | 3 ++-
 py/utils.py         | 5 ++++-
 src/hooks/config.ts | 8 ++++++++
 src/hooks/model.ts  | 1 +
 src/i18n.ts         | 2 ++
 6 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/__init__.py b/__init__.py
index ff4612e..2935af2 100644
--- a/__init__.py
+++ b/__init__.py
@@ -130,7 +130,7 @@ async def list_model_types(request):
     Scan all models and read their information.
     """
     try:
-        result = utils.resolve_model_base_paths()
+        result = utils.resolve_model_base_paths(request)
         return web.json_response({"success": True, "data": result})
     except Exception as e:
         error_msg = f"Read models failed: {str(e)}"
diff --git a/py/config.py b/py/config.py
index 46471ce..f8e74ea 100644
--- a/py/config.py
+++ b/py/config.py
@@ -12,7 +12,8 @@ setting_key = {
         "max_task_count": "ModelManager.Download.MaxTaskCount",
     },
     "scan": {
-        "include_hidden_files": "ModelManager.Scan.IncludeHiddenFiles"
+        "include_hidden_files": "ModelManager.Scan.IncludeHiddenFiles",
+        "exclude_scan_types": "ModelManager.Scan.excludeScanTypes",
     },
 }
 
diff --git a/py/utils.py b/py/utils.py
index 275e636..680a588 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -116,10 +116,13 @@ def download_web_distribution(version: str):
         print_error(f"An unexpected error occurred: {e}")
 
 
-def resolve_model_base_paths():
+def resolve_model_base_paths(request):
     folders = list(folder_paths.folder_names_and_paths.keys())
     model_base_paths = {}
     folder_black_list = ["configs", "custom_nodes"]
+    custom_folders = get_setting_value(request, "scan.exclude_scan_types", "")
+    custom_black_list = [f.strip() for f in custom_folders.split(",") if f.strip()]
+    folder_black_list.extend(custom_black_list)
     for folder in folders:
         if folder in folder_black_list:
             continue
diff --git a/src/hooks/config.ts b/src/hooks/config.ts
index 0ed9f31..be84229 100644
--- a/src/hooks/config.ts
+++ b/src/hooks/config.ts
@@ -190,6 +190,14 @@ function useAddConfigSettings(store: import('hooks/store').StoreProvider) {
       },
     })
 
+    app.ui?.settings.addSetting({
+      id: 'ModelManager.Scan.excludeScanTypes',
+      category: [t('modelManager'), t('setting.scan'), 'ExcludeScanTypes'],
+      name: t('setting.excludeScanTypes'),
+      defaultValue: undefined,
+      type: 'text',
+    })
+
     app.ui?.settings.addSetting({
       id: 'ModelManager.Scan.IncludeHiddenFiles',
       category: [t('modelManager'), t('setting.scan'), 'IncludeHiddenFiles'],
diff --git a/src/hooks/model.ts b/src/hooks/model.ts
index 63c42d7..875c80e 100644
--- a/src/hooks/model.ts
+++ b/src/hooks/model.ts
@@ -55,6 +55,7 @@ export const useModels = defineStore('models', (store) => {
 
   const refreshAllModels = async (force = false) => {
     const forceRefresh = force ? refreshFolders() : Promise.resolve()
+    models.value = {}
     return forceRefresh.then(() =>
       Promise.allSettled(Object.keys(folders.value).map(refreshModels)),
     )
diff --git a/src/i18n.ts b/src/i18n.ts
index e5cedd3..d5088f4 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -45,6 +45,7 @@ const messages = {
       scanMissing: 'Download missing information or preview',
       scanAll: "Override all models' information and preview",
       includeHiddenFiles: 'Include hidden files(start with .)',
+      excludeScanTypes: 'Exclude scan types (separate with commas)',
     },
   },
   zh: {
@@ -90,6 +91,7 @@ const messages = {
       scanMissing: '下载缺失的信息或预览图片',
       scanAll: '覆盖所有模型信息和预览图片',
       includeHiddenFiles: '包含隐藏文件(以 . 开头的文件或文件夹)',
+      excludeScanTypes: '排除扫描类型(使用英文逗号隔开)',
     },
   },
 }