mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-25 23:25:43 -03:00
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
This commit is contained in:
@@ -5,6 +5,7 @@ from .routes.lora_routes import LoraRoutes
|
|||||||
from .routes.api_routes import ApiRoutes
|
from .routes.api_routes import ApiRoutes
|
||||||
from .services.lora_scanner import LoraScanner
|
from .services.lora_scanner import LoraScanner
|
||||||
from .services.file_monitor import LoraFileMonitor
|
from .services.file_monitor import LoraFileMonitor
|
||||||
|
from .services.lora_cache import LoraCache
|
||||||
|
|
||||||
class LoraManager:
|
class LoraManager:
|
||||||
"""Main entry point for LoRA Manager plugin"""
|
"""Main entry point for LoRA Manager plugin"""
|
||||||
@@ -45,8 +46,8 @@ class LoraManager:
|
|||||||
async def _schedule_cache_init(cls, scanner: LoraScanner):
|
async def _schedule_cache_init(cls, scanner: LoraScanner):
|
||||||
"""Schedule cache initialization in the running event loop"""
|
"""Schedule cache initialization in the running event loop"""
|
||||||
try:
|
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:
|
except Exception as e:
|
||||||
print(f"LoRA Manager: Error scheduling cache initialization: {e}")
|
print(f"LoRA Manager: Error scheduling cache initialization: {e}")
|
||||||
|
|
||||||
@@ -54,6 +55,15 @@ class LoraManager:
|
|||||||
async def _initialize_cache(cls, scanner: LoraScanner):
|
async def _initialize_cache(cls, scanner: LoraScanner):
|
||||||
"""Initialize cache in background"""
|
"""Initialize cache in background"""
|
||||||
try:
|
try:
|
||||||
|
# 设置初始缓存占位
|
||||||
|
scanner._cache = LoraCache(
|
||||||
|
raw_data=[],
|
||||||
|
sorted_by_name=[],
|
||||||
|
sorted_by_date=[],
|
||||||
|
folders=[]
|
||||||
|
)
|
||||||
|
|
||||||
|
# 分阶段加载缓存
|
||||||
await scanner.get_cached_data(force_refresh=True)
|
await scanner.get_cached_data(force_refresh=True)
|
||||||
print("LoRA Manager: Cache initialization completed")
|
print("LoRA Manager: Cache initialization completed")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -42,6 +42,15 @@ class LoraScanner:
|
|||||||
"""Get cached LoRA data, refresh if needed"""
|
"""Get cached LoRA data, refresh if needed"""
|
||||||
async with self._initialization_lock:
|
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():
|
if self._initialization_task and not self._initialization_task.done():
|
||||||
try:
|
try:
|
||||||
@@ -120,12 +129,18 @@ class LoraScanner:
|
|||||||
"""Scan all LoRA directories and return metadata"""
|
"""Scan all LoRA directories and return metadata"""
|
||||||
all_loras = []
|
all_loras = []
|
||||||
|
|
||||||
|
# 分目录异步扫描
|
||||||
|
scan_tasks = []
|
||||||
for loras_root in config.loras_roots:
|
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:
|
try:
|
||||||
loras = await self._scan_directory(loras_root)
|
loras = await task
|
||||||
all_loras.extend(loras)
|
all_loras.extend(loras)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error scanning directory {loras_root}: {e}")
|
logger.error(f"Error scanning directory: {e}")
|
||||||
|
|
||||||
return all_loras
|
return all_loras
|
||||||
|
|
||||||
@@ -133,18 +148,33 @@ class LoraScanner:
|
|||||||
"""Scan a single directory for LoRA files"""
|
"""Scan a single directory for LoRA files"""
|
||||||
loras = []
|
loras = []
|
||||||
|
|
||||||
for root, _, files in os.walk(root_path):
|
# 使用异步安全的目录遍历方式
|
||||||
for filename in (f for f in files if f.endswith('.safetensors')):
|
async def scan_recursive(path: str):
|
||||||
try:
|
try:
|
||||||
file_path = os.path.join(root, filename).replace(os.sep, "/")
|
with os.scandir(path) as it:
|
||||||
lora_data = await self._process_lora_file(file_path, root_path)
|
entries = list(it) # 同步获取目录条目
|
||||||
if lora_data:
|
for entry in entries:
|
||||||
loras.append(lora_data)
|
if entry.is_file() and entry.name.endswith('.safetensors'):
|
||||||
except Exception as e:
|
file_path = entry.path.replace(os.sep, "/")
|
||||||
logger.error(f"Error processing {filename}: {e}")
|
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
|
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:
|
async def _process_lora_file(self, file_path: str, root_path: str) -> Dict:
|
||||||
"""Process a single LoRA file and return its metadata"""
|
"""Process a single LoRA file and return its metadata"""
|
||||||
# Try loading existing metadata
|
# Try loading existing metadata
|
||||||
|
|||||||
Reference in New Issue
Block a user