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

Author: psychedelicious

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",

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 (
     <ErrorBoundary onReset={handleReset} FallbackComponent={AppErrorBoundaryFallback}>

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

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<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>
+  );
+};

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: <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>
+    );
+  }
+};

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 = () => {
         </Button>
       </Flex>
       <Flex flexDir="column" layerStyle="second" p={4} gap={4} borderRadius="base" w="full" h="full">
+        <HFToken />
         <ModelListNavigation />
         <ModelList />
       </Flex>

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<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;

invokeai/frontend/web/src/services/api/index.ts

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