feat(nodes): add enable, disable, status to invocation cache

- New routes to clear, enable, disable and get the status of the cache
- Status includes hits, misses, size, max size, enabled
- Add client cache queries and mutations, abstracted into hooks
- Add invocation cache status area (next to queue status) w/ buttons

Author: psychedelicious

invokeai/app/api/routers/app_info.py

@@ -7,6 +7,7 @@
 from pydantic import BaseModel, Field
 
 from invokeai.app.invocations.upscale import ESRGAN_MODELS
+from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
 from invokeai.backend.image_util.invisible_watermark import InvisibleWatermark
 from invokeai.backend.image_util.patchmatch import PatchMatch
 from invokeai.backend.image_util.safety_checker import SafetyChecker
@@ -113,3 +114,33 @@ async def set_log_level(
 async def clear_invocation_cache() -> None:
     """Clears the invocation cache"""
     ApiDependencies.invoker.services.invocation_cache.clear()
+
+
+@app_router.put(
+    "/invocation_cache/enable",
+    operation_id="enable_invocation_cache",
+    responses={200: {"description": "The operation was successful"}},
+)
+async def enable_invocation_cache() -> None:
+    """Clears the invocation cache"""
+    ApiDependencies.invoker.services.invocation_cache.enable()
+
+
+@app_router.put(
+    "/invocation_cache/disable",
+    operation_id="disable_invocation_cache",
+    responses={200: {"description": "The operation was successful"}},
+)
+async def disable_invocation_cache() -> None:
+    """Clears the invocation cache"""
+    ApiDependencies.invoker.services.invocation_cache.disable()
+
+
+@app_router.get(
+    "/invocation_cache/status",
+    operation_id="get_invocation_cache_status",
+    responses={200: {"model": InvocationCacheStatus}},
+)
+async def get_invocation_cache_status() -> InvocationCacheStatus:
+    """Clears the invocation cache"""
+    return ApiDependencies.invoker.services.invocation_cache.get_status()

invokeai/app/services/invocation_cache/invocation_cache_base.py

@@ -2,6 +2,7 @@
 from typing import Optional, Union
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
+from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
 
 
 class InvocationCacheBase(ABC):
@@ -32,7 +33,7 @@ def save(self, key: Union[int, str], invocation_output: BaseInvocationOutput) ->
 
     @abstractmethod
     def delete(self, key: Union[int, str]) -> None:
-        """Deleteds an invocation output from the cache"""
+        """Deletes an invocation output from the cache"""
         pass
 
     @abstractmethod
@@ -44,3 +45,18 @@ def clear(self) -> None:
     def create_key(self, invocation: BaseInvocation) -> int:
         """Gets the key for the invocation's cache item"""
         pass
+
+    @abstractmethod
+    def disable(self) -> None:
+        """Disables the cache, overriding the max cache size"""
+        pass
+
+    @abstractmethod
+    def enable(self) -> None:
+        """Enables the cache, letting the the max cache size take effect"""
+        pass
+
+    @abstractmethod
+    def get_status(self) -> InvocationCacheStatus:
+        """Returns the status of the cache"""
+        pass

invokeai/app/services/invocation_cache/invocation_cache_common.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, Field
+
+
+class InvocationCacheStatus(BaseModel):
+    size: int = Field(description="The current size of the invocation cache")
+    hits: int = Field(description="The number of cache hits")
+    misses: int = Field(description="The number of cache misses")
+    enabled: bool = Field(description="Whether the invocation cache is enabled")
+    max_size: int = Field(description="The maximum size of the invocation cache")

invokeai/app/services/invocation_cache/invocation_cache_memory.py

@@ -3,18 +3,25 @@
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput
 from invokeai.app.services.invocation_cache.invocation_cache_base import InvocationCacheBase
+from invokeai.app.services.invocation_cache.invocation_cache_common import InvocationCacheStatus
 from invokeai.app.services.invoker import Invoker
 
 
 class MemoryInvocationCache(InvocationCacheBase):
     __cache: dict[Union[int, str], tuple[BaseInvocationOutput, str]]
     __max_cache_size: int
+    __disabled: bool
+    __hits: int
+    __misses: int
     __cache_ids: Queue
     __invoker: Invoker
 
     def __init__(self, max_cache_size: int = 0) -> None:
         self.__cache = dict()
         self.__max_cache_size = max_cache_size
+        self.__disabled = False
+        self.__hits = 0
+        self.__misses = 0
         self.__cache_ids = Queue()
 
     def start(self, invoker: Invoker) -> None:
@@ -25,15 +32,17 @@ def start(self, invoker: Invoker) -> None:
         self.__invoker.services.latents.on_deleted(self._delete_by_match)
 
     def get(self, key: Union[int, str]) -> Optional[BaseInvocationOutput]:
-        if self.__max_cache_size == 0:
+        if self.__max_cache_size == 0 or self.__disabled:
             return
 
         item = self.__cache.get(key, None)
         if item is not None:
+            self.__hits += 1
             return item[0]
+        self.__misses += 1
 
     def save(self, key: Union[int, str], invocation_output: BaseInvocationOutput) -> None:
-        if self.__max_cache_size == 0:
+        if self.__max_cache_size == 0 or self.__disabled:
             return
 
         if key not in self.__cache:
@@ -47,24 +56,41 @@ def save(self, key: Union[int, str], invocation_output: BaseInvocationOutput) ->
                     pass
 
     def delete(self, key: Union[int, str]) -> None:
-        if self.__max_cache_size == 0:
+        if self.__max_cache_size == 0 or self.__disabled:
             return
 
         if key in self.__cache:
             del self.__cache[key]
 
     def clear(self, *args, **kwargs) -> None:
-        if self.__max_cache_size == 0:
+        if self.__max_cache_size == 0 or self.__disabled:
             return
 
         self.__cache.clear()
         self.__cache_ids = Queue()
+        self.__misses = 0
+        self.__hits = 0
 
     def create_key(self, invocation: BaseInvocation) -> int:
         return hash(invocation.json(exclude={"id"}))
 
+    def disable(self) -> None:
+        self.__disabled = True
+
+    def enable(self) -> None:
+        self.__disabled = False
+
+    def get_status(self) -> InvocationCacheStatus:
+        return InvocationCacheStatus(
+            hits=self.__hits,
+            misses=self.__misses,
+            enabled=not self.__disabled,
+            size=len(self.__cache),
+            max_size=self.__max_cache_size,
+        )
+
     def _delete_by_match(self, to_match: str) -> None:
-        if self.__max_cache_size == 0:
+        if self.__max_cache_size == 0 or self.__disabled:
             return
 
         keys_to_delete = set()

invokeai/app/services/session_queue/session_queue_common.py

@@ -162,15 +162,15 @@ class SessionQueueItemWithoutGraph(BaseModel):
     session_id: str = Field(
         description="The ID of the session associated with this queue item. The session doesn't exist in graph_executions until the queue item is executed."
     )
-    field_values: Optional[list[NodeFieldValue]] = Field(
-        default=None, description="The field values that were used for this queue item"
-    )
-    queue_id: str = Field(description="The id of the queue with which this item is associated")
     error: Optional[str] = Field(default=None, description="The error message if this queue item errored")
     created_at: Union[datetime.datetime, str] = Field(description="When this queue item was created")
     updated_at: Union[datetime.datetime, str] = Field(description="When this queue item was updated")
     started_at: Optional[Union[datetime.datetime, str]] = Field(description="When this queue item was started")
     completed_at: Optional[Union[datetime.datetime, str]] = Field(description="When this queue item was completed")
+    queue_id: str = Field(description="The id of the queue with which this item is associated")
+    field_values: Optional[list[NodeFieldValue]] = Field(
+        default=None, description="The field values that were used for this queue item"
+    )
 
     @classmethod
     def from_dict(cls, queue_item_dict: dict) -> "SessionQueueItemDTO":

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

@@ -264,6 +264,22 @@
         "graphQueued": "Graph queued",
         "graphFailedToQueue": "Failed to queue graph"
     },
