feat(ui): control image auto-process

This commit is contained in:
psychedelicious 2023-06-02 21:30:21 +10:00
parent fa290aff8d
commit 72b4371804
11 changed files with 395 additions and 45 deletions

View File

@ -71,7 +71,7 @@ import { addStagingAreaImageSavedListener } from './listeners/stagingAreaImageSa
import { addCommitStagingAreaImageListener } from './listeners/addCommitStagingAreaImageListener';
import { addImageCategoriesChangedListener } from './listeners/imageCategoriesChanged';
import { addControlNetImageProcessedListener } from './listeners/controlNetImageProcessed';
import { addControlNetProcessorParamsChangedListener } from './listeners/controlNetProcessorParamsChanged';
import { addControlNetAutoProcessListener } from './listeners/controlNetProcessorParamsChanged';
export const listenerMiddleware = createListenerMiddleware();
@ -178,4 +178,4 @@ addImageCategoriesChangedListener();
// ControlNet
addControlNetImageProcessedListener();
addControlNetProcessorParamsChangedListener();
addControlNetAutoProcessListener();

View File

@ -15,7 +15,10 @@ const moduleLog = log.child({ namespace: 'controlNet' });
export const addControlNetImageProcessedListener = () => {
startAppListening({
actionCreator: controlNetImageProcessed,
effect: async (action, { dispatch, getState, take }) => {
effect: async (
action,
{ dispatch, getState, take, unsubscribe, subscribe }
) => {
const { controlNetId } = action.payload;
const controlNet = getState().controlNet.controlNets[controlNetId];

View File

@ -2,6 +2,7 @@ import { startAppListening } from '..';
import { log } from 'app/logging/useLogger';
import { controlNetImageProcessed } from 'features/controlNet/store/actions';
import {
controlNetImageChanged,
controlNetProcessorParamsChanged,
controlNetProcessorTypeChanged,
} from 'features/controlNet/store/controlNetSlice';
@ -13,10 +14,11 @@ const moduleLog = log.child({ namespace: 'controlNet' });
*
* The network request is debounced by 1 second.
*/
export const addControlNetProcessorParamsChangedListener = () => {
export const addControlNetAutoProcessListener = () => {
startAppListening({
predicate: (action) =>
controlNetProcessorParamsChanged.match(action) ||
controlNetImageChanged.match(action) ||
controlNetProcessorTypeChanged.match(action),
effect: async (
action,
@ -35,11 +37,19 @@ export const addControlNetProcessorParamsChangedListener = () => {
const { controlNetId } = action.payload;
if (!state.controlNet.controlNets[controlNetId].controlImage) {
moduleLog.trace(
{ data: { controlNetId } },
'No ControlNet image to auto-process'
);
return;
}
// Cancel any in-progress instances of this listener
cancelActiveListeners();
// Delay before starting actual work
await delay(1000);
await delay(300);
dispatch(controlNetImageProcessed({ controlNetId }));
},

View File

@ -1,18 +1,21 @@
import { memo, useCallback } from 'react';
import { memo, useCallback, useState } from 'react';
import { ImageDTO } from 'services/api';
import {
ControlNet,
controlNetImageChanged,
controlNetProcessedImageChanged,
controlNetRemoved,
controlNetSelector,
} from '../store/controlNetSlice';
import { useAppDispatch } from 'app/store/storeHooks';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import ParamControlNetModel from './parameters/ParamControlNetModel';
import ParamControlNetWeight from './parameters/ParamControlNetWeight';
import ParamControlNetBeginStepPct from './parameters/ParamControlNetBeginStepPct';
import ParamControlNetEndStepPct from './parameters/ParamControlNetEndStepPct';
import {
Box,
Flex,
Spinner,
Tab,
TabList,
TabPanel,
@ -22,10 +25,15 @@ import {
import IAISelectableImage from './parameters/IAISelectableImage';
import IAIButton from 'common/components/IAIButton';
import { FaUndo } from 'react-icons/fa';
import { TbSquareToggle } from 'react-icons/tb';
import ParamControlNetProcessorSelect from './parameters/ParamControlNetProcessorSelect';
import ControlNetProcessorComponent from './ControlNetProcessorComponent';
import { useIsApplicationReady } from 'features/system/hooks/useIsApplicationReady';
import ControlNetPreprocessButton from './ControlNetPreprocessButton';
import IAIIconButton from 'common/components/IAIIconButton';
import ParamControlNetBeginEnd from './parameters/ParamControlNetBeginEnd';
import ControlNetImagePreview from './ControlNetImagePreview';
import { createSelector } from '@reduxjs/toolkit';
import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
type ControlNetProps = {
controlNet: ControlNet;
@ -45,15 +53,6 @@ const ControlNet = (props: ControlNetProps) => {
processorNode,
} = props.controlNet;
const dispatch = useAppDispatch();
const isReady = useIsApplicationReady();
const handleControlImageChanged = useCallback(
(controlImage: ImageDTO) => {
dispatch(controlNetImageChanged({ controlNetId, controlImage }));
},
[controlNetId, dispatch]
);
const handleReset = useCallback(() => {
dispatch(
controlNetProcessedImageChanged({
@ -63,21 +62,16 @@ const ControlNet = (props: ControlNetProps) => {
);
}, [controlNetId, dispatch]);
const handleControlImageReset = useCallback(() => {
dispatch(controlNetImageChanged({ controlNetId, controlImage: null }));
}, [controlNetId, dispatch]);
const handleControlNetRemoved = useCallback(() => {
dispatch(controlNetRemoved(controlNetId));
}, [controlNetId, dispatch]);
return (
<Flex sx={{ flexDir: 'column', gap: 3 }}>
<IAISelectableImage
image={processedControlImage || controlImage}
onChange={handleControlImageChanged}
onReset={handleControlImageReset}
resetIconSize="sm"
<ControlNetImagePreview
controlNetId={controlNetId}
controlImage={controlImage}
processedControlImage={processedControlImage}
/>
<ParamControlNetModel controlNetId={controlNetId} model={model} />
<Tabs
@ -105,12 +99,9 @@ const ControlNet = (props: ControlNetProps) => {
controlNetId={controlNetId}
weight={weight}
/>
<ParamControlNetBeginStepPct
<ParamControlNetBeginEnd
controlNetId={controlNetId}
beginStepPct={beginStepPct}
/>
<ParamControlNetEndStepPct
controlNetId={controlNetId}
endStepPct={endStepPct}
/>
</TabPanel>

View File

@ -0,0 +1,168 @@
import { memo, useCallback, useState } from 'react';
import { ImageDTO } from 'services/api';
import {
controlNetImageChanged,
controlNetSelector,
} from '../store/controlNetSlice';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { Box, Flex, Spinner } from '@chakra-ui/react';
import IAISelectableImage from './parameters/IAISelectableImage';
import { TbSquareToggle } from 'react-icons/tb';
import IAIIconButton from 'common/components/IAIIconButton';
import { createSelector } from '@reduxjs/toolkit';
import { defaultSelectorOptions } from 'app/store/util/defaultMemoizeOptions';
import { AnimatePresence, motion } from 'framer-motion';
const selector = createSelector(
controlNetSelector,
(controlNet) => {
const { isProcessingControlImage } = controlNet;
return { isProcessingControlImage };
},
defaultSelectorOptions
);
type Props = {
controlNetId: string;
controlImage: ImageDTO | null;
processedControlImage: ImageDTO | null;
};
const ControlNetImagePreview = (props: Props) => {
const { controlNetId, controlImage, processedControlImage } = props;
const dispatch = useAppDispatch();
const { isProcessingControlImage } = useAppSelector(selector);
const [shouldShowProcessedImage, setShouldShowProcessedImage] =
useState(true);
const handleControlImageChanged = useCallback(
(controlImage: ImageDTO) => {
dispatch(controlNetImageChanged({ controlNetId, controlImage }));
},
[controlNetId, dispatch]
);
const handleControlImageReset = useCallback(() => {
dispatch(controlNetImageChanged({ controlNetId, controlImage: null }));
}, [controlNetId, dispatch]);
const shouldShowProcessedImageBackdrop =
Number(controlImage?.width) > Number(processedControlImage?.width) ||
Number(controlImage?.height) > Number(processedControlImage?.height);
return (
<Box sx={{ position: 'relative', aspectRatio: '1/1' }}>
<IAISelectableImage
image={controlImage}
onChange={handleControlImageChanged}
onReset={handleControlImageReset}
isDropDisabled={Boolean(processedControlImage)}
fallback={<ProcessedImageFallback />}
withResetIcon
resetIconSize="sm"
/>
<AnimatePresence>
{controlImage &&
processedControlImage &&
shouldShowProcessedImage &&
!isProcessingControlImage && (
<motion.div
initial={{
opacity: 0,
}}
animate={{
opacity: 1,
transition: { duration: 0.1 },
}}
exit={{
opacity: 0,
transition: { duration: 0.1 },
}}
>
<Box
sx={{
position: 'absolute',
w: 'full',
h: 'full',
top: 0,
insetInlineStart: 0,
}}
>
{shouldShowProcessedImageBackdrop && (
<Box
sx={{
w: 'full',
h: 'full',
bg: 'base.900',
opacity: 0.7,
}}
/>
)}
<Box
sx={{
position: 'absolute',
top: 0,
insetInlineStart: 0,
w: 'full',
h: 'full',
}}
>
<IAISelectableImage
image={processedControlImage}
onChange={handleControlImageChanged}
onReset={handleControlImageReset}
withResetIcon
resetIconSize="sm"
fallback={<ProcessedImageFallback />}
/>
</Box>
</Box>
</motion.div>
)}
</AnimatePresence>
{isProcessingControlImage && (
<Box
sx={{
position: 'absolute',
top: 0,
insetInlineStart: 0,
w: 'full',
h: 'full',
}}
>
<ProcessedImageFallback />
</Box>
)}
{processedControlImage && !isProcessingControlImage && (
<Box sx={{ position: 'absolute', bottom: 0, insetInlineEnd: 0, p: 2 }}>
<IAIIconButton
aria-label="Hide Preview"
icon={<TbSquareToggle />}
size="sm"
onMouseOver={() => setShouldShowProcessedImage(false)}
onMouseOut={() => setShouldShowProcessedImage(true)}
/>
</Box>
)}
</Box>
);
};
export default memo(ControlNetImagePreview);
const ProcessedImageFallback = () => (
<Flex
sx={{
bg: 'base.900',
opacity: 0.7,
w: 'full',
h: 'full',
alignItems: 'center',
justifyContent: 'center',
borderRadius: 'base',
}}
>
<Spinner size="xl" />
</Flex>
);

View File

@ -12,7 +12,7 @@ import IAIIconButton from 'common/components/IAIIconButton';
import ImageMetadataOverlay from 'common/components/ImageMetadataOverlay';
import { useGetUrl } from 'common/util/getUrl';
import { AnimatePresence, motion } from 'framer-motion';
import { SyntheticEvent } from 'react';
import { ReactElement, SyntheticEvent } from 'react';
import { memo, useRef } from 'react';
import { FaImage, FaTimes } from 'react-icons/fa';
import { ImageDTO } from 'services/api';
@ -26,14 +26,29 @@ type IAISelectableImageProps = {
onReset?: () => void;
onError?: (event: SyntheticEvent<HTMLImageElement>) => void;
resetIconSize?: IconButtonProps['size'];
withResetIcon?: boolean;
withMetadataOverlay?: boolean;
isDropDisabled?: boolean;
fallback?: ReactElement;
};
const IAISelectableImage = (props: IAISelectableImageProps) => {
const { image, onChange, onReset, onError, resetIconSize = 'md' } = props;
const {
image,
onChange,
onReset,
onError,
resetIconSize = 'md',
withResetIcon = false,
withMetadataOverlay = false,
isDropDisabled = false,
fallback = <ImageFallback />,
} = props;
const droppableId = useRef(uuidv4());
const { getUrl } = useGetUrl();
const { isOver, setNodeRef, active } = useDroppable({
id: droppableId.current,
disabled: isDropDisabled,
data: {
handleDrop: onChange,
},
@ -54,6 +69,7 @@ const IAISelectableImage = (props: IAISelectableImageProps) => {
<Flex
sx={{
w: 'full',
h: 'full',
position: 'relative',
alignItems: 'center',
justifyContent: 'center',
@ -62,15 +78,15 @@ const IAISelectableImage = (props: IAISelectableImageProps) => {
<Image
src={getUrl(image.image_url)}
fallbackStrategy="beforeLoadOrError"
fallback={<ImageFallback />}
fallback={fallback}
onError={onError}
draggable={false}
sx={{
borderRadius: 'base',
}}
/>
<ImageMetadataOverlay image={image} />
{onReset && (
{withMetadataOverlay && <ImageMetadataOverlay image={image} />}
{onReset && withResetIcon && (
<Box
sx={{
position: 'absolute',

View File

@ -0,0 +1,123 @@
import {
FormControl,
FormLabel,
HStack,
RangeSlider,
RangeSliderFilledTrack,
RangeSliderMark,
RangeSliderThumb,
RangeSliderTrack,
Tooltip,
} from '@chakra-ui/react';
import { useAppDispatch } from 'app/store/storeHooks';
import IAIIconButton from 'common/components/IAIIconButton';
import {
controlNetBeginStepPctChanged,
controlNetEndStepPctChanged,
} from 'features/controlNet/store/controlNetSlice';
import { memo, useCallback } from 'react';
import { useTranslation } from 'react-i18next';
import { BiReset } from 'react-icons/bi';
type Props = {
controlNetId: string;
beginStepPct: number;
endStepPct: number;
};
const formatPct = (v: number) => `${Math.round(v * 100)}%`;
const ParamControlNetBeginEnd = (props: Props) => {
const { controlNetId, beginStepPct, endStepPct } = props;
const dispatch = useAppDispatch();
const { t } = useTranslation();
const handleStepPctChanged = useCallback(
(v: number[]) => {
dispatch(
controlNetBeginStepPctChanged({ controlNetId, beginStepPct: v[0] })
);
dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct: v[1] }));
},
[controlNetId, dispatch]
);
const handleStepPctReset = useCallback(() => {
dispatch(controlNetBeginStepPctChanged({ controlNetId, beginStepPct: 0 }));
dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct: 1 }));
}, [controlNetId, dispatch]);
return (
<FormControl>
<FormLabel>Begin & End Step %</FormLabel>
<HStack w="100%" gap={2} alignItems="center">
<RangeSlider
aria-label={['Begin Step %', 'End Step %']}
value={[beginStepPct, endStepPct]}
onChange={handleStepPctChanged}
min={0}
max={1}
step={0.01}
minStepsBetweenThumbs={5}
>
<RangeSliderTrack>
<RangeSliderFilledTrack />
</RangeSliderTrack>
<Tooltip label={formatPct(beginStepPct)} placement="top" hasArrow>
<RangeSliderThumb index={0} />
</Tooltip>
<Tooltip label={formatPct(endStepPct)} placement="top" hasArrow>
<RangeSliderThumb index={1} />
</Tooltip>
<RangeSliderMark
value={0}
sx={{
fontSize: 'xs',
fontWeight: '500',
color: 'base.200',
insetInlineStart: '0 !important',
insetInlineEnd: 'unset !important',
mt: 1.5,
}}
>
0%
</RangeSliderMark>
<RangeSliderMark
value={0.5}
sx={{
fontSize: 'xs',
fontWeight: '500',
color: 'base.200',
mt: 1.5,
}}
>
50%
</RangeSliderMark>
<RangeSliderMark
value={1}
sx={{
fontSize: 'xs',
fontWeight: '500',
color: 'base.200',
insetInlineStart: 'unset !important',
insetInlineEnd: '0 !important',
mt: 1.5,
}}
>
100%
</RangeSliderMark>
</RangeSlider>
<IAIIconButton
size="sm"
aria-label={t('accessibility.reset')}
tooltip={t('accessibility.reset')}
icon={<BiReset />}
onClick={handleStepPctReset}
/>
</HStack>
</FormControl>
);
};
export default memo(ParamControlNetBeginEnd);

View File

@ -1,6 +1,9 @@
import { useAppDispatch } from 'app/store/storeHooks';
import IAISlider from 'common/components/IAISlider';
import { controlNetBeginStepPctChanged } from 'features/controlNet/store/controlNetSlice';
import {
controlNetBeginStepPctChanged,
controlNetEndStepPctChanged,
} from 'features/controlNet/store/controlNetSlice';
import { memo, useCallback } from 'react';
type ParamControlNetBeginStepPctProps = {
@ -21,9 +24,20 @@ const ParamControlNetBeginStepPct = (
[controlNetId, dispatch]
);
const handleBeginStepPctReset = () => {
const handleBeginStepPctReset = useCallback(() => {
dispatch(controlNetBeginStepPctChanged({ controlNetId, beginStepPct: 0 }));
};
}, [controlNetId, dispatch]);
const handleEndStepPctChanged = useCallback(
(endStepPct: number) => {
dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct }));
},
[controlNetId, dispatch]
);
const handleEndStepPctReset = useCallback(() => {
dispatch(controlNetEndStepPctChanged({ controlNetId, endStepPct: 0 }));
}, [controlNetId, dispatch]);
return (
<IAISlider

View File

@ -8,6 +8,7 @@ import {
RequiredControlNetProcessorNode,
} from './types';
import { CONTROLNET_PROCESSORS } from './constants';
import { controlNetImageProcessed } from './actions';
export const CONTROLNET_MODELS = [
'lllyasviel/sd-controlnet-canny',
@ -52,12 +53,14 @@ export type ControlNetState = {
controlNets: Record<string, ControlNet>;
isEnabled: boolean;
shouldAutoProcess: boolean;
isProcessingControlImage: boolean;
};
export const initialControlNetState: ControlNetState = {
controlNets: {},
isEnabled: false,
shouldAutoProcess: true,
isProcessingControlImage: false,
};
export const controlNetSlice = createSlice({
@ -107,6 +110,9 @@ export const controlNetSlice = createSlice({
const { controlNetId, controlImage } = action.payload;
state.controlNets[controlNetId].controlImage = controlImage;
state.controlNets[controlNetId].processedControlImage = null;
if (state.shouldAutoProcess && controlImage !== null) {
state.isProcessingControlImage = true;
}
},
isControlNetImageProcessedToggled: (
state,
@ -128,6 +134,7 @@ export const controlNetSlice = createSlice({
const { controlNetId, processedControlImage } = action.payload;
state.controlNets[controlNetId].processedControlImage =
processedControlImage;
state.isProcessingControlImage = false;
},
controlNetModelChanged: (
state,
@ -190,6 +197,15 @@ export const controlNetSlice = createSlice({
state.shouldAutoProcess = !state.shouldAutoProcess;
},
},
extraReducers: (builder) => {
builder.addCase(controlNetImageProcessed, (state, action) => {
if (
state.controlNets[action.payload.controlNetId].controlImage !== null
) {
state.isProcessingControlImage = true;
}
});
},
});
export const {

View File

@ -8,7 +8,7 @@ import { isEqual } from 'lodash-es';
import { gallerySelector } from '../store/gallerySelectors';
import ImageMetadataViewer from './ImageMetaDataViewer/ImageMetadataViewer';
import NextPrevImageButtons from './NextPrevImageButtons';
import { DragEvent, memo, useCallback } from 'react';
import { DragEvent, memo, useCallback, useEffect, useState } from 'react';
import { systemSelector } from 'features/system/store/systemSelectors';
import ImageFallbackSpinner from './ImageFallbackSpinner';
import ImageMetadataOverlay from 'common/components/ImageMetadataOverlay';
@ -55,6 +55,7 @@ const CurrentImagePreview = () => {
const { getUrl } = useGetUrl();
const toaster = useAppToaster();
const dispatch = useAppDispatch();
const [isLoaded, setIsLoaded] = useState(false);
const { attributes, listeners, setNodeRef } = useDraggable({
id: `currentImage_${image?.image_name}`,
@ -74,11 +75,15 @@ const CurrentImagePreview = () => {
}
}, [dispatch, toaster, shouldFetchImages]);
useEffect(() => {
setIsLoaded(false);
}, [image]);
return (
<Flex
sx={{
width: '100%',
height: '100%',
width: 'full',
height: 'full',
position: 'relative',
alignItems: 'center',
justifyContent: 'center',
@ -91,8 +96,8 @@ const CurrentImagePreview = () => {
height={progressImage.height}
sx={{
objectFit: 'contain',
maxWidth: '100%',
maxHeight: '100%',
maxWidth: 'full',
maxHeight: 'full',
height: 'auto',
position: 'absolute',
borderRadius: 'base',
@ -124,8 +129,11 @@ const CurrentImagePreview = () => {
touchAction: 'none',
}}
onError={handleError}
onLoad={() => {
setIsLoaded(true);
}}
/>
<ImageMetadataOverlay image={image} />
{isLoaded && <ImageMetadataOverlay image={image} />}
</Flex>
)
)}

View File

@ -14,7 +14,8 @@ const ImageFallbackSpinner = (props: ImageFallbackSpinnerProps) => {
justifyContent: 'center',
position: 'absolute',
color: 'base.400',
minH: 40,
minH: 36,
minW: 36,
}}
>
<Spinner size={size} {...rest} />