mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-22 13:42:12 -03:00
112 lines
4.3 KiB
Python
112 lines
4.3 KiB
Python
import os
|
|
import logging
|
|
import asyncio
|
|
from typing import List, Dict, Optional, Set
|
|
import folder_paths # type: ignore
|
|
|
|
from ..utils.models import CheckpointMetadata
|
|
from ..config import config
|
|
from .model_scanner import ModelScanner
|
|
from .model_hash_index import ModelHashIndex
|
|
from .service_registry import ServiceRegistry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class CheckpointScanner(ModelScanner):
|
|
"""Service for scanning and managing checkpoint files"""
|
|
|
|
_instance = None
|
|
_lock = asyncio.Lock()
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super().__new__(cls)
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
if not hasattr(self, '_initialized'):
|
|
# Define supported file extensions
|
|
file_extensions = {'.safetensors', '.ckpt', '.pt', '.pth', '.sft', '.gguf'}
|
|
super().__init__(
|
|
model_type="checkpoint",
|
|
model_class=CheckpointMetadata,
|
|
file_extensions=file_extensions,
|
|
hash_index=ModelHashIndex()
|
|
)
|
|
self._initialized = True
|
|
|
|
@classmethod
|
|
async def get_instance(cls):
|
|
"""Get singleton instance with async support"""
|
|
async with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
return cls._instance
|
|
|
|
def get_model_roots(self) -> List[str]:
|
|
"""Get checkpoint root directories"""
|
|
return config.base_models_roots
|
|
|
|
async def scan_all_models(self) -> List[Dict]:
|
|
"""Scan all checkpoint directories and return metadata"""
|
|
all_checkpoints = []
|
|
|
|
# Create scan tasks for each directory
|
|
scan_tasks = []
|
|
for root in self.get_model_roots():
|
|
task = asyncio.create_task(self._scan_directory(root))
|
|
scan_tasks.append(task)
|
|
|
|
# Wait for all tasks to complete
|
|
for task in scan_tasks:
|
|
try:
|
|
checkpoints = await task
|
|
all_checkpoints.extend(checkpoints)
|
|
except Exception as e:
|
|
logger.error(f"Error scanning checkpoint directory: {e}")
|
|
|
|
return all_checkpoints
|
|
|
|
async def _scan_directory(self, root_path: str) -> List[Dict]:
|
|
"""Scan a directory for checkpoint files"""
|
|
checkpoints = []
|
|
original_root = root_path
|
|
|
|
async def scan_recursive(path: str, visited_paths: set):
|
|
try:
|
|
real_path = os.path.realpath(path)
|
|
if real_path in visited_paths:
|
|
logger.debug(f"Skipping already visited path: {path}")
|
|
return
|
|
visited_paths.add(real_path)
|
|
|
|
with os.scandir(path) as it:
|
|
entries = list(it)
|
|
for entry in entries:
|
|
try:
|
|
if entry.is_file(follow_symlinks=True):
|
|
# Check if file has supported extension
|
|
ext = os.path.splitext(entry.name)[1].lower()
|
|
if ext in self.file_extensions:
|
|
file_path = entry.path.replace(os.sep, "/")
|
|
await self._process_single_file(file_path, original_root, checkpoints)
|
|
await asyncio.sleep(0)
|
|
elif entry.is_dir(follow_symlinks=True):
|
|
# For directories, continue scanning with original path
|
|
await scan_recursive(entry.path, visited_paths)
|
|
except Exception as e:
|
|
logger.error(f"Error processing entry {entry.path}: {e}")
|
|
except Exception as e:
|
|
logger.error(f"Error scanning {path}: {e}")
|
|
|
|
await scan_recursive(root_path, set())
|
|
return checkpoints
|
|
|
|
async def _process_single_file(self, file_path: str, root_path: str, checkpoints: list):
|
|
"""Process a single checkpoint file and add to results"""
|
|
try:
|
|
result = await self._process_model_file(file_path, root_path)
|
|
if result:
|
|
checkpoints.append(result)
|
|
except Exception as e:
|
|
logger.error(f"Error processing {file_path}: {e}") |