feat(ui): add scale param to upscaling tab

Author: psychedelicious

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

@@ -1649,11 +1649,13 @@
         "structure": "Structure",
         "toInstall": "to install",
         "upscaleModel": "Upcale Model",
+        "scale": "Scale",
         "visit": "Visit",
         "warningNoMainModel": "a model",
         "warningNoTile": "a {{base_model}} tile controlnet required by this feature",
         "warningNoTileOrUpscaleModel": "an upscaler model and {{base_model}} tile controlnet required by this feature",
         "warningNoUpscaleModel": "an upscaler model required by this feature",
+        "upscalingFromTo": "Upscaling from {{from}} to {{to}}"
     },
     "ui": {
         "tabs": {

invokeai/frontend/web/src/features/nodes/util/graph/buildMultidiffusionUpscaleGraph.ts

@@ -2,7 +2,6 @@ import type { RootState } from 'app/store/store';
 import { fetchModelConfigWithTypeGuard } from 'features/metadata/util/modelFetchingHelpers';
 import type { GraphType } from 'features/nodes/util/graph/generation/Graph';
 import { Graph } from 'features/nodes/util/graph/generation/Graph';
-import type { ImageDTO } from 'services/api/types';
 import { isNonRefinerMainModelConfig, isSpandrelImageToImageModelConfig } from 'services/api/types';
 import { assert } from 'tsafe';
 
@@ -15,7 +14,6 @@ import {
   NEGATIVE_CONDITIONING,
   NOISE,
   POSITIVE_CONDITIONING,
-  RESIZE,
   SDXL_MODEL_LOADER,
   SPANDREL,
   TILED_MULTI_DIFFUSION_DENOISE_LATENTS,
@@ -26,27 +24,17 @@ import { addLoRAs } from './generation/addLoRAs';
 import { addSDXLLoRas } from './generation/addSDXLLoRAs';
 import { getBoardField, getSDXLStylePrompts } from './graphBuilderUtils';
 
-const UPSCALE_SCALE = 2;
-
-export const getOutputImageSize = (initialImage: ImageDTO) => {
-  return {
-    width: ((initialImage.width * UPSCALE_SCALE) / 8) * 8,
-    height: ((initialImage.height * UPSCALE_SCALE) / 8) * 8,
-  };
-};
-
 export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promise<GraphType> => {
   const { model, cfgScale: cfg_scale, scheduler, steps, vaePrecision, seed, vae } = state.generation;
   const { positivePrompt, negativePrompt } = state.controlLayers.present;
-  const { upscaleModel, upscaleInitialImage, sharpness, structure, creativity, tileControlnetModel } = state.upscale;
+  const { upscaleModel, upscaleInitialImage, sharpness, structure, creativity, tileControlnetModel, scale } =
+    state.upscale;
 
   assert(model, 'No model found in state');
   assert(upscaleModel, 'No upscale model found in state');
   assert(upscaleInitialImage, 'No initial image found in state');
   assert(tileControlnetModel, 'Tile controlnet is required');
 
-  const { width: outputWidth, height: outputHeight } = getOutputImageSize(upscaleInitialImage);
-
   const g = new Graph();
 
   const unsharpMaskNode1 = g.addNode({
@@ -61,7 +49,8 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
     id: SPANDREL,
     type: 'spandrel_image_to_image',
     image_to_image_model: upscaleModel,
-    tile_size: 500,
+    fit_to_multiple_of_8: true,
+    scale,
   });
 
   g.addEdge(unsharpMaskNode1, 'image', upscaleNode, 'image');
@@ -75,24 +64,14 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
 
   g.addEdge(upscaleNode, 'image', unsharpMaskNode2, 'image');
 
-  const resizeNode = g.addNode({
-    id: RESIZE,
-    type: 'img_resize',
-    width: outputWidth,
-    height: outputHeight,
-    resample_mode: 'lanczos',
-  });
-
-  g.addEdge(unsharpMaskNode2, 'image', resizeNode, 'image');
-
   const noiseNode = g.addNode({
     id: NOISE,
     type: 'noise',
     seed,
   });
 
-  g.addEdge(resizeNode, 'width', noiseNode, 'width');
-  g.addEdge(resizeNode, 'height', noiseNode, 'height');
+  g.addEdge(unsharpMaskNode2, 'width', noiseNode, 'width');
+  g.addEdge(unsharpMaskNode2, 'height', noiseNode, 'height');
 
   const i2lNode = g.addNode({
     id: IMAGE_TO_LATENTS,
@@ -101,7 +80,7 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
     tiled: true,
   });
 
-  g.addEdge(resizeNode, 'image', i2lNode, 'image');
+  g.addEdge(unsharpMaskNode2, 'image', i2lNode, 'image');
 
   const l2iNode = g.addNode({
     type: 'l2i',
@@ -160,8 +139,6 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
 
     g.upsertMetadata({
       cfg_scale,
-      height: outputHeight,
-      width: outputWidth,
       positive_prompt: positivePrompt,
       negative_prompt: negativePrompt,
       positive_style_prompt: positiveStylePrompt,
@@ -204,8 +181,6 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
 
     g.upsertMetadata({
       cfg_scale,
-      height: outputHeight,
-      width: outputWidth,
       positive_prompt: positivePrompt,
       negative_prompt: negativePrompt,
       model: Graph.getModelMetadataField(modelConfig),
@@ -221,6 +196,8 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
   }
 
   g.setMetadataReceivingNode(l2iNode);
+  g.addEdgeToMetadata(upscaleNode, 'width', 'width');
+  g.addEdgeToMetadata(upscaleNode, 'height', 'height');
 
   let vaeNode;
   if (vae) {
@@ -252,7 +229,7 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
     end_step_percent: (structure + 10) * 0.025 + 0.3,
   });
 
-  g.addEdge(resizeNode, 'image', controlnetNode1, 'image');
+  g.addEdge(unsharpMaskNode2, 'image', controlnetNode1, 'image');
 
   const controlnetNode2 = g.addNode({
     id: 'controlnet_2',
@@ -265,7 +242,7 @@ export const buildMultidiffusionUpscsaleGraph = async (state: RootState): Promis
     end_step_percent: 0.85,
   });
 
-  g.addEdge(resizeNode, 'image', controlnetNode2, 'image');
+  g.addEdge(unsharpMaskNode2, 'image', controlnetNode2, 'image');
 
   const collectNode = g.addNode({
     id: CONTROL_NET_COLLECT,

invokeai/frontend/web/src/features/parameters/store/upscaleSlice.ts

@@ -12,6 +12,7 @@ interface UpscaleState {
   structure: number;
   creativity: number;
   tileControlnetModel: ControlNetModelConfig | null;
+  scale: number;
 }
 
 const initialUpscaleState: UpscaleState = {
@@ -22,6 +23,7 @@ const initialUpscaleState: UpscaleState = {
   structure: 0,
   creativity: 0,
   tileControlnetModel: null,
+  scale: 4,
 };
 
 export const upscaleSlice = createSlice({
@@ -46,6 +48,9 @@ export const upscaleSlice = createSlice({
     tileControlnetModelChanged: (state, action: PayloadAction<ControlNetModelConfig | null>) => {
       state.tileControlnetModel = action.payload;
     },
+    scaleChanged: (state, action: PayloadAction<number>) => {
+      state.scale = action.payload;
+    },
   },
 });
 
@@ -56,6 +61,7 @@ export const {
   creativityChanged,
   sharpnessChanged,
   tileControlnetModelChanged,
+  scaleChanged,
 } = upscaleSlice.actions;
 
 export const selectUpscalelice = (state: RootState) => state.upscale;

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

@@ -0,0 +1,39 @@
+import { CompositeSlider, FormControl, FormLabel } from '@invoke-ai/ui-library';
+import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { scaleChanged } from 'features/parameters/store/upscaleSlice';
+import { memo, useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+
+const marks = [2, 4, 8, 16];
+
+const formatValue = (val: number) => `${val}x`;
+
+export const UpscaleScaleSlider = memo(() => {
+  const dispatch = useAppDispatch();
+  const { t } = useTranslation();
+  const scale = useAppSelector((s) => s.upscale.scale);
+
+  const onChange = useCallback(
+    (val: number) => {
+      dispatch(scaleChanged(val));
+    },
+    [dispatch]
+  );
+
+  return (
+    <FormControl orientation="vertical" gap={0}>
+      <FormLabel m={0}>{t('upscaling.scale')}</FormLabel>
+      <CompositeSlider
+        min={2}
+        max={16}
+        value={scale}
+        onChange={onChange}
+        marks={marks}
+        formatValue={formatValue}
+        withThumbTooltip
+      />
+    </FormControl>
+  );
+});
+
+UpscaleScaleSlider.displayName = 'UpscaleScaleSlider';

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

@@ -6,20 +6,30 @@ import ParamSharpness from 'features/parameters/components/Upscale/ParamSharpnes
 import ParamSpandrelModel from 'features/parameters/components/Upscale/ParamSpandrelModel';
 import ParamStructure from 'features/parameters/components/Upscale/ParamStructure';
 import { selectUpscalelice } from 'features/parameters/store/upscaleSlice';
+import { UpscaleScaleSlider } from 'features/settingsAccordions/components/UpscaleSettingsAccordion/UpscaleScaleSlider';
 import { useExpanderToggle } from 'features/settingsAccordions/hooks/useExpanderToggle';
 import { useStandaloneAccordionToggle } from 'features/settingsAccordions/hooks/useStandaloneAccordionToggle';
 import { memo } from 'react';
 import { useTranslation } from 'react-i18next';
 
 import { MultidiffusionWarning } from './MultidiffusionWarning';
 import { UpscaleInitialImage } from './UpscaleInitialImage';
-import { UpscaleSizeDetails } from './UpscaleSizeDetails';
 
-const selector = createMemoizedSelector([selectUpscalelice], (upscale) => {
+const selector = createMemoizedSelector([selectUpscalelice], (upscaleSlice) => {
+  const { upscaleModel, upscaleInitialImage, scale } = upscaleSlice;
+
   const badges: string[] = [];
 
-  if (upscale.upscaleModel) {
-    badges.push(upscale.upscaleModel.name);
+  if (upscaleModel) {
+    badges.push(upscaleModel.name);
+  }
+
+  if (upscaleInitialImage) {
+    // Output height and width are scaled and rounded down to the nearest multiple of 8
+    const outputWidth = Math.floor((upscaleInitialImage.width * scale) / 8) * 8;
+    const outputHeight = Math.floor((upscaleInitialImage.height * scale) / 8) * 8;
+
+    badges.push(`${outputWidth}×${outputHeight}`);
   }
 
   return { badges };
@@ -43,9 +53,9 @@ export const UpscaleSettingsAccordion = memo(() => {
       <Flex pt={4} px={4} w="full" h="full" flexDir="column" data-testid="image-settings-accordion">
         <Flex gap={4}>
           <UpscaleInitialImage />
-          <Flex direction="column" w="full" alignItems="center" gap={4}>
+          <Flex direction="column" w="full" alignItems="center" gap={2}>
             <ParamSpandrelModel />
-            <UpscaleSizeDetails />
+            <UpscaleScaleSlider />
             <MultidiffusionWarning />
           </Flex>
         </Flex>