feat(ui): set HF token in MM tab

- Display a toast on UI launch if the HF token is invalid
- Show form in MM if token is invalid or unable to be verified, let user set the token via this form
This commit is contained in:
psychedelicious 2024-03-20 09:24:10 +11:00
parent 3f6f8199f6
commit bdb52cfcf7
8 changed files with 212 additions and 1 deletions

View File

@ -638,6 +638,14 @@
"huggingFacePlaceholder": "owner/model-name",
"huggingFaceRepoID": "HuggingFace Repo ID",
"huggingFaceHelper": "If multiple models are found in this repo, you will be prompted to select one to install.",
"hfToken": "HuggingFace Token",
"hfTokenHelperText": "A HF token is required to use checkpoint models. Click here to create or get your token.",
"hfTokenInvalid": "Invalid or Missing HF Token",
"hfTokenInvalidErrorMessage": "Invalid or missing HuggingFace token.",
"hfTokenInvalidErrorMessage2": "Update it in the ",
"hfTokenUnableToVerify": "Unable to Verify HF Token",
"hfTokenUnableToVerifyErrorMessage": "Unable to verify HuggingFace token. This is likely due to a network error. Please try again later.",
"hfTokenSaved": "HF Token Saved",
"imageEncoderModelId": "Image Encoder Model ID",
"installQueue": "Install Queue",
"inplaceInstall": "In-place install",

View File

