feat(ui): standardize and clean up workflow loading hooks and logic

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

@@ -4,7 +4,7 @@ import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useAssertSingleton } from 'common/hooks/useAssertSingleton';
 import { useWorkflowLibraryModal } from 'features/nodes/store/workflowLibraryModal';
 import { selectWorkflowIsTouched, workflowModeChanged } from 'features/nodes/store/workflowSlice';
-import { useGetAndLoadLibraryWorkflow } from 'features/workflowLibrary/hooks/useGetAndLoadLibraryWorkflow';
+import { useLoadWorkflowFromLibrary } from 'features/workflowLibrary/hooks/useLoadWorkflowFromLibrary';
 import { atom } from 'nanostores';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -15,7 +15,7 @@ const cleanup = () => $workflowToLoad.set(null);
 export const useLoadWorkflow = () => {
   const dispatch = useAppDispatch();
   const workflowLibraryModal = useWorkflowLibraryModal();
-  const { getAndLoadWorkflow } = useGetAndLoadLibraryWorkflow();
+  const loadWorkflowFromLibrary = useLoadWorkflowFromLibrary();
 
   const isTouched = useAppSelector(selectWorkflowIsTouched);
 
@@ -25,11 +25,14 @@ export const useLoadWorkflow = () => {
       return;
     }
     const { workflowId, mode } = workflow;
-    await getAndLoadWorkflow(workflowId);
-    dispatch(workflowModeChanged(mode));
+    await loadWorkflowFromLibrary(workflowId, {
+      onSuccess: () => {
+        dispatch(workflowModeChanged(mode));
+      },
+    });
     cleanup();
     workflowLibraryModal.close();
-  }, [dispatch, getAndLoadWorkflow, workflowLibraryModal]);
+  }, [dispatch, loadWorkflowFromLibrary, workflowLibraryModal]);
 
   const loadWithDialog = useCallback(
     (workflowId: string, mode: 'view' | 'edit') => {

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

@@ -15,7 +15,7 @@ import {
 } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { graphToWorkflow } from 'features/nodes/util/workflow/graphToWorkflow';
-import { useLoadWorkflow } from 'features/workflowLibrary/hooks/useLoadWorkflow';
+import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useValidateAndLoadWorkflow';
 import { atom } from 'nanostores';
 import type { ChangeEvent } from 'react';
 import { useCallback, useState } from 'react';
@@ -37,16 +37,16 @@ export const useLoadWorkflowFromGraphModal = () => {
 
 export const LoadWorkflowFromGraphModal = () => {
   const { t } = useTranslation();
-  const _loadWorkflow = useLoadWorkflow();
+  const validateAndLoadWorkflow = useValidateAndLoadWorkflow();
   const { isOpen, onClose } = useLoadWorkflowFromGraphModal();
   const [graphRaw, setGraphRaw] = useState<string>('');
-  const [workflowRaw, setWorkflowRaw] = useState<string>('');
+  const [unvalidatedWorkflow, setUnvalidatedWorkflow] = useState<string>('');
   const [shouldAutoLayout, setShouldAutoLayout] = useState(true);
   const onChangeGraphRaw = useCallback((e: ChangeEvent<HTMLTextAreaElement>) => {
     setGraphRaw(e.target.value);
   }, []);
   const onChangeWorkflowRaw = useCallback((e: ChangeEvent<HTMLTextAreaElement>) => {
-    setWorkflowRaw(e.target.value);
+    setUnvalidatedWorkflow(e.target.value);
   }, []);
   const onChangeShouldAutoLayout = useCallback((e: ChangeEvent<HTMLInputElement>) => {
     setShouldAutoLayout(e.target.checked);
@@ -54,12 +54,12 @@ export const LoadWorkflowFromGraphModal = () => {
   const parse = useCallback(() => {
     const graph = JSON.parse(graphRaw);
     const workflow = graphToWorkflow(graph, shouldAutoLayout);
-    setWorkflowRaw(JSON.stringify(workflow, null, 2));
+    setUnvalidatedWorkflow(JSON.stringify(workflow, null, 2));
   }, [graphRaw, shouldAutoLayout]);
-  const loadWorkflow = useCallback(() => {
-    _loadWorkflow({ workflow: workflowRaw, graph: null });
+  const loadWorkflow = useCallback(async () => {
+    await validateAndLoadWorkflow(unvalidatedWorkflow);
     onClose();
-  }, [_loadWorkflow, onClose, workflowRaw]);
+  }, [validateAndLoadWorkflow, onClose, unvalidatedWorkflow]);
   return (
     <Modal isOpen={isOpen} onClose={onClose} isCentered useInert={false}>
       <ModalOverlay />
@@ -95,7 +95,7 @@ export const LoadWorkflowFromGraphModal = () => {
             <FormLabel>{t('nodes.workflow')}</FormLabel>
             <Textarea
               h="full"
-              value={workflowRaw}
+              value={unvalidatedWorkflow}
               fontFamily="monospace"
               whiteSpace="pre-wrap"
               overflowWrap="normal"

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

@@ -1,33 +1,29 @@
 import { Button } from '@invoke-ai/ui-library';
 import { useWorkflowLibraryModal } from 'features/nodes/store/workflowLibraryModal';
-import { saveWorkflowAs } from 'features/workflowLibrary/components/SaveWorkflowAsDialog';
 import { useLoadWorkflowFromFile } from 'features/workflowLibrary/hooks/useLoadWorkflowFromFile';
-import { memo, useCallback, useRef } from 'react';
+import { memo, useCallback } from 'react';
 import { useDropzone } from 'react-dropzone';
 import { useTranslation } from 'react-i18next';
 import { PiUploadSimpleBold } from 'react-icons/pi';
 
 export const UploadWorkflowButton = memo(() => {
   const { t } = useTranslation();
-  const resetRef = useRef<() => void>(null);
   const workflowLibraryModal = useWorkflowLibraryModal();
 
-  const loadWorkflowFromFile = useLoadWorkflowFromFile({
-    resetRef,
-    onSuccess: (workflow) => {
-      workflowLibraryModal.close();
-      saveWorkflowAs(workflow);
-    },
-  });
+  const loadWorkflowFromFile = useLoadWorkflowFromFile();
 
   const onDropAccepted = useCallback(
-    (files: File[]) => {
-      if (!files[0]) {
+    ([file]: File[]) => {
+      if (!file) {
         return;
       }
-      loadWorkflowFromFile(files[0]);
+      loadWorkflowFromFile(file, {
+        onSuccess: () => {
+          workflowLibraryModal.close();
+        },
+      });
     },
-    [loadWorkflowFromFile]
+    [loadWorkflowFromFile, workflowLibraryModal]
   );
 
   const { getInputProps, getRootProps } = useDropzone({
@@ -36,6 +32,7 @@ export const UploadWorkflowButton = memo(() => {
     noDrag: true,
     multiple: false,
   });
+
   return (
     <>
       <Button

invokeai/frontend/web/src/features/workflowLibrary/components/WorkflowLibraryMenu/UploadWorkflowMenuItem.tsx

@@ -1,32 +1,29 @@
 import { MenuItem } from '@invoke-ai/ui-library';
 import { useWorkflowLibraryModal } from 'features/nodes/store/workflowLibraryModal';
-import { saveWorkflowAs } from 'features/workflowLibrary/components/SaveWorkflowAsDialog';
 import { useLoadWorkflowFromFile } from 'features/workflowLibrary/hooks/useLoadWorkflowFromFile';
-import { memo, useCallback, useRef } from 'react';
+import { memo, useCallback } from 'react';
 import { useDropzone } from 'react-dropzone';
 import { useTranslation } from 'react-i18next';
 import { PiUploadSimpleBold } from 'react-icons/pi';
 
 const UploadWorkflowMenuItem = () => {
   const { t } = useTranslation();
-  const resetRef = useRef<() => void>(null);
   const workflowLibraryModal = useWorkflowLibraryModal();
-  const loadWorkflowFromFile = useLoadWorkflowFromFile({
-    resetRef,
-    onSuccess: (workflow) => {
-      workflowLibraryModal.close();
-      saveWorkflowAs(workflow);
-    },
-  });
+
+  const loadWorkflowFromFile = useLoadWorkflowFromFile();
 
   const onDropAccepted = useCallback(
-    (files: File[]) => {
-      if (!files[0]) {
+    ([file]: File[]) => {
+      if (!file) {
         return;
       }
-      loadWorkflowFromFile(files[0]);
+      loadWorkflowFromFile(file, {
+        onSuccess: () => {
+          workflowLibraryModal.close();
+        },
+      });
     },
-    [loadWorkflowFromFile]
+    [loadWorkflowFromFile, workflowLibraryModal]
   );
 
   const { getRootProps, getInputProps } = useDropzone({
@@ -35,6 +32,7 @@ const UploadWorkflowMenuItem = () => {
     noDrag: true,
     multiple: false,
   });
+
   return (
     <MenuItem as="button" icon={<PiUploadSimpleBold />} {...getRootProps()}>
       {t('workflows.uploadWorkflow')}

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

@@ -1,44 +0,0 @@
-import { toast } from 'features/toast/toast';
-import { useLoadWorkflow } from 'features/workflowLibrary/hooks/useLoadWorkflow';
-import { useCallback } from 'react';
-import { useTranslation } from 'react-i18next';
-import { useLazyGetImageWorkflowQuery } from 'services/api/endpoints/images';
-
-type UseGetAndLoadEmbeddedWorkflowOptions = {
-  onSuccess?: () => void;
-  onError?: () => void;
-};
-
-export const useGetAndLoadEmbeddedWorkflow = (options?: UseGetAndLoadEmbeddedWorkflowOptions) => {
-  const { t } = useTranslation();
-  const [_getAndLoadEmbeddedWorkflow, result] = useLazyGetImageWorkflowQuery();
-  const loadWorkflow = useLoadWorkflow();
-  const getAndLoadEmbeddedWorkflow = useCallback(
-    async (imageName: string) => {
-      try {
-        const { data } = await _getAndLoadEmbeddedWorkflow(imageName);
-        if (data) {
-          loadWorkflow(data);
-          // No toast - the listener for this action does that after the workflow is loaded
-          options?.onSuccess && options?.onSuccess();
-        } else {
-          toast({
-            id: 'PROBLEM_RETRIEVING_WORKFLOW',
-            title: t('toast.problemRetrievingWorkflow'),
-            status: 'error',
-          });
-        }
-      } catch {
-        toast({
-          id: 'PROBLEM_RETRIEVING_WORKFLOW',
-          title: t('toast.problemRetrievingWorkflow'),
-          status: 'error',
-        });
-        options?.onError && options?.onError();
-      }
-    },
-    [_getAndLoadEmbeddedWorkflow, loadWorkflow, options, t]
-  );
-
-  return [getAndLoadEmbeddedWorkflow, result] as const;
-};

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

@@ -1,47 +0,0 @@
-import { useToast } from '@invoke-ai/ui-library';
-import { useLoadWorkflow } from 'features/workflowLibrary/hooks/useLoadWorkflow';
-import { useCallback } from 'react';
-import { useTranslation } from 'react-i18next';
-import { useLazyGetWorkflowQuery, useUpdateOpenedAtMutation, workflowsApi } from 'services/api/endpoints/workflows';
-
-type UseGetAndLoadLibraryWorkflowOptions = {
-  onSuccess?: () => void;
-  onError?: () => void;
-};
-
-type UseGetAndLoadLibraryWorkflowReturn = {
-  getAndLoadWorkflow: (workflow_id: string) => Promise<void>;
-  getAndLoadWorkflowResult: ReturnType<typeof useLazyGetWorkflowQuery>[1];
-};
-
-type UseGetAndLoadLibraryWorkflow = (arg?: UseGetAndLoadLibraryWorkflowOptions) => UseGetAndLoadLibraryWorkflowReturn;
-
-export const useGetAndLoadLibraryWorkflow: UseGetAndLoadLibraryWorkflow = (arg) => {
-  const toast = useToast();
-  const { t } = useTranslation();
-  const loadWorkflow = useLoadWorkflow();
-  const [getWorkflow, getAndLoadWorkflowResult] = useLazyGetWorkflowQuery();
-  const [updateOpenedAt] = useUpdateOpenedAtMutation();
-  const getAndLoadWorkflow = useCallback(
-    async (workflow_id: string) => {
-      try {
-        const { workflow } = await getWorkflow(workflow_id).unwrap();
-        // This action expects a stringified workflow, instead of updating the routes and services we will just stringify it here
-        await loadWorkflow({ workflow: JSON.stringify(workflow), graph: null });
-        updateOpenedAt({ workflow_id });
-        // No toast - the listener for this action does that after the workflow is loaded
-        arg?.onSuccess && arg.onSuccess();
-      } catch {
-        toast({
-          id: `AUTH_ERROR_TOAST_${workflowsApi.endpoints.getWorkflow.name}`,
-          title: t('toast.problemRetrievingWorkflow'),
-          status: 'error',
-        });
-        arg?.onError && arg.onError();
-      }
-    },
-    [getWorkflow, loadWorkflow, updateOpenedAt, arg, toast, t]
-  );
-
-  return { getAndLoadWorkflow, getAndLoadWorkflowResult };
-};

invokeai/frontend/web/src/features/workflowLibrary/hooks/useLoadWorkflowFromFile.tsx

@@ -1,46 +1,44 @@
 import { useAppDispatch } from 'app/store/storeHooks';
 import type { WorkflowV3 } from 'features/nodes/types/workflow';
-import { useLoadWorkflow } from 'features/workflowLibrary/hooks/useLoadWorkflow';
+import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useValidateAndLoadWorkflow';
 import { workflowLoadedFromFile } from 'features/workflowLibrary/store/actions';
-import type { RefObject } from 'react';
 import { useCallback } from 'react';
-import { assert } from 'tsafe';
 
-type useLoadWorkflowFromFileOptions = {
-  resetRef: RefObject<() => void>;
-  onSuccess?: (workflow: WorkflowV3) => void;
-};
-
-type UseLoadWorkflowFromFile = (options: useLoadWorkflowFromFileOptions) => (file: File | null) => void;
-
-export const useLoadWorkflowFromFile: UseLoadWorkflowFromFile = ({ resetRef, onSuccess }) => {
+export const useLoadWorkflowFromFile = () => {
   const dispatch = useAppDispatch();
-  const loadWorkflow = useLoadWorkflow();
+  const validatedAndLoadWorkflow = useValidateAndLoadWorkflow();
   const loadWorkflowFromFile = useCallback(
-    (file: File | null) => {
-      if (!file) {
-        return;
-      }
+    (
+      file: File,
+      options: {
+        onSuccess?: (workflow: WorkflowV3) => void;
+        onError?: () => void;
+      } = {}
+    ) => {
       const reader = new FileReader();
       reader.onload = async () => {
         const rawJSON = reader.result;
-
+        const { onSuccess, onError } = options;
         try {
-          const workflow = await loadWorkflow({ workflow: String(rawJSON), graph: null });
-          assert(workflow !== null);
+          const unvalidatedWorkflow = JSON.parse(rawJSON as string);
+          const validatedWorkflow = await validatedAndLoadWorkflow(unvalidatedWorkflow);
+
+          if (!validatedWorkflow) {
+            reader.abort();
+            onError?.();
+            return;
+          }
           dispatch(workflowLoadedFromFile());
-          onSuccess && onSuccess(workflow);
-        } catch (e) {
-          reader.abort();
+          onSuccess?.(validatedWorkflow);
+        } catch {
+          // This is catching the error from the parsing the JSON file
+          onError?.();
         }
       };
 
       reader.readAsText(file);
-
-      // Reset the file picker internal state so that the same file can be loaded again
-      resetRef.current?.();
     },
-    [resetRef, loadWorkflow, dispatch, onSuccess]
+    [validatedAndLoadWorkflow, dispatch]
   );
 
   return loadWorkflowFromFile;

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

@@ -0,0 +1,62 @@
+import type { WorkflowV3 } from 'features/nodes/types/workflow';
+import { graphToWorkflow } from 'features/nodes/util/workflow/graphToWorkflow';
+import { toast } from 'features/toast/toast';
+import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useValidateAndLoadWorkflow';
+import { useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+import { useLazyGetImageWorkflowQuery } from 'services/api/endpoints/images';
+import type { NonNullableGraph } from 'services/api/types';
+import { assert } from 'tsafe';
+
+export const useLoadWorkflowFromImage = () => {
+  const { t } = useTranslation();
+  const [getWorkflowAndGraphFromImage, result] = useLazyGetImageWorkflowQuery();
+  const validateAndLoadWorkflow = useValidateAndLoadWorkflow();
+  const getAndLoadEmbeddedWorkflow = useCallback(
+    async (
+      imageName: string,
+      options: {
+        onSuccess?: (workflow: WorkflowV3) => void;
+        onError?: () => void;
+      } = {}
+    ) => {
+      const { onSuccess, onError } = options;
+      try {
+        const { workflow, graph } = await getWorkflowAndGraphFromImage(imageName).unwrap();
+
+        // Images may have a workflow and/or a graph. We can load either into the workflow editor, but we prefer the
+        // workflow.
+        const unvalidatedWorkflow = workflow
+          ? JSON.parse(workflow)
+          : graph
+            ? graphToWorkflow(JSON.parse(graph) as NonNullableGraph, true)
+            : null;
+
+        assert(unvalidatedWorkflow !== null, 'No workflow or graph provided');
+
+        const validatedWorkflow = await validateAndLoadWorkflow(unvalidatedWorkflow);
+
+        if (!validatedWorkflow) {
+          onError?.();
+          return;
+        }
+
+        onSuccess?.(validatedWorkflow);
+      } catch {
+        // This is catching:
+        // - the error from the getWorkflowAndGraphFromImage query
+        // - the error from parsing the workflow or graph
+        toast({
+          id: 'PROBLEM_RETRIEVING_WORKFLOW',
+          title: t('toast.problemRetrievingWorkflow'),
+          status: 'error',
+        });
+        onError?.();
+        return;
+      }
+    },
+    [getWorkflowAndGraphFromImage, validateAndLoadWorkflow, t]
+  );
+
+  return [getAndLoadEmbeddedWorkflow, result] as const;
+};

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

@@ -0,0 +1,48 @@
+import { useToast } from '@invoke-ai/ui-library';
+import type { WorkflowV3 } from 'features/nodes/types/workflow';
+import { useValidateAndLoadWorkflow } from 'features/workflowLibrary/hooks/useValidateAndLoadWorkflow';
+import { useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+import { useLazyGetWorkflowQuery, useUpdateOpenedAtMutation, workflowsApi } from 'services/api/endpoints/workflows';
+
+export const useLoadWorkflowFromLibrary = () => {
+  const toast = useToast();
+  const { t } = useTranslation();
+  const validateAndLoadWorkflow = useValidateAndLoadWorkflow();
+  const [getWorkflow] = useLazyGetWorkflowQuery();
+  const [updateOpenedAt] = useUpdateOpenedAtMutation();
+  const loadWorkflowFromLibrary = useCallback(
+    async (
+      workflowId: string,
+      options: {
+        onSuccess?: (workflow: WorkflowV3) => void;
+        onError?: () => void;
+      } = {}
+    ) => {
+      const { onSuccess, onError } = options;
+      try {
+        const res = await getWorkflow(workflowId).unwrap();
+
+        const validatedWorkflow = await validateAndLoadWorkflow(res.workflow);
+
+        if (!validatedWorkflow) {
+          onError?.();
+          return;
+        }
+        updateOpenedAt({ workflow_id: workflowId });
+        onSuccess?.(validatedWorkflow);
+      } catch {
+        // This is catching the error from the getWorkflow query
+        toast({
+          id: `AUTH_ERROR_TOAST_${workflowsApi.endpoints.getWorkflow.name}`,
+          title: t('toast.problemRetrievingWorkflow'),
+          status: 'error',
+        });
+        onError?.();
+      }
+    },
+    [getWorkflow, validateAndLoadWorkflow, updateOpenedAt, toast, t]
+  );
+
+  return loadWorkflowFromLibrary;
+};

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

@@ -4,51 +4,40 @@ import { $nodeExecutionStates } from 'features/nodes/hooks/useNodeExecutionState
 import { workflowLoaded } from 'features/nodes/store/actions';
 import { $templates } from 'features/nodes/store/nodesSlice';
 import { $needsFit } from 'features/nodes/store/reactFlowInstance';
-import type { Templates } from 'features/nodes/store/types';
 import { WorkflowMigrationError, WorkflowVersionError } from 'features/nodes/types/error';
 import type { WorkflowV3 } from 'features/nodes/types/workflow';
-import { graphToWorkflow } from 'features/nodes/util/workflow/graphToWorkflow';
 import { validateWorkflow } from 'features/nodes/util/workflow/validateWorkflow';
 import { toast } from 'features/toast/toast';
 import { useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { serializeError } from 'serialize-error';
 import { checkBoardAccess, checkImageAccess, checkModelAccess } from 'services/api/hooks/accessChecks';
-import type { GraphAndWorkflowResponse, NonNullableGraph } from 'services/api/types';
 import { z } from 'zod';
 import { fromZodError } from 'zod-validation-error';
 
 const log = logger('workflows');
 
-const getWorkflowFromStringifiedWorkflowOrGraph = async (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 await validateWorkflow({
-      workflow: parsed,
-      templates,
-      checkImageAccess,
-      checkBoardAccess,
-      checkModelAccess,
-    });
-  } 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 await validateWorkflow({ workflow, templates, checkImageAccess, checkBoardAccess, checkModelAccess });
-  } else {
-    throw new Error('No workflow or graph provided');
-  }
-};
-
-export const useLoadWorkflow = () => {
+export const useValidateAndLoadWorkflow = () => {
   const { t } = useTranslation();
   const dispatch = useAppDispatch();
-  const loadWorkflow = useCallback(
-    async (data: GraphAndWorkflowResponse): Promise<WorkflowV3 | null> => {
+  const validateAndLoadWorkflow = useCallback(
+    /**
+     * Validate and load a workflow into the editor.
+     *
+     * The unvalidated workflow should be a JS object. Do not pass a raw JSON string.
+     *
+     * This function catches all errors. It toasts and logs on success and error.
+     */
+    async (unvalidatedWorkflow: unknown): Promise<WorkflowV3 | null> => {
       try {
         const templates = $templates.get();
-        const { workflow, warnings } = await getWorkflowFromStringifiedWorkflowOrGraph(data, templates);
+        const { workflow, warnings } = await validateWorkflow({
+          workflow: unvalidatedWorkflow,
+          templates,
+          checkImageAccess,
+          checkBoardAccess,
+          checkModelAccess,
+        });
 
         $nodeExecutionStates.set({});
         dispatch(workflowLoaded(workflow));
@@ -119,5 +108,5 @@ export const useLoadWorkflow = () => {
     [dispatch, t]
   );
 
-  return loadWorkflow;
+  return validateAndLoadWorkflow;
 };