[backport core/1.43] refactor: extract missing model refresh pipeline (#11887)

Backport of #11751 to `core/1.43`

Automatically created by backport workflow.

┆Issue is synchronized with this [Notion
page](https://www.notion.so/PR-11887-backport-core-1-43-refactor-extract-missing-model-refresh-pipeline-3566d73d365081f687aafe82b655f135)
by [Unito](https://www.unito.io)

Co-authored-by: jaeone94 <89377375+jaeone94@users.noreply.github.com>
Co-authored-by: DrJKL <DrJKL0424@gmail.com>
Co-authored-by: Alexander Brown <drjkl@comfy.org>
This commit is contained in:
Comfy Org PR Bot
2026-05-04 15:15:33 +09:00
committed by GitHub
parent e8b5e92c48
commit e1bdaa5778
10 changed files with 1303 additions and 589 deletions

View File

@@ -1,4 +1,5 @@
import { createPinia, setActivePinia } from 'pinia'
import { createTestingPinia } from '@pinia/testing'
import { setActivePinia } from 'pinia'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import type {
@@ -6,10 +7,7 @@ import type {
LGraphCanvas,
LGraphNode
} from '@/lib/litegraph/src/litegraph'
import type {
ComfyWorkflowJSON,
ModelFile
} from '@/platform/workflow/validation/schemas/workflowSchema'
import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
import { ComfyApp } from './app'
import { createNode } from '@/utils/litegraphUtil'
import {
@@ -22,14 +20,13 @@ import {
} from '@/composables/usePaste'
import { getWorkflowDataFromFile } from '@/scripts/metadata/parser'
import { useMissingModelStore } from '@/platform/missingModel/missingModelStore'
import type { LoadedComfyWorkflow } from '@/platform/workflow/management/stores/comfyWorkflow'
import type { MissingModelCandidate } from '@/platform/missingModel/types'
const {
mockToastStore,
mockExtensionService,
mockNodeOutputStore,
mockWorkspaceWorkflow
mockWorkspaceWorkflow,
mockRefreshMissingModelPipeline
} = vi.hoisted(() => ({
mockToastStore: {
addAlert: vi.fn(),
@@ -44,8 +41,9 @@ const {
refreshNodeOutputs: vi.fn()
},
mockWorkspaceWorkflow: {
activeWorkflow: null as unknown
}
activeWorkflow: null
},
mockRefreshMissingModelPipeline: vi.fn()
}))
vi.mock('@/utils/litegraphUtil', () => ({
@@ -88,6 +86,11 @@ vi.mock('@/stores/workspaceStore', () => ({
}))
}))
vi.mock('@/platform/missingModel/missingModelPipeline', () => ({
refreshMissingModelPipeline: mockRefreshMissingModelPipeline,
runMissingModelPipeline: vi.fn()
}))
function createMockNode(options: { [K in keyof LGraphNode]?: any } = {}) {
return {
id: 1,
@@ -115,16 +118,6 @@ function createTestFile(name: string, type: string): File {
return new File([''], name, { type })
}
type ComfyAppMissingModelPipelineTarget = {
runMissingModelPipeline: (
graphData: ComfyWorkflowJSON,
options?: { silent?: boolean; missingNodeTypes?: string[] }
) => Promise<{
missingModels: ModelFile[]
confirmedCandidates: MissingModelCandidate[]
}>
}
function createWorkflowGraphData(): ComfyWorkflowJSON {
return {
last_node_id: 0,
@@ -143,7 +136,7 @@ describe('ComfyApp', () => {
let mockCanvas: LGraphCanvas
beforeEach(() => {
setActivePinia(createPinia())
setActivePinia(createTestingPinia({ stubActions: false }))
vi.clearAllMocks()
app = new ComfyApp()
mockCanvas = createMockCanvas() as LGraphCanvas
@@ -187,104 +180,32 @@ describe('ComfyApp', () => {
})
describe('refreshMissingModels', () => {
function mockRefreshMissingModelsApp(
graphData: ComfyWorkflowJSON,
candidates: MissingModelCandidate[] = []
) {
mockWorkspaceWorkflow.activeWorkflow = null
Reflect.set(app, 'rootGraphInternal', {
it('delegates to the app-independent missing model refresh pipeline', async () => {
const graph = {
nodes: [],
serialize: vi.fn(() => graphData)
})
serialize: vi.fn(() => createWorkflowGraphData())
}
const result = {
missingModels: [],
confirmedCandidates: []
}
Reflect.set(app, 'rootGraphInternal', graph)
vi.spyOn(app, 'reloadNodeDefs').mockResolvedValue()
const appWithPrivate =
app as unknown as ComfyAppMissingModelPipelineTarget
const pipelineSpy = vi
.spyOn(appWithPrivate, 'runMissingModelPipeline')
.mockResolvedValue({
missingModels: [],
confirmedCandidates: []
})
useMissingModelStore().missingModelCandidates = candidates
return pipelineSpy
}
mockRefreshMissingModelPipeline.mockResolvedValue(result)
it('reuses active workflow model metadata when refreshing the current graph', async () => {
const graphData = createWorkflowGraphData()
const activeModels = [
{
name: 'embedded.safetensors',
url: 'https://example.com/embedded.safetensors',
directory: 'checkpoints'
}
]
const pipelineSpy = mockRefreshMissingModelsApp(graphData, [
{
nodeId: '1',
nodeType: 'CheckpointLoaderSimple',
widgetName: 'ckpt_name',
name: 'candidate.safetensors',
url: 'https://example.com/candidate.safetensors',
directory: 'checkpoints',
isMissing: true,
isAssetSupported: true
}
])
mockWorkspaceWorkflow.activeWorkflow = {
activeState: { models: activeModels }
} as LoadedComfyWorkflow
await expect(app.refreshMissingModels({ silent: false })).resolves.toBe(
result
)
await app.refreshMissingModels({ silent: false })
expect(mockRefreshMissingModelPipeline).toHaveBeenCalledWith({
graph,
reloadNodeDefs: expect.any(Function),
missingModelStore: useMissingModelStore(),
silent: false
})
await mockRefreshMissingModelPipeline.mock.calls[0][0].reloadNodeDefs()
expect(app.reloadNodeDefs).toHaveBeenCalled()
expect(pipelineSpy).toHaveBeenCalledWith(
expect.objectContaining({ models: activeModels }),
{ silent: false }
)
})
it('falls back to current missing model metadata when workflow state has no models', async () => {
const graphData = createWorkflowGraphData()
const pipelineSpy = mockRefreshMissingModelsApp(graphData, [
{
nodeId: '1',
nodeType: 'CheckpointLoaderSimple',
widgetName: 'ckpt_name',
name: 'candidate.safetensors',
url: 'https://example.com/candidate.safetensors',
directory: 'checkpoints',
hash: 'abc123',
hashType: 'sha256',
isMissing: true,
isAssetSupported: true
},
{
nodeId: '2',
nodeType: 'CheckpointLoaderSimple',
widgetName: 'ckpt_name',
name: 'missing-url.safetensors',
directory: 'checkpoints',
isMissing: true,
isAssetSupported: true
}
])
await app.refreshMissingModels()
expect(pipelineSpy).toHaveBeenCalledWith(
expect.objectContaining({
models: [
{
name: 'candidate.safetensors',
url: 'https://example.com/candidate.safetensors',
directory: 'checkpoints',
hash: 'abc123',
hash_type: 'sha256'
}
]
}),
{ silent: true }
)
})
})

View File

@@ -34,13 +34,13 @@ import { useWorkflowValidation } from '@/platform/workflow/validation/composable
import type {
ComfyApiWorkflow,
ComfyWorkflowJSON,
ModelFile,
NodeId
} from '@/platform/workflow/validation/schemas/workflowSchema'
import {
isSubgraphDefinition,
collectSubgraphDefinitions,
buildSubgraphExecutionPaths
} from '@/platform/workflow/validation/schemas/workflowSchema'
} from '@/platform/workflow/core/utils/workflowFlattening'
import type { FlattenableWorkflowNode } from '@/platform/workflow/core/utils/workflowFlattening'
import type {
ExecutionErrorWsMessage,
NodeError,
@@ -73,7 +73,6 @@ import { useNodeOutputStore } from '@/stores/nodeOutputStore'
import { useJobPreviewStore } from '@/stores/jobPreviewStore'
import { KeyComboImpl } from '@/platform/keybindings/keyCombo'
import { useKeybindingStore } from '@/platform/keybindings/keybindingStore'
import { useModelStore } from '@/stores/modelStore'
import { SYSTEM_NODE_DEFS, useNodeDefStore } from '@/stores/nodeDefStore'
import { useNodeReplacementStore } from '@/platform/nodeReplacement/nodeReplacementStore'
@@ -87,12 +86,11 @@ import type { NodeExecutionId } from '@/types/nodeIdentification'
import { graphToPrompt } from '@/utils/executionUtil'
import { getCnrIdFromProperties } from '@/platform/nodeReplacement/cnrIdUtil'
import { rescanAndSurfaceMissingNodes } from '@/platform/nodeReplacement/missingNodeScan'
import type { MissingModelCandidate } from '@/platform/missingModel/types'
import {
scanAllModelCandidates,
enrichWithEmbeddedMetadata,
verifyAssetSupportedCandidates
} from '@/platform/missingModel/missingModelScan'
refreshMissingModelPipeline,
runMissingModelPipeline
} from '@/platform/missingModel/missingModelPipeline'
import type { MissingModelPipelineResult } from '@/platform/missingModel/missingModelPipeline'
import { useMissingModelStore } from '@/platform/missingModel/missingModelStore'
import { useMissingMediaStore } from '@/platform/missingMedia/missingMediaStore'
import type { MissingMediaCandidate } from '@/platform/missingMedia/types'
@@ -100,8 +98,6 @@ import {
scanAllMediaCandidates,
verifyCloudMediaCandidates
} from '@/platform/missingMedia/missingMediaScan'
import { assetService } from '@/platform/assets/services/assetService'
import { useModelToNodeStore } from '@/stores/modelToNodeStore'
import { anyItemOverlapsRect } from '@/utils/mathUtil'
import {
@@ -154,11 +150,6 @@ import {
pasteVideoNodes
} from '@/composables/usePaste'
interface MissingModelPipelineOptions {
missingNodeTypes?: MissingNodeType[]
silent?: boolean
}
export const ANIM_PREVIEW_WIDGET = '$$comfy_animation_preview'
export function sanitizeNodeName(string: string) {
@@ -1225,7 +1216,7 @@ export class ComfyApp {
// Collect missing node types from all nodes (root + subgraphs)
const collectMissingNodes = (
nodes: ComfyWorkflowJSON['nodes'],
nodes: readonly FlattenableWorkflowNode[],
pathPrefix: string = '',
displayName: string = ''
) => {
@@ -1270,21 +1261,21 @@ export class ComfyApp {
}
collectMissingNodes(graphData.nodes)
const subgraphDefs = graphData.definitions?.subgraphs ?? []
const subgraphDefs = collectSubgraphDefinitions(
graphData.definitions?.subgraphs ?? []
)
const subgraphContainerIdMap = buildSubgraphExecutionPaths(
graphData.nodes,
subgraphDefs
)
for (const subgraph of subgraphDefs) {
if (isSubgraphDefinition(subgraph)) {
const paths = subgraphContainerIdMap.get(subgraph.id) ?? []
for (const pathPrefix of paths) {
collectMissingNodes(
subgraph.nodes,
pathPrefix,
subgraph.name || subgraph.id
)
}
const paths = subgraphContainerIdMap.get(subgraph.id) ?? []
for (const pathPrefix of paths) {
collectMissingNodes(
subgraph.nodes,
pathPrefix,
subgraph.name || subgraph.id
)
}
}
@@ -1454,7 +1445,10 @@ export class ComfyApp {
)
if (!skipAssetScans) {
await this.runMissingModelPipeline(graphData, {
await runMissingModelPipeline({
graph: this.rootGraph,
graphData,
missingModelStore: useMissingModelStore(),
missingNodeTypes: activeMissingNodeTypes,
silent: silentAssetErrors
})
@@ -1477,201 +1471,14 @@ export class ComfyApp {
}
}
private async runMissingModelPipeline(
graphData: ComfyWorkflowJSON,
{ missingNodeTypes, silent = false }: MissingModelPipelineOptions = {}
): Promise<{
missingModels: ModelFile[]
confirmedCandidates: MissingModelCandidate[]
}> {
const missingModelStore = useMissingModelStore()
const controller = missingModelStore.createVerificationAbortController()
const getDirectory = (nodeType: string) =>
useModelToNodeStore().getCategoryForNodeType(nodeType)
const candidates = scanAllModelCandidates(
this.rootGraph,
isCloud
? (nodeType, widgetName) =>
assetService.shouldUseAssetBrowser(nodeType, widgetName)
: () => false,
getDirectory
)
const modelStore = useModelStore()
await modelStore.loadModelFolders()
const enrichedAll = await enrichWithEmbeddedMetadata(
candidates,
graphData,
async (name, directory) => {
const folder = await modelStore.getLoadedModelFolder(directory)
const models = folder?.models
return !!(
models && Object.values(models).some((m) => m.file_name === name)
)
},
isCloud
? (nodeType, widgetName) =>
assetService.shouldUseAssetBrowser(nodeType, widgetName)
: undefined
)
// Drop candidates whose enclosing subgraph is muted/bypassed. Per-node
// scans only checked each node's own mode; the cascade from an
// inactive container to its interior happens here.
// Asymmetric on purpose: a candidate dropped here is not resurrected if
// the user un-bypasses the container mid-verification. The realtime
// mode-change path (handleNodeModeChange → scanAndAddNodeErrors) is
// responsible for surfacing errors after an un-bypass.
const enrichedCandidates = enrichedAll.filter(
(c) =>
c.nodeId == null ||
isAncestorPathActive(this.rootGraph, String(c.nodeId))
)
const missingModels: ModelFile[] = enrichedCandidates
.filter((c) => c.isMissing === true && c.url)
.map((c) => ({
name: c.name,
url: c.url ?? '',
directory: c.directory ?? '',
hash: c.hash,
hash_type: c.hashType
}))
const confirmedCandidates = enrichedCandidates.filter(
(c) => c.isMissing === true
)
const activeWf = useWorkspaceStore().workflow.activeWorkflow
updatePendingWarnings(activeWf, {
...(missingNodeTypes ? { missingNodeTypes } : {}),
missingModelCandidates: confirmedCandidates
})
if (enrichedCandidates.length) {
if (isCloud) {
void verifyAssetSupportedCandidates(
enrichedCandidates,
controller.signal
)
.then(() => {
if (controller.signal.aborted) return
// Re-check ancestor: user may have bypassed a container
// while verification was in flight.
const confirmed = enrichedCandidates.filter((c) =>
isMissingCandidateActive(this.rootGraph, c)
)
useExecutionErrorStore().surfaceMissingModels(confirmed, { silent })
this.cacheModelCandidates(activeWf, confirmed)
})
.catch((err) => {
console.warn(
'[Missing Model Pipeline] Asset verification failed:',
err
)
useToastStore().add({
severity: 'warn',
summary: st(
'toastMessages.missingModelVerificationFailed',
'Failed to verify missing models. Some models may not be shown in the Errors tab.'
),
life: 5000
})
})
} else {
const confirmed = enrichedCandidates.filter((c) => c.isMissing === true)
if (!confirmed.length) {
useExecutionErrorStore().surfaceMissingModels([], { silent })
this.cacheModelCandidates(activeWf, [])
} else {
void api
.getFolderPaths()
.then((paths) => {
if (controller.signal.aborted) return
missingModelStore.setFolderPaths(paths)
})
.catch((err) => {
console.warn(
'[Missing Model Pipeline] Failed to fetch folder paths:',
err
)
})
.finally(() => {
if (controller.signal.aborted) return
useExecutionErrorStore().surfaceMissingModels(confirmed, {
silent
})
this.cacheModelCandidates(activeWf, confirmed)
})
void Promise.allSettled(
confirmed
.filter((c) => c.url)
.map(async (c) => {
const { fetchModelMetadata } =
await import('@/platform/missingModel/missingModelDownload')
const metadata = await fetchModelMetadata(c.url!)
if (!controller.signal.aborted && metadata.fileSize !== null) {
missingModelStore.setFileSize(c.url!, metadata.fileSize)
}
})
)
}
}
} else {
useExecutionErrorStore().surfaceMissingModels([], { silent })
this.cacheModelCandidates(activeWf, [])
}
return { missingModels, confirmedCandidates }
}
async refreshMissingModels(options: { silent?: boolean } = {}): Promise<{
missingModels: ModelFile[]
confirmedCandidates: MissingModelCandidate[]
}> {
await this.reloadNodeDefs()
const graphData = this.rootGraph.serialize() as unknown as ComfyWorkflowJSON
const activeWorkflowState =
useWorkspaceStore().workflow.activeWorkflow?.activeState
const currentModelMetadata =
useMissingModelStore()
.missingModelCandidates?.filter(
(
candidate
): candidate is MissingModelCandidate & {
url: string
directory: string
} => !!candidate.url && !!candidate.directory
)
.map((candidate) => ({
name: candidate.name,
url: candidate.url,
directory: candidate.directory,
hash: candidate.hash,
hash_type: candidate.hashType
})) ?? []
const models = activeWorkflowState?.models?.length
? activeWorkflowState.models
: currentModelMetadata
return this.runMissingModelPipeline(
models.length ? { ...graphData, models } : graphData,
{
silent: options.silent ?? true
}
)
}
private cacheModelCandidates(
wf: ComfyWorkflow | null,
confirmed: MissingModelCandidate[]
) {
if (!wf) return
updatePendingWarnings(wf, {
missingModelCandidates: confirmed
async refreshMissingModels(
options: { silent?: boolean } = {}
): Promise<MissingModelPipelineResult> {
return refreshMissingModelPipeline({
graph: this.rootGraph,
reloadNodeDefs: () => this.reloadNodeDefs(),
missingModelStore: useMissingModelStore(),
silent: options.silent ?? true
})
}