mirror of
https://github.com/Comfy-Org/ComfyUI_frontend.git
synced 2026-05-24 14:45:36 +00:00
[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:
@@ -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 }
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@@ -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
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user