feat(ui): max upscale pixels config (#4765)

* feat(ui): max upscale pixels config

Add `maxUpscalePixels: number` to the app config. The number should be the *total* number of pixels eg `maxUpscalePixels: 4096 * 4096`.

If not provided, any size image may be upscaled.

If the config is provided, users will see be advised if their image is too large for either model, or told to switch to an x2 model if it's only too large for x4.

The message is via tooltip in the popover and via toast if the user uses the hotkey to upscale.

* feat(ui): "mayUpscale" -> "isAllowedToUpscale"

Author: psychedelicious

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

@@ -1015,7 +1015,11 @@
         "variationAmount": "Variation Amount",
         "variations": "Variations",
         "vSymmetryStep": "V Symmetry Step",
-        "width": "Width"
+        "width": "Width",
+        "isAllowedToUpscale": {
+            "useX2Model": "Image is too large to upscale with x4 model, use x2 model",
+            "tooLarge": "Image is too large to upscale, select smaller image"
+        }
     },
     "dynamicPrompts": {
         "combinatorial": "Combinatorial Generation",

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

@@ -6,8 +6,10 @@ import { addToast } from 'features/system/store/systemSlice';
 import { t } from 'i18next';
 import { queueApi } from 'services/api/endpoints/queue';
 import { startAppListening } from '..';
+import { ImageDTO } from 'services/api/types';
+import { createIsAllowedToUpscaleSelector } from 'features/parameters/hooks/useIsAllowedToUpscale';
 
-export const upscaleRequested = createAction<{ image_name: string }>(
+export const upscaleRequested = createAction<{ imageDTO: ImageDTO }>(
   `upscale/upscaleRequested`
 );
 
@@ -17,8 +19,28 @@ export const addUpscaleRequestedListener = () => {
     effect: async (action, { dispatch, getState }) => {
       const log = logger('session');
 
-      const { image_name } = action.payload;
+      const { imageDTO } = action.payload;
+      const { image_name } = imageDTO;
       const state = getState();
+
+      const { isAllowedToUpscale, detailTKey } =
+        createIsAllowedToUpscaleSelector(imageDTO)(state);
+
+      // if we can't upscale, show a toast and return
+      if (!isAllowedToUpscale) {
+        log.error(
+          { imageDTO },
+          t(detailTKey ?? 'parameters.isAllowedToUpscale.tooLarge') // should never coalesce
+        );
+        dispatch(
+          addToast({
+            title: t(detailTKey ?? 'parameters.isAllowedToUpscale.tooLarge'), // should never coalesce
+            status: 'error',
+          })
+        );
+        return;
+      }
+
       const { esrganModelName } = state.postprocessing;
       const { autoAddBoardId } = state.gallery;
 

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

@@ -56,6 +56,7 @@ export type AppConfig = {
   canRestoreDeletedImagesFromBin: boolean;
   nodesAllowlist: string[] | undefined;
   nodesDenylist: string[] | undefined;
+  maxUpscalePixels?: number;
   sd: {
     defaultModel?: string;
     disabledControlNetModels: string[];

invokeai/frontend/web/src/features/gallery/components/CurrentImage/CurrentImageButtons.tsx

@@ -181,7 +181,7 @@ const CurrentImageButtons = (props: CurrentImageButtonsProps) => {
     if (!imageDTO) {
       return;
     }
-    dispatch(upscaleRequested({ image_name: imageDTO.image_name }));
+    dispatch(upscaleRequested({ imageDTO }));
   }, [dispatch, imageDTO]);
 
   const handleDelete = useCallback(() => {

invokeai/frontend/web/src/features/nodes/components/flow/Flow.tsx

@@ -30,8 +30,8 @@ import {
   connectionEnded,
   connectionMade,
   connectionStarted,
-  edgeChangeStarted,
   edgeAdded,
+  edgeChangeStarted,
   edgeDeleted,
   edgesChanged,
   edgesDeleted,

invokeai/frontend/web/src/features/parameters/components/Parameters/Upscale/ParamUpscaleSettings.tsx

@@ -4,6 +4,7 @@ import { useAppDispatch } from 'app/store/storeHooks';
 import IAIButton from 'common/components/IAIButton';
 import IAIIconButton from 'common/components/IAIIconButton';
 import IAIPopover from 'common/components/IAIPopover';
+import { useIsAllowedToUpscale } from 'features/parameters/hooks/useIsAllowedToUpscale';
 import { useIsQueueMutationInProgress } from 'features/queue/hooks/useIsQueueMutationInProgress';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
@@ -19,14 +20,15 @@ const ParamUpscalePopover = (props: Props) => {
   const inProgress = useIsQueueMutationInProgress();
   const { t } = useTranslation();
   const { isOpen, onOpen, onClose } = useDisclosure();
+  const { isAllowedToUpscale, detail } = useIsAllowedToUpscale(imageDTO);
 
   const handleClickUpscale = useCallback(() => {
     onClose();
-    if (!imageDTO) {
+    if (!imageDTO || !isAllowedToUpscale) {
       return;
     }
-    dispatch(upscaleRequested({ image_name: imageDTO.image_name }));
-  }, [dispatch, imageDTO, onClose]);
+    dispatch(upscaleRequested({ imageDTO }));
+  }, [dispatch, imageDTO, isAllowedToUpscale, onClose]);
 
   return (
     <IAIPopover
@@ -49,8 +51,9 @@ const ParamUpscalePopover = (props: Props) => {
       >
         <ParamESRGANModel />
         <IAIButton
+          tooltip={detail}
           size="sm"
-          isDisabled={!imageDTO || inProgress}
+          isDisabled={!imageDTO || inProgress || !isAllowedToUpscale}
           onClick={handleClickUpscale}
         >
           {t('parameters.upscaleImage')}

invokeai/frontend/web/src/features/parameters/hooks/useIsAllowedToUpscale.ts

@@ -0,0 +1,100 @@
+import { createSelector } from '@reduxjs/toolkit';
+import { stateSelector } from 'app/store/store';
+import { useAppSelector } from 'app/store/storeHooks';
+import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
+import { useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+import { ImageDTO } from 'services/api/types';
+
+const getUpscaledPixels = (imageDTO?: ImageDTO, maxUpscalePixels?: number) => {
+  if (!imageDTO) {
+    return;
+  }
+  if (!maxUpscalePixels) {
+    return;
+  }
+  const { width, height } = imageDTO;
+  const x4 = height * 4 * width * 4;
+  const x2 = height * 2 * width * 2;
+  return { x4, x2 };
+};
+
+const getIsAllowedToUpscale = (
+  upscaledPixels?: ReturnType<typeof getUpscaledPixels>,
+  maxUpscalePixels?: number
+) => {
+  if (!upscaledPixels || !maxUpscalePixels) {
+    return { x4: true, x2: true };
+  }
+  const isAllowedToUpscale = { x4: false, x2: false };
+  if (upscaledPixels.x4 <= maxUpscalePixels) {
+    isAllowedToUpscale.x4 = true;
+  }
+  if (upscaledPixels.x2 <= maxUpscalePixels) {
+    isAllowedToUpscale.x2 = true;
+  }
+
+  return isAllowedToUpscale;
+};
+
+const getDetailTKey = (
+  isAllowedToUpscale?: ReturnType<typeof getIsAllowedToUpscale>,
+  scaleFactor?: number
+) => {
+  if (!isAllowedToUpscale || !scaleFactor) {
+    return;
+  }
+
+  if (isAllowedToUpscale.x4 && isAllowedToUpscale.x2) {
+    return;
+  }
+
+  if (!isAllowedToUpscale.x2 && !isAllowedToUpscale.x4) {
+    return 'parameters.isAllowedToUpscale.tooLarge';
+  }
+
+  if (!isAllowedToUpscale.x4 && isAllowedToUpscale.x2 && scaleFactor === 4) {
+    return 'parameters.isAllowedToUpscale.useX2Model';
+  }
+
+  return;
+};
+
+export const createIsAllowedToUpscaleSelector = (imageDTO?: ImageDTO) =>
+  createSelector(
+    stateSelector,
+    ({ postprocessing, config }) => {
+      const { esrganModelName } = postprocessing;
+      const { maxUpscalePixels } = config;
+
+      const upscaledPixels = getUpscaledPixels(imageDTO, maxUpscalePixels);
+      const isAllowedToUpscale = getIsAllowedToUpscale(
+        upscaledPixels,
+        maxUpscalePixels
+      );
+      const scaleFactor = esrganModelName.includes('x2') ? 2 : 4;
+      const detailTKey = getDetailTKey(isAllowedToUpscale, scaleFactor);
+      return {
+        isAllowedToUpscale:
+          scaleFactor === 2 ? isAllowedToUpscale.x2 : isAllowedToUpscale.x4,
+        detailTKey,
+      };
+    },
+    defaultSelectorOptions
+  );
+
+export const useIsAllowedToUpscale = (imageDTO?: ImageDTO) => {
+  const { t } = useTranslation();
+  const selectIsAllowedToUpscale = useMemo(
+    () => createIsAllowedToUpscaleSelector(imageDTO),
+    [imageDTO]
+  );
+  const { isAllowedToUpscale, detailTKey } = useAppSelector(
+    selectIsAllowedToUpscale
+  );
+
+  return {
+    isAllowedToUpscale,
+    detail: detailTKey ? t(detailTKey) : undefined,
+  };
+};