Adds pagination & refresh on re-connect to gallery

This commit is contained in:
psychedelicious 2022-09-26 11:50:47 +10:00
parent db537f154e
commit 460dc897ad
13 changed files with 263 additions and 139 deletions

View File

@ -103,8 +103,8 @@ socketio = SocketIO(
engineio_logger=engineio_logger, engineio_logger=engineio_logger,
max_http_buffer_size=max_http_buffer_size, max_http_buffer_size=max_http_buffer_size,
cors_allowed_origins=cors_allowed_origins, cors_allowed_origins=cors_allowed_origins,
ping_interval=(50,50), ping_interval=(50, 50),
ping_timeout=60 ping_timeout=60,
) )
@ -188,17 +188,50 @@ def handle_request_capabilities():
socketio.emit("systemConfig", config) socketio.emit("systemConfig", config)
@socketio.on("requestAllImages") @socketio.on("requestImages")
def handle_request_all_images(): def handle_request_images(page=1, offset=0, last_mtime=None):
print(f">> All images requested") chunk_size = 50
paths = list(filter(os.path.isfile, glob.glob(result_path + "*.png")))
paths.sort(key=lambda x: os.path.getmtime(x)) if last_mtime:
print(f">> Latest images requested")
else:
print(
f">> Page {page} of images requested (page size {chunk_size} offset {offset})"
)
paths = glob.glob(os.path.join(result_path, "*.png"))
sorted_paths = sorted(paths, key=lambda x: os.path.getmtime(x), reverse=True)
if last_mtime:
image_paths = filter(lambda x: os.path.getmtime(x) > last_mtime, sorted_paths)
else:
image_paths = sorted_paths[
slice(chunk_size * (page - 1) + offset, chunk_size * page + offset)
]
page = page + 1
image_array = [] image_array = []
for path in paths:
for path in image_paths:
metadata = retrieve_metadata(path) metadata = retrieve_metadata(path)
image_array.append({"url": path, "metadata": metadata["sd-metadata"]}) image_array.append(
socketio.emit("galleryImages", {"images": image_array}) {
eventlet.sleep(0) "url": path,
"mtime": os.path.getmtime(path),
"metadata": metadata["sd-metadata"],
}
)
socketio.emit(
"galleryImages",
{
"images": image_array,
"nextPage": page,
"offset": offset,
"onlyNewImages": True if last_mtime else False,
},
)
@socketio.on("generateImage") @socketio.on("generateImage")
@ -277,6 +310,7 @@ def handle_run_esrgan_event(original_image, esrgan_parameters):
"esrganResult", "esrganResult",
{ {
"url": os.path.relpath(path), "url": os.path.relpath(path),
"mtime": os.path.getmtime(path),
"metadata": metadata, "metadata": metadata,
}, },
) )
@ -345,6 +379,7 @@ def handle_run_gfpgan_event(original_image, gfpgan_parameters):
"gfpganResult", "gfpganResult",
{ {
"url": os.path.relpath(path), "url": os.path.relpath(path),
"mtime": os.path.mtime(path),
"metadata": metadata, "metadata": metadata,
}, },
) )
@ -644,7 +679,11 @@ def generate_images(generation_parameters, esrgan_parameters, gfpgan_parameters)
step_index += 1 step_index += 1
socketio.emit( socketio.emit(
"intermediateResult", "intermediateResult",
{"url": os.path.relpath(path), "metadata": generation_parameters}, {
"url": os.path.relpath(path),
"mtime": os.path.getmtime(path),
"metadata": generation_parameters,
},
) )
socketio.emit("progressUpdate", progress) socketio.emit("progressUpdate", progress)
eventlet.sleep(0) eventlet.sleep(0)
@ -737,7 +776,11 @@ def generate_images(generation_parameters, esrgan_parameters, gfpgan_parameters)
socketio.emit( socketio.emit(
"generationResult", "generationResult",
{"url": os.path.relpath(path), "metadata": metadata}, {
"url": os.path.relpath(path),
"mtime": os.path.getmtime(path),
"metadata": metadata,
},
) )
eventlet.sleep(0) eventlet.sleep(0)

