(ui): add image resolution badge to initial upscale image

Author: Mary Hipp

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

@@ -0,0 +1,28 @@
+import { createMemoizedSelector } from 'app/store/createMemoizedSelector';
+import { useAppSelector } from 'app/store/storeHooks';
+import { selectUpscalelice } from 'features/parameters/store/upscaleSlice';
+import { selectConfigSlice } from 'features/system/store/configSlice';
+import { useMemo } from 'react';
+import type { ImageDTO } from 'services/api/types';
+
+
+export const createIsTooLargeToUpscaleSelector = (imageDTO?: ImageDTO) =>
+  createMemoizedSelector(selectUpscalelice, selectConfigSlice, (upscale, config) => {
+    const { upscaleModel, scale } = upscale;
+    const { maxUpscalePixels } = config;
+
+    if (!maxUpscalePixels || !upscaleModel || !imageDTO) {
+      return false
+    }
+
+    const upscaledPixels = imageDTO.width * scale * imageDTO.height * scale;
+    console.log({ upscaledPixels })
+    console.log({ maxUpscalePixels })
+    return upscaledPixels > maxUpscalePixels
+
+  });
+
+export const useIsTooLargeToUpscale = (imageDTO?: ImageDTO) => {
+  const selectIsTooLargeToUpscale = useMemo(() => createIsTooLargeToUpscaleSelector(imageDTO), [imageDTO]);
+  return useAppSelector(selectIsTooLargeToUpscale);
+};

invokeai/frontend/web/src/features/settingsAccordions/components/UpscaleSettingsAccordion/UpscaleInitialImage.tsx

@@ -1,4 +1,4 @@
-import { Flex } from '@invoke-ai/ui-library';
+import { Flex, Text } from '@invoke-ai/ui-library';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import IAIDndImage from 'common/components/IAIDndImage';
 import IAIDndImageIcon from 'common/components/IAIDndImageIcon';
@@ -41,13 +41,30 @@ export const UpscaleInitialImage = () => {
           postUploadAction={postUploadAction}
         />
         {imageDTO && (
-          <Flex position="absolute" flexDir="column" top={1} insetInlineEnd={1} gap={1}>
-            <IAIDndImageIcon
-              onClick={onReset}
-              icon={<PiArrowCounterClockwiseBold size={16} />}
-              tooltip={t('controlnet.resetControlImage')}
-            />
-          </Flex>
+          <>
+            <Flex position="absolute" flexDir="column" top={1} insetInlineEnd={1} gap={1}>
+              <IAIDndImageIcon
+                onClick={onReset}
+                icon={<PiArrowCounterClockwiseBold size={16} />}
+                tooltip={t('controlnet.resetControlImage')}
+              />
+            </Flex>
+            <Text
+              position="absolute"
+              background="base.900"
+              color="base.50"
+              fontSize="sm"
+              fontWeight="semibold"
+              bottom={0}
+              left={0}
+              opacity={0.7}
+              px={2}
+              lineHeight={1.25}
+              borderTopEndRadius="base"
+              borderBottomStartRadius="base"
+              pointerEvents="none"
+            >{`${imageDTO.width}x${imageDTO.height}`}</Text>
+          </>
         )}
       </Flex>
     </Flex>

invokeai/frontend/web/src/features/settingsAccordions/components/UpscaleSettingsAccordion/UpscaleWarning.tsx

@@ -6,16 +6,19 @@ import { setActiveTab } from 'features/ui/store/uiSlice';
 import { useCallback, useEffect, useMemo } from 'react';
 import { Trans, useTranslation } from 'react-i18next';
 import { useControlNetModels } from 'services/api/hooks/modelsByType';
+import { useIsTooLargeToUpscale } from '../../../parameters/hooks/useIsTooLargeToUpscale';
 
 export const UpscaleWarning = () => {
   const { t } = useTranslation();
   const model = useAppSelector((s) => s.generation.model);
   const upscaleModel = useAppSelector((s) => s.upscale.upscaleModel);
   const tileControlnetModel = useAppSelector((s) => s.upscale.tileControlnetModel);
+  const upscaleInitialImage = useAppSelector((s) => s.upscale.upscaleInitialImage);
   const dispatch = useAppDispatch();
   const [modelConfigs, { isLoading }] = useControlNetModels();
   const disabledTabs = useAppSelector((s) => s.config.disabledTabs);
   const shouldShowButton = useMemo(() => !disabledTabs.includes('models'), [disabledTabs]);
+  const isTooLargeToUpscale = useIsTooLargeToUpscale(upscaleInitialImage || undefined);
 
   useEffect(() => {
     const validModel = modelConfigs.find((cnetModel) => {
@@ -24,7 +27,7 @@ export const UpscaleWarning = () => {
     dispatch(tileControlnetModelChanged(validModel || null));
   }, [model?.base, modelConfigs, dispatch]);
 
-  const warnings = useMemo(() => {
+  const modelWarnings = useMemo(() => {
     const _warnings: string[] = [];
     if (!model) {
       _warnings.push(t('upscaling.mainModelDesc'));
@@ -35,33 +38,43 @@ export const UpscaleWarning = () => {
     if (!upscaleModel) {
       _warnings.push(t('upscaling.upscaleModelDesc'));
     }
-
     return _warnings;
   }, [model, tileControlnetModel, upscaleModel, t]);
 
+  const otherWarnings = useMemo(() => {
+    console.log({ isTooLargeToUpscale });
+    const _warnings: string[] = [];
+    if (isTooLargeToUpscale) {
+      _warnings.push(t('upscaling.outputTooLarge'));
+    }
+    return _warnings;
+  }, [isTooLargeToUpscale]);
+
   const handleGoToModelManager = useCallback(() => {
     dispatch(setActiveTab('models'));
     $installModelsTab.set(3);
   }, [dispatch]);
 
-  if (!warnings.length || isLoading || !shouldShowButton) {
+  if ((!modelWarnings.length && !otherWarnings.length) || isLoading || !shouldShowButton) {
     return null;
   }
 
   return (
     <Flex bg="error.500" borderRadius="base" padding={4} direction="column" fontSize="sm" gap={2}>
-      <Text>
-        <Trans
-          i18nKey="upscaling.missingModelsWarning"
-          components={{
-            LinkComponent: (
-              <Button size="sm" flexGrow={0} variant="link" color="base.50" onClick={handleGoToModelManager} />
-            ),
-          }}
-        />
-      </Text>
+      {!!modelWarnings.length && (
+        <Text>
+          <Trans
+            i18nKey="upscaling.missingModelsWarning"
+            components={{
+              LinkComponent: (
+                <Button size="sm" flexGrow={0} variant="link" color="base.50" onClick={handleGoToModelManager} />
+              ),
+            }}
+          />
+        </Text>
+      )}
       <UnorderedList>
-        {warnings.map((warning) => (
+        {[...modelWarnings, ...otherWarnings].map((warning) => (
           <ListItem key={warning}>{warning}</ListItem>
         ))}
       </UnorderedList>