alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-09-20 18:51:26 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(ui): show live scan progress and ETA for cache refresh

Browse Source 
Broadcast typed scan_progress messages over /ws/fetch-progress from the
manual refresh/rebuild paths of ModelScanner and RecipeScanner, and
render percent, processed/total, current file name and an EMA-smoothed
ETA in the loading overlay. Hardcoded refresh strings move to i18n
(common.scanProgress); WS connection failure falls back to the previous
static loading behavior.
This commit is contained in:
Will Miao
2026-09-03 11:38:27 +08:00
parent da71985c3e
commit 14da8a6f17
19 changed files with 1668 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/services/model_scanner.py
+137 -16
View File
@@ -66,6 +66,14 @@ def _is_hidden_relative_path(rel_path: str) -> bool:
# requests (modal open + autocomplete) do not re-walk the model roots.
ALL_FOLDERS_CACHE_TTL_SECONDS = 5.0
# Maps a scanner model type to the manager page type used in progress
# broadcasts (e.g. 'lora' -> 'loras').
PAGE_TYPE_MAP = {
'lora': 'loras',
'checkpoint': 'checkpoints',
'embedding': 'embeddings',
}
def _is_pending_delete_path(path: str) -> bool:
"""Return True when any path component is the pending-delete staging dir."""
@@ -149,6 +157,38 @@ class ModelScanner:
# Register this service
asyncio.create_task(self._register_service())
@property
def page_type(self) -> str:
"""Manager page type used in progress broadcasts (e.g. 'loras')."""
return PAGE_TYPE_MAP.get(self.model_type, self.model_type)
async def _broadcast_scan_progress(
self,
status: str,
stage: str,
progress: int,
full_rebuild: bool,
**extra: Any,
) -> None:
"""Broadcast manual-refresh scan progress on the generic WS channel.
Best-effort only: broadcast failures must never affect the scan itself.
"""
payload: Dict[str, Any] = {
'type': 'scan_progress',
'status': status,
'model_type': self.model_type,
'pageType': self.page_type,
'stage': stage,
'full_rebuild': full_rebuild,
'progress': progress,
}
payload.update(extra)
try:
await ws_manager.broadcast(payload)
except Exception as exc: # pragma: no cover - defensive logging
logger.error(f"Error broadcasting scan progress for {self.model_type}: {exc}")
@property
def cache_version(self) -> int:
"""Monotonic version counter for the in-memory cache.
@@ -434,12 +474,7 @@ class ModelScanner:
self._is_initializing = True
# Determine the page type based on model type
page_type_map = {
'lora': 'loras',
'checkpoint': 'checkpoints',
'embedding': 'embeddings'
}
page_type = page_type_map.get(self.model_type, self.model_type)
page_type = self.page_type
# First, try to load from cache
await ws_manager.broadcast_init_progress({
@@ -804,7 +839,7 @@ class ModelScanner:
last_progress_time = time.time()
last_progress_percent = 0
async def progress_callback(processed_files: int, expected_total: int) -> None:
async def progress_callback(processed_files: int, expected_total: int, current_name: str = '') -> None:
nonlocal last_progress_time, last_progress_percent
if expected_total <= 0:
@@ -871,32 +906,84 @@ class ModelScanner:
async def _initialize_cache(self) -> None:
"""Initialize or refresh the cache"""
self._is_initializing = True # Set flag
last_progress_percent = 0
try:
start_time = time.time()
await self._broadcast_scan_progress('started', 'scan_folders', 0, True)
# Manually trigger a symlink rescan during a full rebuild.
# This ensures that any new symlink mappings are correctly picked up.
config.rebuild_symlink_cache()
# Determine the page type based on model type
# Count files in a thread so the event loop stays responsive
loop = asyncio.get_running_loop()
total_files = await loop.run_in_executor(None, self._count_model_files)
await self._broadcast_scan_progress(
'processing', 'count_models', 1, True,
processed=0, total=total_files,
)
last_progress_time = time.time()
async def progress_callback(processed_files: int, expected_total: int, current_name: str = '') -> None:
nonlocal last_progress_time, last_progress_percent
if expected_total <= 0:
return
current_time = time.time()
progress_percent = min(99, int(1 + (processed_files / expected_total) * 98))
if progress_percent <= last_progress_percent:
return
if current_time - last_progress_time <= 0.5 and processed_files != expected_total:
return
last_progress_percent = progress_percent
last_progress_time = current_time
await self._broadcast_scan_progress(
'processing', 'process_models', progress_percent, True,
processed=processed_files, total=expected_total,
current_name=current_name,
)
# Scan for new data
scan_result = await self._gather_model_data()
scan_result = await self._gather_model_data(
total_files=total_files,
progress_callback=progress_callback,
)
if not self.is_cancelled():
await self._broadcast_scan_progress('finalizing', 'finalizing', 99, True)
await self._apply_scan_result(scan_result)
await self._save_persistent_cache(scan_result)
await self._sync_download_history(scan_result.raw_data, source='scan')
await self._broadcast_scan_progress(
'completed', 'finalizing', 100, True,
elapsed_seconds=time.time() - start_time,
)
logger.info(
f"{self.model_type.capitalize()} Scanner: Cache initialization completed in {time.time() - start_time:.2f} seconds, "
f"found {len(scan_result.raw_data)} models"
)
else:
await self._broadcast_scan_progress(
'cancelled', 'process_models', last_progress_percent, True,
elapsed_seconds=time.time() - start_time,
)
logger.info(
f"{self.model_type.capitalize()} Scanner: Cache initialization cancelled "
f"after {time.time() - start_time:.2f} seconds"
)
except Exception as e:
logger.error(f"{self.model_type.capitalize()} Scanner: Error initializing cache: {e}")
await self._broadcast_scan_progress(
'error', 'process_models', last_progress_percent, True,
error=str(e),
)
# Ensure cache is at least an empty structure on error
if self._cache is None:
self._cache = ModelCache(
@@ -914,6 +1001,8 @@ class ModelScanner:
try:
start_time = time.time()
logger.info(f"{self.model_type.capitalize()} Scanner: Starting fast cache reconciliation...")
await self._broadcast_scan_progress('started', 'reconcile_scan', 0, False)
# Get current cached file paths
cached_paths = {item['file_path'] for item in self._cache.raw_data}
@@ -987,6 +1076,10 @@ class ModelScanner:
await asyncio.sleep(0)
if self.is_cancelled():
logger.info(f"{self.model_type.capitalize()} Scanner: Reconcile scan cancelled")
await self._broadcast_scan_progress(
'cancelled', 'reconcile_scan', 0, False,
elapsed_seconds=time.time() - start_time,
)
return
# Process new files in batches
@@ -994,10 +1087,14 @@ class ModelScanner:
if new_files:
logger.info(f"{self.model_type.capitalize()} Scanner: Found {len(new_files)} new files to process")
batch_size = 50
for i in range(0, len(new_files), batch_size):
total_new = len(new_files)
processed_new = 0
last_progress_time = time.time()
for i in range(0, total_new, batch_size):
batch = new_files[i:i+batch_size]
for path in batch:
logger.info(f"{self.model_type.capitalize()} Scanner: Processing {path}")
processed_new += 1
try:
# Find the appropriate root path for this file
root_path = None
@@ -1053,9 +1150,24 @@ class ModelScanner:
logger.error(f"Could not determine root path for {path}")
except Exception as e:
logger.error(f"Error adding {path} to cache: {e}")
current_time = time.time()
if current_time - last_progress_time > 0.5 or processed_new == total_new:
last_progress_time = current_time
await self._broadcast_scan_progress(
'processing', 'process_new',
min(99, int(1 + (processed_new / total_new) * 98)), False,
processed=processed_new, total=total_new,
current_name=os.path.basename(path),
)
if self.is_cancelled():
logger.info(f"{self.model_type.capitalize()} Scanner: Reconcile processing cancelled")
await self._broadcast_scan_progress(
'cancelled', 'process_new',
min(99, int(1 + (processed_new / total_new) * 98)), False,
elapsed_seconds=time.time() - start_time,
)
return
# Find missing files (in cache but not in filesystem)
@@ -1121,8 +1233,17 @@ class ModelScanner:
await self._persist_current_cache()
logger.info(f"{self.model_type.capitalize()} Scanner: Cache reconciliation completed in {time.time() - start_time:.2f} seconds. Added {total_added}, removed {total_removed} models.")
await self._broadcast_scan_progress(
'completed', 'process_new', 100, False,
added=total_added, removed=total_removed,
elapsed_seconds=time.time() - start_time,
)
except Exception as e:
logger.error(f"{self.model_type.capitalize()} Scanner: Error reconciling cache: {e}", exc_info=True)
await self._broadcast_scan_progress(
'error', 'reconcile_scan', 0, False,
error=str(e),
)
finally:
self._is_initializing = False # Unset flag
self.bump_cache_version()
@@ -1498,7 +1619,7 @@ class ModelScanner:
self,
*,
total_files: int = 0,
progress_callback: Optional[Callable[[int, int], Awaitable[None]]] = None
progress_callback: Optional[Callable[[int, int, str], Awaitable[None]]] = None
) -> CacheBuildResult:
"""Collect metadata for all model files."""
@@ -1510,11 +1631,11 @@ class ModelScanner:
processed_real_files: Set[str] = set()
visited_real_dirs: Set[str] = set()
async def handle_progress() -> None:
async def handle_progress(current_name: str = '') -> None:
if progress_callback is None:
return
try:
await progress_callback(processed_files, total_files)
await progress_callback(processed_files, total_files, current_name)
except Exception as exc: # pragma: no cover - defensive logging
logger.error(f"Error reporting progress for {self.model_type}: {exc}")
@@ -1580,7 +1701,7 @@ class ModelScanner:
for tag in result.get('tags') or []:
tags_count[tag] = tags_count.get(tag, 0) + 1
await handle_progress()
await handle_progress(entry.name)
await asyncio.sleep(0)
if self.is_cancelled():
return
py/services/recipe_scanner.py
+102 -4
View File
@@ -1753,7 +1753,36 @@ class RecipeScanner:
# Mark initialization as complete regardless of outcome
self._is_initializing = False
def _initialize_recipe_cache_sync(self):
async def _broadcast_scan_progress(
self,
status: str,
stage: str,
progress: int,
full_rebuild: bool,
**extra: Any,
) -> None:
"""Broadcast manual-refresh scan progress on the generic WS channel.
Mirrors ``ModelScanner._broadcast_scan_progress`` so the recipes page
can reuse the same frontend contract. Best-effort only: broadcast
failures must never affect the scan itself.
"""
payload: Dict[str, Any] = {
'type': 'scan_progress',
'status': status,
'model_type': 'recipe',
'pageType': 'recipes',
'stage': stage,
'full_rebuild': full_rebuild,
'progress': progress,
}
payload.update(extra)
try:
await ws_manager.broadcast(payload)
except Exception as exc: # pragma: no cover - defensive logging
logger.error(f"Error broadcasting scan progress for recipe: {exc}")
def _initialize_recipe_cache_sync(self, report_progress: bool = False):
"""Synchronous version of recipe cache initialization for thread pool execution.
Uses persistent cache for fast startup when available:
@@ -1761,8 +1790,14 @@ class RecipeScanner:
2. Reconcile with filesystem (check mtime/size for changes)
3. Fall back to full directory scan if cache miss or reconciliation fails
4. Persist results for next startup
Args:
report_progress: When True (manual force-refresh only), broadcast
scan_progress messages during the full directory scan. Startup
initialization leaves this False and behaves as before.
"""
loop = None
scan_start_time: Optional[float] = None
try:
# Ensure cache exists to avoid None reference errors
if self._cache is None:
@@ -1844,7 +1879,17 @@ class RecipeScanner:
# Fall back to full directory scan
logger.info("Recipe cache miss: performing full directory scan")
recipes, json_paths = self._full_directory_scan_sync(recipes_dir)
if report_progress:
scan_start_time = time.time()
# Broadcast from the worker thread via its own event loop,
# mirroring ModelScanner._initialize_cache_sync.
loop.run_until_complete(
self._broadcast_scan_progress('started', 'scan_folders', 0, True)
)
recipes, json_paths = self._full_directory_scan_sync(
recipes_dir,
progress_loop=loop if report_progress else None,
)
self._json_path_map = json_paths
# Update cache with the collected data
@@ -1858,12 +1903,30 @@ class RecipeScanner:
recipes, json_paths, self._cache.image_id_map
)
if report_progress:
loop.run_until_complete(
self._broadcast_scan_progress(
'completed', 'finalizing', 100, True,
elapsed_seconds=time.time() - (scan_start_time or time.time()),
total=len(recipes),
)
)
return self._cache
except Exception as e:
logger.error(f"Error in thread-based recipe cache initialization: {e}")
import traceback
traceback.print_exc(file=sys.stderr)
if report_progress and loop is not None:
try:
loop.run_until_complete(
self._broadcast_scan_progress(
'error', 'process_models', 0, True, error=str(e)
)
)
except Exception: # pragma: no cover - defensive logging
logger.error("Error broadcasting recipe scan failure", exc_info=True)
return self._cache if hasattr(self, "_cache") else None
finally:
# Clean up the event loop
@@ -2017,12 +2080,16 @@ class RecipeScanner:
return updated
def _full_directory_scan_sync(
self, recipes_dir: str
self,
recipes_dir: str,
progress_loop: Optional[asyncio.AbstractEventLoop] = None,
) -> Tuple[List[Dict[str, Any]], Dict[str, str]]:
"""Perform a full synchronous directory scan for recipes.
Args:
recipes_dir: Path to the recipes directory.
progress_loop: When set (manual force-refresh only), broadcast
scan_progress messages through this thread-local event loop.
Returns:
Tuple of (recipes list, json_paths dict).
@@ -2037,6 +2104,17 @@ class RecipeScanner:
if file.lower().endswith(".recipe.json"):
recipe_files.append(os.path.join(root, file))
total_files = len(recipe_files)
if progress_loop is not None:
progress_loop.run_until_complete(
self._broadcast_scan_progress(
'processing', 'count_models', 1, True,
processed=0, total=total_files,
)
)
last_progress_time = time.time()
# Process each recipe file
for i, recipe_path in enumerate(recipe_files):
recipe_data = self._load_recipe_file_sync(recipe_path)
@@ -2044,6 +2122,23 @@ class RecipeScanner:
recipe_id = str(recipe_data.get("id", ""))
recipes.append(recipe_data)
json_paths[recipe_id] = recipe_path
if progress_loop is not None and total_files > 0:
processed = i + 1
current_time = time.time()
# Throttle to one update per 0.5s; always send the final one.
if (
processed == total_files
or current_time - last_progress_time > 0.5
):
last_progress_time = current_time
progress_percent = min(99, int(1 + (processed / total_files) * 98))
progress_loop.run_until_complete(
self._broadcast_scan_progress(
'processing', 'process_models', progress_percent, True,
processed=processed, total=total_files,
current_name=os.path.basename(recipe_path),
)
)
# Periodically release GIL so the event loop thread can run
if i % 100 == 0:
time.sleep(0)
@@ -2613,11 +2708,14 @@ class RecipeScanner:
start_time = time.time()
# Run the heavy lifting in a thread pool – same path
# used by initialize_in_background().
# used by initialize_in_background(). Pass
# report_progress=True so manual refreshes broadcast
# scan_progress updates; startup init keeps it off.
loop = asyncio.get_event_loop()
cache = await loop.run_in_executor(
None,
self._initialize_recipe_cache_sync,
True,
)
if cache is not None:
self._cache = cache