diff --git a/__init__.py b/__init__.py index 0880c43..b14464b 100644 --- a/__init__.py +++ b/__init__.py @@ -183,7 +183,7 @@ class ModelResolverExtension: read_json_safe, write_json_atomic, ) - from .core.type_utils import first_non_empty, to_int, to_bool + from .core.type_utils import first_non_empty, to_int, to_bool, MODEL_EXTENSIONS, format_size_bytes from .core.settings import ( TEMPLATE_KEY_ALIASES, bool_setting as resolver_bool_setting, @@ -404,10 +404,8 @@ class ModelResolverExtension: ) if not has_perfect_match: - filename = ( + filename = get_filename_from_path( missing.get("original_path", "") - .split("/")[-1] - .split("\\")[-1] ) # 0. Check workflow URL first (highest priority - directly from workflow) @@ -1148,9 +1146,7 @@ class ModelResolverExtension: if not existing: loaded_models.append( { - "name": name.split("/")[-1].split("\\")[ - -1 - ], + "name": get_filename_from_path(name), "category": directory or "checkpoints", "node_id": node.get("id"), "widget_index": None, @@ -1485,18 +1481,7 @@ class ModelResolverExtension: expected = _os.path.basename(str(expected_filename or "")).lower() if expected and basename == expected: return True - model_extensions = { - ".safetensors", - ".ckpt", - ".pt", - ".pth", - ".bin", - ".gguf", - ".onnx", - ".pb", - ".pkl", - ".pickle", - } + model_extensions = MODEL_EXTENSIONS return _os.path.splitext(basename)[1].lower() in model_extensions def collect_result_download_urls(result): @@ -2131,21 +2116,7 @@ class ModelResolverExtension: return text def format_log_size(value): - if value is None or value == "": - return None - if isinstance(value, str): - return value - try: - size = float(value) - except (TypeError, ValueError): - return str(value) - - units = ("B", "KB", "MB", "GB", "TB") - unit_index = 0 - while size >= 1024 and unit_index < len(units) - 1: - size /= 1024 - unit_index += 1 - return f"{size:.1f}{units[unit_index]}" + return format_size_bytes(value, include_space=False) def format_log_fields(**fields): parts = [] diff --git a/core/downloader.py b/core/downloader.py index ad6d04b..01a7d79 100644 --- a/core/downloader.py +++ b/core/downloader.py @@ -24,7 +24,7 @@ log = create_module_logger(__name__) from .resolver import normalize_sha256 from .path_utils import is_path_within, get_path_identity, write_json_atomic, read_json_safe, get_comfy_root_path, calculate_file_sha256 as _calculate_file_sha256 -from .type_utils import as_dict, as_list, first_non_empty +from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes try: import folder_paths @@ -632,17 +632,6 @@ def write_lora_manager_metadata( return None -def format_bytes(bytes_value: int) -> str: - """Format bytes to human readable string (e.g., 1.5 GB).""" - if bytes_value == 0: - return "0 B" - k = 1024 - sizes = ["B", "KB", "MB", "GB", "TB"] - i = 0 - while bytes_value >= k and i < len(sizes) - 1: - bytes_value /= k - i += 1 - return f"{bytes_value:.1f} {sizes[i]}" # Imported from .settings diff --git a/core/type_utils.py b/core/type_utils.py index 2bca978..296e1c4 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -55,6 +55,8 @@ MODEL_EXTENSIONS = { ".sft", ".onnx", ".gguf", + ".pb", + ".pickle", } @@ -422,5 +424,31 @@ def check_credential_http( } +def format_size_bytes(bytes_value: Any, include_space: bool = True) -> Optional[str]: + """Format bytes to human readable string (e.g., 1.5 GB or 1.5GB).""" + if bytes_value is None or bytes_value == "": + return None + if isinstance(bytes_value, str) and not bytes_value.strip().replace(".", "", 1).isdigit(): + return bytes_value + try: + val = float(bytes_value) + except (TypeError, ValueError): + return str(bytes_value) + + if val <= 0: + return f"0{ ' B' if include_space else 'B' }" + + k = 1024 + sizes = ["B", "KB", "MB", "GB", "TB"] + i = 0 + while val >= k and i < len(sizes) - 1: + val /= k + i += 1 + + space = " " if include_space else "" + return f"{val:.1f}{space}{sizes[i]}" + + + diff --git a/core/workflow_analyzer.py b/core/workflow_analyzer.py index 19ddcd4..fd9a4c6 100644 --- a/core/workflow_analyzer.py +++ b/core/workflow_analyzer.py @@ -21,21 +21,7 @@ except ImportError: log.warn("Model Resolver: folder_paths not available yet - will retry later") -# Common model file extensions -MODEL_EXTENSIONS = { - ".ckpt", - ".pt", - ".pt2", - ".bin", - ".pth", - ".safetensors", - ".pkl", - ".sft", - ".onnx", - ".gguf", -} - -from .type_utils import URN_TYPE_MAP +from .type_utils import URN_TYPE_MAP, MODEL_EXTENSIONS URN_REGEX = re.compile(r"^urn:air:([^:]+):([^:]+):([^:]+):(\d+)@(\d+)$") diff --git a/web/resolver/actions/queue_methods.js b/web/resolver/actions/queue_methods.js index 2a4c90c..06bf059 100644 --- a/web/resolver/actions/queue_methods.js +++ b/web/resolver/actions/queue_methods.js @@ -2428,7 +2428,7 @@ export const queueMethods = { category: selection.category || sourceMissing?.category || resolvedModel.category || 'unknown', original_path: resolvedRelativePath, name: resolvedRelativePath, - filename: resolvedModel.filename || resolvedRelativePath.split('/').pop()?.split('\\').pop() || resolvedRelativePath, + filename: resolvedModel.filename || this.getFilenameFromPath(resolvedRelativePath) || resolvedRelativePath, relative_path: resolvedRelativePath, full_path: resolvedModel.path || selection.resolved_path || sourceMissing?.full_path || '', path: resolvedModel.path || selection.resolved_path || '', diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index 7eeef56..322de1c 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/web/resolver/actions/resolve_download_methods.js @@ -239,7 +239,7 @@ export const resolveDownloadMethods = { const state = this.searchResultCache.get(this.getMissingSearchKey(missing)); const results = state?.results || {}; - const filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + const filename = this.getFilenameFromPath(missing.original_path); const first = (value) => Array.isArray(value) ? value[0] : value; const candidates = [ { @@ -447,8 +447,7 @@ export const resolveDownloadMethods = { return downloadedFilename || missing.civitai_info?.expected_filename || missing.download_source?.filename - || missing.original_path?.split('/').pop()?.split('\\').pop() - || ''; + || this.getFilenameFromPath(missing.original_path); }, getDownloadFolderContext(progress = {}, info = {}) { @@ -1178,7 +1177,7 @@ export const resolveDownloadMethods = { const minRefreshFeedback = new Promise(resolve => setTimeout(resolve, 420)); const refreshAnimation = this.startRefreshButtonAnimation?.(button); const body = this.contentElement?.querySelector(`#local-matches-body-${this.getMissingModelDomKey(missing)}`); - const displayName = targetFilename.split('/').pop()?.split('\\').pop() || targetFilename; + const displayName = this.getFilenameFromPath(targetFilename) || targetFilename; const missingKey = this.getMissingModelKey(missing); const currentMissing = (this.missingModels || []).find(item => this.getMissingModelKey(item) === missingKey) || missing; const previousMatches = this.cloneLocalMatches?.(currentMissing.matches || missing.matches || []) || []; @@ -1500,7 +1499,7 @@ export const resolveDownloadMethods = { return; } - const originalFilename = missing.original_path?.split('/').pop()?.split('\\').pop() || 'model.safetensors'; + const originalFilename = this.getFilenameFromPath(missing.original_path) || 'model.safetensors'; const filename = source.filename || originalFilename; const targetSelection = this.getDownloadTargetSelection(missing, source.directory || missing.category || 'checkpoints'); const downloadBtn = this.contentElement?.querySelector(`#${this.getDownloadButtonElementId(missing)}`); @@ -1836,7 +1835,7 @@ export const resolveDownloadMethods = { * Search online for a model */ async searchOnline(missing, { workflowKey = this.getWorkflowScopedQueueKey(), forceSearch = false } = {}) { - let filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + let filename = this.getFilenameFromPath(missing.original_path); let category = missing.category || this.getNodeTypeDownloadCategory?.(missing.node_type) || ''; const state = this.getSearchStateForWorkflow(workflowKey, missing); const missingSearchKey = this.getMissingSearchKey(missing); @@ -2551,7 +2550,7 @@ export const resolveDownloadMethods = { addRow(knownDownloadRow); if (popular) { - const popularFilename = popular.filename || missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + const popularFilename = popular.filename || this.getFilenameFromPath(missing.original_path); const missingCategory = this.getMissingDownloadCategory?.(missing, 'checkpoints') || missing.category || 'checkpoints'; const popularSize = popular.size || ( @@ -2616,7 +2615,7 @@ export const resolveDownloadMethods = { } if (civarchiveResult && civarchiveResult.download_url) { - const archiveFilename = civarchiveResult.filename || missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + const archiveFilename = civarchiveResult.filename || this.getFilenameFromPath(missing.original_path); const archiveName = civarchiveResult.name || archiveFilename || 'Model'; const archiveCategory = this.getSourceResultDownloadCategory?.( civarchiveResult, @@ -2652,7 +2651,7 @@ export const resolveDownloadMethods = { } if (loraManagerArchiveResult && loraManagerArchiveResult.download_url) { - const archiveFilename = loraManagerArchiveResult.filename || missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + const archiveFilename = loraManagerArchiveResult.filename || this.getFilenameFromPath(missing.original_path); const archiveName = loraManagerArchiveResult.name || archiveFilename; const archiveCategory = this.getSourceResultDownloadCategory?.( loraManagerArchiveResult, diff --git a/web/resolver/search/download_target_methods.js b/web/resolver/search/download_target_methods.js index a1f94b7..711445a 100644 --- a/web/resolver/search/download_target_methods.js +++ b/web/resolver/search/download_target_methods.js @@ -1060,8 +1060,7 @@ export const downloadTargetMethods = { const filename = merged.downloadFilename || merged.filename || merged.file_name - || missing.original_path?.split('/').pop()?.split('\\').pop() - || ''; + || this.getFilenameFromPath(missing.original_path); const modelName = merged.model_name || merged.model || merged.name || filename.replace(/\.[^.]+$/, '') || ''; const creator = merged.creator || (merged.creator_username ? { username: merged.creator_username } : null) @@ -1117,8 +1116,7 @@ export const downloadTargetMethods = { || merged.filename || merged.file_name || pathMetadata.filename - || missing.original_path?.split('/').pop()?.split('\\').pop() - || ''; + || this.getFilenameFromPath(missing.original_path); const modelName = merged.model_name || merged.model || merged.name diff --git a/web/resolver/search/search_panel.js b/web/resolver/search/search_panel.js index c76d957..fdbba59 100644 --- a/web/resolver/search/search_panel.js +++ b/web/resolver/search/search_panel.js @@ -1677,7 +1677,7 @@ export const searchPanelMethods = { getDownloadSourceTableRow(missing, downloadSource = {}) { if (!downloadSource?.url) return null; - const originalFilename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + const originalFilename = this.getFilenameFromPath(missing.original_path); const downloadFilename = downloadSource.filename || originalFilename || 'model'; const isFromWorkflow = downloadSource.url_source === 'workflow'; const source = downloadSource.source || (isFromWorkflow ? 'workflow' : 'online'); diff --git a/web/resolver/utils/render_format_methods.js b/web/resolver/utils/render_format_methods.js index 645ded0..4296661 100644 --- a/web/resolver/utils/render_format_methods.js +++ b/web/resolver/utils/render_format_methods.js @@ -77,6 +77,11 @@ export const renderFormatMethods = { return `${label}`; }, + getFilenameFromPath(path) { + if (!path) return ''; + return path.split(/[\/\\]/).pop() || path; + }, + /** * Format a filename with smart truncation * @param {string} path - Full path or filename @@ -87,7 +92,7 @@ export const renderFormatMethods = { if (!path) return { display: 'Unknown', full: 'Unknown' }; // Extract just the filename from path - const filename = path.split(/[\/\\]/).pop() || path; + const filename = this.getFilenameFromPath(path); if (filename.length <= maxLength) { return { display: filename, full: filename }; diff --git a/web/resolver/views/missing_browser_methods.js b/web/resolver/views/missing_browser_methods.js index 1ca7561..d69d849 100644 --- a/web/resolver/views/missing_browser_methods.js +++ b/web/resolver/views/missing_browser_methods.js @@ -8,7 +8,7 @@ import { startSplitterDrag } from "../utils/splitter_drag.js"; export const missingBrowserMethods = { getMissingFilename(missing = {}) { - return missing.original_path?.split('/').pop()?.split('\\').pop() || missing.name || 'Missing model'; + return this.getFilenameFromPath(missing.original_path) || missing.name || 'Missing model'; }, getMissingLocateTarget(missing = {}) { @@ -419,7 +419,7 @@ export const missingBrowserMethods = { normalizeResolvedWorkflowModel(model = {}) { const originalPath = model.original_path || model.name || model.filename || ''; - const filename = originalPath.split('/').pop()?.split('\\').pop() || originalPath || 'Resolved model'; + const filename = this.getFilenameFromPath(originalPath) || originalPath || 'Resolved model'; const fullPath = model.full_path || model.path || ''; const relativePath = model.relative_path || originalPath || filename; const category = model.category || 'unknown'; @@ -2107,7 +2107,7 @@ export const missingBrowserMethods = { html += `
`; html += `
Download
`; - const filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + const filename = this.getFilenameFromPath(missing.original_path); const downloadSource = missing.download_source; const urnDownloadId = `urn-download-${missing.node_id}-${missing.widget_index}`;