From bdb52cfcf7fdbe6cda7ca5fee44cd909a0f41eed Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Wed, 20 Mar 2024 09:24:10 +1100 Subject: [PATCH] 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 --- invokeai/frontend/web/public/locales/en.json | 8 ++ .../frontend/web/src/app/components/App.tsx | 2 + .../frontend/web/src/app/types/invokeai.ts | 3 +- .../modelManagerV2/components/HFToken.tsx | 81 +++++++++++++++++ .../modelManagerV2/hooks/useHFLoginToast.tsx | 89 +++++++++++++++++++ .../modelManagerV2/subpanels/ModelManager.tsx | 2 + .../web/src/services/api/endpoints/models.ts | 27 ++++++ .../frontend/web/src/services/api/index.ts | 1 + 8 files changed, 212 insertions(+), 1 deletion(-) create mode 100644 invokeai/frontend/web/src/features/modelManagerV2/components/HFToken.tsx create mode 100644 invokeai/frontend/web/src/features/modelManagerV2/hooks/useHFLoginToast.tsx diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 0cfaf912f0..733f558819 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -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", diff --git a/invokeai/frontend/web/src/app/components/App.tsx b/invokeai/frontend/web/src/app/components/App.tsx index ae1a762f61..13231e6aeb 100644 --- a/invokeai/frontend/web/src/app/components/App.tsx +++ b/invokeai/frontend/web/src/app/components/App.tsx @@ -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 ( diff --git a/invokeai/frontend/web/src/app/types/invokeai.ts b/invokeai/frontend/web/src/app/types/invokeai.ts index 1e40b7074d..4982dbb83f 100644 --- a/invokeai/frontend/web/src/app/types/invokeai.ts +++ b/invokeai/frontend/web/src/app/types/invokeai.ts @@ -25,7 +25,8 @@ export type AppFeature = | 'prependQueue' | 'invocationCache' | 'bulkDownload' - | 'starterModels'; + | 'starterModels' + | 'hfToken'; /** * A disable-able Stable Diffusion feature diff --git a/invokeai/frontend/web/src/features/modelManagerV2/components/HFToken.tsx b/invokeai/frontend/web/src/features/modelManagerV2/components/HFToken.tsx new file mode 100644 index 0000000000..86ca7e128b --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/components/HFToken.tsx @@ -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) => { + 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 ( + + + {t('modelManager.hfToken')} + + + + + + + + {error} + + + ); +}; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/hooks/useHFLoginToast.tsx b/invokeai/frontend/web/src/features/modelManagerV2/hooks/useHFLoginToast.tsx new file mode 100644 index 0000000000..972dfaccb3 --- /dev/null +++ b/invokeai/frontend/web/src/features/modelManagerV2/hooks/useHFLoginToast.tsx @@ -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: , + 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 ( + + {t('modelManager.hfTokenInvalidErrorMessage')}{' '} + {t('modelManager.hfTokenInvalidErrorMessage2')} + + + ); + } + + if (token_status === 'unknown') { + return ( + + {t('modelManager.hfTokenUnableToErrorMessage')}{' '} + + + ); + } +}; diff --git a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManager.tsx b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManager.tsx index ed75f86078..51c2ab0f7b 100644 --- a/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManager.tsx +++ b/invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManager.tsx @@ -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 = () => { + diff --git a/invokeai/frontend/web/src/services/api/endpoints/models.ts b/invokeai/frontend/web/src/services/api/endpoints/models.ts index 520208bf5b..a4c756c1c6 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/models.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/models.ts @@ -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({ query: () => buildModelsUrl('starter_models'), }), + getHFTokenStatus: build.query({ + query: () => buildModelsUrl('hf_login'), + providesTags: ['HFTokenStatus'], + }), + setHFToken: build.mutation({ + 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; diff --git a/invokeai/frontend/web/src/services/api/index.ts b/invokeai/frontend/web/src/services/api/index.ts index 63b34a08ad..32ae2319d2 100644 --- a/invokeai/frontend/web/src/services/api/index.ts +++ b/invokeai/frontend/web/src/services/api/index.ts @@ -12,6 +12,7 @@ export const tagTypes = [ 'Board', 'BoardImagesTotal', 'BoardAssetsTotal', + 'HFTokenStatus', 'Image', 'ImageNameList', 'ImageList',