diff --git a/__init__.py b/__init__.py index e4f62cc..463615a 100644 --- a/__init__.py +++ b/__init__.py @@ -59,9 +59,20 @@ class JobProgressTracker: self.progress.pop(pid, None) self.cancelled.difference_update(expired) - def update(self, progress_id, **payload): + def update(self, progress_id, source=None, stage=None, message=None, percent=None, status=None, **payload): if not progress_id: return + if source is not None: + payload["source"] = source + if stage is not None: + payload["stage"] = stage + if message is not None: + payload["message"] = message + if percent is not None: + payload["percent"] = percent + if status is not None: + payload["status"] = status + now = time.time() with self.lock: current = self.progress.get(progress_id, {}) @@ -126,7 +137,7 @@ class ModelResolverExtension: def __init__(self): self.routes_setup = False self.logger = create_module_logger(__name__) - self.analysis_progress = {} + self.analysis_progress = JobProgressTracker("Analyzing...") self.search_tracker = JobProgressTracker("Searching...") self.hash_tracker = JobProgressTracker("Preparing hash calculation...") self.search_result_timestamps = {} @@ -349,24 +360,24 @@ class ModelResolverExtension: ) if analysis_id: - self.analysis_progress[analysis_id] = { - "status": "starting", - "stage": "starting", - "message": "Starting analysis...", - "current": 0, - "total": 0, - } + self.analysis_progress.update( + analysis_id, + status="starting", + stage="starting", + message="Starting analysis...", + current=0, + total=0, + ) def update_analysis_progress(payload): if not analysis_id: return - self.analysis_progress[analysis_id] = { - **self.analysis_progress.get(analysis_id, {}), - **payload, - "status": "running" - if payload.get("stage") != "completed" - else "completed", - } + status = "running" if payload.get("stage") != "completed" else "completed" + self.analysis_progress.update( + analysis_id, + status=status, + **payload + ) # Analyze and find matches result = await asyncio.to_thread( @@ -497,25 +508,26 @@ class ModelResolverExtension: # when user clicks "Search Online" button, not automatically if analysis_id: - self.analysis_progress[analysis_id] = { - **self.analysis_progress.get(analysis_id, {}), - "status": "completed", - "stage": "completed", - "message": "Analysis complete", - "current": result.get("total_missing", 0), - "total": result.get("total_missing", 0), - } + self.analysis_progress.update( + analysis_id, + status="completed", + stage="completed", + message="Analysis complete", + current=result.get("total_missing", 0), + total=result.get("total_missing", 0), + ) return web.json_response(result) except Exception as e: if "analysis_id" in locals() and analysis_id: - self.analysis_progress[analysis_id] = { - "status": "error", - "stage": "error", - "message": str(e), - "current": 0, - "total": 0, - } + self.analysis_progress.update( + analysis_id, + status="error", + stage="error", + message=str(e), + current=0, + total=0, + ) self.logger.error(f"Model Resolver analyze error: {e}", exc_info=True) return web.json_response({"error": str(e)}, status=500) @@ -529,6 +541,7 @@ class ModelResolverExtension: {"error": "Analysis ID is required"}, status=400 ) + self.analysis_progress.cleanup() progress = self.analysis_progress.get(analysis_id) if not progress: return web.json_response( @@ -647,17 +660,7 @@ class ModelResolverExtension: return web.json_response({"success": True}) - def cleanup_hash_progress(max_age_seconds=300): - self.hash_tracker.cleanup(max_age_seconds) - - def update_hash_progress(progress_id, **payload): - self.hash_tracker.update(progress_id, **payload) - - def is_hash_progress_cancelled(progress_id): - return self.hash_tracker.is_cancelled(progress_id) - - def mark_hash_progress_cancelled(progress_id): - return self.hash_tracker.mark_cancelled(progress_id, "Stopping hash calculation...") + # cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed class HashCalculationCancelled(Exception): pass @@ -743,7 +746,7 @@ class ModelResolverExtension: last_update = 0.0 chunk_size = 1024 * 1024 * 4 - update_hash_progress( + self.hash_tracker.update( progress_id, status="running", stage="hashing", @@ -759,12 +762,12 @@ class ModelResolverExtension: continue sha256_hash.update(chunk) bytes_read += len(chunk) - if is_hash_progress_cancelled(progress_id): + if self.hash_tracker.is_cancelled(progress_id): percent = 0 if total_bytes <= 0 else min( 98, (bytes_read / total_bytes) * 98, ) - update_hash_progress( + self.hash_tracker.update( progress_id, status="cancelled", stage="cancelled", @@ -780,7 +783,7 @@ class ModelResolverExtension: 98, (bytes_read / total_bytes) * 98, ) - update_hash_progress( + self.hash_tracker.update( progress_id, status="running", stage="hashing", @@ -850,7 +853,7 @@ class ModelResolverExtension: {"error": error}, status=404 ) - cleanup_hash_progress() + self.hash_tracker.cleanup() progress_id = f"hash_{uuid.uuid4().hex}" self.hash_tracker.update( progress_id, @@ -867,9 +870,9 @@ class ModelResolverExtension: normalized_path, progress_id=progress_id, ) - if is_hash_progress_cancelled(progress_id): + if self.hash_tracker.is_cancelled(progress_id): raise HashCalculationCancelled() - update_hash_progress( + self.hash_tracker.update( progress_id, status="running", stage="metadata", @@ -881,7 +884,7 @@ class ModelResolverExtension: metadata_path, sha256, ) - update_hash_progress( + self.hash_tracker.update( progress_id, status="done", stage="done", @@ -894,7 +897,7 @@ class ModelResolverExtension: metadata_updated=metadata_updated, ) except HashCalculationCancelled: - update_hash_progress( + self.hash_tracker.update( progress_id, status="cancelled", stage="cancelled", @@ -902,7 +905,7 @@ class ModelResolverExtension: ) except Exception as exc: self.logger.exception(f"Hash calculation failed for {normalized_path}: {exc}") - update_hash_progress( + self.hash_tracker.update( progress_id, status="error", stage="error", @@ -924,7 +927,7 @@ class ModelResolverExtension: async def calculate_file_hash_progress_route(request): """Return progress for a background SHA256 calculation.""" progress_id = request.match_info.get("progress_id", "").strip() - cleanup_hash_progress() + self.hash_tracker.cleanup() progress = self.hash_tracker.get(progress_id) if not progress: return web.json_response( @@ -941,7 +944,7 @@ class ModelResolverExtension: return web.json_response( {"error": "progress_id is required"}, status=400 ) - cancelled = mark_hash_progress_cancelled(progress_id) + cancelled = self.hash_tracker.mark_cancelled(progress_id, "Stopping hash calculation...") return web.json_response( { "success": True, @@ -1979,43 +1982,14 @@ class ModelResolverExtension: dedupe_local_base_directories, ) - def cleanup_search_progress(max_age_seconds=300): - self.search_tracker.cleanup(max_age_seconds) - - def is_search_progress_cancelled(progress_id): - return self.search_tracker.is_cancelled(progress_id) - - def mark_search_progress_cancelled(progress_id, source=""): - res = self.search_tracker.mark_cancelled(progress_id, "Cancelled") - if res and source: - self.search_tracker.update(progress_id, source=source) - return res - - def update_search_progress( - progress_id, - source="", - stage="running", - message="Searching...", - percent=None, - status="running", - **extra, - ): - self.search_tracker.update( - progress_id, - source=source, - stage=stage, - message=message, - percent=percent, - status=status, - **extra - ) + # cleanup_search_progress, is_search_progress_cancelled, mark_search_progress_cancelled, and update_search_progress removed @routes.get("/model_resolver/search-progress/{progress_id}") @json_api_endpoint("search-progress") async def get_search_progress_route(request): """Return live progress for an in-flight source search.""" progress_id = request.match_info.get("progress_id", "") - cleanup_search_progress() + self.search_tracker.cleanup() progress = self.search_tracker.get(progress_id) if not progress: return web.json_response({"exists": False}) @@ -2030,7 +2004,7 @@ class ModelResolverExtension: return web.json_response( {"error": "Progress ID is required"}, status=400 ) - marked = mark_search_progress_cancelled(progress_id) + marked = self.search_tracker.mark_cancelled(progress_id, "Cancelled") return web.json_response( {"success": True, "cancelled": marked, "progress_id": progress_id} ) @@ -2045,7 +2019,7 @@ class ModelResolverExtension: pass def raise_if_search_cancelled(source=""): - if is_search_progress_cancelled(progress_id): + if self.search_tracker.is_cancelled(progress_id): raise SearchCancelled("Search cancelled") def format_log_value(value): @@ -2212,7 +2186,7 @@ class ModelResolverExtension: ) force_search = to_bool(data.get("force_search"), False) - update_search_progress( + self.search_tracker.update( progress_id, progress_source, "starting", @@ -2222,7 +2196,7 @@ class ModelResolverExtension: raise_if_search_cancelled(progress_source) if force_search: - update_search_progress( + self.search_tracker.update( progress_id, progress_source, "cache", @@ -2444,7 +2418,7 @@ class ModelResolverExtension: continue seen_hash_sources.add(hash_source_key) - update_search_progress( + self.search_tracker.update( progress_id, progress_source, "local_hash", @@ -2525,7 +2499,7 @@ class ModelResolverExtension: except (TypeError, ValueError): percent = None - update_search_progress( + self.search_tracker.update( progress_id, source_key, stage, @@ -2549,7 +2523,7 @@ class ModelResolverExtension: if initial_message is None: initial_message = f"Querying {source_key.capitalize()}" raise_if_search_cancelled(source_key) - update_search_progress( + self.search_tracker.update( progress_id, source_key, initial_stage, @@ -2573,7 +2547,7 @@ class ModelResolverExtension: "lora_manager_archive": "LoRA Manager archive checked", } done_msg = done_messages.get(source_key, f"{source_key} checked") - update_search_progress( + self.search_tracker.update( progress_id, source_key, "done", @@ -2603,7 +2577,7 @@ class ModelResolverExtension: if not res and base_model_context: raise_if_search_cancelled(source_key) - update_search_progress( + self.search_tracker.update( progress_id, source_key, "any_model", @@ -2635,7 +2609,7 @@ class ModelResolverExtension: popular_info = get_popular_model_url(filename) log_search_result("popular", popular_info) - update_search_progress( + self.search_tracker.update( progress_id, "local", "model_list", @@ -2726,7 +2700,7 @@ class ModelResolverExtension: version_id_val = data.get("version_id") if model_id_val and version_id_val: - update_search_progress( + self.search_tracker.update( progress_id, "civitai", "urn", @@ -2735,7 +2709,7 @@ class ModelResolverExtension: ) model_info = resolve_urn(model_id_val, version_id_val) if model_info: - update_search_progress( + self.search_tracker.update( progress_id, "civitai", "file", @@ -2802,7 +2776,7 @@ class ModelResolverExtension: }, ) elif category: - update_search_progress( + self.search_tracker.update( progress_id, "civitai", "fallback", @@ -2882,7 +2856,7 @@ class ModelResolverExtension: model_id_val = data.get("model_id") version_id_val = data.get("version_id") if model_id_val and version_id_val: - update_search_progress( + self.search_tracker.update( progress_id, "civarchive", "urn", @@ -2935,7 +2909,7 @@ class ModelResolverExtension: def handle_civarchive_error(e): error_message = f"CivArchive search failed: {e}" self.logger.warning(error_message) - update_search_progress( + self.search_tracker.update( progress_id, "civarchive", "error", @@ -3022,7 +2996,7 @@ class ModelResolverExtension: "Search sources async " + format_log_fields(count=len(search_tasks)) ) - update_search_progress( + self.search_tracker.update( progress_id, progress_source, "running", @@ -3054,7 +3028,7 @@ class ModelResolverExtension: found=results["found"], ) ) - update_search_progress( + self.search_tracker.update( progress_id, progress_source, "completed", @@ -3066,7 +3040,7 @@ class ModelResolverExtension: return web.json_response(results) except SearchCancelled: - update_search_progress( + self.search_tracker.update( progress_id if "progress_id" in locals() else "", progress_source if "progress_source" in locals() else "", "cancelled", @@ -3094,7 +3068,7 @@ class ModelResolverExtension: ) except Exception as e: - update_search_progress( + self.search_tracker.update( progress_id if "progress_id" in locals() else "", progress_source if "progress_source" in locals() else "", "error", diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 488b713..b6da9d3 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -38,6 +38,7 @@ from ..type_utils import ( fetch_remote_file_size_cached, clear_remote_size_cache, extract_trained_words, + extract_file_size, ) from ..progress import report_progress, get_progress_reporter from ..log_system import create_module_logger @@ -89,35 +90,8 @@ def is_civarchive_available() -> bool: return True -def _extract_file_size_bytes(file_info: Dict[str, Any]) -> Optional[int]: - if not isinstance(file_info, dict): - return None +# _extract_file_size_bytes removed in favor of extract_file_size from type_utils - for key in ("sizeKB", "size_kb"): - val = file_info.get(key) - if val is not None and val != "": - try: - return int(float(val) * 1024) - except (TypeError, ValueError): - pass - - for key in ("sizeBytes", "size_bytes", "fileSize", "file_size", "bytes", "size"): - val = file_info.get(key) - if val is not None and val != "": - size = parse_size_to_bytes(val) - if size: - return size - - mirrors = file_info.get("mirrors") or [] - if not isinstance(mirrors, list): - mirrors = [mirrors] - for mirror in mirrors: - if isinstance(mirror, dict): - size = _extract_file_size_bytes(mirror) - if size: - return size - - return None def _normalize_payload(payload: Dict[str, Any]) -> Dict[str, Any]: @@ -628,7 +602,7 @@ def _resolve_file_size_bytes( file_info: Dict[str, Any], download_urls: Optional[List[str]] = None, ) -> Optional[int]: - size = _extract_file_size_bytes(file_info) + size = extract_file_size(file_info) if size: return size @@ -800,8 +774,8 @@ def _merge_top_file_mirrors( mirror_url = mirror.get("url") if mirror_url and not any(existing.get("url") == mirror_url for existing in mirrors): mirrors.append(mirror) - if not _extract_file_size_bytes(target): - target_size_bytes = _extract_file_size_bytes(top_file) + if not extract_file_size(target): + target_size_bytes = extract_file_size(top_file) if target_size_bytes: target["sizeBytes"] = target_size_bytes @@ -940,7 +914,7 @@ def _normalize_archive_file(file_info: Dict[str, Any], model_id: Optional[int], "id": transformed.get("id"), "name": transformed.get("name"), "type": transformed.get("type"), - "size": _extract_file_size_bytes(transformed), + "size": extract_file_size(transformed), "download_url": download_urls[0] if download_urls else transformed.get("downloadUrl"), "download_urls": download_urls, "primary": bool(transformed.get("primary")), diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 8abe5e0..f42399a 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -18,7 +18,7 @@ from ..log_system import create_module_logger log = create_module_logger(__name__) from ..path_utils import write_json_atomic, METADATA_DIR, read_json_safe -from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size, fetch_remote_file_size_cached, clear_remote_size_cache +from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size, fetch_remote_file_size_cached, clear_remote_size_cache, extract_file_size from ..matcher import build_filename_search_queries HF_API_URL = "https://huggingface.co/api" @@ -329,22 +329,8 @@ def get_huggingface_download_url(repo: str, filename: str, branch: str = "main") -def _extract_file_info_size(file_info: Dict[str, Any]) -> Optional[int]: - if not isinstance(file_info, dict): - return None +# _extract_file_info_size removed in favor of extract_file_size from type_utils - for key in ("size", "sizeBytes", "size_bytes", "fileSize", "file_size", "bytes"): - size = _parse_size_header(file_info.get(key)) - if size: - return size - - lfs_info = file_info.get("lfs") if isinstance(file_info.get("lfs"), dict) else {} - for key in ("size", "sizeBytes", "size_bytes"): - size = _parse_size_header(lfs_info.get(key)) - if size: - return size - - return None def _normalize_huggingface_size_probe_url(url: str) -> Optional[str]: @@ -404,7 +390,7 @@ def _build_huggingface_result( or lfs_info.get("oid") ) download_url = get_huggingface_download_url(repo_id, file_path) - size = _extract_file_info_size(file_info) + size = extract_file_size(file_info) if not size: size = _fetch_remote_file_size_bytes(download_url, headers=headers) return { @@ -1052,7 +1038,7 @@ def get_repo_files(repo: str, token: Optional[str] = None) -> List[Dict[str, Any (".safetensors", ".ckpt", ".pt", ".bin", ".pth", ".onnx", ".gguf") ): url = get_huggingface_download_url(repo, path) - size = _extract_file_info_size(item) + size = extract_file_size(item) if not size: size = _fetch_remote_file_size_bytes(url, headers=headers) files.append( diff --git a/core/type_utils.py b/core/type_utils.py index a292001..80e8e94 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -809,4 +809,51 @@ def clear_remote_size_cache() -> None: _remote_size_cache.clear() +def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]: + """ + Extracts file size in bytes from various metadata formats (HuggingFace, CivitAI, CivArchive). + """ + if not isinstance(file_info, dict): + return None + + # 1. Check direct sizeKB or size_kb keys (e.g. from CivArchive or CivitAI) + for key in ("sizeKB", "size_kb"): + val = file_info.get(key) + if val is not None and val != "": + try: + return int(float(val) * 1024) + except (TypeError, ValueError): + pass + + # 2. Check direct size bytes keys + for key in ("sizeBytes", "size_bytes", "fileSize", "file_size", "bytes", "size"): + val = file_info.get(key) + if val is not None and val != "": + size = parse_size_header(val) + if size: + return size + + # 3. Check nested LFS info (common in HuggingFace API metadata) + lfs_info = file_info.get("lfs") + if isinstance(lfs_info, dict): + for key in ("size", "sizeBytes", "size_bytes"): + size = parse_size_header(lfs_info.get(key)) + if size: + return size + + # 4. Check nested CivArchive mirrors list + mirrors = file_info.get("mirrors") + if mirrors: + if not isinstance(mirrors, list): + mirrors = [mirrors] + for mirror in mirrors: + if isinstance(mirror, dict): + size = extract_file_size(mirror) + if size: + return size + + return None + + + diff --git a/web/resolver/resolver_dialog.js b/web/resolver/resolver_dialog.js index e43fa25..cb81702 100644 --- a/web/resolver/resolver_dialog.js +++ b/web/resolver/resolver_dialog.js @@ -256,6 +256,9 @@ export class ResolverManagerDialog extends ComfyDialog { if (response.status === 204) { return null; } + if (options.raw) { + return response; + } return await response.json(); } catch (error) { console.error(`Model Resolver: ${errorContext} failed:`, error); diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index b56b8c4..158ac8b 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -876,15 +876,11 @@ export const optionsMethods = { try { if (exportBackendLogsBtn) exportBackendLogsBtn.disabled = true; setExportLogsStatus(exportBackendLogsStatus, 'Preparing backend logs...'); - const response = await api.fetchApi('/model_resolver/logs/backend/export'); - if (!response.ok) { - let errorMsg = `Server returned ${response.status}: ${response.statusText}`; - try { - const errData = await response.json(); - if (errData?.error) errorMsg = errData.error; - } catch (_) { } - throw new Error(errorMsg); - } + const response = await this.fetchJson( + '/model_resolver/logs/backend/export', + { raw: true }, + 'Export backend logs' + ); const blob = await response.blob(); const filename = parseDownloadFilename( response.headers.get('Content-Disposition'), @@ -894,9 +890,7 @@ export const optionsMethods = { setExportLogsStatus(exportBackendLogsStatus, 'Backend logs exported.'); this.showNotification('Backend logs exported', 'success'); } catch (error) { - console.error('Model Resolver: backend log export failed:', error); setExportLogsStatus(exportBackendLogsStatus, error.message || 'Backend log export failed.'); - this.showNotification('Backend log export failed', 'error'); } finally { if (exportBackendLogsBtn) exportBackendLogsBtn.disabled = false; } diff --git a/web/resolver/views/tabs_loaded_methods.js b/web/resolver/views/tabs_loaded_methods.js index 3a571cd..d02145a 100644 --- a/web/resolver/views/tabs_loaded_methods.js +++ b/web/resolver/views/tabs_loaded_methods.js @@ -137,6 +137,20 @@ export const tabsLoadedMethods = { } }, + getModelNameAndStrength(model) { + const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; + const name = this.stripModelExtension(fullName); + const strength = model.strength !== null && model.strength !== undefined + ? model.strength.toFixed(2) + : null; + return { name, strength }; + }, + + getModelToken(model, category) { + const { name, strength } = this.getModelNameAndStrength(model); + return `<${this.getCategoryTokenName(category)}:${name}:${strength || '1.00'}>`; + }, + displayLoadedModels(container, data) { const loadedModels = data.loaded_models || []; const total = data.total || 0; @@ -168,21 +182,11 @@ export const tabsLoadedMethods = { const activeCount = Object.values(byCategory).reduce((sum, cat) => sum + cat.active.length, 0); const inactiveCount = Object.values(byCategory).reduce((sum, cat) => sum + cat.inactive.length, 0); - const stripModelExtension = (value) => this.stripModelExtension(value || 'Unknown'); - const buildCategoryStrings = (filter) => { const result = {}; for (const [category, modelsObj] of Object.entries(byCategory)) { - const displayCat = this.getCategoryTokenName(category); const models = filter === 'active' ? modelsObj.active : filter === 'inactive' ? modelsObj.inactive : [...modelsObj.active, ...modelsObj.inactive]; - const parts = models.map(model => { - const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - const name = stripModelExtension(fullName); - const strength = model.strength !== null && model.strength !== undefined - ? model.strength.toFixed(2) - : '1.00'; - return `<${displayCat}:${name}:${strength}>`; - }); + const parts = models.map(model => this.getModelToken(model, category)); result[category] = parts.join(' '); } return Object.values(result).join(' '); @@ -193,12 +197,6 @@ export const tabsLoadedMethods = { const allString = buildCategoryStrings('all'); const copyIcon = getSvgIcon('copy', 'currentColor', 'mr-copy-btn-icon'); - const modelToken = (model, category) => { - const fullName = model.name || model.original_path?.split(/[\/\\]/).pop() || 'Unknown'; - const name = stripModelExtension(fullName); - const strength = model.strength !== null && model.strength !== undefined ? model.strength.toFixed(2) : '1.00'; - return `<${this.getCategoryTokenName(category)}:${name}:${strength}>`; - }; const copyButtonHtml = (text, extraClass = '') => `