@ -11,6 +11,7 @@ import { useGlobalHotkeys } from 'common/hooks/useGlobalHotkeys';
import ChangeBoardModal from 'features/changeBoardModal/components/ChangeBoardModal';
import DeleteImageModal from 'features/deleteImageModal/components/DeleteImageModal';
import { DynamicPromptsModal } from 'features/dynamicPrompts/components/DynamicPromptsPreviewModal';
import { useHFLoginToast } from 'features/modelManagerV2/hooks/useHFLoginToast';
import { useStarterModelsToast } from 'features/modelManagerV2/hooks/useStarterModelsToast';
import { configChanged } from 'features/system/store/configSlice';
import { languageSelector } from 'features/system/store/systemSelectors';
@ -70,6 +71,7 @@ const App = ({ config = DEFAULT_CONFIG, selectedImage }: Props) => {
}, [dispatch]);
useStarterModelsToast();
useHFLoginToast()
return (
<ErrorBoundary onReset={handleReset} FallbackComponent={AppErrorBoundaryFallback}>

View File

@ -25,7 +25,8 @@ export type AppFeature =
| 'prependQueue'
| 'invocationCache'
| 'bulkDownload'
| 'starterModels';
| 'starterModels'
| 'hfToken';
/**
* A disable-able Stable Diffusion feature

View File

@ -0,0 +1,81 @@
import {
Button,
ExternalLink,
Flex,
FormControl,
FormErrorMessage,
FormHelperText,
FormLabel,
Input,
useToast,
} from '@invoke-ai/ui-library';
import { skipToken } from '@reduxjs/toolkit/query';
import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
import type { ChangeEvent } from 'react';
import { useCallback, useMemo, useState } from 'react';
import { useTranslation } from 'react-i18next';
import { useGetHFTokenStatusQuery, useSetHFTokenMutation } from 'services/api/endpoints/models';
export const HFToken = () => {
const { t } = useTranslation();
const isEnabled = useFeatureStatus('hfToken').isFeatureEnabled;
const [token, setToken] = useState('');
const { currentData } = useGetHFTokenStatusQuery(isEnabled ? undefined : skipToken);
const [trigger, { isLoading }] = useSetHFTokenMutation();
const toast = useToast();
const onChange = useCallback((e: ChangeEvent<HTMLInputElement>) => {
setToken(e.target.value);
}, []);
const onClick = useCallback(() => {
trigger({ token })
.unwrap()
.then((res) => {
if (res === 'valid') {
setToken('');
toast({
title: t('modelManager.hfTokenSaved'),
status: 'success',
duration: 3000,
});
}
});
}, [t, toast, token, trigger]);
const error = useMemo(() => {
if (!currentData || isLoading) {
return null;
}
if (currentData === 'invalid') {
return t('modelManager.hfTokenInvalidErrorMessage');
}
if (currentData === 'unknown') {
return t('modelManager.hfTokenUnableToVerifyErrorMessage');
}
return null;
}, [currentData, isLoading, t]);
if (!currentData || currentData === 'valid') {
return null;
}
return (
<Flex borderRadius="base" w="full">
<FormControl isInvalid={Boolean(error)} orientation="vertical">
<FormLabel>{t('modelManager.hfToken')}</FormLabel>
<Flex gap={3} alignItems="center" w="full">
<Input type="password" value={token} onChange={onChange} />
<Button onClick={onClick} size="sm" isDisabled={token.trim().length === 0} isLoading={isLoading}>
{t('common.save')}
</Button>
</Flex>
<FormHelperText>
<ExternalLink
label={t('modelManager.hfTokenHelperText')}
href="https://huggingface.co/settings/tokens"
/>
</FormHelperText>
<FormErrorMessage>{error}</FormErrorMessage>
</FormControl>
</Flex>
);
};

View File

@ -0,0 +1,89 @@
import { Button, Text, useToast } from '@invoke-ai/ui-library';
import { skipToken } from '@reduxjs/toolkit/query';
import { useAppDispatch } from 'app/store/storeHooks';
import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
import { setActiveTab } from 'features/ui/store/uiSlice';
import { t } from 'i18next';
import { useCallback, useEffect, useState } from 'react';
import { useTranslation } from 'react-i18next';
import { useGetHFTokenStatusQuery } from 'services/api/endpoints/models';
import type { S } from 'services/api/types';
const FEATURE_ID = 'hfToken';
const getTitle = (token_status: S['HFTokenStatus']) => {
switch (token_status) {
case 'invalid':
return t('modelManager.hfTokenInvalid');
case 'unknown':
return t('modelManager.hfTokenUnableToVerify');
}
};
export const useHFLoginToast = () => {
const { t } = useTranslation();
const isEnabled = useFeatureStatus(FEATURE_ID).isFeatureEnabled;
const [didToast, setDidToast] = useState(false);
const { data } = useGetHFTokenStatusQuery(isEnabled ? undefined : skipToken);
const toast = useToast();
useEffect(() => {
if (toast.isActive(FEATURE_ID)) {
if (data === 'valid') {
setDidToast(true);
toast.close(FEATURE_ID);
}
return;
}
if (data && data !== 'valid' && !didToast && isEnabled) {
const title = getTitle(data);
toast({
id: FEATURE_ID,
title,
description: <ToastDescription token_status={data} />,
status: 'info',
isClosable: true,
duration: null,
onCloseComplete: () => setDidToast(true),
});
}
}, [data, didToast, isEnabled, t, toast]);
};
type Props = {
token_status: S['HFTokenStatus'];
};
const ToastDescription = ({ token_status }: Props) => {
const { t } = useTranslation();
const dispatch = useAppDispatch();
const toast = useToast();
const onClick = useCallback(() => {
dispatch(setActiveTab('modelManager'));
toast.close(FEATURE_ID);
}, [dispatch, toast]);
if (token_status === 'invalid') {
return (
<Text fontSize="md">
{t('modelManager.hfTokenInvalidErrorMessage')}{' '}
{t('modelManager.hfTokenInvalidErrorMessage2')}
<Button onClick={onClick} variant="link" color="base.50" flexGrow={0}>
{t('modelManager.modelManager')}.
</Button>
</Text>
);
}
if (token_status === 'unknown') {
return (
<Text fontSize="md">
{t('modelManager.hfTokenUnableToErrorMessage')}{' '}
<Button onClick={onClick} variant="link" color="base.50" flexGrow={0}>
{t('modelManager.modelManager')}.
</Button>
</Text>
);
}
};

View File

@ -1,5 +1,6 @@
import { Button, Flex, Heading, Spacer } from '@invoke-ai/ui-library';
import { useAppDispatch } from 'app/store/storeHooks';
import { HFToken } from 'features/modelManagerV2/components/HFToken';
import { SyncModelsButton } from 'features/modelManagerV2/components/SyncModels/SyncModelsButton';
import { setSelectedModelKey } from 'features/modelManagerV2/store/modelManagerV2Slice';
import { useCallback } from 'react';
@ -27,6 +28,7 @@ export const ModelManager = () => {
</Button>
</Flex>
<Flex flexDir="column" layerStyle="second" p={4} gap={4} borderRadius="base" w="full" h="full">
<HFToken />
<ModelListNavigation />
<ModelList />
</Flex>

View File

@ -27,6 +27,15 @@ type GetModelConfigsResponse = NonNullable<
paths['/api/v2/models/']['get']['responses']['200']['content']['application/json']
>;
type GetHFTokenStatusResponse =
paths['/api/v2/models/hf_login']['get']['responses']['200']['content']['application/json'];
type SetHFTokenResponse = NonNullable<
paths['/api/v2/models/hf_login']['post']['responses']['200']['content']['application/json']
>;
type SetHFTokenArg = NonNullable<
paths['/api/v2/models/hf_login']['post']['requestBody']['content']['application/json']
>;
export type GetStarterModelsResponse =
paths['/api/v2/models/starter_models']['get']['responses']['200']['content']['application/json'];
@ -265,6 +274,22 @@ export const modelsApi = api.injectEndpoints({
getStarterModels: build.query<GetStarterModelsResponse, void>({
query: () => buildModelsUrl('starter_models'),
}),
getHFTokenStatus: build.query<GetHFTokenStatusResponse, void>({
query: () => buildModelsUrl('hf_login'),
providesTags: ['HFTokenStatus'],
}),
setHFToken: build.mutation<SetHFTokenResponse, SetHFTokenArg>({
query: (body) => ({ url: buildModelsUrl('hf_login'), method: 'POST', body }),
invalidatesTags: ['HFTokenStatus'],
onQueryStarted: async (_, { dispatch, queryFulfilled }) => {
try {
const { data } = await queryFulfilled;
dispatch(modelsApi.util.updateQueryData('getHFTokenStatus', undefined, () => data));
} catch {
// no-op
}
},
}),
}),
});
@ -284,4 +309,6 @@ export const {
useCancelModelInstallMutation,
usePruneCompletedModelInstallsMutation,
useGetStarterModelsQuery,
useGetHFTokenStatusQuery,
useSetHFTokenMutation,
} = modelsApi;

View File

@ -12,6 +12,7 @@ export const tagTypes = [
'Board',
'BoardImagesTotal',
'BoardAssetsTotal',
'HFTokenStatus',
'Image',
'ImageNameList',
'ImageList',