From 78f16ad651da2131af1c578da48f2133549f322e Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Tue, 4 Feb 2025 19:03:09 +0800
Subject: [PATCH] Optimize LoRA scanning and caching with asynchronous
 improvements

- Implement asynchronous directory scanning with concurrent tasks
- Add fallback mechanism for uninitialized cache
- Improve error handling and logging during file processing
- Enhance file scanning with safer async recursive traversal
---
 lora_manager.py          | 14 +++++++++--
 services/lora_scanner.py | 54 +++++++++++++++++++++++++++++++---------
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/lora_manager.py b/lora_manager.py
index 1afa1a1c..d1b0d2ce 100644
--- a/lora_manager.py
+++ b/lora_manager.py
@@ -5,6 +5,7 @@ from .routes.lora_routes import LoraRoutes
 from .routes.api_routes import ApiRoutes
 from .services.lora_scanner import LoraScanner
 from .services.file_monitor import LoraFileMonitor
+from .services.lora_cache import LoraCache
 
 class LoraManager:
     """Main entry point for LoRA Manager plugin"""
@@ -45,8 +46,8 @@ class LoraManager:
     async def _schedule_cache_init(cls, scanner: LoraScanner):
         """Schedule cache initialization in the running event loop"""
         try:
-            # Create the initialization task
-            asyncio.create_task(cls._initialize_cache(scanner))
+            # 创建低优先级的初始化任务
+            asyncio.create_task(cls._initialize_cache(scanner), name='lora_cache_init')
         except Exception as e:
             print(f"LoRA Manager: Error scheduling cache initialization: {e}")
     
@@ -54,6 +55,15 @@ class LoraManager:
     async def _initialize_cache(cls, scanner: LoraScanner):
         """Initialize cache in background"""
         try:
+            # 设置初始缓存占位
+            scanner._cache = LoraCache(
+                raw_data=[],
+                sorted_by_name=[],
+                sorted_by_date=[],
+                folders=[]
+            )
+            
+            # 分阶段加载缓存
             await scanner.get_cached_data(force_refresh=True)
             print("LoRA Manager: Cache initialization completed")
         except Exception as e:
diff --git a/services/lora_scanner.py b/services/lora_scanner.py
index a63aad15..92c875ea 100644
--- a/services/lora_scanner.py
+++ b/services/lora_scanner.py
@@ -42,6 +42,15 @@ class LoraScanner:
         """Get cached LoRA data, refresh if needed"""
         async with self._initialization_lock:
             
+            # 如果缓存未初始化但需要响应请求，返回空缓存
+            if self._cache is None and not force_refresh:
+                return LoraCache(
+                    raw_data=[],
+                    sorted_by_name=[],
+                    sorted_by_date=[],
+                    folders=[]
+                )
+
             # 如果正在初始化，等待完成
             if self._initialization_task and not self._initialization_task.done():
                 try:
@@ -120,12 +129,18 @@ class LoraScanner:
         """Scan all LoRA directories and return metadata"""
         all_loras = []
         
+        # 分目录异步扫描
+        scan_tasks = []
         for loras_root in config.loras_roots:
+            task = asyncio.create_task(self._scan_directory(loras_root))
+            scan_tasks.append(task)
+            
+        for task in scan_tasks:
             try:
-                loras = await self._scan_directory(loras_root)
+                loras = await task
                 all_loras.extend(loras)
             except Exception as e:
-                logger.error(f"Error scanning directory {loras_root}: {e}")
+                logger.error(f"Error scanning directory: {e}")
                 
         return all_loras
 
@@ -133,18 +148,33 @@ class LoraScanner:
         """Scan a single directory for LoRA files"""
         loras = []
         
-        for root, _, files in os.walk(root_path):
-            for filename in (f for f in files if f.endswith('.safetensors')):
-                try:
-                    file_path = os.path.join(root, filename).replace(os.sep, "/")
-                    lora_data = await self._process_lora_file(file_path, root_path)
-                    if lora_data:
-                        loras.append(lora_data)
-                except Exception as e:
-                    logger.error(f"Error processing {filename}: {e}")
-                    
+        # 使用异步安全的目录遍历方式
+        async def scan_recursive(path: str):
+            try:
+                with os.scandir(path) as it:
+                    entries = list(it)  # 同步获取目录条目
+                    for entry in entries:
+                        if entry.is_file() and entry.name.endswith('.safetensors'):
+                            file_path = entry.path.replace(os.sep, "/")
+                            await self._process_single_file(file_path, root_path, loras)
+                            await asyncio.sleep(0)  # 释放事件循环
+                        elif entry.is_dir():
+                            await scan_recursive(entry.path)
+            except Exception as e:
+                logger.error(f"Error scanning {path}: {e}")
+
+        await scan_recursive(root_path)
         return loras
 
+    async def _process_single_file(self, file_path: str, root_path: str, loras: list):
+        """处理单个文件并添加到结果列表"""
+        try:
+            result = await self._process_lora_file(file_path, root_path)
+            if result:
+                loras.append(result)
+        except Exception as e:
+            logger.error(f"Error processing {file_path}: {e}")
+
     async def _process_lora_file(self, file_path: str, root_path: str) -> Dict:
         """Process a single LoRA file and return its metadata"""
         # Try loading existing metadata