From 02be06baf64edf9cdc1f8b3bcfe4c0c7ba636f8b Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Thu, 14 May 2026 12:12:01 -0500 Subject: [PATCH 1/2] Improve image collection chaining --- invokeai/app/invocations/primitives.py | 19 ++- .../panels/TopPanel/UpdateNodesButton.tsx | 10 +- .../nodes/util/node/nodeUpdate.test.ts | 122 ++++++++++++++++++ .../features/nodes/util/node/nodeUpdate.ts | 31 ++++- .../nodes/util/workflow/validateWorkflow.ts | 7 +- .../frontend/web/src/services/api/schema.ts | 8 +- tests/app/invocations/test_image.py | 32 +++++ 7 files changed, 221 insertions(+), 8 deletions(-) create mode 100644 invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.test.ts diff --git a/invokeai/app/invocations/primitives.py b/invokeai/app/invocations/primitives.py index 7ec6c3dc149..6249de0cd8e 100644 --- a/invokeai/app/invocations/primitives.py +++ b/invokeai/app/invocations/primitives.py @@ -279,15 +279,28 @@ def invoke(self, context: InvocationContext) -> ImageOutput: title="Image Collection Primitive", tags=["primitives", "image", "collection"], category="primitives", - version="1.0.1", + version="1.0.2", ) class ImageCollectionInvocation(BaseInvocation): """A collection of image primitive values""" - collection: list[ImageField] = InputField(description="The collection of image values") + collection: Optional[list[ImageField]] = InputField( + default=None, + description="An optional image collection to append to", + input=Input.Connection, + title="Collection", + ui_order=0, + ) + images: Optional[list[ImageField]] = InputField( + default=None, + description="The images to append to the collection", + input=Input.Direct, + title="Images", + ui_order=1, + ) def invoke(self, context: InvocationContext) -> ImageCollectionOutput: - return ImageCollectionOutput(collection=self.collection) + return ImageCollectionOutput(collection=[*(self.collection or []), *(self.images or [])]) # endregion diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/UpdateNodesButton.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/UpdateNodesButton.tsx index 5072d3ae7f3..007cdcf9c6e 100644 --- a/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/UpdateNodesButton.tsx +++ b/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/UpdateNodesButton.tsx @@ -3,7 +3,7 @@ import { logger } from 'app/logging/logger'; import { useAppStore } from 'app/store/storeHooks'; import { useGetNodesNeedUpdate } from 'features/nodes/hooks/useGetNodesNeedUpdate'; import { $templates, nodesChanged } from 'features/nodes/store/nodesSlice'; -import { selectNodes } from 'features/nodes/store/selectors'; +import { selectEdges, selectNodes } from 'features/nodes/store/selectors'; import { NodeUpdateError } from 'features/nodes/types/error'; import { isInvocationNode } from 'features/nodes/types/invocation'; import { getNeedsUpdate, updateNode } from 'features/nodes/util/node/nodeUpdate'; @@ -20,6 +20,7 @@ const useUpdateNodes = () => { const updateNodes = useCallback(() => { const nodes = selectNodes(store.getState()); + const edges = selectEdges(store.getState()); const templates = $templates.get(); let unableToUpdateCount = 0; @@ -35,7 +36,12 @@ const useUpdateNodes = () => { return; } try { - const updatedNode = updateNode(node, template); + const connectedInputNames = new Set( + edges.flatMap((edge) => + edge.type === 'default' && edge.target === node.id && edge.targetHandle ? [edge.targetHandle] : [] + ) + ); + const updatedNode = updateNode(node, template, { connectedInputNames }); store.dispatch( nodesChanged([ { type: 'remove', id: updatedNode.id }, diff --git a/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.test.ts b/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.test.ts new file mode 100644 index 00000000000..94c62b6b768 --- /dev/null +++ b/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.test.ts @@ -0,0 +1,122 @@ +import type { InvocationTemplate } from 'features/nodes/types/invocation'; +import { buildInvocationNode } from 'features/nodes/util/node/buildInvocationNode'; +import { updateNode } from 'features/nodes/util/node/nodeUpdate'; +import { describe, expect, it } from 'vitest'; + +const imageCollectionOutput = { + collection: { + fieldKind: 'output', + name: 'collection', + title: 'Collection', + description: 'The output images', + type: { + name: 'ImageField', + cardinality: 'COLLECTION', + batch: false, + }, + ui_hidden: false, + }, +} satisfies InvocationTemplate['outputs']; + +const oldImageCollectionTemplate = { + title: 'Image Collection Primitive', + type: 'image_collection', + version: '1.0.1', + tags: ['primitives', 'image', 'collection'], + description: 'A collection of image primitive values', + outputType: 'image_collection_output', + inputs: { + collection: { + name: 'collection', + title: 'Collection', + required: false, + description: 'The collection of image values', + fieldKind: 'input', + input: 'any', + ui_hidden: false, + type: { + name: 'ImageField', + cardinality: 'COLLECTION', + batch: false, + }, + default: undefined, + }, + }, + outputs: imageCollectionOutput, + useCache: true, + nodePack: 'invokeai', + classification: 'stable', + category: 'primitives', +} satisfies InvocationTemplate; + +const currentImageCollectionTemplate = { + ...oldImageCollectionTemplate, + version: '1.0.2', + inputs: { + collection: { + name: 'collection', + title: 'Collection', + required: false, + description: 'An optional image collection to append to', + fieldKind: 'input', + input: 'connection', + ui_hidden: false, + type: { + name: 'ImageField', + cardinality: 'COLLECTION', + batch: false, + }, + default: [], + }, + images: { + name: 'images', + title: 'Images', + required: false, + description: 'The images to append to the collection', + fieldKind: 'input', + input: 'direct', + ui_hidden: false, + type: { + name: 'ImageField', + cardinality: 'COLLECTION', + batch: false, + }, + default: undefined, + }, + }, +} satisfies InvocationTemplate; + +describe('updateNode', () => { + it('moves old image_collection direct collection values to the new images field', () => { + const node = buildInvocationNode({ x: 0, y: 0 }, oldImageCollectionTemplate); + const images = [{ image_name: 'first' }, { image_name: 'second' }]; + const collectionInput = node.data.inputs.collection; + if (!collectionInput) { + throw new Error('Expected collection input'); + } + collectionInput.value = images; + + const updated = updateNode(node, currentImageCollectionTemplate); + + expect(updated.data.version).toBe('1.0.2'); + expect(updated.data.inputs.images?.value).toEqual(images); + expect(updated.data.inputs.collection?.value).toEqual([]); + }); + + it('does not move old image_collection direct collection values when collection is connected', () => { + const node = buildInvocationNode({ x: 0, y: 0 }, oldImageCollectionTemplate); + const images = [{ image_name: 'stale' }]; + const collectionInput = node.data.inputs.collection; + if (!collectionInput) { + throw new Error('Expected collection input'); + } + collectionInput.value = images; + + const updated = updateNode(node, currentImageCollectionTemplate, { + connectedInputNames: new Set(['collection']), + }); + + expect(updated.data.inputs.images?.value).toBeUndefined(); + expect(updated.data.inputs.collection?.value).toEqual([]); + }); +}); diff --git a/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.ts b/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.ts index 59eac81b33b..157295dfe2e 100644 --- a/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.ts +++ b/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.ts @@ -7,6 +7,10 @@ import { zParsedSemver } from 'features/nodes/types/semver'; import { buildInvocationNode } from './buildInvocationNode'; +type UpdateNodeOptions = { + connectedInputNames?: Set; +}; + export const getNeedsUpdate = (data: InvocationNodeData, template: InvocationTemplate): boolean => { if (data.type !== template.type) { return true; @@ -29,6 +33,26 @@ const getMayUpdateNode = (node: InvocationNode, template: InvocationTemplate): b return satisfies(node.data.version, `^${templateMajor}`); }; +const migrateImageCollectionInputValues = (node: InvocationNode, options?: UpdateNodeOptions) => { + if (node.data.type !== 'image_collection') { + return; + } + + const collection = node.data.inputs.collection; + const images = node.data.inputs.images; + if (!collection || !images || !Array.isArray(collection.value)) { + return; + } + if (Array.isArray(images.value) && images.value.length > 0) { + return; + } + + if (!options?.connectedInputNames?.has('collection')) { + images.value = collection.value; + } + collection.value = []; +}; + /** * Updates a node to the latest version of its template: * - Create a new node data object with the latest version of the template. @@ -40,7 +64,11 @@ const getMayUpdateNode = (node: InvocationNode, template: InvocationTemplate): b * @param template The invocation template to update to. * @throws {NodeUpdateError} If the node is not an invocation node. */ -export const updateNode = (node: InvocationNode, template: InvocationTemplate): InvocationNode => { +export const updateNode = ( + node: InvocationNode, + template: InvocationTemplate, + options?: UpdateNodeOptions +): InvocationNode => { const mayUpdate = getMayUpdateNode(node, template); if (!mayUpdate || node.data.type !== template.type) { @@ -56,6 +84,7 @@ export const updateNode = (node: InvocationNode, template: InvocationTemplate): const clone = deepClone(node); clone.data.version = template.version; defaultsDeep(clone, defaults); // mutates! + migrateImageCollectionInputValues(clone, options); // Remove any fields that are not in the template clone.data.inputs = pick(clone.data.inputs, keys(defaults.data.inputs)); diff --git a/invokeai/frontend/web/src/features/nodes/util/workflow/validateWorkflow.ts b/invokeai/frontend/web/src/features/nodes/util/workflow/validateWorkflow.ts index 448214defe3..d5f0737d336 100644 --- a/invokeai/frontend/web/src/features/nodes/util/workflow/validateWorkflow.ts +++ b/invokeai/frontend/web/src/features/nodes/util/workflow/validateWorkflow.ts @@ -81,7 +81,12 @@ export const validateWorkflow = async (args: ValidateWorkflowArgs): Promise + edge.type === 'default' && edge.target === node.id && edge.targetHandle ? [edge.targetHandle] : [] + ) + ); + const updatedNode = updateNode(node, template, { connectedInputNames }); node.data = updatedNode.data; } catch { const message = t('nodes.unableToUpdateNode', { diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index e9de97db96d..b9ee06c454a 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -13760,10 +13760,16 @@ export type components = { use_cache?: boolean; /** * Collection - * @description The collection of image values + * @description An optional image collection to append to * @default null */ collection?: components["schemas"]["ImageField"][] | null; + /** + * Images + * @description The images to append to the collection + * @default null + */ + images?: components["schemas"]["ImageField"][] | null; /** * type * @default image_collection diff --git a/tests/app/invocations/test_image.py b/tests/app/invocations/test_image.py index 7c0036d3db3..05d2cd32b76 100644 --- a/tests/app/invocations/test_image.py +++ b/tests/app/invocations/test_image.py @@ -8,6 +8,7 @@ from PIL import Image, ImageFilter from invokeai.app.invocations.image import ImageField, OklabUnsharpMaskInvocation, OklchImageHueAdjustmentInvocation +from invokeai.app.invocations.primitives import ImageCollectionInvocation from invokeai.backend.image_util.color_conversion import ( linear_srgb_from_oklab, linear_srgb_from_oklch, @@ -47,6 +48,37 @@ def _max_abs_diff_uint8(left: Image.Image, right: Image.Image) -> int: return int(numpy.abs(left_arr - right_arr).max()) +def test_image_collection_invocation_preserves_existing_collection_values() -> None: + images = [ImageField(image_name="first"), ImageField(image_name="second")] + + output = ImageCollectionInvocation(collection=images).invoke(MagicMock()) + + assert output.collection == images + + +def test_image_collection_invocation_appends_direct_images_after_chained_collection() -> None: + chained_images = [ImageField(image_name="chained")] + direct_images = [ImageField(image_name="direct_1"), ImageField(image_name="direct_2")] + + output = ImageCollectionInvocation(collection=chained_images, images=direct_images).invoke(MagicMock()) + + assert output.collection == [*chained_images, *direct_images] + + +def test_image_collection_invocation_supports_empty_direct_images() -> None: + chained_images = [ImageField(image_name="chained")] + + output = ImageCollectionInvocation(collection=chained_images, images=None).invoke(MagicMock()) + + assert output.collection == chained_images + + +def test_image_collection_invocation_outputs_empty_collection_when_inputs_are_empty() -> None: + output = ImageCollectionInvocation(collection=None, images=None).invoke(MagicMock()) + + assert output.collection == [] + + def test_oklab_unsharp_mask_invocation_preserves_alpha_and_sharpens_lightness_only() -> None: input_image = Image.new("RGBA", (3, 1)) input_image.putdata( From 81f969552debcba85a38643c69f06f379d77f576 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Fri, 15 May 2026 05:33:25 -0500 Subject: [PATCH 2/2] chore: openapi.json --- invokeai/frontend/web/openapi.json | 33 +++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 4a36cf31e2b..71ac7fa9854 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -32336,11 +32336,34 @@ } ], "default": null, - "description": "The collection of image values", + "description": "An optional image collection to append to", "field_kind": "input", - "input": "any", - "orig_required": true, - "title": "Collection" + "input": "connection", + "orig_default": null, + "orig_required": false, + "title": "Collection", + "ui_order": 0 + }, + "images": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/ImageField" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "description": "The images to append to the collection", + "field_kind": "input", + "input": "direct", + "orig_default": null, + "orig_required": false, + "title": "Images", + "ui_order": 1 }, "type": { "const": "image_collection", @@ -32354,7 +32377,7 @@ "tags": ["primitives", "image", "collection"], "title": "Image Collection Primitive", "type": "object", - "version": "1.0.1", + "version": "1.0.2", "output": { "$ref": "#/components/schemas/ImageCollectionOutput" }