+    "invocationCache": {
+        "invocationCache": "Invocation Cache",
+        "cacheSize": "Cache Size",
+        "maxCacheSize": "Max Cache Size",
+        "hits": "Cache Hits",
+        "misses": "Cache Misses",
+        "clear": "Clear",
+        "clearSucceeded": "Invocation Cache Cleared",
+        "clearFailed": "Problem Clearing Invocation Cache",
+        "enable": "Enable",
+        "enableSucceeded": "Invocation Cache Enabled",
+        "enableFailed": "Problem Enabling Invocation Cache",
+        "disable": "Disable",
+        "disableSucceeded": "Invocation Cache Disabled",
+        "disableFailed": "Problem Disabling Invocation Cache"
+    },
     "gallery": {
         "allImagesLoaded": "All Images Loaded",
         "assets": "Assets",

invokeai/frontend/web/src/app/types/invokeai.ts

@@ -21,7 +21,8 @@ export type AppFeature =
   | 'multiselect'
   | 'pauseQueue'
   | 'resumeQueue'
-  | 'prependQueue';
+  | 'prependQueue'
+  | 'invocationCache';
 
 /**
  * A disable-able Stable Diffusion feature

invokeai/frontend/web/src/features/queue/components/ClearInvocationCacheButton.tsx

@@ -0,0 +1,22 @@
+import IAIButton from 'common/components/IAIButton';
+import { memo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { useClearInvocationCache } from '../hooks/useClearInvocationCache';
+
+const ClearInvocationCacheButton = () => {
+  const { t } = useTranslation();
+  const { clearInvocationCache, isDisabled, isLoading } =
+    useClearInvocationCache();
+
+  return (
+    <IAIButton
+      isDisabled={isDisabled}
+      isLoading={isLoading}
+      onClick={clearInvocationCache}
+    >
+      {t('invocationCache.clear')}
+    </IAIButton>
+  );
+};
+
+export default memo(ClearInvocationCacheButton);

invokeai/frontend/web/src/features/queue/components/InvocationCacheStatus.tsx

@@ -0,0 +1,45 @@
+import { ButtonGroup } from '@chakra-ui/react';
+import { memo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { useGetInvocationCacheStatusQuery } from 'services/api/endpoints/appInfo';
+import ClearInvocationCacheButton from './ClearInvocationCacheButton';
+import ToggleInvocationCacheButton from './ToggleInvocationCacheButton';
+import StatusStatGroup from './common/StatusStatGroup';
+import StatusStatItem from './common/StatusStatItem';
+
+const InvocationCacheStatus = () => {
+  const { data: cacheStatus } = useGetInvocationCacheStatusQuery(undefined, {
+    pollingInterval: 5000,
+  });
+  const { t } = useTranslation();
+  return (
+    <StatusStatGroup>
+      <StatusStatItem
+        isDisabled={!cacheStatus?.enabled}
+        label={t('invocationCache.cacheSize')}
+        value={cacheStatus?.size ?? 0}
+      />
+      <StatusStatItem
+        isDisabled={!cacheStatus?.enabled}
+        label={t('invocationCache.hits')}
+        value={cacheStatus?.hits ?? 0}
+      />
+      <StatusStatItem
+        isDisabled={!cacheStatus?.enabled}
+        label={t('invocationCache.misses')}
+        value={cacheStatus?.misses ?? 0}
+      />
+      <StatusStatItem
+        isDisabled={!cacheStatus?.enabled}
+        label={t('invocationCache.maxCacheSize')}
+        value={cacheStatus?.max_size ?? 0}
+      />
+      <ButtonGroup w={24} orientation="vertical" size="xs">
+        <ClearInvocationCacheButton />
+        <ToggleInvocationCacheButton />
+      </ButtonGroup>
+    </StatusStatGroup>
+  );
+};
+
+export default memo(InvocationCacheStatus);

invokeai/frontend/web/src/features/queue/components/QueueStatus.tsx

@@ -1,38 +1,39 @@
-import { Stat, StatGroup, StatLabel, StatNumber } from '@chakra-ui/react';
 import { memo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useGetQueueStatusQuery } from 'services/api/endpoints/queue';
+import StatusStatGroup from './common/StatusStatGroup';
+import StatusStatItem from './common/StatusStatItem';
 
 const QueueStatus = () => {
   const { data: queueStatus } = useGetQueueStatusQuery();
   const { t } = useTranslation();
   return (
-    <StatGroup alignItems="center" justifyContent="center" w="full" h="full">
-      <Stat w={24}>
-        <StatLabel>{t('queue.in_progress')}</StatLabel>
-        <StatNumber>{queueStatus?.queue.in_progress ?? 0}</StatNumber>
-      </Stat>
-      <Stat w={24}>
-        <StatLabel>{t('queue.pending')}</StatLabel>
-        <StatNumber>{queueStatus?.queue.pending ?? 0}</StatNumber>
-      </Stat>
-      <Stat w={24}>
-        <StatLabel>{t('queue.completed')}</StatLabel>
-        <StatNumber>{queueStatus?.queue.completed ?? 0}</StatNumber>
-      </Stat>
-      <Stat w={24}>
-        <StatLabel>{t('queue.failed')}</StatLabel>
-        <StatNumber>{queueStatus?.queue.failed ?? 0}</StatNumber>
-      </Stat>
-      <Stat w={24}>
-        <StatLabel>{t('queue.canceled')}</StatLabel>
-        <StatNumber>{queueStatus?.queue.canceled ?? 0}</StatNumber>
-      </Stat>
-      <Stat w={24}>
-        <StatLabel>{t('queue.total')}</StatLabel>
-        <StatNumber>{queueStatus?.queue.total}</StatNumber>
-      </Stat>
-    </StatGroup>
+    <StatusStatGroup>
+      <StatusStatItem
+        label={t('queue.in_progress')}
+        value={queueStatus?.queue.in_progress ?? 0}
+      />
+      <StatusStatItem
+        label={t('queue.pending')}
+        value={queueStatus?.queue.pending ?? 0}
+      />
+      <StatusStatItem
+        label={t('queue.completed')}
+        value={queueStatus?.queue.completed ?? 0}
+      />
+      <StatusStatItem
+        label={t('queue.failed')}
+        value={queueStatus?.queue.failed ?? 0}
+      />
+      <StatusStatItem
+        label={t('queue.canceled')}
+        value={queueStatus?.queue.canceled ?? 0}
+      />
+      <StatusStatItem
+        label={t('queue.total')}
+        value={queueStatus?.queue.total ?? 0}
+      />
+    </StatusStatGroup>
   );
 };