alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-24 06:32:12 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat: Enhance initialization component with progress tracking and UI improvements

Browse Source
This commit is contained in:
Will Miao
2025-04-13 12:58:38 +08:00
parent a043b487bd
commit 0b11e6e6d0
5 changed files with 414 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
py/services/model_scanner.py
View File

@@ -62,11 +62,16 @@ class ModelScanner:
# Set initializing flag to true
self._is_initializing = True
# Determine the page type based on model type
page_type = 'loras' if self.model_type == 'lora' else 'checkpoints'
# First, count all model files to track progress
await ws_manager.broadcast_init_progress({
'stage': 'scan_folders',
'progress': 0,
'details': f"Scanning {self.model_type} folders..."
'details': f"Scanning {self.model_type} folders...",
'scanner_type': self.model_type,
'pageType': page_type
})
# Count files in a separate thread to avoid blocking
@@ -79,7 +84,9 @@ class ModelScanner:
await ws_manager.broadcast_init_progress({
'stage': 'count_models',
'progress': 1, # Changed from 10 to 1
'details': f"Found {total_files} {self.model_type} files"
'details': f"Found {total_files} {self.model_type} files",
'scanner_type': self.model_type,
'pageType': page_type
})
start_time = time.time()
@@ -88,14 +95,17 @@ class ModelScanner:
await loop.run_in_executor(
None, # Use default thread pool
self._initialize_cache_sync, # Run synchronous version in thread
total_files # Pass the total file count for progress reporting
total_files, # Pass the total file count for progress reporting
page_type # Pass the page type for progress reporting
)
# Send final progress update
await ws_manager.broadcast_init_progress({
'stage': 'finalizing',
'progress': 99, # Changed from 95 to 99
'details': f"Finalizing {self.model_type} cache..."
'details': f"Finalizing {self.model_type} cache...",
'scanner_type': self.model_type,
'pageType': page_type
})
logger.info(f"{self.model_type.capitalize()} cache initialized in {time.time() - start_time:.2f} seconds. Found {len(self._cache.raw_data)} models")
@@ -106,7 +116,9 @@ class ModelScanner:
'stage': 'finalizing',
'progress': 100,
'status': 'complete',
'details': f"Completed! Found {len(self._cache.raw_data)} {self.model_type} files."
'details': f"Completed! Found {len(self._cache.raw_data)} {self.model_type} files.",
'scanner_type': self.model_type,
'pageType': page_type
})
except Exception as e:
@@ -154,7 +166,7 @@ class ModelScanner:
return total_files
def _initialize_cache_sync(self, total_files=0):
def _initialize_cache_sync(self, total_files=0, page_type='loras'):
"""Synchronous version of cache initialization for thread pool execution"""
try:
# Create a new event loop for this thread
@@ -222,7 +234,9 @@ class ModelScanner:
await ws_manager.broadcast_init_progress({
'stage': 'process_models',
'progress': progress_percent,
'details': f"Processing {self.model_type} files: {processed_files}/{total_files}"
'details': f"Processing {self.model_type} files: {processed_files}/{total_files}",
'scanner_type': self.model_type,
'pageType': page_type
})
elif entry.is_dir(follow_symlinks=True):
@@ -299,6 +313,9 @@ class ModelScanner:
# Clear existing tags count
self._tags_count = {}
# Determine the page type based on model type
page_type = 'loras' if self.model_type == 'lora' else 'checkpoints'
# Scan for new data
raw_data = await self.scan_all_models()

8
py/services/websocket_manager.py
View File

@@ -55,6 +55,14 @@ class WebSocketManager:
if not self._init_websockets:
return
# Ensure data has all required fields
if 'stage' not in data:
data['stage'] = 'processing'
if 'progress' not in data:
data['progress'] = 0
if 'details' not in data:
data['details'] = 'Processing...'
for ws in self._init_websockets:
try:
await ws.send_json(data)