File diff suppressed because one or more lines are too long

View File

@ -4,7 +4,7 @@
<meta charset="UTF-8" /> <meta charset="UTF-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" /> <meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>InvokeAI Stable Diffusion Dream Server</title> <title>InvokeAI Stable Diffusion Dream Server</title>
<script type="module" crossorigin src="/assets/index.48fa0a78.js"></script> <script type="module" crossorigin src="/assets/index.66192cce.js"></script>
<link rel="stylesheet" href="/assets/index.447eb2a9.css"> <link rel="stylesheet" href="/assets/index.447eb2a9.css">
</head> </head>
<body> <body>

View File

@ -10,15 +10,13 @@ import PromptInput from '../features/options/PromptInput';
import LogViewer from '../features/system/LogViewer'; import LogViewer from '../features/system/LogViewer';
import Loading from '../Loading'; import Loading from '../Loading';
import { useAppDispatch } from './store'; import { useAppDispatch } from './store';
import { requestAllImages, requestSystemConfig } from './socketio/actions'; import { requestSystemConfig } from './socketio/actions';
const App = () => { const App = () => {
const dispatch = useAppDispatch(); const dispatch = useAppDispatch();
const [isReady, setIsReady] = useState<boolean>(false); const [isReady, setIsReady] = useState<boolean>(false);
// Load images from the gallery once
useEffect(() => { useEffect(() => {
dispatch(requestAllImages());
dispatch(requestSystemConfig()); dispatch(requestSystemConfig());
setIsReady(true); setIsReady(true);
}, [dispatch]); }, [dispatch]);

View File

@ -107,6 +107,7 @@ export declare type Metadata = SystemConfig & {
export declare type Image = { export declare type Image = {
uuid: string; uuid: string;
url: string; url: string;
mtime: number;
metadata: Metadata; metadata: Metadata;
}; };
@ -148,6 +149,7 @@ export declare type SystemConfigResponse = SystemConfig;
export declare type ImageResultResponse = { export declare type ImageResultResponse = {
url: string; url: string;
mtime: number;
metadata: Metadata; metadata: Metadata;
}; };
@ -157,7 +159,10 @@ export declare type ErrorResponse = {
}; };
export declare type GalleryImagesResponse = { export declare type GalleryImagesResponse = {
images: Array<{ url: string; metadata: Metadata }>; images: Array<Omit<Image, 'uuid'>>;
nextPage: number;
offset: number;
onlyNewImages: boolean;
}; };
export declare type ImageUrlAndUuidResponse = { export declare type ImageUrlAndUuidResponse = {

View File

@ -12,8 +12,11 @@ export const generateImage = createAction<undefined>('socketio/generateImage');
export const runESRGAN = createAction<InvokeAI.Image>('socketio/runESRGAN'); export const runESRGAN = createAction<InvokeAI.Image>('socketio/runESRGAN');
export const runGFPGAN = createAction<InvokeAI.Image>('socketio/runGFPGAN'); export const runGFPGAN = createAction<InvokeAI.Image>('socketio/runGFPGAN');
export const deleteImage = createAction<InvokeAI.Image>('socketio/deleteImage'); export const deleteImage = createAction<InvokeAI.Image>('socketio/deleteImage');
export const requestAllImages = createAction<undefined>( export const requestImages = createAction<undefined>(
'socketio/requestAllImages' 'socketio/requestImages'
);
export const requestNewImages = createAction<undefined>(
'socketio/requestNewImages'
); );
export const cancelProcessing = createAction<undefined>( export const cancelProcessing = createAction<undefined>(
'socketio/cancelProcessing' 'socketio/cancelProcessing'
@ -23,4 +26,6 @@ export const uploadInitialImage = createAction<File>(
); );
export const uploadMaskImage = createAction<File>('socketio/uploadMaskImage'); export const uploadMaskImage = createAction<File>('socketio/uploadMaskImage');
export const requestSystemConfig = createAction<undefined>('socketio/requestSystemConfig'); export const requestSystemConfig = createAction<undefined>(
'socketio/requestSystemConfig'
);

View File

@ -83,8 +83,17 @@ const makeSocketIOEmitters = (
const { url, uuid } = imageToDelete; const { url, uuid } = imageToDelete;
socketio.emit('deleteImage', url, uuid); socketio.emit('deleteImage', url, uuid);
}, },
emitRequestAllImages: () => { emitRequestImages: () => {
socketio.emit('requestAllImages'); const { nextPage, offset } = getState().gallery;
socketio.emit('requestImages', nextPage, offset);
},
emitRequestNewImages: () => {
const { nextPage, offset, images } = getState().gallery;
if (images.length > 0) {
socketio.emit('requestImages', nextPage, offset, images[0].mtime);
} else {
socketio.emit('requestImages', nextPage, offset);
}
}, },
emitCancelProcessing: () => { emitCancelProcessing: () => {
socketio.emit('cancel'); socketio.emit('cancel');
@ -96,8 +105,8 @@ const makeSocketIOEmitters = (
socketio.emit('uploadMaskImage', file, file.name); socketio.emit('uploadMaskImage', file, file.name);
}, },
emitRequestSystemConfig: () => { emitRequestSystemConfig: () => {
socketio.emit('requestSystemConfig') socketio.emit('requestSystemConfig');
} },
}; };
}; };

View File

@ -14,10 +14,10 @@ import {
} from '../../features/system/systemSlice'; } from '../../features/system/systemSlice';
import { import {
addGalleryImages,
addImage, addImage,
clearIntermediateImage, clearIntermediateImage,
removeImage, removeImage,
setGalleryImages,
setIntermediateImage, setIntermediateImage,
} from '../../features/gallery/gallerySlice'; } from '../../features/gallery/gallerySlice';
@ -25,6 +25,7 @@ import {
setInitialImagePath, setInitialImagePath,
setMaskPath, setMaskPath,
} from '../../features/options/optionsSlice'; } from '../../features/options/optionsSlice';
import { requestNewImages } from './actions';
/** /**
* Returns an object containing listener callbacks for socketio events. * Returns an object containing listener callbacks for socketio events.
@ -43,6 +44,7 @@ const makeSocketIOListeners = (
try { try {
dispatch(setIsConnected(true)); dispatch(setIsConnected(true));
dispatch(setCurrentStatus('Connected')); dispatch(setCurrentStatus('Connected'));
dispatch(requestNewImages());
} catch (e) { } catch (e) {
console.error(e); console.error(e);
} }
@ -53,7 +55,6 @@ const makeSocketIOListeners = (
onDisconnect: () => { onDisconnect: () => {
try { try {
dispatch(setIsConnected(false)); dispatch(setIsConnected(false));
dispatch(setIsProcessing(false));
dispatch(setCurrentStatus('Disconnected')); dispatch(setCurrentStatus('Disconnected'));
dispatch( dispatch(
@ -72,13 +73,14 @@ const makeSocketIOListeners = (
*/ */
onGenerationResult: (data: InvokeAI.ImageResultResponse) => { onGenerationResult: (data: InvokeAI.ImageResultResponse) => {
try { try {
const { url, metadata } = data; const { url, mtime, metadata } = data;
const newUuid = uuidv4(); const newUuid = uuidv4();
dispatch( dispatch(
addImage({ addImage({
uuid: newUuid, uuid: newUuid,
url, url,
mtime,
metadata: metadata, metadata: metadata,
}) })
); );
@ -99,11 +101,12 @@ const makeSocketIOListeners = (
onIntermediateResult: (data: InvokeAI.ImageResultResponse) => { onIntermediateResult: (data: InvokeAI.ImageResultResponse) => {
try { try {
const uuid = uuidv4(); const uuid = uuidv4();
const { url, metadata } = data; const { url, metadata, mtime } = data;
dispatch( dispatch(
setIntermediateImage({ setIntermediateImage({
uuid, uuid,
url, url,
mtime,
metadata, metadata,
}) })
); );
@ -123,12 +126,13 @@ const makeSocketIOListeners = (
*/ */
onESRGANResult: (data: InvokeAI.ImageResultResponse) => { onESRGANResult: (data: InvokeAI.ImageResultResponse) => {
try { try {
const { url, metadata } = data; const { url, metadata, mtime } = data;
dispatch( dispatch(
addImage({ addImage({
uuid: uuidv4(), uuid: uuidv4(),
url, url,
mtime,
metadata, metadata,
}) })
); );
@ -149,12 +153,13 @@ const makeSocketIOListeners = (
*/ */
onGFPGANResult: (data: InvokeAI.ImageResultResponse) => { onGFPGANResult: (data: InvokeAI.ImageResultResponse) => {
try { try {
const { url, metadata } = data; const { url, metadata, mtime } = data;
dispatch( dispatch(
addImage({ addImage({
uuid: uuidv4(), uuid: uuidv4(),
url, url,
mtime,
metadata, metadata,
}) })
); );
@ -209,16 +214,26 @@ const makeSocketIOListeners = (
* Callback to run when we receive a 'galleryImages' event. * Callback to run when we receive a 'galleryImages' event.
*/ */
onGalleryImages: (data: InvokeAI.GalleryImagesResponse) => { onGalleryImages: (data: InvokeAI.GalleryImagesResponse) => {
const { images } = data; const { images, nextPage, offset } = data;
/**
* the logic here ideally would be in the reducer but we have a side effect:
* generating a uuid. so the logic needs to be here, outside redux.
*/
// Generate a UUID for each image
const preparedImages = images.map((image): InvokeAI.Image => { const preparedImages = images.map((image): InvokeAI.Image => {
const { url, metadata } = image; const { url, metadata, mtime } = image;
return { return {
uuid: uuidv4(), uuid: uuidv4(),
url, url,
mtime,
metadata, metadata,
}; };
}); });
dispatch(setGalleryImages(preparedImages));
dispatch(addGalleryImages({ images: preparedImages, nextPage, offset }));
dispatch( dispatch(
addLogEntry({ addLogEntry({
timestamp: dateFormat(new Date(), 'isoDateTime'), timestamp: dateFormat(new Date(), 'isoDateTime'),

View File

@ -53,7 +53,8 @@ export const socketioMiddleware = () => {
emitRunESRGAN, emitRunESRGAN,
emitRunGFPGAN, emitRunGFPGAN,
emitDeleteImage, emitDeleteImage,
emitRequestAllImages, emitRequestImages,
emitRequestNewImages,
emitCancelProcessing, emitCancelProcessing,
emitUploadInitialImage, emitUploadInitialImage,
emitUploadMaskImage, emitUploadMaskImage,
@ -142,11 +143,17 @@ export const socketioMiddleware = () => {
break; break;
} }
case 'socketio/requestAllImages': { case 'socketio/requestImages': {
emitRequestAllImages(); emitRequestImages();
break; break;
} }
case 'socketio/requestNewImages': {
emitRequestNewImages();
break;
}
case 'socketio/cancelProcessing': { case 'socketio/cancelProcessing': {
emitCancelProcessing(); emitCancelProcessing();
break; break;

View File

@ -1,5 +1,6 @@
import { Center, Flex, Text } from '@chakra-ui/react'; import { Button, Center, Flex, Text } from '@chakra-ui/react';
import { RootState } from '../../app/store'; import { requestImages } from '../../app/socketio/actions';
import { RootState, useAppDispatch } from '../../app/store';
import { useAppSelector } from '../../app/store'; import { useAppSelector } from '../../app/store';
import HoverableImage from './HoverableImage'; import HoverableImage from './HoverableImage';
@ -10,7 +11,7 @@ const ImageGallery = () => {
const { images, currentImageUuid } = useAppSelector( const { images, currentImageUuid } = useAppSelector(
(state: RootState) => state.gallery (state: RootState) => state.gallery
); );
const dispatch = useAppDispatch();
/** /**
* I don't like that this needs to rerender whenever the current image is changed. * I don't like that this needs to rerender whenever the current image is changed.
* What if we have a large number of images? I suppose pagination (planned) will * What if we have a large number of images? I suppose pagination (planned) will
@ -19,15 +20,22 @@ const ImageGallery = () => {
* TODO: Refactor if performance complaints, or after migrating to new API which supports pagination. * TODO: Refactor if performance complaints, or after migrating to new API which supports pagination.
*/ */
const handleClickLoadMore = () => {
dispatch(requestImages());
};
return images.length ? ( return images.length ? (
<Flex gap={2} wrap="wrap" pb={2}> <Flex direction={'column'} gap={2} pb={2}>
{[...images].reverse().map((image) => { <Flex gap={2} wrap="wrap">
const { uuid } = image; {images.map((image) => {
const isSelected = currentImageUuid === uuid; const { uuid } = image;
return ( const isSelected = currentImageUuid === uuid;
<HoverableImage key={uuid} image={image} isSelected={isSelected} /> return (
); <HoverableImage key={uuid} image={image} isSelected={isSelected} />
})} );
})}
</Flex>
<Button onClick={handleClickLoadMore}>Load more...</Button>
</Flex> </Flex>
) : ( ) : (
<Center height={'100%'} position={'relative'}> <Center height={'100%'} position={'relative'}>

View File

@ -8,11 +8,15 @@ export interface GalleryState {
currentImageUuid: string; currentImageUuid: string;
images: Array<InvokeAI.Image>; images: Array<InvokeAI.Image>;
intermediateImage?: InvokeAI.Image; intermediateImage?: InvokeAI.Image;
nextPage: number;
offset: number;
} }
const initialState: GalleryState = { const initialState: GalleryState = {
currentImageUuid: '', currentImageUuid: '',
images: [], images: [],
nextPage: 1,
offset: 0,
}; };
export const gallerySlice = createSlice({ export const gallerySlice = createSlice({
@ -50,7 +54,7 @@ export const gallerySlice = createSlice({
* Clamp the new index to ensure it is valid.. * Clamp the new index to ensure it is valid..
*/ */
const newCurrentImageIndex = clamp( const newCurrentImageIndex = clamp(
imageToDeleteIndex - 1, imageToDeleteIndex,
0, 0,
newImages.length - 1 newImages.length - 1
); );
@ -67,10 +71,11 @@ export const gallerySlice = createSlice({
state.images = newImages; state.images = newImages;
}, },
addImage: (state, action: PayloadAction<InvokeAI.Image>) => { addImage: (state, action: PayloadAction<InvokeAI.Image>) => {
state.images.push(action.payload); state.images.unshift(action.payload);
state.currentImageUuid = action.payload.uuid; state.currentImageUuid = action.payload.uuid;
state.intermediateImage = undefined; state.intermediateImage = undefined;
state.currentImage = action.payload; state.currentImage = action.payload;
state.offset += 1
}, },
setIntermediateImage: (state, action: PayloadAction<InvokeAI.Image>) => { setIntermediateImage: (state, action: PayloadAction<InvokeAI.Image>) => {
state.intermediateImage = action.payload; state.intermediateImage = action.payload;
@ -78,13 +83,24 @@ export const gallerySlice = createSlice({
clearIntermediateImage: (state) => { clearIntermediateImage: (state) => {
state.intermediateImage = undefined; state.intermediateImage = undefined;
}, },
setGalleryImages: (state, action: PayloadAction<Array<InvokeAI.Image>>) => { addGalleryImages: (
const newImages = action.payload; state,
if (newImages.length) { action: PayloadAction<{
const newCurrentImage = newImages[newImages.length - 1]; images: Array<InvokeAI.Image>;
state.images = newImages; nextPage: number;
offset: number;
}>
) => {
const { images, nextPage, offset } = action.payload;
if (images.length) {
const newCurrentImage = images[0];
state.images = state.images
.concat(images)
.sort((a, b) => b.mtime - a.mtime);
state.currentImage = newCurrentImage; state.currentImage = newCurrentImage;
state.currentImageUuid = newCurrentImage.uuid; state.currentImageUuid = newCurrentImage.uuid;
state.nextPage = nextPage;
state.offset = offset;
} }
}, },
}, },
@ -95,7 +111,7 @@ export const {
clearIntermediateImage, clearIntermediateImage,
removeImage, removeImage,
setCurrentImage, setCurrentImage,
setGalleryImages, addGalleryImages,
setIntermediateImage, setIntermediateImage,
} = gallerySlice.actions; } = gallerySlice.actions;

View File

@ -32,8 +32,16 @@ import { cloneElement, ReactElement } from 'react';
const systemSelector = createSelector( const systemSelector = createSelector(
(state: RootState) => state.system, (state: RootState) => state.system,
(system: SystemState) => { (system: SystemState) => {
const { shouldDisplayInProgress, shouldConfirmOnDelete, shouldDisplayGuides } = system; const {
return { shouldDisplayInProgress, shouldConfirmOnDelete, shouldDisplayGuides }; shouldDisplayInProgress,
shouldConfirmOnDelete,
shouldDisplayGuides,
} = system;
return {
shouldDisplayInProgress,
shouldConfirmOnDelete,
shouldDisplayGuides,
};
}, },
{ {
memoizeOptions: { resultEqualityCheck: isEqual }, memoizeOptions: { resultEqualityCheck: isEqual },
@ -64,8 +72,11 @@ const SettingsModal = ({ children }: SettingsModalProps) => {
onClose: onRefreshModalClose, onClose: onRefreshModalClose,
} = useDisclosure(); } = useDisclosure();
const { shouldDisplayInProgress, shouldConfirmOnDelete, shouldDisplayGuides } = const {
useAppSelector(systemSelector); shouldDisplayInProgress,
shouldConfirmOnDelete,
shouldDisplayGuides,
} = useAppSelector(systemSelector);
const dispatch = useAppDispatch(); const dispatch = useAppDispatch();

View File

@ -1,7 +1,7 @@
import { createSlice } from '@reduxjs/toolkit'; import { createSlice } from '@reduxjs/toolkit';
import type { PayloadAction } from '@reduxjs/toolkit'; import type { PayloadAction } from '@reduxjs/toolkit';
import { ExpandedIndex } from '@chakra-ui/react'; import { ExpandedIndex } from '@chakra-ui/react';
import * as InvokeAI from '../../app/invokeai' import * as InvokeAI from '../../app/invokeai';
export type LogLevel = 'info' | 'warning' | 'error'; export type LogLevel = 'info' | 'warning' | 'error';
@ -15,7 +15,9 @@ export interface Log {
[index: number]: LogEntry; [index: number]: LogEntry;
} }
export interface SystemState extends InvokeAI.SystemStatus, InvokeAI.SystemConfig { export interface SystemState
extends InvokeAI.SystemStatus,
InvokeAI.SystemConfig {
shouldDisplayInProgress: boolean; shouldDisplayInProgress: boolean;
log: Array<LogEntry>; log: Array<LogEntry>;
shouldShowLogViewer: boolean; shouldShowLogViewer: boolean;
@ -31,7 +33,6 @@ export interface SystemState extends InvokeAI.SystemStatus, InvokeAI.SystemConfi
totalIterations: number; totalIterations: number;
currentStatus: string; currentStatus: string;
currentStatusHasSteps: boolean; currentStatusHasSteps: boolean;
shouldDisplayGuides: boolean; shouldDisplayGuides: boolean;
} }
@ -51,7 +52,7 @@ const initialSystemState = {
totalSteps: 0, totalSteps: 0,
currentIteration: 0, currentIteration: 0,
totalIterations: 0, totalIterations: 0,
currentStatus: '', currentStatus: 'Disconnected',
currentStatusHasSteps: false, currentStatusHasSteps: false,
model: '', model: '',
model_id: '', model_id: '',
@ -107,6 +108,12 @@ export const systemSlice = createSlice({
}, },
setIsConnected: (state, action: PayloadAction<boolean>) => { setIsConnected: (state, action: PayloadAction<boolean>) => {
state.isConnected = action.payload; state.isConnected = action.payload;
state.isProcessing = false;
state.currentStep = 0;
state.totalSteps = 0;
state.currentIteration = 0;
state.totalIterations = 0;
state.currentStatusHasSteps = false;
}, },
setSocketId: (state, action: PayloadAction<string>) => { setSocketId: (state, action: PayloadAction<string>) => {
state.socketId = action.payload; state.socketId = action.payload;