Skip to content

Commit

Permalink
Maryhipp/option fetch metadata from api (#4491)
Browse files Browse the repository at this point in the history
## What type of PR is this? (check all applicable)

- [ ] Refactor
- [x] Feature
- [ ] Bug Fix
- [ ] Optimization
- [ ] Documentation Update
- [ ] Community Node Submission


## Description

Adds a configuration option to fetch metadata and workflows from api
isntead of the image file. Needed for commercial.
  • Loading branch information
blessedcoolant authored Sep 8, 2023
2 parents cf83dde + 501cb4c commit b700809
Show file tree
Hide file tree
Showing 6 changed files with 73 additions and 21 deletions.
1 change: 1 addition & 0 deletions invokeai/frontend/web/src/app/types/invokeai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ export type AppConfig = {
* Whether or not we should update image urls when image loading errors
*/
shouldUpdateImagesOnConnect: boolean;
shouldFetchMetadataFromApi: boolean;
disabledTabs: InvokeTabName[];
disabledFeatures: AppFeature[];
disabledSDFeatures: SDFeature[];
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import {
setShouldShowImageDetails,
setShouldShowProgressInViewer,
} from 'features/ui/store/uiSlice';
import { memo, useCallback } from 'react';
import { memo, useCallback, useMemo } from 'react';
import { useHotkeys } from 'react-hotkeys-hook';
import { useTranslation } from 'react-i18next';
import {
Expand All @@ -49,7 +49,7 @@ import SingleSelectionMenuItems from '../ImageContextMenu/SingleSelectionMenuIte

const currentImageButtonsSelector = createSelector(
[stateSelector, activeTabNameSelector],
({ gallery, system, ui }, activeTabName) => {
({ gallery, system, ui, config }, activeTabName) => {
const { isProcessing, isConnected, shouldConfirmOnDelete, progressImage } =
system;

Expand All @@ -59,6 +59,8 @@ const currentImageButtonsSelector = createSelector(
shouldShowProgressInViewer,
} = ui;

const { shouldFetchMetadataFromApi } = config;

const lastSelectedImage = gallery.selection[gallery.selection.length - 1];

return {
Expand All @@ -72,6 +74,7 @@ const currentImageButtonsSelector = createSelector(
shouldHidePreview,
shouldShowProgressInViewer,
lastSelectedImage,
shouldFetchMetadataFromApi,
};
},
{
Expand All @@ -92,6 +95,7 @@ const CurrentImageButtons = (props: CurrentImageButtonsProps) => {
shouldShowImageDetails,
lastSelectedImage,
shouldShowProgressInViewer,
shouldFetchMetadataFromApi,
} = useAppSelector(currentImageButtonsSelector);

const isUpscalingEnabled = useFeatureStatus('upscaling').isFeatureEnabled;
Expand All @@ -106,8 +110,16 @@ const CurrentImageButtons = (props: CurrentImageButtonsProps) => {
lastSelectedImage?.image_name ?? skipToken
);

const getMetadataArg = useMemo(() => {
if (lastSelectedImage) {
return { image: lastSelectedImage, shouldFetchMetadataFromApi };
} else {
return skipToken;
}
}, [lastSelectedImage, shouldFetchMetadataFromApi]);

const { metadata, workflow, isLoading } = useGetImageMetadataFromFileQuery(
lastSelectedImage ?? skipToken,
getMetadataArg,
{
selectFromResult: (res) => ({
isLoading: res.isFetching,
Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { Flex, MenuItem, Spinner } from '@chakra-ui/react';
import { useAppToaster } from 'app/components/Toaster';
import { useAppDispatch } from 'app/store/storeHooks';
import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
import { setInitialCanvasImage } from 'features/canvas/store/canvasSlice';
import {
imagesToChangeSelected,
Expand Down Expand Up @@ -34,6 +34,7 @@ import {
import { ImageDTO } from 'services/api/types';
import { sentImageToCanvas, sentImageToImg2Img } from '../../store/actions';
import { workflowLoadRequested } from 'features/nodes/store/actions';
import { configSelector } from '../../../system/store/configSelectors';

type SingleSelectionMenuItemsProps = {
imageDTO: ImageDTO;
Expand All @@ -48,9 +49,10 @@ const SingleSelectionMenuItems = (props: SingleSelectionMenuItemsProps) => {
const toaster = useAppToaster();

const isCanvasEnabled = useFeatureStatus('unifiedCanvas').isFeatureEnabled;
const { shouldFetchMetadataFromApi } = useAppSelector(configSelector);

const { metadata, workflow, isLoading } = useGetImageMetadataFromFileQuery(
imageDTO,
{ image: imageDTO, shouldFetchMetadataFromApi },
{
selectFromResult: (res) => ({
isLoading: res.isFetching,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ import { useGetImageMetadataFromFileQuery } from 'services/api/endpoints/images'
import { ImageDTO } from 'services/api/types';
import DataViewer from './DataViewer';
import ImageMetadataActions from './ImageMetadataActions';
import { useAppSelector } from '../../../../app/store/storeHooks';
import { configSelector } from '../../../system/store/configSelectors';

type ImageMetadataViewerProps = {
image: ImageDTO;
Expand All @@ -27,12 +29,17 @@ const ImageMetadataViewer = ({ image }: ImageMetadataViewerProps) => {
// dispatch(setShouldShowImageDetails(false));
// });

const { metadata, workflow } = useGetImageMetadataFromFileQuery(image, {
selectFromResult: (res) => ({
metadata: res?.currentData?.metadata,
workflow: res?.currentData?.workflow,
}),
});
const { shouldFetchMetadataFromApi } = useAppSelector(configSelector);

const { metadata, workflow } = useGetImageMetadataFromFileQuery(
{ image, shouldFetchMetadataFromApi },
{
selectFromResult: (res) => ({
metadata: res?.currentData?.metadata,
workflow: res?.currentData?.workflow,
}),
}
);

return (
<Flex
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import { merge } from 'lodash-es';

export const initialConfigState: AppConfig = {
shouldUpdateImagesOnConnect: false,
shouldFetchMetadataFromApi: false,
disabledTabs: [],
disabledFeatures: ['lightbox', 'faceRestore', 'batches'],
disabledSDFeatures: [
Expand Down
49 changes: 39 additions & 10 deletions invokeai/frontend/web/src/services/api/endpoints/images.ts
Original file line number Diff line number Diff line change
@@ -1,14 +1,20 @@
import { EntityState, Update } from '@reduxjs/toolkit';
import { fetchBaseQuery } from '@reduxjs/toolkit/dist/query';
import { PatchCollection } from '@reduxjs/toolkit/dist/query/core/buildThunks';
import {
ASSETS_CATEGORIES,
BoardId,
IMAGE_CATEGORIES,
IMAGE_LIMIT,
} from 'features/gallery/store/types';
import {
ImageMetadataAndWorkflow,
zCoreMetadata,
} from 'features/nodes/types/types';
import { getMetadataAndWorkflowFromImageBlob } from 'features/nodes/util/getMetadataAndWorkflowFromImageBlob';
import { keyBy } from 'lodash-es';
import { ApiFullTagDescription, LIST_TAG, api } from '..';
import { $authToken, $projectId } from '../client';
import { components, paths } from '../schema';
import {
DeleteBoardResult,
Expand All @@ -27,9 +33,6 @@ import {
imagesSelectors,
} from '../util';
import { boardsApi } from './boards';
import { ImageMetadataAndWorkflow } from 'features/nodes/types/types';
import { fetchBaseQuery } from '@reduxjs/toolkit/dist/query';
import { $authToken, $projectId } from '../client';

export const imagesApi = api.injectEndpoints({
endpoints: (build) => ({
Expand Down Expand Up @@ -117,8 +120,16 @@ export const imagesApi = api.injectEndpoints({
],
keepUnusedDataFor: 86400, // 24 hours
}),
getImageMetadataFromFile: build.query<ImageMetadataAndWorkflow, ImageDTO>({
queryFn: async (args: ImageDTO, api, extraOptions) => {
getImageMetadataFromFile: build.query<
ImageMetadataAndWorkflow,
{ image: ImageDTO; shouldFetchMetadataFromApi: boolean }
>({
queryFn: async (
args: { image: ImageDTO; shouldFetchMetadataFromApi: boolean },
api,
extraOptions,
fetchWithBaseQuery
) => {
const authToken = $authToken.get();
const projectId = $projectId.get();
const customBaseQuery = fetchBaseQuery({
Expand All @@ -139,17 +150,35 @@ export const imagesApi = api.injectEndpoints({
});

const response = await customBaseQuery(
args.image_url,
args.image.image_url,
api,
extraOptions
);
const data = await getMetadataAndWorkflowFromImageBlob(
const blobData = await getMetadataAndWorkflowFromImageBlob(
response.data as Blob
);
return { data };

let metadata = blobData.metadata;

if (args.shouldFetchMetadataFromApi) {
const metadataResponse = await fetchWithBaseQuery(
`images/i/${args.image.image_name}/metadata`
);
if (metadataResponse.data) {
const metadataResult = zCoreMetadata.safeParse(
// eslint-disable-next-line @typescript-eslint/no-explicit-any
(metadataResponse.data as any)?.metadata
);
if (metadataResult.success) {
metadata = metadataResult.data;
}
}
}

return { data: { ...blobData, metadata } };
},
providesTags: (result, error, image_dto) => [
{ type: 'ImageMetadataFromFile', id: image_dto.image_name },
providesTags: (result, error, { image }) => [
{ type: 'ImageMetadataFromFile', id: image.image_name },
],
keepUnusedDataFor: 86400, // 24 hours
}),
Expand Down

0 comments on commit b700809

Please sign in to comment.