From 5f3c515323b27a95bbeaed6da97cf10e996f099a Mon Sep 17 00:00:00 2001
From: Will Miao <13051207myq@gmail.com>
Date: Sat, 12 Apr 2025 19:06:17 +0800
Subject: [PATCH] feat: Enhance checkpoint handling by initializing paths and
 adding static routes

---
 py/config.py       | 35 ++++++++++++++++++++++++++++++++++-
 py/lora_manager.py | 23 ++++++++++++++++-------
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/py/config.py b/py/config.py
index b0ca1f4c..f5ce927f 100644
--- a/py/config.py
+++ b/py/config.py
@@ -17,6 +17,7 @@ class Config:
         # 静态路由映射字典, target to route mapping
         self._route_mappings = {}
         self.loras_roots = self._init_lora_paths()
+        self.checkpoints_roots = self._init_checkpoint_paths()
         self.temp_directory = folder_paths.get_temp_directory()
         # 在初始化时扫描符号链接
         self._scan_symbolic_links()
@@ -39,9 +40,12 @@ class Config:
             return False
 
     def _scan_symbolic_links(self):
-        """扫描所有 LoRA 根目录中的符号链接"""
+        """扫描所有 LoRA 和 Checkpoint 根目录中的符号链接"""
         for root in self.loras_roots:
             self._scan_directory_links(root)
+        
+        for root in self.checkpoints_roots:
+            self._scan_directory_links(root)
 
     def _scan_directory_links(self, root: str):
         """递归扫描目录中的符号链接"""
@@ -115,6 +119,35 @@ class Config:
         
         return paths
 
+    def _init_checkpoint_paths(self) -> List[str]:
+        """Initialize and validate checkpoint paths from ComfyUI settings"""
+        # Get checkpoint paths from folder_paths
+        checkpoint_paths = folder_paths.get_folder_paths("checkpoints")
+        diffusion_paths = folder_paths.get_folder_paths("diffusers")
+        unet_paths = folder_paths.get_folder_paths("unet")
+        
+        # Combine all checkpoint-related paths
+        all_paths = checkpoint_paths + diffusion_paths + unet_paths
+        
+        # Filter and normalize paths
+        paths = sorted(set(path.replace(os.sep, "/") 
+                for path in all_paths 
+                if os.path.exists(path)), key=lambda p: p.lower())
+        
+        print("Found checkpoint roots:", paths)
+        
+        if not paths:
+            logger.warning("No valid checkpoint folders found in ComfyUI configuration")
+            return []
+        
+        # 初始化路径映射，与 LoRA 路径处理方式相同
+        for path in paths:
+            real_path = os.path.normpath(os.path.realpath(path)).replace(os.sep, '/')
+            if real_path != path:
+                self.add_path_mapping(path, real_path)
+        
+        return paths
+
     def get_preview_static_url(self, preview_path: str) -> str:
         """Convert local preview path to static URL"""
         if not preview_path:
diff --git a/py/lora_manager.py b/py/lora_manager.py
index 3e9a3bbf..fae7e910 100644
--- a/py/lora_manager.py
+++ b/py/lora_manager.py
@@ -38,11 +38,8 @@ class LoraManager:
             config.add_route_mapping(real_root, preview_path)
             added_targets.add(real_root)
         
-        # Get checkpoint scanner instance
-        checkpoint_scanner = asyncio.run(ServiceRegistry.get_checkpoint_scanner())
-        
         # Add static routes for each checkpoint root
-        for idx, root in enumerate(checkpoint_scanner.get_model_roots(), start=1):
+        for idx, root in enumerate(config.checkpoints_roots, start=1):
             preview_path = f'/checkpoints_static/root{idx}/preview'
             
             real_root = root
@@ -60,16 +57,28 @@ class LoraManager:
             added_targets.add(real_root)
         
         # Add static routes for symlink target paths
-        link_idx = 1
+        link_idx = {
+            'lora': 1,
+            'checkpoint': 1
+        }
         
         for target_path, link_path in config._path_mappings.items():
             if target_path not in added_targets:
-                route_path = f'/loras_static/link_{link_idx}/preview'
+                # Determine if this is a checkpoint or lora link based on path
+                is_checkpoint = any(cp_root in link_path for cp_root in config.checkpoints_roots)
+                is_checkpoint = is_checkpoint or any(cp_root in target_path for cp_root in config.checkpoints_roots)
+                
+                if is_checkpoint:
+                    route_path = f'/checkpoints_static/link_{link_idx["checkpoint"]}/preview'
+                    link_idx["checkpoint"] += 1
+                else:
+                    route_path = f'/loras_static/link_{link_idx["lora"]}/preview'
+                    link_idx["lora"] += 1
+                
                 app.router.add_static(route_path, target_path)
                 logger.info(f"Added static route for link target {route_path} -> {target_path}")
                 config.add_route_mapping(target_path, route_path)
                 added_targets.add(target_path)
-                link_idx += 1
         
         # Add static route for plugin assets
         app.router.add_static('/loras_static', config.static_path)