fix(workflow): recall generation graphs from images

invokeai/frontend/webv2/src/features/workflow/core/buildGraph.ts

@@ -21,7 +21,7 @@ import {
   type WorkflowGeneratorResolutions,
 } from './batch';
 import { CALL_SAVED_WORKFLOW_DYNAMIC_FIELD_PREFIX } from './callSavedWorkflow';
-import { createWorkflowId } from './document';
+import { createWorkflowId, nodeAcceptsExtraInputs } from './document';
 import { getWorkflowFieldInvalidReason, isDirectInputField, isWorkflowCollectionItemValid } from './fields';
 import {
   createForLoopValidationReason,
@@ -285,7 +285,27 @@ export const compileProjectGraph = (
     const workflowInputs: Record<string, unknown> = {};
 
     for (const instance of Object.values(node.data.inputs)) {
-      const inputTemplate = node.data.dynamicInputTemplates?.[instance.name] ?? template.inputs[instance.name];
+      const dynamicInputs = node.data.dynamicInputTemplates;
+      const inputTemplate =
+        dynamicInputs && Object.hasOwn(dynamicInputs, instance.name)
+          ? dynamicInputs[instance.name]
+          : Object.hasOwn(template.inputs, instance.name)
+            ? template.inputs[instance.name]
+            : undefined;
+
+      if (
+        !inputTemplate &&
+        nodeAcceptsExtraInputs(node.data.type) &&
+        !['id', 'type', 'is_intermediate', 'use_cache'].includes(instance.name) &&
+        instance.value !== undefined
+      ) {
+        Object.defineProperty(graphNode, instance.name, {
+          configurable: true,
+          enumerable: true,
+          value: structuredClone(instance.value),
+          writable: true,
+        });
+      }
 
       // An unset board shows as Auto, so it resolves like one.
       if (!inputTemplate || (instance.value === undefined && inputTemplate.type.name !== 'BoardField')) {

invokeai/frontend/webv2/src/features/workflow/core/document.ts

@@ -443,13 +443,53 @@ export const getNodeUpdateStatus = (
 /** Node types whose undeclared inputs the backend accepts (`extra='allow'`), so an update must carry them over. */
 const EXTRA_INPUT_NODE_TYPES = new Set(['core_metadata']);
 
+export const nodeAcceptsExtraInputs = (type: string): boolean => EXTRA_INPUT_NODE_TYPES.has(type);
+
 interface InvocationNodeUpdate {
   node: WorkflowInvocationNode;
   droppedInputNames: string[];
   /** Inputs whose exposed form elements now belong to another input (`collection` → `images`). */
   renamedInputs: Record<string, string>;
 }
 
+/** Unversioned execution graphs need the same value migration as older authored workflows. */
+export const migrateImageCollectionInputValues = (
+  inputs: Record<string, WorkflowFieldInstance>,
+  template: InvocationTemplate,
+  isCollectionConnected: boolean
+): void => {
+  const collection = inputs.collection;
+  const images = inputs.images;
+
+  if (
+    images &&
+    collection &&
+    Array.isArray(collection.value) &&
+    !(Array.isArray(images.value) && images.value.length > 0) &&
+    !isCollectionConnected
+  ) {
+    inputs.images = { ...images, value: collection.value };
+    inputs.collection = { ...collection, value: structuredClone(template.inputs.collection?.default) };
+  }
+};
+
+export const migrateMiniMaxH3NumFramesValue = (
+  inputs: Record<string, WorkflowFieldInstance>,
+  template: InvocationTemplate
+): void => {
+  const frames = inputs.num_frames;
+  const frameTemplate = template.inputs.num_frames;
+
+  if (frames && frameTemplate?.options && typeof frames.value === 'number') {
+    const choice = String(frames.value);
+
+    inputs.num_frames = {
+      ...frames,
+      value: frameTemplate.options.includes(choice) ? choice : frameTemplate.default,
+    };
+  }
+};
+
 /**
  * Fresh template defaults underneath the node's own values: existing labels, values, and node settings win, new
  * inputs start at their defaults, inputs the template no longer declares are dropped (dynamic and extra inputs
@@ -479,7 +519,7 @@ const updateInvocationNodeToTemplate = (
     }
 
     if (
-      EXTRA_INPUT_NODE_TYPES.has(node.data.type) ||
+      nodeAcceptsExtraInputs(node.data.type) ||
       node.data.dynamicInputTemplates?.[name] !== undefined ||
       name.startsWith(CALL_SAVED_WORKFLOW_DYNAMIC_FIELD_PREFIX)
     ) {
@@ -494,41 +534,16 @@ const updateInvocationNodeToTemplate = (
   // `image_collection` 1.0.2 moved its direct list from `collection` to `images`; a connected `collection` keeps
   // feeding the node and is left alone.
   if (node.data.type === 'image_collection' && sourceVersion && compareVersions(sourceVersion, [1, 0, 2]) < 0) {
-    const collection = node.data.inputs.collection;
-    const images = inputs.images;
-
-    if (images) {
+    if (inputs.images) {
       renamedInputs.collection = 'images';
-
-      if (
-        collection &&
-        Array.isArray(collection.value) &&
-        !(Array.isArray(images.value) && images.value.length > 0) &&
-        !connectedInputNames.has('collection')
-      ) {
-        inputs.images = { ...images, value: collection.value };
-
-        if (inputs.collection) {
-          inputs.collection = { ...inputs.collection, value: fresh.data.inputs.collection?.value };
-        }
-      }
     }
+    migrateImageCollectionInputValues(inputs, template, connectedInputNames.has('collection'));
   }
 
   // `minimax_h3_denoise` 1.3.0 turned the free frame count into a choice list; a count still on the grid becomes
   // that choice, anything else falls back to the default.
   if (node.data.type === 'minimax_h3_denoise' && sourceVersion && compareVersions(sourceVersion, [1, 3, 0]) < 0) {
-    const frames = inputs.num_frames;
-    const frameTemplate = template.inputs.num_frames;
-
-    if (frames && frameTemplate?.options && typeof frames.value === 'number') {
-      const choice = String(frames.value);
-
-      inputs.num_frames = {
-        ...frames,
-        value: frameTemplate.options.includes(choice) ? choice : frameTemplate.default,
-      };
-    }
+    migrateMiniMaxH3NumFramesValue(inputs, template);
   }
 
   return {

invokeai/frontend/webv2/src/features/workflow/core/graphToDocument.test.ts

@@ -2,9 +2,11 @@ import { describe, expect, it } from 'vitest';
 
 import type { FieldInputTemplate, InvocationTemplate } from './types';
 
-import { compileProjectGraph } from './buildGraph';
-import { previewGraphToDocument, type PreviewGraphLike } from './graphToDocument';
+import { compileProjectGraph, getProjectGraphReadiness } from './buildGraph';
+import { updateWorkflowNodes } from './document';
+import { previewGraphToDocument, storedGraphToDocument, type PreviewGraphLike } from './graphToDocument';
 import { isInvocationNode } from './types';
+import { parseWorkflowJson, serializeWorkflowJson } from './workflowJson';
 
 const input = (name: string, overrides: Partial<FieldInputTemplate> = {}): FieldInputTemplate => ({
   default: undefined,
@@ -87,7 +89,186 @@ const graph: PreviewGraphLike = {
   ],
 };
 
+describe('storedGraphToDocument', () => {
+  it('loads unversioned frame counts as current editable choices', () => {
+    const frameTemplates = {
+      minimax_h3_denoise: {
+        ...template('minimax_h3_denoise', {
+          num_frames: input('num_frames', {
+            default: '124',
+            options: ['90', '124'],
+            type: { batch: false, cardinality: 'SINGLE', name: 'EnumField' },
+          }),
+        }),
+        version: '1.4.0',
+      },
+    };
+    const { document } = storedGraphToDocument(
+      {
+        nodes: {
+          valid: { id: 'valid', type: 'minimax_h3_denoise', num_frames: 90 },
+          invalid: { id: 'invalid', type: 'minimax_h3_denoise', num_frames: 100 },
+          current: { id: 'current', type: 'minimax_h3_denoise', num_frames: '90' },
+        },
+        edges: [],
+      },
+      frameTemplates
+    );
+    const loaded = updateWorkflowNodes(parseWorkflowJson(serializeWorkflowJson(document)).document, frameTemplates);
+    expect(
+      getProjectGraphReadiness(loaded.document, { status: 'loaded', error: null, templates: frameTemplates })
+    ).toMatchObject({ canInvoke: true, reasons: [] });
+    const compiled = compileProjectGraph(loaded.document, frameTemplates).backendGraph;
+    expect(compiled.nodes.valid?.num_frames).toBe('90');
+    expect(compiled.nodes.invalid?.num_frames).toBe('124');
+    expect(compiled.nodes.current?.num_frames).toBe('90');
+  });
+
+  it('loads old image lists as editable images without replacing current lists or connected collections', () => {
+    const imageList = { batch: false, cardinality: 'COLLECTION' as const, name: 'ImageField' };
+    const collectionTemplates = {
+      image_collection: {
+        ...template('image_collection', {
+          collection: input('collection', { input: 'connection', type: imageList }),
+          images: input('images', { default: [], type: imageList }),
+        }),
+        outputs: { collection: { description: '', name: 'collection', title: 'Collection', type: imageList } },
+        version: '1.0.2',
+      },
+    };
+    const oldImages = [{ image_name: 'old.png' }];
+    const currentImages = [{ image_name: 'current.png' }];
+    const edge = {
+      source: { node_id: 'legacy', field: 'collection' },
+      destination: { node_id: 'connected', field: 'collection' },
+    };
+    const { document } = storedGraphToDocument(
+      {
+        nodes: {
+          legacy: { id: 'legacy', type: 'image_collection', collection: oldImages },
+          current: { id: 'current', type: 'image_collection', collection: oldImages, images: currentImages },
+          connected: { id: 'connected', type: 'image_collection', collection: oldImages },
+        },
+        edges: [edge],
+      },
+      collectionTemplates
+    );
+    const loaded = updateWorkflowNodes(
+      parseWorkflowJson(serializeWorkflowJson(document)).document,
+      collectionTemplates
+    );
+    expect(
+      getProjectGraphReadiness(loaded.document, { status: 'loaded', error: null, templates: collectionTemplates })
+    ).toMatchObject({ canInvoke: true, reasons: [] });
+    const compiled = compileProjectGraph(loaded.document, collectionTemplates).backendGraph;
+    expect(compiled.nodes.legacy?.images).toEqual(oldImages);
+    expect(compiled.nodes.legacy).not.toHaveProperty('collection');
+    expect(compiled.nodes.current?.images).toEqual(currentImages);
+    expect(compiled.nodes.connected?.images).toEqual([]);
+    expect(compiled.nodes.connected).not.toHaveProperty('collection');
+    expect(compiled.edges).toEqual([{ ...edge, type: 'default' }]);
+  });
+
+  it('recalls generation values, flags, metadata extras and edges through serialization and compilation', () => {
+    const storedTemplates = { ...templates, core_metadata: template('core_metadata', {}) };
+    const stored = {
+      id: 'generated',
+      nodes: {
+        seed: { id: 'seed', type: 'integer', value: 42, use_cache: false },
+        denoise: { id: 'denoise', type: 'denoise', steps: 12, is_intermediate: false, unrelated: 'drop' },
+        metadata: {
+          id: 'metadata',
+          type: 'core_metadata',
+          is_intermediate: true,
+          loras: [{ model: { key: 'lora-1' }, weight: 0.75 }],
+          z_image_seed_variance_enabled: true,
+        },
+      },
+      edges: [
+        { source: { node_id: 'seed', field: 'value' }, destination: { node_id: 'denoise', field: 'seed' } },
+        { source: { node_id: 'seed', field: 'value' }, destination: { node_id: 'metadata', field: 'custom_seed' } },
+      ],
+    };
+    const { document, skippedNodeTypes } = storedGraphToDocument(stored, storedTemplates);
+    expect(skippedNodeTypes).toEqual([]);
+    const recalled = parseWorkflowJson(serializeWorkflowJson(document)).document;
+    const compiled = compileProjectGraph(recalled, storedTemplates).backendGraph;
+
+    expect(compiled.nodes.seed).toMatchObject({ value: 42, use_cache: false });
+    expect(compiled.nodes.denoise).toMatchObject({ steps: 12, is_intermediate: false });
+    expect(compiled.nodes.denoise).not.toHaveProperty('unrelated');
+    expect(compiled.nodes.metadata).toMatchObject({
+      loras: [{ model: { key: 'lora-1' }, weight: 0.75 }],
+      z_image_seed_variance_enabled: true,
+      is_intermediate: true,
+    });
+    expect(compiled.edges).toEqual(stored.edges.map((edge) => ({ ...edge, type: 'default' })));
+    expect(recalled.nodes[0]?.position).not.toEqual(recalled.nodes[1]?.position);
+
+    const metadata = recalled.nodes.find((node) => node.id === 'metadata');
+    if (!metadata || !isInvocationNode(metadata)) {
+      throw new Error('Missing metadata node');
+    }
+    for (const name of ['id', 'type', 'is_intermediate', 'use_cache']) {
+      metadata.data.inputs[name] = { name, label: '', value: 'injected' };
+    }
+    expect(compileProjectGraph(recalled, storedTemplates).backendGraph.nodes.metadata).toMatchObject({
+      id: 'metadata',
+      type: 'core_metadata',
+      is_intermediate: true,
+      use_cache: true,
+    });
+  });
+
+  it('reports unknown nodes and omits their connections', () => {
+    const { document, skippedNodeTypes } = storedGraphToDocument(
+      {
+        nodes: { seed: { id: 'seed', type: 'integer', value: 42 }, absent: { id: 'absent', type: 'missing' } },
+        edges: [{ source: { node_id: 'seed', field: 'value' }, destination: { node_id: 'absent', field: 'seed' } }],
+      },
+      templates
+    );
+    expect(skippedNodeTypes).toEqual(['missing']);
+    expect(document.nodes.map((node) => node.id)).toEqual(['seed']);
+    expect(document.edges).toEqual([]);
+  });
+
+  it.each([null, {}, { nodes: [], edges: [] }, { nodes: { seed: { type: 'integer' } }, edges: [] }])(
+    'rejects malformed stored graphs: %j',
+    (raw) => {
+      expect(() => storedGraphToDocument(raw, templates)).toThrow();
+    }
+  );
+});
+
 describe('previewGraphToDocument', () => {
+  it('treats prototype-named metadata fields as own values without changing object prototypes', () => {
+    const metadataTemplates = { core_metadata: template('core_metadata', {}) };
+    const prototypeBefore = Object.getOwnPropertyDescriptors(Object.prototype);
+    const { document } = previewGraphToDocument(
+      {
+        nodes: [
+          {
+            id: 'metadata',
+            type: 'core_metadata',
+            inputs: JSON.parse('{"__proto__":{"marker":42},"constructor":"custom"}'),
+          },
+        ],
+        edges: [],
+      },
+      metadataTemplates
+    );
+    const node = document.nodes[0];
+    expect(node && isInvocationNode(node) && Object.hasOwn(node.data.inputs, '__proto__')).toBe(true);
+
+    const compiled = compileProjectGraph(document, metadataTemplates).backendGraph.nodes.metadata;
+    expect(JSON.parse(JSON.stringify(compiled))).toMatchObject(
+      JSON.parse('{"__proto__":{"marker":42},"constructor":"custom"}')
+    );
+    expect(Object.getPrototypeOf(compiled)).toBe(Object.prototype);
+    expect(Object.getOwnPropertyDescriptors(Object.prototype)).toEqual(prototypeBefore);
+  });
+
   it('converts a preview graph into an editable document, skipping unknown node types', () => {
     const { document, skippedNodeTypes } = previewGraphToDocument(graph, templates);