Improve LoRA cache initialization with async locking and background task

This commit is contained in:
Will Miao
2025-02-02 23:51:36 +08:00
parent 8e8b80ddcf
commit 9aa0dcfb2b
3 changed files with 79 additions and 25 deletions

View File

@@ -1,7 +1,9 @@
import asyncio
from server import PromptServer # type: ignore
from .config import config
from .routes.lora_routes import LoraRoutes
from .routes.api_routes import ApiRoutes
from .services.lora_scanner import LoraScanner
class LoraManager:
"""Main entry point for LoRA Manager plugin"""
@@ -20,5 +22,29 @@ class LoraManager:
app.router.add_static('/loras_static', config.static_path)
# Setup feature routes
routes = LoraRoutes()
api_routes = ApiRoutes()
LoraRoutes.setup_routes(app)
ApiRoutes.setup_routes(app)
ApiRoutes.setup_routes(app)
# Schedule cache initialization using the application's startup handler
app.on_startup.append(lambda app: cls._schedule_cache_init(routes.scanner))
@classmethod
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))
except Exception as e:
print(f"LoRA Manager: Error scheduling cache initialization: {e}")
@classmethod
async def _initialize_cache(cls, scanner: LoraScanner):
"""Initialize cache in background"""
try:
await scanner.get_cached_data(force_refresh=True)
print("LoRA Manager: Cache initialization completed")
except Exception as e:
print(f"LoRA Manager: Error initializing cache: {e}")