diff --git a/core/settings.py b/core/settings.py index ea18797..3127a54 100644 --- a/core/settings.py +++ b/core/settings.py @@ -238,6 +238,9 @@ def normalize_download_path_mode(value: Any, auto_fill_subfolder: Any = None) -> def normalize_settings(settings: Optional[Mapping[str, Any]]) -> Dict[str, Any]: data = dict(settings or {}) + data["auto_refresh_comfy_models_after_apply"] = bool_setting( + data.get("auto_refresh_comfy_models_after_apply"), True + ) data["download_path_mode"] = normalize_download_path_mode( data.get("download_path_mode"), data.get("auto_fill_subfolder") ) diff --git a/web/resolver/actions/queue_methods.js b/web/resolver/actions/queue_methods.js index 1950508..98b8012 100644 --- a/web/resolver/actions/queue_methods.js +++ b/web/resolver/actions/queue_methods.js @@ -1814,6 +1814,7 @@ export const queueMethods = { const data = await response.json(); if (data.success) { const optimisticData = this.getOptimisticAnalysisDataAfterApply(appliedSelections); + await this.refreshComfyModelCatalogAfterApply?.(data.workflow, applyResolutions); await this.updateWorkflowInComfyUI(data.workflow); this.rememberAppliedResolvedSelections(appliedSelections); diff --git a/web/resolver/actions/resolve_download_methods.js b/web/resolver/actions/resolve_download_methods.js index db175bf..5dd5e46 100644 --- a/web/resolver/actions/resolve_download_methods.js +++ b/web/resolver/actions/resolve_download_methods.js @@ -58,6 +58,8 @@ export const resolveDownloadMethods = { log.debug('Resolve response: success=', data.success, ' missing count:', data.workflow?.nodes?.length); if (data.success) { + await this.refreshComfyModelCatalogAfterApply?.(data.workflow, resolutions); + // Update workflow in ComfyUI await this.updateWorkflowInComfyUI(data.workflow); @@ -149,6 +151,8 @@ export const resolveDownloadMethods = { const resolveData = await resolveResponse.json(); if (resolveData.success) { + await this.refreshComfyModelCatalogAfterApply?.(resolveData.workflow, resolutions); + // Update workflow in ComfyUI await this.updateWorkflowInComfyUI(resolveData.workflow); diff --git a/web/resolver/resolver_dialog.js b/web/resolver/resolver_dialog.js index aab1ce5..7b80f30 100644 --- a/web/resolver/resolver_dialog.js +++ b/web/resolver/resolver_dialog.js @@ -98,6 +98,7 @@ export class ResolverManagerDialog extends ComfyDialog { this._workflowRefreshGeneration = 0; this._workflowRefreshExpectedRoute = null; this._workflowRefreshPreviousSignature = null; + this._comfyModelCatalogRefreshPromise = null; this._boundHandleViewportResize = () => this.scheduleModalViewportClamp(true); // Create backdrop overlay for click-outside-to-close diff --git a/web/resolver/search/download_target_methods.js b/web/resolver/search/download_target_methods.js index 34daf22..e92a54c 100644 --- a/web/resolver/search/download_target_methods.js +++ b/web/resolver/search/download_target_methods.js @@ -2379,6 +2379,7 @@ export const downloadTargetMethods = { hf_use_brave_fallback: localStorage.getItem('ModelResolver.hfUseBraveFallback') !== 'false', auto_fill_base_model: localStorage.getItem('ModelResolver.autoFillBaseModel') !== 'false', auto_fill_subfolder: localStorage.getItem('ModelResolver.autoFillSubfolder') !== 'false', + auto_refresh_comfy_models_after_apply: localStorage.getItem('ModelResolver.autoRefreshComfyModelsAfterApply') !== 'false', frontend_logs_enabled: storedFrontendLogsEnabled === null ? frontendLogger.enabled !== false : storedFrontendLogsEnabled !== 'false', @@ -2448,6 +2449,8 @@ export const downloadTargetMethods = { localStorage.setItem('ModelResolver.autoFillBaseModel', data.auto_fill_base_model ? 'true' : 'false'); if (data.auto_fill_subfolder !== undefined) localStorage.setItem('ModelResolver.autoFillSubfolder', data.auto_fill_subfolder ? 'true' : 'false'); + if (data.auto_refresh_comfy_models_after_apply !== undefined) + localStorage.setItem('ModelResolver.autoRefreshComfyModelsAfterApply', data.auto_refresh_comfy_models_after_apply ? 'true' : 'false'); if (data.download_path_mode !== undefined) localStorage.setItem('ModelResolver.downloadPathMode', this.normalizeDownloadPathMode(data.download_path_mode)); if (data.download_path_templates !== undefined) diff --git a/web/resolver/shell/workflow_update_methods.js b/web/resolver/shell/workflow_update_methods.js index bb895cf..03a0419 100644 --- a/web/resolver/shell/workflow_update_methods.js +++ b/web/resolver/shell/workflow_update_methods.js @@ -3,6 +3,272 @@ import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; export const workflowUpdateMethods = { + async refreshComfyModelCatalogAfterApply(workflow = null, resolutions = []) { + if (!this.isAutoRefreshComfyModelsAfterApplyEnabled()) { + return false; + } + + if (!this.shouldRefreshComfyModelCatalogForApply(workflow, resolutions)) { + return false; + } + + if (this._comfyModelCatalogRefreshPromise) { + return this._comfyModelCatalogRefreshPromise; + } + + this._comfyModelCatalogRefreshPromise = (async () => { + try { + if (typeof app?.refreshComboInNodes === 'function') { + await app.refreshComboInNodes(); + return true; + } + + const nodeDefs = await this.fetchComfyNodeDefs(); + this.applyComfyNodeDefs(nodeDefs); + this.refreshGraphComboWidgetsFromNodeDefs(nodeDefs); + app?.graph?.setDirtyCanvas?.(true, true); + return true; + } catch (error) { + console.warn('Model Resolver: could not refresh ComfyUI model catalog:', error); + return false; + } finally { + this._comfyModelCatalogRefreshPromise = null; + } + })(); + + return this._comfyModelCatalogRefreshPromise; + }, + + isAutoRefreshComfyModelsAfterApplyEnabled() { + const tokens = this.getStoredTokens?.(); + return tokens?.auto_refresh_comfy_models_after_apply !== false; + }, + + async fetchComfyNodeDefs() { + const response = await api.fetchApi(`/object_info?model_resolver_refresh=${Date.now()}`); + if (!response.ok) { + throw new Error(`object_info refresh failed: ${response.status}`); + } + + return await response.json(); + }, + + shouldRefreshComfyModelCatalogForApply(workflow = null, resolutions = []) { + const targets = this.getComfyCatalogApplyTargets(workflow, resolutions); + if (!targets.length) return false; + + return targets.some(target => !this.isComfyCatalogTargetAvailable(target)); + }, + + getComfyCatalogApplyTargets(workflow = null, resolutions = []) { + if (!workflow || !Array.isArray(resolutions) || !resolutions.length) return []; + + const targets = []; + for (const resolution of resolutions) { + const node = this.findWorkflowNodeForResolution(workflow, resolution); + if (!node) continue; + + const widgetIndex = Number(resolution?.widget_index); + if (!Number.isInteger(widgetIndex)) continue; + + const value = this.getAppliedWidgetValueFromNode(node, widgetIndex, resolution); + if (!value) continue; + + const nodeType = node.comfyClass || node.type || resolution.node_type || ''; + const widgetName = this.getWorkflowNodeWidgetName(node, widgetIndex) + || this.getGraphNodeWidgetName(node, widgetIndex, resolution) + || this.getComfyWidgetNameByIndex(nodeType, widgetIndex); + if (!nodeType || !widgetName) continue; + + targets.push({ nodeType, widgetName, value }); + } + + return targets; + }, + + findWorkflowNodeForResolution(workflow = {}, resolution = {}) { + const nodeId = String(resolution?.node_id ?? ''); + if (!nodeId) return null; + + const isTopLevel = resolution?.is_top_level !== false; + if (isTopLevel) { + return (workflow.nodes || []).find(node => String(node?.id) === nodeId) || null; + } + + const subgraphId = String(resolution?.subgraph_id || ''); + const subgraphs = workflow.definitions?.subgraphs || []; + for (const subgraph of subgraphs) { + if (subgraphId && String(subgraph?.id) !== subgraphId) continue; + const node = (subgraph?.nodes || []).find(item => String(item?.id) === nodeId); + if (node) return node; + } + + return null; + }, + + getAppliedWidgetValueFromNode(node = {}, widgetIndex = -1, resolution = {}) { + const widgetsValues = Array.isArray(node.widgets_values) ? node.widgets_values : []; + const rawValue = widgetsValues[widgetIndex]; + const nestedKey = resolution?.nested_key; + if (nestedKey && rawValue && typeof rawValue === 'object' && !Array.isArray(rawValue)) { + return String(rawValue[nestedKey] || '').trim(); + } + + if (typeof rawValue === 'string') return rawValue.trim(); + return ''; + }, + + getWorkflowNodeWidgetName(node = {}, widgetIndex = -1) { + const widget = Array.isArray(node.widgets) ? node.widgets[widgetIndex] : null; + return String(widget?.name || widget?.label || widget?.widget || '').trim(); + }, + + getGraphNodeWidgetName(workflowNode = {}, widgetIndex = -1, resolution = {}) { + if (resolution?.is_top_level === false) return ''; + + const nodeId = String(resolution?.node_id ?? workflowNode?.id ?? ''); + const graphNode = (app?.graph?.nodes || []).find(node => String(node?.id) === nodeId); + const widget = graphNode?.widgets?.[widgetIndex]; + return String(widget?.name || widget?.label || '').trim(); + }, + + getComfyWidgetNameByIndex(nodeType = '', widgetIndex = -1) { + const nodeDef = this.getCurrentComfyNodeDef(nodeType); + if (!nodeDef || widgetIndex < 0) return ''; + + const names = []; + const input = nodeDef.input || {}; + for (const sectionName of ['required', 'optional']) { + const section = input[sectionName]; + if (!section || typeof section !== 'object') continue; + for (const [name, spec] of Object.entries(section)) { + if (this.isComfyWidgetInputSpec(spec)) { + names.push(name); + } + } + } + + return names[widgetIndex] || ''; + }, + + getCurrentComfyNodeDef(nodeType = '') { + if (!nodeType) return null; + + return app?.nodeDefs?.[nodeType] + || globalThis?.LiteGraph?.registered_node_types?.[nodeType]?.nodeData + || globalThis?.LiteGraph?.registered_node_types?.[nodeType]?.prototype?.nodeData + || null; + }, + + isComfyWidgetInputSpec(inputSpec) { + if (!Array.isArray(inputSpec)) return false; + if (Array.isArray(inputSpec[0])) return true; + const type = String(inputSpec[0] || '').toUpperCase(); + return ['STRING', 'INT', 'FLOAT', 'BOOLEAN'].includes(type); + }, + + isComfyCatalogTargetAvailable(target = {}) { + const values = this.getCurrentComfyCatalogValues(target.nodeType, target.widgetName); + if (!values) return true; + + const wanted = this.normalizeComfyCatalogValue(target.value); + return values.some(value => this.normalizeComfyCatalogValue(value) === wanted); + }, + + getCurrentComfyCatalogValues(nodeType = '', widgetName = '') { + const nodeDef = this.getCurrentComfyNodeDef(nodeType); + const values = this.getComfyComboValuesFromSpec( + this.getComfyWidgetInputSpec(nodeDef, widgetName) + ); + if (values) return values; + + const graphNodes = app?.graph?.nodes || []; + for (const node of graphNodes) { + const type = node?.comfyClass || node?.type; + if (type !== nodeType || !Array.isArray(node.widgets)) continue; + + const widget = node.widgets.find(item => item?.name === widgetName); + const widgetValues = widget?.options?.values; + if (Array.isArray(widgetValues)) return widgetValues; + } + + return null; + }, + + normalizeComfyCatalogValue(value = '') { + return String(value || '').trim().replace(/\\/g, '/'); + }, + + applyComfyNodeDefs(nodeDefs = {}) { + if (!nodeDefs || typeof nodeDefs !== 'object') return; + + try { + if (app && typeof app === 'object') { + if (!app.nodeDefs || typeof app.nodeDefs !== 'object') { + app.nodeDefs = {}; + } + Object.assign(app.nodeDefs, nodeDefs); + } + } catch (error) { + console.warn('Model Resolver: could not update app.nodeDefs:', error); + } + + const registeredTypes = globalThis?.LiteGraph?.registered_node_types; + if (!registeredTypes || typeof registeredTypes !== 'object') return; + + for (const [nodeType, nodeDef] of Object.entries(nodeDefs)) { + const registered = registeredTypes[nodeType]; + if (!registered) continue; + + try { + registered.nodeData = nodeDef; + registered.prototype.nodeData = nodeDef; + } catch (error) { + console.warn(`Model Resolver: could not update node definition for ${nodeType}:`, error); + } + } + }, + + getComfyWidgetInputSpec(nodeDef = {}, widgetName = '') { + const input = nodeDef?.input; + if (!input || !widgetName) return null; + + return input.required?.[widgetName] + || input.optional?.[widgetName] + || input.hidden?.[widgetName] + || null; + }, + + getComfyComboValuesFromSpec(inputSpec) { + if (!Array.isArray(inputSpec)) return null; + + const values = inputSpec[0]; + return Array.isArray(values) ? values : null; + }, + + refreshGraphComboWidgetsFromNodeDefs(nodeDefs = {}) { + const graphNodes = app?.graph?.nodes; + if (!Array.isArray(graphNodes) || !nodeDefs || typeof nodeDefs !== 'object') return; + + for (const node of graphNodes) { + const nodeType = node?.comfyClass || node?.type; + const nodeDef = nodeDefs[nodeType]; + if (!nodeDef || !Array.isArray(node.widgets)) continue; + + for (const widget of node.widgets) { + const values = this.getComfyComboValuesFromSpec( + this.getComfyWidgetInputSpec(nodeDef, widget?.name) + ); + if (!values) continue; + + if (!widget.options || typeof widget.options !== 'object') { + widget.options = {}; + } + widget.options.values = values; + } + } + }, + /** * Extract model page URL from a download URL * HuggingFace file: https://huggingface.co/Owner/Repo/resolve/main/file.safetensors -> https://huggingface.co/Owner/Repo/blob/main/file.safetensors diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index 3325e8d..7b84c5b 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -296,6 +296,15 @@ export const optionsMethods = { + @@ -607,6 +616,7 @@ export const optionsMethods = { const hfUseBraveFallbackInput = this.contentElement.querySelector('#mr-options-hf-use-brave-fallback'); 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 downloadPathModeInput = this.contentElement.querySelector('#mr-options-download-path-mode'); const defaultRootSelectInputs = Array.from(this.contentElement.querySelectorAll('.mr-options-default-root')); const templatePresetInputs = Array.from(this.contentElement.querySelectorAll('.mr-options-template-preset')); @@ -660,6 +670,7 @@ export const optionsMethods = { hfUseBraveFallbackInput, autoFillBaseModelInput, autoFillSubfolderInput, + autoRefreshComfyModelsInput, downloadPathModeInput, ...defaultRootSelectInputs, ...templatePresetInputs, @@ -1613,6 +1624,7 @@ export const optionsMethods = { hf_use_brave_fallback: Boolean(hfUseBraveFallbackInput?.checked), auto_fill_base_model: Boolean(autoFillBaseModelInput?.checked), auto_fill_subfolder: Boolean(autoFillSubfolderInput?.checked), + auto_refresh_comfy_models_after_apply: Boolean(autoRefreshComfyModelsInput?.checked), download_path_mode: downloadPathMode, download_path_templates: downloadPathTemplates, base_model_path_mappings: baseModelPathMappings, @@ -1637,6 +1649,7 @@ export const optionsMethods = { localStorage.setItem('ModelResolver.hfUseBraveFallback', newSettings.hf_use_brave_fallback ? 'true' : 'false'); localStorage.setItem('ModelResolver.autoFillBaseModel', newSettings.auto_fill_base_model ? 'true' : 'false'); localStorage.setItem('ModelResolver.autoFillSubfolder', newSettings.auto_fill_subfolder ? 'true' : 'false'); + localStorage.setItem('ModelResolver.autoRefreshComfyModelsAfterApply', newSettings.auto_refresh_comfy_models_after_apply ? 'true' : 'false'); localStorage.setItem('ModelResolver.downloadPathMode', newSettings.download_path_mode); localStorage.setItem('ModelResolver.downloadPathTemplates', JSON.stringify(newSettings.download_path_templates)); localStorage.setItem('ModelResolver.baseModelPathMappings', JSON.stringify(newSettings.base_model_path_mappings));