Skip to content

Commit 61ff9ee

Browse files
feat(ui): add button to ref image to recall size & optimize for model
This is useful for FLUX Kontext, where you typically want the generation size to at least roughly match the first ref image size.
1 parent 111408c commit 61ff9ee

File tree

4 files changed

+52
-1
lines changed

4 files changed

+52
-1
lines changed

invokeai/frontend/web/src/features/controlLayers/components/RefImage/RefImageImage.tsx

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,20 @@
11
import { Flex } from '@invoke-ai/ui-library';
22
import { useStore } from '@nanostores/react';
33
import { skipToken } from '@reduxjs/toolkit/query';
4+
import { useAppSelector, useAppStore } from 'app/store/storeHooks';
45
import { UploadImageIconButton } from 'common/hooks/useImageUploadButton';
6+
import { bboxSizeOptimized, bboxSizeRecalled } from 'features/controlLayers/store/canvasSlice';
7+
import { useCanvasIsStaging } from 'features/controlLayers/store/canvasStagingAreaSlice';
8+
import { sizeOptimized, sizeRecalled } from 'features/controlLayers/store/paramsSlice';
59
import type { ImageWithDims } from 'features/controlLayers/store/types';
610
import type { setGlobalReferenceImageDndTarget, setRegionalGuidanceReferenceImageDndTarget } from 'features/dnd/dnd';
711
import { DndDropTarget } from 'features/dnd/DndDropTarget';
812
import { DndImage } from 'features/dnd/DndImage';
913
import { DndImageIcon } from 'features/dnd/DndImageIcon';
14+
import { selectActiveTab } from 'features/ui/store/uiSelectors';
1015
import { memo, useCallback, useEffect } from 'react';
1116
import { useTranslation } from 'react-i18next';
12-
import { PiArrowCounterClockwiseBold } from 'react-icons/pi';
17+
import { PiArrowCounterClockwiseBold, PiRulerBold } from 'react-icons/pi';
1318
import { useGetImageDTOQuery } from 'services/api/endpoints/images';
1419
import type { ImageDTO } from 'services/api/types';
1520
import { $isConnected } from 'services/events/stores';
@@ -29,7 +34,10 @@ export const RefImageImage = memo(
2934
dndTargetData,
3035
}: Props<T>) => {
3136
const { t } = useTranslation();
37+
const store = useAppStore();
3238
const isConnected = useStore($isConnected);
39+
const tab = useAppSelector(selectActiveTab);
40+
const isStaging = useCanvasIsStaging();
3341
const { currentData: imageDTO, isError } = useGetImageDTOQuery(image?.image_name ?? skipToken);
3442
const handleResetControlImage = useCallback(() => {
3543
onChangeImage(null);
@@ -48,6 +56,20 @@ export const RefImageImage = memo(
4856
[onChangeImage]
4957
);
5058

59+
const recallSizeAndOptimize = useCallback(() => {
60+
if (!imageDTO || (tab === 'canvas' && isStaging)) {
61+
return;
62+
}
63+
const { width, height } = imageDTO;
64+
if (tab === 'canvas') {
65+
store.dispatch(bboxSizeRecalled({ width, height }));
66+
store.dispatch(bboxSizeOptimized());
67+
} else if (tab === 'generate') {
68+
store.dispatch(sizeRecalled({ width, height }));
69+
store.dispatch(sizeOptimized());
70+
}
71+
}, [imageDTO, isStaging, store, tab]);
72+
5173
return (
5274
<Flex position="relative" w="full" h="full" alignItems="center" data-error={!imageDTO && !image?.image_name}>
5375
{!imageDTO && (
@@ -69,6 +91,14 @@ export const RefImageImage = memo(
6991
tooltip={t('common.reset')}
7092
/>
7193
</Flex>
94+
<Flex position="absolute" flexDir="column" bottom={2} insetInlineEnd={2} gap={1}>
95+
<DndImageIcon
96+
onClick={recallSizeAndOptimize}
97+
icon={<PiRulerBold size={16} />}
98+
tooltip={t('parameters.useSize')}
99+
isDisabled={!imageDTO || (tab === 'canvas' && isStaging)}
100+
/>
101+
</Flex>
72102
</>
73103
)}
74104
<DndDropTarget dndTarget={dndTarget} dndTargetData={dndTargetData} label={t('gallery.drop')} />

invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1091,6 +1091,15 @@ const slice = createSlice({
10911091

10921092
syncScaledSize(state);
10931093
},
1094+
bboxSizeRecalled: (state, action: PayloadAction<{ width: number; height: number }>) => {
1095+
const { width, height } = action.payload;
1096+
const gridSize = getGridSize(state.bbox.modelBase);
1097+
state.bbox.rect.width = Math.max(roundDownToMultiple(width, gridSize), 64);
1098+
state.bbox.rect.height = Math.max(roundDownToMultiple(height, gridSize), 64);
1099+
state.bbox.aspectRatio.value = state.bbox.rect.width / state.bbox.rect.height;
1100+
state.bbox.aspectRatio.id = 'Free';
1101+
state.bbox.aspectRatio.isLocked = true;
1102+
},
10941103
bboxAspectRatioLockToggled: (state) => {
10951104
state.bbox.aspectRatio.isLocked = !state.bbox.aspectRatio.isLocked;
10961105
syncScaledSize(state);
@@ -1627,6 +1636,7 @@ export const {
16271636
bboxScaledWidthChanged,
16281637
bboxScaledHeightChanged,
16291638
bboxScaleMethodChanged,
1639+
bboxSizeRecalled,
16301640
bboxWidthChanged,
16311641
bboxHeightChanged,
16321642
bboxAspectRatioLockToggled,

invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,15 @@ const slice = createSlice({
241241
},
242242

243243
//#region Dimensions
244+
sizeRecalled: (state, action: PayloadAction<{ width: number; height: number }>) => {
245+
const { width, height } = action.payload;
246+
const gridSize = getGridSize(state.model?.base);
247+
state.dimensions.rect.width = Math.max(roundDownToMultiple(width, gridSize), 64);
248+
state.dimensions.rect.height = Math.max(roundDownToMultiple(height, gridSize), 64);
249+
state.dimensions.aspectRatio.value = state.dimensions.rect.width / state.dimensions.rect.height;
250+
state.dimensions.aspectRatio.id = 'Free';
251+
state.dimensions.aspectRatio.isLocked = true;
252+
},
244253
widthChanged: (state, action: PayloadAction<{ width: number; updateAspectRatio?: boolean; clamp?: boolean }>) => {
245254
const { width, updateAspectRatio, clamp } = action.payload;
246255
const gridSize = getGridSize(state.model?.base);
@@ -429,6 +438,7 @@ export const {
429438
modelChanged,
430439

431440
// Dimensions
441+
sizeRecalled,
432442
widthChanged,
433443
heightChanged,
434444
aspectRatioLockToggled,

invokeai/frontend/web/src/features/dnd/DndImageIcon.tsx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ export const DndImageIcon = memo((props: Props) => {
2727
return (
2828
<IconButton
2929
onClick={onClick}
30+
tooltip={tooltip}
3031
aria-label={tooltip}
3132
icon={icon}
3233
variant="link"

0 commit comments

Comments
 (0)