feat(ui): add canvas snapshot save/restore functionality

invokeai/app/api/routers/client_state.py

@@ -45,6 +45,44 @@ async def set_client_state(
         raise HTTPException(status_code=500, detail="Error setting client state")
 
 
+@client_state_router.get(
+    "/{queue_id}/get_keys_by_prefix",
+    operation_id="get_client_state_keys_by_prefix",
+    response_model=list[str],
+)
+async def get_client_state_keys_by_prefix(
+    current_user: CurrentUserOrDefault,
+    queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"),
+    prefix: str = Query(..., description="Prefix to filter keys by"),
+) -> list[str]:
+    """Gets client state keys matching a prefix for the current user"""
+    try:
+        return ApiDependencies.invoker.services.client_state_persistence.get_keys_by_prefix(
+            current_user.user_id, prefix
+        )
+    except Exception as e:
+        logging.error(f"Error getting client state keys: {e}")
+        raise HTTPException(status_code=500, detail="Error getting client state keys")
+
+
+@client_state_router.post(
+    "/{queue_id}/delete_by_key",
+    operation_id="delete_client_state_by_key",
+    responses={204: {"description": "Client state key deleted"}},
+)
+async def delete_client_state_by_key(
+    current_user: CurrentUserOrDefault,
+    queue_id: str = Path(description="The queue id (ignored, kept for backwards compatibility)"),
+    key: str = Query(..., description="Key to delete"),
+) -> None:
+    """Deletes a specific client state key for the current user"""
+    try:
+        ApiDependencies.invoker.services.client_state_persistence.delete_by_key(current_user.user_id, key)
+    except Exception as e:
+        logging.error(f"Error deleting client state key: {e}")
+        raise HTTPException(status_code=500, detail="Error deleting client state key")
+
+
 @client_state_router.post(
     "/{queue_id}/delete",
     operation_id="delete_client_state",

invokeai/app/services/client_state_persistence/client_state_persistence_base.py

@@ -36,6 +36,31 @@ def get_by_key(self, user_id: str, key: str) -> str | None:
         """
         pass
 
+    @abstractmethod
+    def get_keys_by_prefix(self, user_id: str, prefix: str) -> list[str]:
+        """
+        Get all keys matching a prefix for a user.
+
+        Args:
+            user_id (str): The user ID to get keys for.
+            prefix (str): The prefix to filter keys by.
+
+        Returns:
+            list[str]: A list of keys matching the prefix.
+        """
+        pass
+
+    @abstractmethod
+    def delete_by_key(self, user_id: str, key: str) -> None:
+        """
+        Delete a specific key-value pair for a user.
+
+        Args:
+            user_id (str): The user ID to delete state for.
+            key (str): The key to delete.
+        """
+        pass
+
     @abstractmethod
     def delete(self, user_id: str) -> None:
         """

invokeai/app/services/client_state_persistence/client_state_persistence_sqlite.py

@@ -44,6 +44,28 @@ def get_by_key(self, user_id: str, key: str) -> str | None:
                 return None
             return row[0]
 
+    def get_keys_by_prefix(self, user_id: str, prefix: str) -> list[str]:
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """
+                SELECT key FROM client_state
+                WHERE user_id = ? AND key LIKE ?
+                ORDER BY rowid DESC
+                """,
+                (user_id, f"{prefix}%"),
+            )
+            return [row[0] for row in cursor.fetchall()]
+
+    def delete_by_key(self, user_id: str, key: str) -> None:
+        with self._db.transaction() as cursor:
+            cursor.execute(
+                """
+                DELETE FROM client_state
+                WHERE user_id = ? AND key = ?
+                """,
+                (user_id, key),
+            )
+
     def delete(self, user_id: str) -> None:
         with self._db.transaction() as cursor:
             cursor.execute(

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

@@ -2898,6 +2898,20 @@
             "off": "Off",
             "switchOnStart": "On Start",
             "switchOnFinish": "On Finish"
+        },
+        "snapshot": {
+            "snapshots": "Save or Load Canvas Snapshot",
+            "saveSnapshot": "Save Snapshot",
+            "restoreSnapshot": "Restore Snapshot",
+            "snapshotNamePlaceholder": "Snapshot name",
+            "save": "Save",
+            "delete": "Delete",
+            "snapshotSaved": "Snapshot \"{{name}}\" saved",
+            "snapshotRestored": "Snapshot \"{{name}}\" restored",
+            "snapshotDeleted": "Snapshot \"{{name}}\" deleted",
+            "snapshotSaveFailed": "Failed to save snapshot",
+            "snapshotRestoreFailed": "Failed to restore snapshot",
+            "snapshotDeleteFailed": "Failed to delete snapshot"
         }
     },
     "upscaling": {

invokeai/frontend/web/src/features/controlLayers/components/Toolbar/CanvasToolbar.tsx

@@ -14,6 +14,7 @@ import { CanvasToolbarRedoButton } from 'features/controlLayers/components/Toolb
 import { CanvasToolbarResetViewButton } from 'features/controlLayers/components/Toolbar/CanvasToolbarResetViewButton';
 import { CanvasToolbarSaveToGalleryButton } from 'features/controlLayers/components/Toolbar/CanvasToolbarSaveToGalleryButton';
 import { CanvasToolbarScale } from 'features/controlLayers/components/Toolbar/CanvasToolbarScale';
+import { CanvasToolbarSnapshotMenuButton } from 'features/controlLayers/components/Toolbar/CanvasToolbarSnapshotMenuButton';
 import { CanvasToolbarUndoButton } from 'features/controlLayers/components/Toolbar/CanvasToolbarUndoButton';
 import { useCanvasDeleteLayerHotkey } from 'features/controlLayers/hooks/useCanvasDeleteLayerHotkey';
 import { useCanvasEntityQuickSwitchHotkey } from 'features/controlLayers/hooks/useCanvasEntityQuickSwitchHotkey';
@@ -68,6 +69,7 @@ export const CanvasToolbar = memo(() => {
       <Divider orientation="vertical" />
       <Flex alignItems="center" h="full">
         <CanvasToolbarSaveToGalleryButton />
+        <CanvasToolbarSnapshotMenuButton />
         <CanvasToolbarUndoButton />
         <CanvasToolbarRedoButton />
         <CanvasToolbarNewSessionMenuButton />

invokeai/frontend/web/src/features/controlLayers/components/Toolbar/CanvasToolbarSnapshotMenuButton.tsx

@@ -0,0 +1,173 @@
+import {
+  Flex,
+  IconButton,
+  Input,
+  Menu,
+  MenuButton,
+  MenuDivider,
+  MenuGroup,
+  MenuItem,
+  MenuList,
+  Text,
+} from '@invoke-ai/ui-library';
+import type { SnapshotInfo } from 'features/controlLayers/hooks/useCanvasSnapshots';
+import { useCanvasSnapshots } from 'features/controlLayers/hooks/useCanvasSnapshots';
+import { toast } from 'features/toast/toast';
+import type { ChangeEvent, KeyboardEvent, MouseEvent } from 'react';
+import { memo, useCallback, useState } from 'react';
+import { useTranslation } from 'react-i18next';
+import { PiCameraBold, PiFloppyDiskBold, PiTrashBold } from 'react-icons/pi';
+
+const SnapshotItem = memo(
+  ({
+    snapshot,
+    onRestore,
+    onDelete,
+  }: {
+    snapshot: SnapshotInfo;
+    onRestore: (key: string, name: string) => void;
+    onDelete: (e: MouseEvent, key: string, name: string) => void;
+  }) => {
+    const handleClick = useCallback(() => {
+      onRestore(snapshot.key, snapshot.name);
+    }, [onRestore, snapshot.key, snapshot.name]);
+
+    const handleDelete = useCallback(
+      (e: MouseEvent) => {
+        onDelete(e, snapshot.key, snapshot.name);
+      },
+      [onDelete, snapshot.key, snapshot.name]
+    );
+
+    return (
+      <MenuItem onClick={handleClick}>
+        <Flex w="full" justifyContent="space-between" alignItems="center">
+          <Text fontSize="sm" noOfLines={1}>
+            {snapshot.name}
+          </Text>
+          <IconButton
+            aria-label="Delete"
+            icon={<PiTrashBold />}
+            size="xs"
+            variant="ghost"
+            colorScheme="error"
+            onClick={handleDelete}
+          />
+        </Flex>
+      </MenuItem>
+    );
+  }
+);
+
+SnapshotItem.displayName = 'SnapshotItem';
+
+const getDefaultSnapshotName = (): string => {
+  const now = new Date();
+  const y = now.getFullYear();
+  const mo = String(now.getMonth() + 1).padStart(2, '0');
+  const d = String(now.getDate()).padStart(2, '0');
+  const h = String(now.getHours()).padStart(2, '0');
+  const mi = String(now.getMinutes()).padStart(2, '0');
+  return `${y}/${mo}/${d} ${h}:${mi}`;
+};
+
+export const CanvasToolbarSnapshotMenuButton = memo(() => {
+  const { t } = useTranslation();
+  const { snapshots, saveSnapshot, restoreSnapshot, deleteSnapshot } = useCanvasSnapshots();
+  const [snapshotName, setSnapshotName] = useState('');
+
+  const onNameChange = useCallback((e: ChangeEvent<HTMLInputElement>) => {
+    setSnapshotName(e.target.value);
+  }, []);
+
+  const onSave = useCallback(async () => {
+    const name = snapshotName.trim() || getDefaultSnapshotName();
+    const success = await saveSnapshot(name);
+    if (success) {
+      toast({ title: t('controlLayers.snapshot.snapshotSaved', { name }), status: 'info' });
+      setSnapshotName('');
+    } else {
+      toast({ title: t('controlLayers.snapshot.snapshotSaveFailed'), status: 'error' });
+    }
+  }, [snapshotName, saveSnapshot, t]);
+
+  const onKeyDown = useCallback(
+    (e: KeyboardEvent<HTMLInputElement>) => {
+      if (e.key === 'Enter') {
+        e.preventDefault();
+        e.stopPropagation();
+        onSave();
+      }
+    },
+    [onSave]
+  );
+
+  const onRestore = useCallback(
+    async (key: string, name: string) => {
+      const success = await restoreSnapshot(key);
+      if (success) {
+        toast({ title: t('controlLayers.snapshot.snapshotRestored', { name }), status: 'info' });
+      } else {
+        toast({ title: t('controlLayers.snapshot.snapshotRestoreFailed'), status: 'error' });
+      }
+    },
+    [restoreSnapshot, t]
+  );
+
+  const onDelete = useCallback(
+    async (e: MouseEvent, key: string, name: string) => {
+      e.stopPropagation();
+      const success = await deleteSnapshot(key);
+      if (success) {
+        toast({ title: t('controlLayers.snapshot.snapshotDeleted', { name }), status: 'info' });
+      } else {
+        toast({ title: t('controlLayers.snapshot.snapshotDeleteFailed'), status: 'error' });
+      }
+    },
+    [deleteSnapshot, t]
+  );
+
+  return (
+    <Menu placement="bottom-end" closeOnSelect={false}>
+      <MenuButton
+        as={IconButton}
+        aria-label={t('controlLayers.snapshot.snapshots')}
+        tooltip={t('controlLayers.snapshot.snapshots')}
+        icon={<PiCameraBold />}
+        variant="link"
+        alignSelf="stretch"
+      />
+      <MenuList maxH="60vh" overflowY="auto">
+        <MenuGroup title={t('controlLayers.snapshot.saveSnapshot')}>
+          <Flex px={3} pb={2} gap={2} alignItems="center">
+            <Input
+              size="sm"
+              placeholder={t('controlLayers.snapshot.snapshotNamePlaceholder')}
+              value={snapshotName}
+              onChange={onNameChange}
+              onKeyDown={onKeyDown}
+            />
+            <IconButton
+              aria-label={t('controlLayers.snapshot.save')}
+              icon={<PiFloppyDiskBold />}
+              size="sm"
+              onClick={onSave}
+            />
+          </Flex>
+        </MenuGroup>
+        {snapshots.length > 0 && (
+          <>
+            <MenuDivider />
+            <MenuGroup title={t('controlLayers.snapshot.restoreSnapshot')}>
+              {snapshots.map((snapshot) => (
+                <SnapshotItem key={snapshot.key} snapshot={snapshot} onRestore={onRestore} onDelete={onDelete} />
+              ))}
+            </MenuGroup>
+          </>
+        )}
+      </MenuList>
+    </Menu>
+  );
+});
+
+CanvasToolbarSnapshotMenuButton.displayName = 'CanvasToolbarSnapshotMenuButton';