checkpoint

py/lora_manager.py

@@ -7,10 +7,12 @@ from .routes.api_routes import ApiRoutes
 from .routes.recipe_routes import RecipeRoutes
 from .routes.checkpoints_routes import CheckpointsRoutes
 from .services.lora_scanner import LoraScanner
+from .services.checkpoint_scanner import CheckpointScanner
 from .services.recipe_scanner import RecipeScanner
-from .services.file_monitor import LoraFileMonitor
+from .services.file_monitor import LoraFileMonitor, CheckpointFileMonitor
 from .services.lora_cache import LoraCache
 from .services.recipe_cache import RecipeCache
+from .services.model_cache import ModelCache
 import logging
 
 logger = logging.getLogger(__name__)
@@ -23,7 +25,7 @@ class LoraManager:
         """Initialize and register all routes"""
         app = PromptServer.instance.app
 
-        added_targets = set()  # 用于跟踪已添加的目标路径
+        added_targets = set()  # Track already added target paths
         
         # Add static routes for each lora root
         for idx, root in enumerate(config.loras_roots, start=1):
@@ -35,15 +37,34 @@ class LoraManager:
                     if link == root:
                         real_root = target
                         break
-            # 为原始路径添加静态路由
+            # Add static route for original path
             app.router.add_static(preview_path, real_root)
             logger.info(f"Added static route {preview_path} -> {real_root}")
             
-            # 记录路由映射
+            # Record route mapping
             config.add_route_mapping(real_root, preview_path)
             added_targets.add(real_root)
         
-        # 为符号链接的目标路径添加额外的静态路由
+        # Add static routes for each checkpoint root
+        checkpoint_scanner = CheckpointScanner()
+        for idx, root in enumerate(checkpoint_scanner.get_model_roots(), start=1):
+            preview_path = f'/checkpoints_static/root{idx}/preview'
+            
+            real_root = root
+            if root in config._path_mappings.values():
+                for target, link in config._path_mappings.items():
+                    if link == root:
+                        real_root = target
+                        break
+            # Add static route for original path
+            app.router.add_static(preview_path, real_root)
+            logger.info(f"Added static route {preview_path} -> {real_root}")
+            
+            # Record route mapping
+            config.add_route_mapping(real_root, preview_path)
+            added_targets.add(real_root)
+        
+        # Add static routes for symlink target paths
         link_idx = 1
         
         for target_path, link_path in config._path_mappings.items():
@@ -59,37 +80,47 @@ class LoraManager:
         app.router.add_static('/loras_static', config.static_path)
         
         # Setup feature routes
-        routes = LoraRoutes()
+        lora_routes = LoraRoutes()
         checkpoints_routes = CheckpointsRoutes()
         
         # Setup file monitoring
-        monitor = LoraFileMonitor(routes.scanner, config.loras_roots)
-        monitor.start()
+        lora_monitor = LoraFileMonitor(lora_routes.scanner, config.loras_roots)
+        lora_monitor.start()
         
-        routes.setup_routes(app)
+        checkpoint_monitor = CheckpointFileMonitor(checkpoints_routes.scanner, checkpoints_routes.scanner.get_model_roots())
+        checkpoint_monitor.start()
+        
+        lora_routes.setup_routes(app)
         checkpoints_routes.setup_routes(app)
-        ApiRoutes.setup_routes(app, monitor)
+        ApiRoutes.setup_routes(app, lora_monitor)
         RecipeRoutes.setup_routes(app)
         
-        # Store monitor in app for cleanup
-        app['lora_monitor'] = monitor
+        # Store monitors in app for cleanup
+        app['lora_monitor'] = lora_monitor
+        app['checkpoint_monitor'] = checkpoint_monitor
+
+        logger.info("PromptServer app: ", app)
         
         # Schedule cache initialization using the application's startup handler
