diff --git a/__init__.py b/__init__.py index 9737197..52a91cc 100644 --- a/__init__.py +++ b/__init__.py @@ -187,6 +187,7 @@ class ModelResolverExtension: search_local_matches_by_hash, search_local_matches, ) + from .core.workflow_analyzer import analyze_workflow_models from .core.path_templates import infer_download_path_templates from .core.metadata_audit import audit_metadata_sizes from .core.metadata_builder import ( @@ -673,6 +674,105 @@ class ModelResolverExtension: get_local_model_hash_metadata(path, model=model) ) + @routes.post("/model_resolver/workflow-model-hashes") + @json_api_endpoint("workflow-model-hashes") + async def workflow_model_hashes(request): + """Return hash metadata for existing local models used by a workflow.""" + data = await request.json() + workflow_json = data.get("workflow") + if not isinstance(workflow_json, dict): + return web.json_response( + {"error": "Workflow JSON must be an object"}, status=400 + ) + + settings = await asyncio.to_thread(load_resolver_settings) + if not resolver_bool_setting( + settings.get("workflow_hash_metadata_enabled"), True + ): + return web.json_response( + { + "success": True, + "enabled": False, + "models": [], + "by_node": {}, + "by_path": {}, + "count": 0, + } + ) + + available_models = await asyncio.to_thread(get_model_files, False) + refs = await asyncio.to_thread( + analyze_workflow_models, + workflow_json, + available_models, + ) + + by_node = {} + by_path = {} + models = [] + seen = set() + + for ref in refs: + if not isinstance(ref, dict) or not ref.get("exists"): + continue + full_path = str(ref.get("full_path") or "").strip() + if not full_path: + continue + model_info = { + "path": full_path, + "filename": get_filename_from_path(full_path), + "relative_path": ref.get("original_path") or "", + "category": ref.get("category") or "", + } + metadata = get_local_model_hash_metadata( + full_path, model=model_info + ) + sha256 = normalize_sha256(metadata.get("sha256")) + if not sha256: + continue + + entry = { + "node_id": ref.get("node_id"), + "node_type": ref.get("node_type") or "", + "widget_index": ref.get("widget_index"), + "widget_name": ref.get("widget_name") or "", + "path": ref.get("original_path") or "", + "filename": get_filename_from_path( + ref.get("original_path") or full_path + ), + "category": ref.get("category") or "", + "sha256": sha256, + "size": metadata.get("size") or 0, + } + entry_key = ( + str(entry.get("node_id")), + str(entry.get("widget_index")), + entry.get("path"), + sha256, + ) + if entry_key in seen: + continue + seen.add(entry_key) + models.append(entry) + + path_key = str(entry.get("path") or entry.get("filename") or "") + if path_key: + by_path[path_key] = entry + by_path[get_filename_from_path(path_key)] = entry + node_key = f"{entry.get('node_id')}:{entry.get('widget_index')}" + by_node[node_key] = entry + + return web.json_response( + { + "success": True, + "enabled": True, + "models": models, + "by_node": by_node, + "by_path": by_path, + "count": len(models), + } + ) + @routes.post("/model_resolver/local-matches-by-hash") @json_api_endpoint("local-matches-by-hash") async def local_matches_by_hash(request): diff --git a/core/resolver.py b/core/resolver.py index 0aa079c..5b7884c 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -287,6 +287,98 @@ def get_workflow_url_info_for_filename( return None +def _workflow_hash_entry_candidates(entry: Any) -> List[Dict[str, Any]]: + if not isinstance(entry, dict): + return [] + sha256 = normalize_sha256(extract_sha256_from_metadata(entry)) + if not sha256: + return [] + return [{**entry, "sha256": sha256}] + + +def _collect_workflow_hash_entries(value: Any) -> List[Dict[str, Any]]: + if not isinstance(value, dict): + return [] + + entries: List[Dict[str, Any]] = [] + for key in ( + "model_resolver_hashes", + "anomalous_hashes", + "model_hashes", + "hashes", + ): + raw = value.get(key) + if isinstance(raw, dict): + for nested_key in ("models", "by_path", "by_node"): + nested = raw.get(nested_key) + if isinstance(nested, list): + for item in nested: + entries.extend(_workflow_hash_entry_candidates(item)) + elif isinstance(nested, dict): + for item in nested.values(): + entries.extend(_workflow_hash_entry_candidates(item)) + entries.extend(_workflow_hash_entry_candidates(raw)) + elif isinstance(raw, list): + for item in raw: + entries.extend(_workflow_hash_entry_candidates(item)) + return entries + + +def extract_workflow_hash_metadata(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: + """Return embedded workflow SHA256 metadata indexed by path, filename, and node/widget.""" + if not isinstance(workflow_json, dict): + return {} + + entries: List[Dict[str, Any]] = [] + entries.extend(_collect_workflow_hash_entries(workflow_json)) + extra = workflow_json.get("extra") + if isinstance(extra, dict): + entries.extend(_collect_workflow_hash_entries(extra)) + + index: Dict[str, Dict[str, Any]] = {} + for entry in entries: + sha256 = normalize_sha256(entry.get("sha256")) + if not sha256: + continue + normalized = {**entry, "sha256": sha256} + for candidate in ( + entry.get("path"), + entry.get("original_path"), + entry.get("filename"), + entry.get("name"), + entry.get("file_name"), + ): + text = str(candidate or "").strip() + if not text: + continue + index[text] = normalized + index[get_filename_from_path(text)] = normalized + node_id = entry.get("node_id") + widget_index = entry.get("widget_index") + if node_id is not None and widget_index is not None: + index[f"{node_id}:{widget_index}"] = normalized + return index + + +def get_workflow_hash_info_for_ref( + workflow_hashes: Dict[str, Dict[str, Any]], model_ref: Dict[str, Any] +) -> Optional[Dict[str, Any]]: + if not workflow_hashes: + return None + original_path = str(model_ref.get("original_path") or "") + candidates = [ + f"{model_ref.get('node_id')}:{model_ref.get('widget_index')}", + original_path, + get_filename_from_path(original_path), + model_ref.get("filename") or "", + model_ref.get("name") or "", + ] + for candidate in candidates: + if candidate and candidate in workflow_hashes: + return workflow_hashes[candidate] + return None + + def workflow_has_nodes(workflow_json: Dict[str, Any]) -> bool: """Return True when the active top-level workflow contains nodes.""" if not isinstance(workflow_json, dict): @@ -599,6 +691,12 @@ def get_local_model_hash_metadata( normalized_path = os.path.abspath(os.path.normpath(raw_path)) exists = os.path.exists(normalized_path) + file_size = 0 + if exists and os.path.isfile(normalized_path): + try: + file_size = os.path.getsize(normalized_path) + except OSError: + file_size = 0 model_info: Dict[str, Any] = { **(model if isinstance(model, dict) else {}), "path": normalized_path, @@ -620,6 +718,7 @@ def get_local_model_hash_metadata( "hash_status": last_hash_status, "hashes": hashes, "sha256": hashes[0], + "size": file_size, } return { @@ -628,6 +727,7 @@ def get_local_model_hash_metadata( "hash_status": last_hash_status, "hashes": [], "sha256": "", + "size": file_size, } @@ -879,6 +979,8 @@ def analyze_and_find_matches( # Extract URLs from workflow (node.properties.models + regex) workflow_urls = extract_workflow_urls(workflow_json) log.debug(f"Extracted {len(workflow_urls)} URLs from workflow") + workflow_hashes = extract_workflow_hash_metadata(workflow_json) + log.debug(f"Extracted {len(workflow_hashes)} workflow hash metadata keys") if progress_callback: progress_callback( @@ -944,6 +1046,10 @@ def analyze_and_find_matches( missing["workflow_model_url"] = url_info.get("model_url", "") missing["workflow_directory"] = url_info.get("directory", "") missing["url_source"] = url_info.get("source", "") + hash_info = get_workflow_hash_info_for_ref(workflow_hashes, missing) + if hash_info: + missing["workflow_sha256"] = normalize_sha256(hash_info.get("sha256")) + missing["hash_lookup_source"] = hash_info.get("source") or "workflow_metadata" # Handle URNs: mark for async resolution by frontend # No sync CivitAI calls here - frontend will fetch asynchronously @@ -1049,13 +1155,32 @@ def analyze_and_find_matches( def find_local_matches_for_ref(model_ref: Dict[str, Any]) -> List[Dict[str, Any]]: target_for_matching = get_match_target(model_ref) - if not target_for_matching: - return [] - category = get_match_category(model_ref) + workflow_sha256 = normalize_sha256(model_ref.get("workflow_sha256")) + hash_matches = [] + if workflow_sha256: + hash_matches = search_local_matches_by_hash( + workflow_sha256, + category=category, + max_matches=max_matches_per_model, + force_rescan=False, + ) + hash_matches = [ + { + **match, + "hash_lookup_source": model_ref.get("hash_lookup_source") or "workflow_metadata", + "hash_lookup_filename": get_filename_from_path(model_ref.get("original_path") or ""), + "hash_lookup_sha256": workflow_sha256, + } + for match in hash_matches + ] + + if not target_for_matching: + return hash_matches + cache_key = (target_for_matching, category) if cache_key in local_match_cache: - return local_match_cache[cache_key] + return deduplicate_matches(hash_matches + local_match_cache[cache_key]) matches = find_matches( target_for_matching, @@ -1068,7 +1193,7 @@ def analyze_and_find_matches( active_downloads_by_path, ) local_match_cache[cache_key] = deduplicated_matches - return deduplicated_matches + return deduplicate_matches(hash_matches + deduplicated_matches) # Find matches for each missing model missing_with_matches = [] diff --git a/core/settings.py b/core/settings.py index 3a73895..69e5feb 100644 --- a/core/settings.py +++ b/core/settings.py @@ -163,6 +163,9 @@ def normalize_download_backend(value: Any) -> str: def normalize_settings(settings: Optional[Mapping[str, Any]]) -> Dict[str, Any]: data = dict(settings or {}) + data["workflow_hash_metadata_enabled"] = bool_setting( + data.get("workflow_hash_metadata_enabled"), True + ) data["download_backend"] = normalize_download_backend(data.get("download_backend")) data["aria2c_path"] = str(data.get("aria2c_path") or "").strip() data["aria2_auto_stop_daemon"] = bool_setting( diff --git a/web/resolver/model_resolver.js b/web/resolver/model_resolver.js index 191a738..25dcef2 100644 --- a/web/resolver/model_resolver.js +++ b/web/resolver/model_resolver.js @@ -98,6 +98,10 @@ export class ModelResolver { this.dialog = null; this.isCheckingMissing = false; // Prevent multiple simultaneous checks this.lastCheckedWorkflow = null; // Track to avoid duplicate checks + this.workflowHashMetadataCache = null; + this.workflowHashMetadataSignature = null; + this.workflowHashMetadataRefreshTimer = null; + this.workflowHashMetadataRefreshing = false; } setup = async () => { @@ -117,6 +121,7 @@ export class ModelResolver { // Listen for workflow load events to auto-check for missing models this.setupAutoOpenOnMissingModels(); + this.setupWorkflowHashMetadataInjection(); this.setupActiveWorkflowChangeListeners(); const sidebarRegistered = this.registerSidebarButton(); @@ -509,10 +514,111 @@ export class ModelResolver { } handleActiveWorkflowRouteChange(reason = 'workflow-change') { + this.scheduleWorkflowHashMetadataRefresh(); if (!this.dialog?.isVisible()) return; this.dialog.scheduleActiveWorkflowRefresh(reason); } + isWorkflowHashMetadataEnabled() { + return localStorage.getItem('ModelResolver.workflowHashMetadataEnabled') !== 'false'; + } + + setupWorkflowHashMetadataInjection() { + window.__ModelResolverWorkflowHashMetadataOwner = this; + if (window.__ModelResolverWorkflowHashMetadataPatched) { + this.scheduleWorkflowHashMetadataRefresh(); + return; + } + + if (!app?.graph || typeof app.graph.serialize !== 'function') return; + + const originalSerialize = app.graph.serialize; + app.graph.serialize = function(...args) { + const workflow = originalSerialize.apply(this, args); + window.__ModelResolverWorkflowHashMetadataOwner?.injectWorkflowHashMetadata(workflow); + window.__ModelResolverWorkflowHashMetadataOwner?.scheduleWorkflowHashMetadataRefresh(workflow); + return workflow; + }; + window.__ModelResolverWorkflowHashMetadataPatched = true; + this.scheduleWorkflowHashMetadataRefresh(); + } + + injectWorkflowHashMetadata(workflow) { + if (!workflow || typeof workflow !== 'object') return workflow; + if (!this.isWorkflowHashMetadataEnabled()) return workflow; + + const cache = this.workflowHashMetadataCache; + if (!cache || !Array.isArray(cache.models) || !cache.models.length) return workflow; + + workflow.extra = workflow.extra && typeof workflow.extra === 'object' ? workflow.extra : {}; + workflow.extra.model_resolver_hashes = { + version: 1, + source: 'comfyui-model-resolver', + models: cache.models, + updated_at: cache.updated_at || Date.now() + }; + return workflow; + } + + scheduleWorkflowHashMetadataRefresh(workflow = null) { + if (!this.isWorkflowHashMetadataEnabled()) return; + if (this.workflowHashMetadataRefreshTimer) { + clearTimeout(this.workflowHashMetadataRefreshTimer); + } + this.workflowHashMetadataRefreshTimer = setTimeout(() => { + this.workflowHashMetadataRefreshTimer = null; + this.refreshWorkflowHashMetadata(workflow); + }, 800); + } + + async refreshWorkflowHashMetadata(workflow = null) { + if (this.workflowHashMetadataRefreshing || !this.isWorkflowHashMetadataEnabled()) return; + + const currentWorkflow = workflow || app?.graph?.serialize?.(); + if (!currentWorkflow) return; + + let signature = ''; + try { + signature = JSON.stringify((currentWorkflow.nodes || []).map((node) => [ + node?.id, + node?.type, + node?.widgets_values + ])); + } catch (error) { + signature = ''; + } + if (signature && signature === this.workflowHashMetadataSignature && this.workflowHashMetadataCache) { + return; + } + + this.workflowHashMetadataRefreshing = true; + try { + const response = await api.fetchApi('/model_resolver/workflow-model-hashes', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ workflow: currentWorkflow }) + }); + if (!response.ok) return; + const data = await response.json(); + if (!data?.enabled) { + this.workflowHashMetadataCache = null; + this.workflowHashMetadataSignature = signature; + return; + } + this.workflowHashMetadataCache = { + models: Array.isArray(data.models) ? data.models : [], + by_path: data.by_path || {}, + by_node: data.by_node || {}, + updated_at: Date.now() + }; + this.workflowHashMetadataSignature = signature; + } catch (error) { + log.debug('Model Resolver: workflow hash metadata refresh failed', error); + } finally { + this.workflowHashMetadataRefreshing = false; + } + } + /** * Setup MutationObserver to detect ComfyUI's Missing Models popup and inject our button */ diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index ad49216..84d3788 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -18,6 +18,7 @@ const SETTINGS_MAP = [ { serverKey: 'auto_fill_base_model', localKey: 'ModelResolver.autoFillBaseModel', type: 'boolean', default: true }, { serverKey: 'auto_fill_subfolder', localKey: 'ModelResolver.autoFillSubfolder', type: 'boolean', default: true }, { serverKey: 'auto_refresh_comfy_models_after_apply', localKey: 'ModelResolver.autoRefreshComfyModelsAfterApply', type: 'boolean', default: true }, + { serverKey: 'workflow_hash_metadata_enabled', localKey: 'ModelResolver.workflowHashMetadataEnabled', type: 'boolean', default: true }, { serverKey: 'download_backend', localKey: 'ModelResolver.downloadBackend', type: 'backend', default: 'python' }, { serverKey: 'aria2c_path', localKey: 'ModelResolver.aria2cPath', type: 'string', default: '' }, { serverKey: 'aria2_auto_stop_daemon', localKey: 'ModelResolver.aria2AutoStopDaemon', type: 'boolean', default: true }, @@ -433,6 +434,15 @@ export const optionsMethods = { + @@ -903,6 +913,7 @@ export const optionsMethods = { const autoFillBaseModelInput = this.contentElement.querySelector('#mr-options-auto-fill-base-model'); const autoFillSubfolderInput = this.contentElement.querySelector('#mr-options-auto-fill-subfolder'); const autoRefreshComfyModelsInput = this.contentElement.querySelector('#mr-options-auto-refresh-comfy-models'); + const workflowHashMetadataInput = this.contentElement.querySelector('#mr-options-workflow-hash-metadata'); const downloadBackendInput = this.contentElement.querySelector('#mr-options-download-backend'); const aria2cPathInput = this.contentElement.querySelector('#mr-options-aria2c-path'); const aria2InstallBtn = this.contentElement.querySelector('#mr-options-aria2-install'); @@ -3422,6 +3433,7 @@ export const optionsMethods = { auto_fill_base_model: Boolean(autoFillBaseModelInput?.checked), auto_fill_subfolder: Boolean(autoFillSubfolderInput?.checked), auto_refresh_comfy_models_after_apply: Boolean(autoRefreshComfyModelsInput?.checked), + workflow_hash_metadata_enabled: Boolean(workflowHashMetadataInput?.checked), download_backend: downloadBackend, aria2c_path: aria2cPathInput?.value?.trim() || '', aria2_auto_stop_daemon: Boolean(aria2AutoStopInput?.checked),