feat(ui): store graph in image metadata

The previous super-minimal implementation had a major issue - the saved workflow didn't take into account batched field values. When generating with multiple iterations or dynamic prompts, the same workflow with the first prompt, seed, etc was stored in each image.

As a result, when the batch results in multiple queue items, only one of the images has the correct workflow - the others are mismatched.

To work around this, we can store the _graph_ in the image metadata (alongside the workflow, if generated via workflow editor). When loading a workflow from an image, we can choose to load the workflow or the graph, preferring the workflow.

Internally, we need to update images router image-saving services. The changes are minimal.

To avoid pydantic errors deserializing the graph, when we extract it from the image, we will leave it as stringified JSON and let the frontend's more sophisticated and flexible parsing handle it. The worklow is also changed to just return stringified JSON, so the API is consistent.

invokeai/frontend/web/public/locales/en.json

@@ -880,6 +880,7 @@
         "versionUnknown": " Version Unknown",
         "workflow": "Workflow",
         "graph": "Graph",
+        "noGraph": "No Graph",
         "workflowAuthor": "Author",
         "workflowContact": "Contact",
         "workflowDescription": "Short Description",

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/workflowLoadRequested.ts

@@ -4,31 +4,49 @@ import { parseify } from 'common/util/serialize';
 import { workflowLoaded, workflowLoadRequested } from 'features/nodes/store/actions';
 import { $templates } from 'features/nodes/store/nodesSlice';
 import { $flow } from 'features/nodes/store/reactFlowInstance';
+import type { Templates } from 'features/nodes/store/types';
 import { WorkflowMigrationError, WorkflowVersionError } from 'features/nodes/types/error';
+import { graphToWorkflow } from 'features/nodes/util/workflow/graphToWorkflow';
 import { validateWorkflow } from 'features/nodes/util/workflow/validateWorkflow';
 import { addToast } from 'features/system/store/systemSlice';
 import { makeToast } from 'features/system/util/makeToast';
 import { t } from 'i18next';
+import type { GraphAndWorkflowResponse, NonNullableGraph } from 'services/api/types';
 import { z } from 'zod';
 import { fromZodError } from 'zod-validation-error';
 