-        app.on_startup.append(lambda app: cls._schedule_cache_init(routes.scanner, routes.recipe_scanner))
+        app.on_startup.append(lambda app: cls._schedule_cache_init(
+            lora_routes.scanner, 
+            checkpoints_routes.scanner,
+            lora_routes.recipe_scanner
+        ))
         
         # Add cleanup
         app.on_shutdown.append(cls._cleanup)
         app.on_shutdown.append(ApiRoutes.cleanup)
     
     @classmethod
-    async def _schedule_cache_init(cls, scanner: LoraScanner, recipe_scanner: RecipeScanner):
+    async def _schedule_cache_init(cls, lora_scanner, checkpoint_scanner, recipe_scanner):
         """Schedule cache initialization in the running event loop"""
         try:
-            # 创建低优先级的初始化任务
-            lora_task = asyncio.create_task(cls._initialize_lora_cache(scanner), name='lora_cache_init')
-            
-            # Schedule recipe cache initialization with a delay to let lora scanner initialize first
-            recipe_task = asyncio.create_task(cls._initialize_recipe_cache(recipe_scanner, delay=2), name='recipe_cache_init')
+            # Create low-priority initialization tasks
+            lora_task = asyncio.create_task(cls._initialize_lora_cache(lora_scanner), name='lora_cache_init')
+            checkpoint_task = asyncio.create_task(cls._initialize_checkpoint_cache(checkpoint_scanner), name='checkpoint_cache_init')
+            recipe_task = asyncio.create_task(cls._initialize_recipe_cache(recipe_scanner), name='recipe_cache_init')
+            logger.info("Cache initialization tasks scheduled to run in background")
         except Exception as e:
             logger.error(f"LoRA Manager: Error scheduling cache initialization: {e}")
     
@@ -97,26 +128,45 @@ class LoraManager:
     async def _initialize_lora_cache(cls, scanner: LoraScanner):
         """Initialize lora cache in background"""
         try:
-            # 设置初始缓存占位
+            # Set initial placeholder cache
             scanner._cache = LoraCache(
                 raw_data=[],
                 sorted_by_name=[],
                 sorted_by_date=[],
                 folders=[]
             )
-            
-            # 分阶段加载缓存
-            await scanner.get_cached_data(force_refresh=True)
+            # 使用线程池执行耗时操作
+            loop = asyncio.get_event_loop()
+            await loop.run_in_executor(
+                None,  # 使用默认线程池
+                lambda: scanner.get_cached_data_sync(force_refresh=True)  # 创建同步版本的方法
+            )
+            # Load cache in phases
+            # await scanner.get_cached_data(force_refresh=True)
         except Exception as e:
             logger.error(f"LoRA Manager: Error initializing lora cache: {e}")
     
     @classmethod
-    async def _initialize_recipe_cache(cls, scanner: RecipeScanner, delay: float = 2.0):
-        """Initialize recipe cache in background with a delay"""
+    async def _initialize_checkpoint_cache(cls, scanner: CheckpointScanner):
+        """Initialize checkpoint cache in background"""
         try:
-            # Wait for the specified delay to let lora scanner initialize first
-            await asyncio.sleep(delay)
+            # Set initial placeholder cache
+            scanner._cache = ModelCache(
+                raw_data=[],
+                sorted_by_name=[],
+                sorted_by_date=[],
+                folders=[]
+            )
             
+            # Load cache in phases
+            await scanner.get_cached_data(force_refresh=True)
+        except Exception as e:
+            logger.error(f"LoRA Manager: Error initializing checkpoint cache: {e}")
+    
+    @classmethod
+    async def _initialize_recipe_cache(cls, scanner: RecipeScanner):
+        """Initialize recipe cache in background with a delay"""
+        try:          
             # Set initial empty cache
             scanner._cache = RecipeCache(
                 raw_data=[],
@@ -134,3 +184,6 @@ class LoraManager:
         """Cleanup resources"""
         if 'lora_monitor' in app:
             app['lora_monitor'].stop()
+            
+        if 'checkpoint_monitor' in app:
+            app['checkpoint_monitor'].stop()