mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
fix: Fix app crashing when you upload an incorrect JSON to node editor (#3911)
## What type of PR is this? (check all applicable) - [x] Bug Fix ## Have you discussed this change with the InvokeAI team? - [x] Yes, we feel very passionate about this. ## Description Uploading an incorrect JSON file to the Node Editor would crash the app. While this is a much larger problem that we will tackle while refining the Node Editor, this is a fix that should address 99% of the cases out there. When saving an InvokeAI node graph, there are three primary keys. 1. `nodes` - which has all the node related data. 2. `edges` - which has all the edges related data 3. `viewport` - which has all the viewport related data. So when we load back the JSON, we now check if all three of these keys exist in the retrieved JSON object. While the `viewport` itself is not a mandatory key to repopulate the graph, checking for it will allow us to treat it as an additional check to ensure that the graph was saved from InvokeAI. As a result ... - If you upload an invalid JSON file, the app now warns you that the JSON is invalid. - If you upload a JSON of a graph editor that is not InvokeAI, it simply warns you that you are uploading a non InvokeAI graph. So effectively, you should not be able to load any graph that is not generated by ReactFlow. Here are the edge cases: - What happens if a user maintains the above key structure but tampers with the data inside them? Well tested it. Turns out because we validate and build the graph based on the JSON data, if you tamper with any data that is needed to rebuild that node, it simply will skip that and load the rest of the graph with valid data. - What happens if a user uploads a graph that was made by some other random ReactFlow app? Well, same as above. Because we do not have to parse that in our setup, it simply will skip it and only display what are setup to do. I think that just about covers 99% of the cases where this could go wrong. If there's any other edges cases, can add checks if need be. But can't think of any at the moment. ## Related Tickets & Documents ### Closes - #3893 - #3881 ## [optional] Are there any post deployment tasks we need to perform? Yes. Making @psychedelicious a little bit happier. :P
This commit is contained in:
commit
02618a701d
@ -102,8 +102,7 @@
|
||||
"openInNewTab": "Open in New Tab",
|
||||
"dontAskMeAgain": "Don't ask me again",
|
||||
"areYouSure": "Are you sure?",
|
||||
"imagePrompt": "Image Prompt",
|
||||
"clearNodes": "Are you sure you want to clear all nodes?"
|
||||
"imagePrompt": "Image Prompt"
|
||||
},
|
||||
"gallery": {
|
||||
"generations": "Generations",
|
||||
@ -615,6 +614,11 @@
|
||||
"initialImageNotSetDesc": "Could not load initial image",
|
||||
"nodesSaved": "Nodes Saved",
|
||||
"nodesLoaded": "Nodes Loaded",
|
||||
"nodesNotValidGraph": "Not a valid InvokeAI Node Graph",
|
||||
"nodesNotValidJSON": "Not a valid JSON",
|
||||
"nodesCorruptedGraph": "Cannot load. Graph seems to be corrupted.",
|
||||
"nodesUnrecognizedTypes": "Cannot load. Graph has unrecognized types",
|
||||
"nodesBrokenConnections": "Cannot load. Some connections are broken.",
|
||||
"nodesLoadedFailed": "Failed To Load Nodes",
|
||||
"nodesCleared": "Nodes Cleared"
|
||||
},
|
||||
@ -700,9 +704,10 @@
|
||||
},
|
||||
"nodes": {
|
||||
"reloadSchema": "Reload Schema",
|
||||
"saveNodes": "Save Nodes",
|
||||
"loadNodes": "Load Nodes",
|
||||
"clearNodes": "Clear Nodes",
|
||||
"saveGraph": "Save Graph",
|
||||
"loadGraph": "Load Graph (saved from Node Editor) (Do not copy-paste metadata)",
|
||||
"clearGraph": "Clear Graph",
|
||||
"clearGraphDesc": "Are you sure you want to clear all nodes?",
|
||||
"zoomInNodes": "Zoom In",
|
||||
"zoomOutNodes": "Zoom Out",
|
||||
"fitViewportNodes": "Fit View",
|
||||
|
@ -2,11 +2,11 @@ import { HStack } from '@chakra-ui/react';
|
||||
import CancelButton from 'features/parameters/components/ProcessButtons/CancelButton';
|
||||
import { memo } from 'react';
|
||||
import { Panel } from 'reactflow';
|
||||
import LoadNodesButton from '../ui/LoadNodesButton';
|
||||
import ClearGraphButton from '../ui/ClearGraphButton';
|
||||
import LoadGraphButton from '../ui/LoadGraphButton';
|
||||
import NodeInvokeButton from '../ui/NodeInvokeButton';
|
||||
import ReloadSchemaButton from '../ui/ReloadSchemaButton';
|
||||
import SaveNodesButton from '../ui/SaveNodesButton';
|
||||
import ClearNodesButton from '../ui/ClearNodesButton';
|
||||
import SaveGraphButton from '../ui/SaveGraphButton';
|
||||
|
||||
const TopCenterPanel = () => {
|
||||
return (
|
||||
@ -15,9 +15,9 @@ const TopCenterPanel = () => {
|
||||
<NodeInvokeButton />
|
||||
<CancelButton />
|
||||
<ReloadSchemaButton />
|
||||
<SaveNodesButton />
|
||||
<LoadNodesButton />
|
||||
<ClearNodesButton />
|
||||
<SaveGraphButton />
|
||||
<LoadGraphButton />
|
||||
<ClearGraphButton />
|
||||
</HStack>
|
||||
</Panel>
|
||||
);
|
||||
|
@ -9,17 +9,17 @@ import {
|
||||
Text,
|
||||
useDisclosure,
|
||||
} from '@chakra-ui/react';
|
||||
import { makeToast } from 'features/system/util/makeToast';
|
||||
import { RootState } from 'app/store/store';
|
||||
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
|
||||
import IAIIconButton from 'common/components/IAIIconButton';
|
||||
import { nodeEditorReset } from 'features/nodes/store/nodesSlice';
|
||||
import { addToast } from 'features/system/store/systemSlice';
|
||||
import { memo, useRef, useCallback } from 'react';
|
||||
import { makeToast } from 'features/system/util/makeToast';
|
||||
import { memo, useCallback, useRef } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { FaTrash } from 'react-icons/fa';
|
||||
|
||||
const ClearNodesButton = () => {
|
||||
const ClearGraphButton = () => {
|
||||
const { t } = useTranslation();
|
||||
const dispatch = useAppDispatch();
|
||||
const { isOpen, onOpen, onClose } = useDisclosure();
|
||||
@ -46,8 +46,8 @@ const ClearNodesButton = () => {
|
||||
<>
|
||||
<IAIIconButton
|
||||
icon={<FaTrash />}
|
||||
tooltip={t('nodes.clearNodes')}
|
||||
aria-label={t('nodes.clearNodes')}
|
||||
tooltip={t('nodes.clearGraph')}
|
||||
aria-label={t('nodes.clearGraph')}
|
||||
onClick={onOpen}
|
||||
isDisabled={nodes.length === 0}
|
||||
/>
|
||||
@ -62,11 +62,11 @@ const ClearNodesButton = () => {
|
||||
|
||||
<AlertDialogContent>
|
||||
<AlertDialogHeader fontSize="lg" fontWeight="bold">
|
||||
{t('nodes.clearNodes')}
|
||||
{t('nodes.clearGraph')}
|
||||
</AlertDialogHeader>
|
||||
|
||||
<AlertDialogBody>
|
||||
<Text>{t('common.clearNodes')}</Text>
|
||||
<Text>{t('nodes.clearGraphDesc')}</Text>
|
||||
</AlertDialogBody>
|
||||
|
||||
<AlertDialogFooter>
|
||||
@ -83,4 +83,4 @@ const ClearNodesButton = () => {
|
||||
);
|
||||
};
|
||||
|
||||
export default memo(ClearNodesButton);
|
||||
export default memo(ClearGraphButton);
|
@ -0,0 +1,161 @@
|
||||
import { FileButton } from '@mantine/core';
|
||||
import { useAppDispatch } from 'app/store/storeHooks';
|
||||
import IAIIconButton from 'common/components/IAIIconButton';
|
||||
import { loadFileEdges, loadFileNodes } from 'features/nodes/store/nodesSlice';
|
||||
import { addToast } from 'features/system/store/systemSlice';
|
||||
import { makeToast } from 'features/system/util/makeToast';
|
||||
import i18n from 'i18n';
|
||||
import { memo, useCallback, useRef } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { FaUpload } from 'react-icons/fa';
|
||||
import { useReactFlow } from 'reactflow';
|
||||
|
||||
interface JsonFile {
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
function sanityCheckInvokeAIGraph(jsonFile: JsonFile): {
|
||||
isValid: boolean;
|
||||
message: string;
|
||||
} {
|
||||
// Check if primary keys exist
|
||||
const keys = ['nodes', 'edges', 'viewport'];
|
||||
for (const key of keys) {
|
||||
if (!(key in jsonFile)) {
|
||||
return {
|
||||
isValid: false,
|
||||
message: i18n.t('toast.nodesNotValidGraph'),
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Check if nodes and edges are arrays
|
||||
if (!Array.isArray(jsonFile.nodes) || !Array.isArray(jsonFile.edges)) {
|
||||
return {
|
||||
isValid: false,
|
||||
message: i18n.t('toast.nodesNotValidGraph'),
|
||||
};
|
||||
}
|
||||
|
||||
// Check if data is present in nodes
|
||||
const nodeKeys = ['data', 'type'];
|
||||
const nodeTypes = ['invocation', 'progress_image'];
|
||||
if (jsonFile.nodes.length > 0) {
|
||||
for (const node of jsonFile.nodes) {
|
||||
for (const nodeKey of nodeKeys) {
|
||||
if (!(nodeKey in node)) {
|
||||
return {
|
||||
isValid: false,
|
||||
message: i18n.t('toast.nodesNotValidGraph'),
|
||||
};
|
||||
}
|
||||
if (nodeKey === 'type' && !nodeTypes.includes(node[nodeKey])) {
|
||||
return {
|
||||
isValid: false,
|
||||
message: i18n.t('toast.nodesUnrecognizedTypes'),
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check Edge Object
|
||||
const edgeKeys = ['source', 'sourceHandle', 'target', 'targetHandle'];
|
||||
if (jsonFile.edges.length > 0) {
|
||||
for (const edge of jsonFile.edges) {
|
||||
for (const edgeKey of edgeKeys) {
|
||||
if (!(edgeKey in edge)) {
|
||||
return {
|
||||
isValid: false,
|
||||
message: i18n.t('toast.nodesBrokenConnections'),
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
isValid: true,
|
||||
message: i18n.t('toast.nodesLoaded'),
|
||||
};
|
||||
}
|
||||
|
||||
const LoadGraphButton = () => {
|
||||
const { t } = useTranslation();
|
||||
const dispatch = useAppDispatch();
|
||||
const { fitView } = useReactFlow();
|
||||
|
||||
const uploadedFileRef = useRef<() => void>(null);
|
||||
|
||||
const restoreJSONToEditor = useCallback(
|
||||
(v: File | null) => {
|
||||
if (!v) return;
|
||||
const reader = new FileReader();
|
||||
reader.onload = async () => {
|
||||
const json = reader.result;
|
||||
|
||||
try {
|
||||
const retrievedNodeTree = await JSON.parse(String(json));
|
||||
const { isValid, message } =
|
||||
sanityCheckInvokeAIGraph(retrievedNodeTree);
|
||||
|
||||
if (isValid) {
|
||||
dispatch(loadFileNodes(retrievedNodeTree.nodes));
|
||||
dispatch(loadFileEdges(retrievedNodeTree.edges));
|
||||
fitView();
|
||||
|
||||
dispatch(
|
||||
addToast(makeToast({ title: message, status: 'success' }))
|
||||
);
|
||||
} else {
|
||||
dispatch(
|
||||
addToast(
|
||||
makeToast({
|
||||
title: message,
|
||||
status: 'error',
|
||||
})
|
||||
)
|
||||
);
|
||||
}
|
||||
// Cleanup
|
||||
reader.abort();
|
||||
} catch (error) {
|
||||
if (error) {
|
||||
dispatch(
|
||||
addToast(
|
||||
makeToast({
|
||||
title: t('toast.nodesNotValidJSON'),
|
||||
status: 'error',
|
||||
})
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
reader.readAsText(v);
|
||||
|
||||
// Cleanup
|
||||
uploadedFileRef.current?.();
|
||||
},
|
||||
[fitView, dispatch, t]
|
||||
);
|
||||
return (
|
||||
<FileButton
|
||||
resetRef={uploadedFileRef}
|
||||
accept="application/json"
|
||||
onChange={restoreJSONToEditor}
|
||||
>
|
||||
{(props) => (
|
||||
<IAIIconButton
|
||||
icon={<FaUpload />}
|
||||
tooltip={t('nodes.loadGraph')}
|
||||
aria-label={t('nodes.loadGraph')}
|
||||
{...props}
|
||||
/>
|
||||
)}
|
||||
</FileButton>
|
||||
);
|
||||
};
|
||||
|
||||
export default memo(LoadGraphButton);
|
@ -1,79 +0,0 @@
|
||||
import { FileButton } from '@mantine/core';
|
||||
import { makeToast } from 'features/system/util/makeToast';
|
||||
import { useAppDispatch } from 'app/store/storeHooks';
|
||||
import IAIIconButton from 'common/components/IAIIconButton';
|
||||
import { loadFileEdges, loadFileNodes } from 'features/nodes/store/nodesSlice';
|
||||
import { addToast } from 'features/system/store/systemSlice';
|
||||
import { memo, useCallback, useRef } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { FaUpload } from 'react-icons/fa';
|
||||
import { useReactFlow } from 'reactflow';
|
||||
|
||||
const LoadNodesButton = () => {
|
||||
const { t } = useTranslation();
|
||||
const dispatch = useAppDispatch();
|
||||
const { fitView } = useReactFlow();
|
||||
|
||||
const uploadedFileRef = useRef<() => void>(null);
|
||||
|
||||
const restoreJSONToEditor = useCallback(
|
||||
(v: File | null) => {
|
||||
if (!v) return;
|
||||
const reader = new FileReader();
|
||||
reader.onload = async () => {
|
||||
const json = reader.result;
|
||||
const retrievedNodeTree = await JSON.parse(String(json));
|
||||
|
||||
if (!retrievedNodeTree) {
|
||||
dispatch(
|
||||
addToast(
|
||||
makeToast({
|
||||
title: t('toast.nodesLoadedFailed'),
|
||||
status: 'error',
|
||||
})
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
if (retrievedNodeTree) {
|
||||
dispatch(loadFileNodes(retrievedNodeTree.nodes));
|
||||
dispatch(loadFileEdges(retrievedNodeTree.edges));
|
||||
fitView();
|
||||
|
||||
dispatch(
|
||||
addToast(
|
||||
makeToast({ title: t('toast.nodesLoaded'), status: 'success' })
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
reader.abort();
|
||||
};
|
||||
|
||||
reader.readAsText(v);
|
||||
|
||||
// Cleanup
|
||||
uploadedFileRef.current?.();
|
||||
},
|
||||
[fitView, dispatch, t]
|
||||
);
|
||||
return (
|
||||
<FileButton
|
||||
resetRef={uploadedFileRef}
|
||||
accept="application/json"
|
||||
onChange={restoreJSONToEditor}
|
||||
>
|
||||
{(props) => (
|
||||
<IAIIconButton
|
||||
icon={<FaUpload />}
|
||||
tooltip={t('nodes.loadNodes')}
|
||||
aria-label={t('nodes.loadNodes')}
|
||||
{...props}
|
||||
/>
|
||||
)}
|
||||
</FileButton>
|
||||
);
|
||||
};
|
||||
|
||||
export default memo(LoadNodesButton);
|
@ -6,7 +6,7 @@ import { memo, useCallback } from 'react';
|
||||
import { useTranslation } from 'react-i18next';
|
||||
import { FaSave } from 'react-icons/fa';
|
||||
|
||||
const SaveNodesButton = () => {
|
||||
const SaveGraphButton = () => {
|
||||
const { t } = useTranslation();
|
||||
const editorInstance = useAppSelector(
|
||||
(state: RootState) => state.nodes.editorInstance
|
||||
@ -37,12 +37,12 @@ const SaveNodesButton = () => {
|
||||
<IAIIconButton
|
||||
icon={<FaSave />}
|
||||
fontSize={18}
|
||||
tooltip={t('nodes.saveNodes')}
|
||||
aria-label={t('nodes.saveNodes')}
|
||||
tooltip={t('nodes.saveGraph')}
|
||||
aria-label={t('nodes.saveGraph')}
|
||||
onClick={saveEditorToJSON}
|
||||
isDisabled={nodes.length === 0}
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
||||
export default memo(SaveNodesButton);
|
||||
export default memo(SaveGraphButton);
|
Loading…
Reference in New Issue
Block a user