InvokeAI/invokeai/frontend/web/src/features/nodes/util/node/nodeUpdate.test.ts at 3733ae60deff1e610e2bb42617baa185a323e561 · invoke-ai/InvokeAI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
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([]);
  });
});