+const getWorkflow = (data: GraphAndWorkflowResponse, templates: Templates) => {
+  if (data.workflow) {
+    // Prefer to load the workflow if it's available - it has more information
+    const parsed = JSON.parse(data.workflow);
+    return validateWorkflow(parsed, templates);
+  } else if (data.graph) {
+    // Else we fall back on the graph, using the graphToWorkflow function to convert and do layout
+    const parsed = JSON.parse(data.graph);
+    const workflow = graphToWorkflow(parsed as NonNullableGraph, true);
+    return validateWorkflow(workflow, templates);
+  } else {
+    throw new Error('No workflow or graph provided');
+  }
+};
+
 export const addWorkflowLoadRequestedListener = (startAppListening: AppStartListening) => {
   startAppListening({
     actionCreator: workflowLoadRequested,
     effect: (action, { dispatch }) => {
       const log = logger('nodes');
-      const { workflow, asCopy } = action.payload;
+      const { data, asCopy } = action.payload;
       const nodeTemplates = $templates.get();
 
       try {
-        const { workflow: validatedWorkflow, warnings } = validateWorkflow(workflow, nodeTemplates);
+        const { workflow, warnings } = getWorkflow(data, nodeTemplates);
 
         if (asCopy) {
           // If we're loading a copy, we need to remove the ID so that the backend will create a new workflow
-          delete validatedWorkflow.id;
+          delete workflow.id;
         }
 
-        dispatch(workflowLoaded(validatedWorkflow));
+        dispatch(workflowLoaded(workflow));
         if (!warnings.length) {
           dispatch(
             addToast(

invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataGraphTabContent.tsx

@@ -0,0 +1,34 @@
+import { IAINoContentFallback } from 'common/components/IAIImageFallback';
+import { memo, useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { useDebouncedImageWorkflow } from 'services/api/hooks/useDebouncedImageWorkflow';
+import type { ImageDTO } from 'services/api/types';
+
+import DataViewer from './DataViewer';
+
+type Props = {
+  image: ImageDTO;
+};
+
+const ImageMetadataGraphTabContent = ({ image }: Props) => {
+  const { t } = useTranslation();
+  const { currentData } = useDebouncedImageWorkflow(image);
+  const graph = useMemo(() => {
+    if (currentData?.graph) {
+      try {
+        return JSON.parse(currentData.graph);
+      } catch {
+        return null;
+      }
+    }
+    return null;
+  }, [currentData]);
+
+  if (!graph) {
+    return <IAINoContentFallback label={t('nodes.noGraph')} />;
+  }
+
+  return <DataViewer data={graph} label={t('nodes.graph')} />;
+};
+
+export default memo(ImageMetadataGraphTabContent);

invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataViewer.tsx

@@ -1,6 +1,7 @@
 import { ExternalLink, Flex, Tab, TabList, TabPanel, TabPanels, Tabs, Text } from '@invoke-ai/ui-library';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
+import ImageMetadataGraphTabContent from 'features/gallery/components/ImageMetadataViewer/ImageMetadataGraphTabContent';
 import { useMetadataItem } from 'features/metadata/hooks/useMetadataItem';
 import { handlers } from 'features/metadata/util/handlers';
 import { memo } from 'react';
@@ -52,6 +53,7 @@ const ImageMetadataViewer = ({ image }: ImageMetadataViewerProps) => {
           <Tab>{t('metadata.metadata')}</Tab>
           <Tab>{t('metadata.imageDetails')}</Tab>
           <Tab>{t('metadata.workflow')}</Tab>
+          <Tab>{t('nodes.graph')}</Tab>
         </TabList>
 
         <TabPanels>
@@ -81,6 +83,9 @@ const ImageMetadataViewer = ({ image }: ImageMetadataViewerProps) => {
           <TabPanel>
             <ImageMetadataWorkflowTabContent image={image} />
           </TabPanel>
+          <TabPanel>
+            <ImageMetadataGraphTabContent image={image} />
+          </TabPanel>
         </TabPanels>
       </Tabs>
     </Flex>

invokeai/frontend/web/src/features/gallery/components/ImageMetadataViewer/ImageMetadataWorkflowTabContent.tsx

@@ -1,5 +1,5 @@
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
-import { memo } from 'react';
+import { memo, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useDebouncedImageWorkflow } from 'services/api/hooks/useDebouncedImageWorkflow';
 import type { ImageDTO } from 'services/api/types';
@@ -12,7 +12,17 @@ type Props = {
 
 const ImageMetadataWorkflowTabContent = ({ image }: Props) => {
   const { t } = useTranslation();
-  const { workflow } = useDebouncedImageWorkflow(image);
+  const { currentData } = useDebouncedImageWorkflow(image);
+  const workflow = useMemo(() => {
+    if (currentData?.workflow) {
+      try {
+        return JSON.parse(currentData.workflow);
+      } catch {
+        return null;
+      }
+    }
+    return null;
+  }, [currentData]);
 
   if (!workflow) {
     return <IAINoContentFallback label={t('nodes.noWorkflow')} />;

invokeai/frontend/web/src/features/nodes/store/actions.ts

@@ -1,6 +1,6 @@
 import { createAction, isAnyOf } from '@reduxjs/toolkit';
 import type { WorkflowV3 } from 'features/nodes/types/workflow';
-import type { Graph } from 'services/api/types';
+import type { Graph, GraphAndWorkflowResponse } from 'services/api/types';
 
 const textToImageGraphBuilt = createAction<Graph>('nodes/textToImageGraphBuilt');
 const imageToImageGraphBuilt = createAction<Graph>('nodes/imageToImageGraphBuilt');
@@ -15,7 +15,7 @@ export const isAnyGraphBuilt = isAnyOf(
 );
 
 export const workflowLoadRequested = createAction<{
-  workflow: unknown;
+  data: GraphAndWorkflowResponse;
   asCopy: boolean;
 }>('nodes/workflowLoadRequested');
 

invokeai/frontend/web/src/features/workflowLibrary/components/LoadWorkflowFromGraphModal/LoadWorkflowFromGraphModal.tsx

@@ -58,8 +58,7 @@ export const LoadWorkflowFromGraphModal = () => {
     setWorkflowRaw(JSON.stringify(workflow, null, 2));
   }, [graphRaw, shouldAutoLayout]);
   const loadWorkflow = useCallback(() => {
-    const workflow = JSON.parse(workflowRaw);
-    dispatch(workflowLoadRequested({ workflow, asCopy: true }));
+    dispatch(workflowLoadRequested({ data: { workflow: workflowRaw, graph: null }, asCopy: true }));
     onClose();
   }, [dispatch, onClose, workflowRaw]);
   return (

invokeai/frontend/web/src/features/workflowLibrary/hooks/useGetAndLoadEmbeddedWorkflow.ts

@@ -27,10 +27,17 @@ export const useGetAndLoadEmbeddedWorkflow: UseGetAndLoadEmbeddedWorkflow = ({ o
   const getAndLoadEmbeddedWorkflow = useCallback(
     async (imageName: string) => {
       try {
-        const workflow = await _getAndLoadEmbeddedWorkflow(imageName);
-        dispatch(workflowLoadRequested({ workflow: workflow.data, asCopy: true }));
-        // No toast - the listener for this action does that after the workflow is loaded
-        onSuccess && onSuccess();
+        const { data } = await _getAndLoadEmbeddedWorkflow(imageName);
+        if (data) {
+          dispatch(workflowLoadRequested({ data, asCopy: true }));
+          // No toast - the listener for this action does that after the workflow is loaded
+          onSuccess && onSuccess();
+        } else {
+          toaster({
+            title: t('toast.problemRetrievingWorkflow'),
+            status: 'error',
+          });
+        }
       } catch {
         toaster({
           title: t('toast.problemRetrievingWorkflow'),

invokeai/frontend/web/src/services/api/endpoints/images.ts

@@ -10,6 +10,7 @@ import { keyBy } from 'lodash-es';
 import type { components, paths } from 'services/api/schema';
 import type {
   DeleteBoardResult,
+  GraphAndWorkflowResponse,
   ImageCategory,
   ImageDTO,
   ListImagesArgs,
@@ -122,10 +123,7 @@ export const imagesApi = api.injectEndpoints({
       providesTags: (result, error, image_name) => [{ type: 'ImageMetadata', id: image_name }],
       keepUnusedDataFor: 86400, // 24 hours
     }),
-    getImageWorkflow: build.query<
-      paths['/api/v1/images/i/{image_name}/workflow']['get']['responses']['200']['content']['application/json'],
-      string
-    >({
+    getImageWorkflow: build.query<GraphAndWorkflowResponse, string>({
       query: (image_name) => ({ url: buildImagesUrl(`i/${image_name}/workflow`) }),
       providesTags: (result, error, image_name) => [{ type: 'ImageWorkflow', id: image_name }],
       keepUnusedDataFor: 86400, // 24 hours

invokeai/frontend/web/src/services/api/hooks/useDebouncedImageWorkflow.ts

@@ -9,7 +9,7 @@ export const useDebouncedImageWorkflow = (imageDTO?: ImageDTO | null) => {
 
   const [debouncedImageName] = useDebounce(imageDTO?.has_workflow ? imageDTO.image_name : null, workflowFetchDebounce);
 
-  const { data: workflow, isLoading } = useGetImageWorkflowQuery(debouncedImageName ?? skipToken);
+  const result = useGetImageWorkflowQuery(debouncedImageName ?? skipToken);
 
-  return { workflow, isLoading };
+  return result;
 };

invokeai/frontend/web/src/services/api/types.ts

@@ -16,6 +16,9 @@ export type UpdateBoardArg = paths['/api/v1/boards/{board_id}']['patch']['parame
   changes: paths['/api/v1/boards/{board_id}']['patch']['requestBody']['content']['application/json'];
 };
 
+export type GraphAndWorkflowResponse =
+  paths['/api/v1/images/i/{image_name}/workflow']['get']['responses']['200']['content']['application/json'];
+
 export type BatchConfig =
   paths['/api/v1/queue/{queue_id}/enqueue_batch']['post']['requestBody']['content']['application/json'];