pref: optimize scan cost (#117)

py/manager.py

@@ -1,6 +1,7 @@
 import os
 import folder_paths
 from aiohttp import web
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
 
 from . import utils
@@ -115,43 +116,57 @@ class ModelManager:
     def scan_models(self, folder: str, request):
         result = []
 
+        include_hidden_files = utils.get_setting_value(request, "scan.include_hidden_files", False)
         folders, extensions = folder_paths.folder_names_and_paths[folder]
+
+        def get_file_info(entry: os.DirEntry[str], base_path: str, path_index: int):
+            fullname = entry.path.replace(base_path, "")
+            basename = os.path.splitext(fullname)[0]
+            extension = os.path.splitext(fullname)[1]
+
+            if extension not in folder_paths.supported_pt_extensions:
+                return None
+
+            model_preview = f"/model-manager/preview/{folder}/{path_index}/{basename}.webp"
+
+            stat = entry.stat()
+            return {
+                "fullname": fullname,
+                "basename": basename,
+                "extension": extension,
+                "type": folder,
+                "pathIndex": path_index,
+                "sizeBytes": stat.st_size,
+                "preview": model_preview,
+                "createdAt": round(stat.st_ctime_ns / 1000000),
+                "updatedAt": round(stat.st_mtime_ns / 1000000),
+            }
+
+        def get_all_files_entry(directory: str):
+            files = []
+            with os.scandir(directory) as it:
+                for entry in it:
+                    # Skip hidden files
+                    if not include_hidden_files:
+                        if entry.name.startswith("."):
+                            continue
+                    if entry.is_dir():
+                        files.extend(get_all_files_entry(entry.path))
+                    elif entry.is_file():
+                        files.append(entry)
+            return files
+
         for path_index, base_path in enumerate(folders):
-            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)
-                basename = os.path.splitext(fullname)[0]
-                extension = os.path.splitext(fullname)[1]
-
-                abs_path = utils.join_path(base_path, fullname)
-                file_stats = os.stat(abs_path)
-
-                # Resolve preview
-                image_name = utils.get_model_preview_name(abs_path)
-                image_name = utils.join_path(os.path.dirname(fullname), image_name)
-                abs_image_path = utils.join_path(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/{folder}/{path_index}/{image_name}"
-
-                model_info = {
-                    "fullname": fullname,
-                    "basename": basename,
-                    "extension": extension,
-                    "type": folder,
-                    "pathIndex": path_index,
-                    "sizeBytes": file_stats.st_size,
-                    "preview": model_preview,
-                    "createdAt": round(file_stats.st_ctime_ns / 1000000),
-                    "updatedAt": round(file_stats.st_mtime_ns / 1000000),
-                }
-
-                result.append(model_info)
+            if not os.path.exists(base_path):
+                continue
+            file_entries = get_all_files_entry(base_path)
+            with ThreadPoolExecutor() as executor:
+                futures = {executor.submit(get_file_info, entry, base_path, path_index): entry for entry in file_entries}
+                for future in as_completed(futures):
+                    file_info = future.result()
+                    if file_info is None:
+                        continue
+                    result.append(file_info)
 
         return result