diff --git a/__init__.py b/__init__.py index 5492320..bf26d0b 100644 --- a/__init__.py +++ b/__init__.py @@ -1796,6 +1796,15 @@ class ModelResolverExtension: # Also extract from node.properties.models workflow_nodes = workflow_json.get("nodes", []) nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else [] + node_scopes = { + id(node): { + "subgraph_id": "", + "subgraph_name": "", + "is_top_level": True, + } + for node in nodes + if isinstance(node, dict) + } definitions = workflow_json.get("definitions", {}) subgraphs = ( definitions.get("subgraphs", []) @@ -1807,6 +1816,44 @@ class ModelResolverExtension: subgraph_nodes = subgraph.get("nodes", []) if isinstance(subgraph_nodes, list): nodes.extend(subgraph_nodes) + for node in subgraph_nodes: + if not isinstance(node, dict): + continue + node_scopes[id(node)] = { + "subgraph_id": subgraph.get("id") or "", + "subgraph_name": ( + subgraph.get("name") + or subgraph.get("id") + or "" + ), + "is_top_level": False, + } + + def get_node_scope(node): + return node_scopes.get( + id(node), + { + "subgraph_id": "", + "subgraph_name": "", + "is_top_level": True, + }, + ) + + def node_matches_ref(node, ref): + if not isinstance(node, dict): + return False + if str(node.get("id")) != str(ref.get("node_id")): + return False + scope = get_node_scope(node) + if scope["is_top_level"] != ( + ref.get("is_top_level") is not False + ): + return False + if not scope["is_top_level"]: + return str(scope["subgraph_id"]) == str( + ref.get("subgraph_id") or "" + ) + return True # Collect all loaded models with their values loaded_models = [] @@ -1847,7 +1894,7 @@ class ModelResolverExtension: ): # Find the node in workflow to get strength value for node in nodes: - if str(node.get("id")) == str(node_id): + if node_matches_ref(node, ref): widgets_values = node.get("widgets_values", []) if ( isinstance(widgets_values, list) @@ -1890,6 +1937,25 @@ class ModelResolverExtension: "node_id": node_id, "widget_index": widget_index, "node_type": node_type, + "node_title": ref.get("node_title", ""), + "subgraph_id": ref.get("subgraph_id") or "", + "subgraph_name": ref.get("subgraph_name") or "", + "is_top_level": ref.get("is_top_level") is not False, + "locate_node_id": ref.get("locate_node_id"), + "locate_node_type": ref.get("locate_node_type", ""), + "locate_node_title": ref.get("locate_node_title", ""), + "locate_subgraph_id": ( + ref.get("locate_subgraph_id") or "" + ), + "locate_subgraph_name": ( + ref.get("locate_subgraph_name") or "" + ), + "locate_is_top_level": ( + ref.get("locate_is_top_level") is not False + ), + "locate_via_promoted_widget": ref.get( + "locate_via_promoted_widget", False + ), "exists": exists, "strength": strength, "original_path": original_path, @@ -1932,6 +1998,7 @@ class ModelResolverExtension: if isinstance(model_info, dict): name = model_info.get("name", "") directory = model_info.get("directory", "") + node_scope = get_node_scope(node) if name: # Check if this model is already in loaded_models @@ -1951,6 +2018,14 @@ class ModelResolverExtension: "node_id": node.get("id"), "widget_index": None, "node_type": node_type, + "node_title": node.get("title", ""), + "subgraph_id": node_scope["subgraph_id"], + "subgraph_name": node_scope[ + "subgraph_name" + ], + "is_top_level": node_scope[ + "is_top_level" + ], "exists": True, # Embedded models are loaded "strength": None, "original_path": name, diff --git a/tests/test_downloads_tab_workflow_route.mjs b/tests/test_downloads_tab_workflow_route.mjs index 89a52c5..8c0cf75 100644 --- a/tests/test_downloads_tab_workflow_route.mjs +++ b/tests/test_downloads_tab_workflow_route.mjs @@ -46,6 +46,14 @@ const missingBrowserMethodsSource = fs.readFileSync( path.join(projectRoot, 'web/resolver/views/missing_browser_methods.js'), 'utf8' ); +const tabsLoadedMethodsSource = fs.readFileSync( + path.join(projectRoot, 'web/resolver/views/tabs_loaded_methods.js'), + 'utf8' +); +const resolverDialogSource = fs.readFileSync( + path.join(projectRoot, 'web/resolver/resolver_dialog.js'), + 'utf8' +); const renderFormatMethodsSource = fs.readFileSync( path.join(projectRoot, 'web/resolver/utils/render_format_methods.js'), 'utf8' @@ -328,6 +336,44 @@ test('pending node context model resolves against the current workflow analysis assert.equal(isExistingResolvedModel(resolvedModel), true); }); +test('loaded model context menu locates its exact workflow node and subgraph', () => { + const getLoadedModelContext = eval(`(${extractMethod(tabsLoadedMethodsSource, 'getLoadedModelContext')})`); + const handleContextMenuAction = eval(`(${extractMethod(modelInfoMethodsSource, 'handleContextMenuAction')})`); + const calls = []; + const model = getLoadedModelContext.call({}, { + node_id: 19, + subgraph_id: 'subgraph-a', + is_top_level: false, + }); + const dialog = { + _contextMenuModel: model, + hideContextMenu() {}, + getMissingLocateTarget(value) { + return { + nodeId: value.node_id, + subgraphId: value.subgraph_id, + isTopLevel: value.is_top_level, + }; + }, + locateNodeInGraph(nodeId, options) { + calls.push({ nodeId, options }); + }, + }; + + handleContextMenuAction.call(dialog, 'locateNode'); + + assert.equal(model.context_scope, 'loaded_model'); + assert.deepEqual(calls, [{ + nodeId: 19, + options: { + subgraphId: 'subgraph-a', + isTopLevel: false, + }, + }]); + assert.match(resolverDialogSource, /data-menu-action": "locateNode"/); + assert.match(resolverDialogSource, /textContent: "Locate Node"/); +}); + test('queued workflow model selection survives missing data and completes after analysis', () => { const queueWorkflowModelReferenceSelection = eval(`(${extractMethod(missingBrowserMethodsSource, 'queueWorkflowModelReferenceSelection')})`); const applyPendingWorkflowModelSelection = eval(`(${extractMethod(missingBrowserMethodsSource, 'applyPendingWorkflowModelSelection')})`); diff --git a/web/resolver/resolver_dialog.js b/web/resolver/resolver_dialog.js index df6bbda..272aa71 100644 --- a/web/resolver/resolver_dialog.js +++ b/web/resolver/resolver_dialog.js @@ -149,6 +149,15 @@ export class ResolverManagerDialog extends ComfyDialog { $el("div.mr-context-menu-divider.mr-context-menu-divider-workflow", { "data-menu-divider": "workflow" }), + $el("div.mr-context-menu-item.mr-context-menu-action-locate-node", { + "data-menu-action": "locateNode", + onclick: () => this.handleContextMenuAction('locateNode') + }, [ + $el("span.mr-context-menu-item-icon", { + innerHTML: getSvgIcon('locate', 'currentColor', 'mr-context-menu-item-svg') + }), + $el("span", { textContent: "Locate Node" }) + ]), $el("div.mr-context-menu-item.mr-context-menu-action-switch-workflow", { "data-menu-action": "switchWorkflow", onclick: () => this.handleContextMenuAction('switchWorkflow') diff --git a/web/resolver/views/model_info_methods.js b/web/resolver/views/model_info_methods.js index 9fce1f9..35e3d04 100644 --- a/web/resolver/views/model_info_methods.js +++ b/web/resolver/views/model_info_methods.js @@ -37,6 +37,7 @@ export const modelInfoMethods = { const isDownloadRootContext = model?.context_scope === 'download_root'; const isDownloadQueueContext = model?.context_scope === 'download_queue'; const isDownloadHistoryContext = model?.context_scope === 'download_history'; + const isLoadedModelContext = model?.context_scope === 'loaded_model'; const isLocalModelContext = model?.context_scope === 'local_model' || model?.context_scope === 'local_match'; const isFolderOnlyContext = isDownloadFolderContext || isDownloadRootContext; const isSourceModelContext = !isDownloadTableContext && !isFolderOnlyContext && !isDownloadQueueContext && !isDownloadHistoryContext; @@ -46,15 +47,23 @@ export const modelInfoMethods = { && hasLocalPath && Boolean(model?.missing_key || model?.missing_search_key); const showSwitchWorkflow = (isDownloadQueueContext || isDownloadHistoryContext) && Boolean(this.canSwitchToDownloadWorkflow?.(model)); + const showLocateNode = isLoadedModelContext + && model?.node_id !== undefined + && model?.node_id !== null + && model?.node_id !== ''; this.setContextMenuItemVisible('showInfo', isSourceModelContext); this.setContextMenuItemVisible('showMore', canShowMore); this.setContextMenuItemVisible('source', showSourceLink); this.setContextMenuItemVisible('civitai', false); + this.setContextMenuItemVisible('locateNode', showLocateNode); this.setContextMenuItemVisible('switchWorkflow', showSwitchWorkflow); this.setContextMenuItemVisible('compareHashes', showCompareHashes); this.setContextMenuItemVisible('openFolder', showOpenFolder); this.setContextMenuDividerVisible('source', showSourceLink && (isSourceModelContext || canShowMore)); - this.setContextMenuDividerVisible('workflow', showSwitchWorkflow && (isSourceModelContext || canShowMore || showSourceLink)); + this.setContextMenuDividerVisible( + 'workflow', + (showLocateNode || showSwitchWorkflow) && (isSourceModelContext || canShowMore || showSourceLink) + ); this.setContextMenuDividerVisible('folder', (showCompareHashes || showOpenFolder) && (isSourceModelContext || canShowMore || showSourceLink || showSwitchWorkflow)); this.updateContextMenuSourceItem(sourceLink); @@ -107,6 +116,12 @@ export const modelInfoMethods = { this.compareLocalModelHashesWithCurrentFinding(model); } else if (action === 'switchWorkflow') { this.switchToDownloadWorkflow(model); + } else if (action === 'locateNode') { + const locateTarget = this.getMissingLocateTarget(model); + this.locateNodeInGraph(locateTarget.nodeId, { + subgraphId: locateTarget.subgraphId || '', + isTopLevel: locateTarget.isTopLevel + }); } else if (action === 'showInfo') { this.showModelInfo(model); } else if (action === 'showMore') { diff --git a/web/resolver/views/tabs_loaded_methods.js b/web/resolver/views/tabs_loaded_methods.js index 48ca0e7..4513c73 100644 --- a/web/resolver/views/tabs_loaded_methods.js +++ b/web/resolver/views/tabs_loaded_methods.js @@ -176,6 +176,13 @@ export const tabsLoadedMethods = { return `<${this.getCategoryTokenName(category)}:${name}:${strength || '1.00'}>`; }, + getLoadedModelContext(model = {}) { + return { + ...model, + context_scope: 'loaded_model' + }; + }, + displayLoadedModels(container, data) { const loadedModels = data.loaded_models || []; const total = data.total || 0; @@ -274,7 +281,7 @@ export const tabsLoadedMethods = { for (const model of modelsObj.active) { const { name, strength } = this.getModelNameAndStrength(model); - html += `${this.escapeHtml(name)}${strength !== null ? `${this.escapeHtml(strength)}` : ''}`; + html += `${this.escapeHtml(name)}${strength !== null ? `${this.escapeHtml(strength)}` : ''}`; } html += ``; } @@ -291,7 +298,7 @@ export const tabsLoadedMethods = { for (const model of modelsObj.inactive) { const { name, strength } = this.getModelNameAndStrength(model); - html += `${this.escapeHtml(name)}${strength !== null ? `${this.escapeHtml(strength)}` : ''}`; + html += `${this.escapeHtml(name)}${strength !== null ? `${this.escapeHtml(strength)}` : ''}`; } html += ``; }