Skip to content

Commit 69c5882

Browse files
JPPhotoJPPhoto
authored
Improve image collection chaining (#9168)
* Improve image collection chaining * chore: openapi.json * Fix image collection migration paths * Fix node update knip exports * Fix image collection version migration gate
1 parent 6ca505f commit 69c5882

11 files changed

Lines changed: 696 additions & 108 deletions

File tree

‎invokeai/app/invocations/primitives.py‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -279,15 +279,28 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
279279
title="Image Collection Primitive",
280280
tags=["primitives", "image", "collection"],
281281
category="primitives",
282-
version="1.0.1",
282+
version="1.0.2",
283283
)
284284
class ImageCollectionInvocation(BaseInvocation):
285285
"""A collection of image primitive values"""
286286

287-
collection: list[ImageField] = InputField(description="The collection of image values")
287+
collection: Optional[list[ImageField]] = InputField(
288+
default=None,
289+
description="An optional image collection to append to",
290+
input=Input.Connection,
291+
title="Collection",
292+
ui_order=0,
293+
)
294+
images: Optional[list[ImageField]] = InputField(
295+
default=None,
296+
description="The images to append to the collection",
297+
input=Input.Direct,
298+
title="Images",
299+
ui_order=1,
300+
)
288301

289302
def invoke(self, context: InvocationContext) -> ImageCollectionOutput:
290-
return ImageCollectionOutput(collection=self.collection)
303+
return ImageCollectionOutput(collection=[*(self.collection or []), *(self.images or [])])
291304

292305

293306
# endregion

‎invokeai/frontend/web/openapi.json‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32518,11 +32518,34 @@
3251832518
}
3251932519
],
3252032520
"default": null,
32521-
"description": "The collection of image values",
32521+
"description": "An optional image collection to append to",
3252232522
"field_kind": "input",
32523-
"input": "any",
32524-
"orig_required": true,
32525-
"title": "Collection"
32523+
"input": "connection",
32524+
"orig_default": null,
32525+
"orig_required": false,
32526+
"title": "Collection",
32527+
"ui_order": 0
32528+
},
32529+
"images": {
32530+
"anyOf": [
32531+
{
32532+
"items": {
32533+
"$ref": "#/components/schemas/ImageField"
32534+
},
32535+
"type": "array"
32536+
},
32537+
{
32538+
"type": "null"
32539+
}
32540+
],
32541+
"default": null,
32542+
"description": "The images to append to the collection",
32543+
"field_kind": "input",
32544+
"input": "direct",
32545+
"orig_default": null,
32546+
"orig_required": false,
32547+
"title": "Images",
32548+
"ui_order": 1
3252632549
},
3252732550
"type": {
3252832551
"const": "image_collection",
@@ -32536,7 +32559,7 @@
3253632559
"tags": ["primitives", "image", "collection"],
3253732560
"title": "Image Collection Primitive",
3253832561
"type": "object",
32539-
"version": "1.0.1",
32562+
"version": "1.0.2",
3254032563
"output": {
3254132564
"$ref": "#/components/schemas/ImageCollectionOutput"
3254232565
}

‎invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/UpdateNodesButton.tsx‎

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,9 @@ import { logger } from 'app/logging/logger';
33
import { useAppStore } from 'app/store/storeHooks';
44
import { useGetNodesNeedUpdate } from 'features/nodes/hooks/useGetNodesNeedUpdate';
55
import { $templates, nodesChanged } from 'features/nodes/store/nodesSlice';
6-
import { selectNodes } from 'features/nodes/store/selectors';
7-
import { NodeUpdateError } from 'features/nodes/types/error';
6+
import { selectEdges, selectNodes } from 'features/nodes/store/selectors';
87
import { isInvocationNode } from 'features/nodes/types/invocation';
9-
import { getNeedsUpdate, updateNode } from 'features/nodes/util/node/nodeUpdate';
8+
import { getConnectedInputNames, getNeedsUpdate, updateNode } from 'features/nodes/util/node/nodeUpdate';
109
import { toast } from 'features/toast/toast';
1110
import { memo, useCallback } from 'react';
1211
import { useTranslation } from 'react-i18next';
@@ -20,6 +19,7 @@ const useUpdateNodes = () => {
2019

2120
const updateNodes = useCallback(() => {
2221
const nodes = selectNodes(store.getState());
22+
const edges = selectEdges(store.getState());
2323
const templates = $templates.get();
2424

2525
let unableToUpdateCount = 0;
@@ -35,17 +35,16 @@ const useUpdateNodes = () => {
3535
return;
3636
}
3737
try {
38-
const updatedNode = updateNode(node, template);
38+
const connectedInputNames = getConnectedInputNames(node.id, edges);
39+
const updatedNode = updateNode(node, template, { connectedInputNames });
3940
store.dispatch(
4041
nodesChanged([
4142
{ type: 'remove', id: updatedNode.id },
4243
{ type: 'add', item: updatedNode },
4344
])
4445
);
45-
} catch (e) {
46-
if (e instanceof NodeUpdateError) {
47-
unableToUpdateCount++;
48-
}
46+
} catch {
47+
unableToUpdateCount++;
4948
}
5049
});
5150

‎invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.test.ts‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
import type { InvocationTemplate } from 'features/nodes/types/invocation';
2+
import { buildInvocationNode } from 'features/nodes/util/node/buildInvocationNode';
3+
import { updateNode } from 'features/nodes/util/node/nodeUpdate';
4+
import { describe, expect, it } from 'vitest';
5+
6+
const imageCollectionOutput = {
7+
collection: {
8+
fieldKind: 'output',
9+
name: 'collection',
10+
title: 'Collection',
11+
description: 'The output images',
12+
type: {
13+
name: 'ImageField',
14+
cardinality: 'COLLECTION',
15+
batch: false,
16+
},
17+
ui_hidden: false,
18+
},
19+
} satisfies InvocationTemplate['outputs'];
20+
21+
const oldImageCollectionTemplate = {
22+
title: 'Image Collection Primitive',
23+
type: 'image_collection',
24+
version: '1.0.1',
25+
tags: ['primitives', 'image', 'collection'],
26+
description: 'A collection of image primitive values',
27+
outputType: 'image_collection_output',
28+
inputs: {
29+
collection: {
30+
name: 'collection',
31+
title: 'Collection',
32+
required: false,
33+
description: 'The collection of image values',
34+
fieldKind: 'input',
35+
input: 'any',
36+
ui_hidden: false,
37+
type: {
38+
name: 'ImageField',
39+
cardinality: 'COLLECTION',
40+
batch: false,
41+
},
42+
default: undefined,
43+
},
44+
},
45+
outputs: imageCollectionOutput,
46+
useCache: true,
47+
nodePack: 'invokeai',
48+
classification: 'stable',
49+
category: 'primitives',
50+
} satisfies InvocationTemplate;
51+
52+
const oldestImageCollectionTemplate = {
53+
...oldImageCollectionTemplate,
54+
version: '1.0.0',
55+
} satisfies InvocationTemplate;
56+
57+
const currentImageCollectionTemplate = {
58+
...oldImageCollectionTemplate,
59+
version: '1.0.2',
60+
inputs: {
61+
collection: {
62+
name: 'collection',
63+
title: 'Collection',
64+
required: false,
65+
description: 'An optional image collection to append to',
66+
fieldKind: 'input',
67+
input: 'connection',
68+
ui_hidden: false,
69+
type: {
70+
name: 'ImageField',
71+
cardinality: 'COLLECTION',
72+
batch: false,
73+
},
74+
default: undefined,
75+
},
76+
images: {
77+
name: 'images',
78+
title: 'Images',
79+
required: false,
80+
description: 'The images to append to the collection',
81+
fieldKind: 'input',
82+
input: 'direct',
83+
ui_hidden: false,
84+
type: {
85+
name: 'ImageField',
86+
cardinality: 'COLLECTION',
87+
batch: false,
88+
},
89+
default: undefined,
90+
},
91+
},
92+
} satisfies InvocationTemplate;
93+
94+
describe('updateNode', () => {
95+
it('moves old image_collection direct collection values to the new images field', () => {
96+
const node = buildInvocationNode({ x: 0, y: 0 }, oldImageCollectionTemplate);
97+
const images = [{ image_name: 'first' }, { image_name: 'second' }];
98+
const collectionInput = node.data.inputs.collection;
99+
if (!collectionInput) {
100+
throw new Error('Expected collection input');
101+
}
102+
collectionInput.value = images;
103+
104+
const updated = updateNode(node, currentImageCollectionTemplate, { connectedInputNames: new Set() });
105+
106+
expect(updated.data.version).toBe('1.0.2');
107+
expect(updated.data.inputs.images?.value).toEqual(images);
108+
expect(updated.data.inputs.collection?.value).toEqual([]);
109+
});
110+
111+
it('moves 1.0.0 image_collection direct collection values to the new images field', () => {
112+
const node = buildInvocationNode({ x: 0, y: 0 }, oldestImageCollectionTemplate);
113+
const images = [{ image_name: 'first' }];
114+
const collectionInput = node.data.inputs.collection;
115+
if (!collectionInput) {
116+
throw new Error('Expected collection input');
117+
}
118+
collectionInput.value = images;
119+
120+
const updated = updateNode(node, currentImageCollectionTemplate, { connectedInputNames: new Set() });
121+
122+
expect(updated.data.version).toBe('1.0.2');
123+
expect(updated.data.inputs.images?.value).toEqual(images);
124+
expect(updated.data.inputs.collection?.value).toEqual([]);
125+
});
126+
127+
it('preserves old image_collection direct collection values when collection is connected', () => {
128+
const node = buildInvocationNode({ x: 0, y: 0 }, oldImageCollectionTemplate);
129+
const images = [{ image_name: 'stale' }];
130+
const collectionInput = node.data.inputs.collection;
131+
if (!collectionInput) {
132+
throw new Error('Expected collection input');
133+
}
134+
collectionInput.value = images;
135+
136+
const updated = updateNode(node, currentImageCollectionTemplate, {
137+
connectedInputNames: new Set(['collection']),
138+
});
139+
140+
expect(updated.data.inputs.images?.value).toBeUndefined();
141+
expect(updated.data.inputs.collection?.value).toEqual(images);
142+
});
143+
});

‎invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.ts‎

Lines changed: 56 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,32 @@
11
import { deepClone } from 'common/util/deepClone';
2-
import { satisfies } from 'compare-versions';
2+
import { compare, satisfies } from 'compare-versions';
33
import { defaultsDeep, keys, pick } from 'es-toolkit/compat';
44
import { NodeUpdateError } from 'features/nodes/types/error';
55
import type { InvocationNode, InvocationNodeData, InvocationTemplate } from 'features/nodes/types/invocation';
66
import { zParsedSemver } from 'features/nodes/types/semver';
77

88
import { buildInvocationNode } from './buildInvocationNode';
99

10+
type ConnectedInputEdge = { type?: string; target: string; targetHandle?: string | null };
11+
12+
type UpdateNodeOptions = {
13+
connectedInputNames: Set<string>;
14+
};
15+
16+
export const getConnectedInputNames = (nodeId: string, edges: ConnectedInputEdge[]): Set<string> =>
17+
new Set(
18+
edges.flatMap((edge) =>
19+
edge.type === 'default' && edge.target === nodeId && edge.targetHandle ? [edge.targetHandle] : []
20+
)
21+
);
22+
23+
export const getUpdatedFieldName = (node: InvocationNode, fieldName: string): string => {
24+
if (node.data.type === 'image_collection' && fieldName === 'collection' && node.data.inputs.images) {
25+
return 'images';
26+
}
27+
return fieldName;
28+
};
29+
1030
export const getNeedsUpdate = (data: InvocationNodeData, template: InvocationTemplate): boolean => {
1131
if (data.type !== template.type) {
1232
return true;
@@ -29,6 +49,34 @@ const getMayUpdateNode = (node: InvocationNode, template: InvocationTemplate): b
2949
return satisfies(node.data.version, `^${templateMajor}`);
3050
};
3151

52+
export const migrateImageCollectionInputValues = (
53+
node: InvocationNode,
54+
options: UpdateNodeOptions & { sourceVersion?: string }
55+
) => {
56+
if (node.data.type !== 'image_collection') {
57+
return;
58+
}
59+
if (options.sourceVersion && compare(options.sourceVersion, '1.0.2', '>=')) {
60+
return;
61+
}
62+
63+
const collection = node.data.inputs.collection;
64+
const images = node.data.inputs.images;
65+
if (!collection || !images || !Array.isArray(collection.value)) {
66+
return;
67+
}
68+
if (Array.isArray(images.value) && images.value.length > 0) {
69+
return;
70+
}
71+
72+
if (options.connectedInputNames.has('collection')) {
73+
return;
74+
}
75+
76+
images.value = collection.value;
77+
collection.value = [];
78+
};
79+
3280
/**
3381
* Updates a node to the latest version of its template:
3482
* - Create a new node data object with the latest version of the template.
@@ -40,7 +88,11 @@ const getMayUpdateNode = (node: InvocationNode, template: InvocationTemplate): b
4088
* @param template The invocation template to update to.
4189
* @throws {NodeUpdateError} If the node is not an invocation node.
4290
*/
43-
export const updateNode = (node: InvocationNode, template: InvocationTemplate): InvocationNode => {
91+
export const updateNode = (
92+
node: InvocationNode,
93+
template: InvocationTemplate,
94+
options: UpdateNodeOptions
95+
): InvocationNode => {
4496
const mayUpdate = getMayUpdateNode(node, template);
4597

4698
if (!mayUpdate || node.data.type !== template.type) {
@@ -54,8 +106,10 @@ export const updateNode = (node: InvocationNode, template: InvocationTemplate):
54106
// being valid. We rely on the template's major version to be majorly incremented if this kind of
55107
// merge would result in an invalid node.
56108
const clone = deepClone(node);
109+
const sourceVersion = clone.data.version;
57110
clone.data.version = template.version;
58111
defaultsDeep(clone, defaults); // mutates!
112+
migrateImageCollectionInputValues(clone, { ...options, sourceVersion });
59113

60114
// Remove any fields that are not in the template
61115
clone.data.inputs = pick(clone.data.inputs, keys(defaults.data.inputs));

0 commit comments

Comments
 (0)