feat(ui,api): support for HF tokens in UI, handle Unauthorized and Forbidden errors

invokeai/frontend/web/src/features/modelManagerV2/hooks/useHFForbiddenToast.tsx

@@ -0,0 +1,49 @@
+import { Button, ExternalLink, Text, useToast } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
+import { useAppDispatch } from 'app/store/storeHooks';
+import { setActiveTab } from 'features/ui/store/uiSlice';
+import { atom } from 'nanostores';
+import { useCallback, useEffect } from 'react';
+import { useTranslation } from 'react-i18next';
+
+const FEATURE_ID = 'hfToken';
+const TOAST_ID = 'hfForbidden';
+
+/**
+ * Tracks whether or not the HF Login toast is showing
+ */
+export const $isHFForbiddenToastOpen = atom<{ isEnabled: boolean; source?: string }>({ isEnabled: false });
+
+export const useHFForbiddenToast = () => {
+  const { t } = useTranslation();
+  const toast = useToast();
+  const isHFForbiddenToastOpen = useStore($isHFForbiddenToastOpen);
+
+  useEffect(() => {
+    if (!isHFForbiddenToastOpen.isEnabled) {
+      toast.close(TOAST_ID);
+      return
+    }
+
+
+    if (isHFForbiddenToastOpen.isEnabled) {
+      toast({
+        id: TOAST_ID,
+        title: t('modelManager.hfForbidden'),
+        description: (
+          <Text fontSize="md">
+            {t('modelManager.hfForbiddenErrorMessage')}
+            <ExternalLink
+              label={isHFForbiddenToastOpen.source || ''}
+              href={`https://huggingface.co/${isHFForbiddenToastOpen.source}`}
+            />
+          </Text>
+        ),
+        status: 'error',
+        isClosable: true,
+        duration: null,
+        onCloseComplete: () => $isHFForbiddenToastOpen.set({ isEnabled: false }),
+      });
+    }
+  }, [isHFForbiddenToastOpen, t]);
+};

invokeai/frontend/web/src/features/modelManagerV2/hooks/useHFLoginToast.tsx

@@ -0,0 +1,93 @@
+import { Button, Text, useToast } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
+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 { atom } from 'nanostores';
+import { useCallback, useEffect } 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 TOAST_ID = 'hfTokenLogin';
+/**
+ * Tracks whether or not the HF Login toast is showing
+ */
+export const $isHFLoginToastOpen = atom<boolean>(false);
+
+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 isEnabled = useFeatureStatus(FEATURE_ID);
+  const { data } = useGetHFTokenStatusQuery(isEnabled ? undefined : skipToken);
+  const toast = useToast();
+  const isHFLoginToastOpen = useStore($isHFLoginToastOpen);
+
+  useEffect(() => {
+    if (!isHFLoginToastOpen) {
+      toast.close(TOAST_ID);
+      return
+    }
+
+    if (isHFLoginToastOpen && data) {
+      const title = getTitle(data);
+      toast({
+        id: TOAST_ID,
+        title,
+        description: <ToastDescription token_status={data} />,
+        status: 'error',
+        isClosable: true,
+        duration: null,
+        onCloseComplete: () => $isHFLoginToastOpen.set(false),
+      });
+    }
+  }, [isHFLoginToastOpen, data]);
+};
+
+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('models'));
+    toast.close(FEATURE_ID);
+  }, [dispatch, toast]);
+
+  if (token_status === 'invalid') {
+    return (
+      <Text fontSize="md">
+        {t('modelManager.hfTokenInvalidErrorMessage')} {t('modelManager.hfTokenRequired')}{' '}
+        {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/AddModelPanel/HuggingFaceFolder/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';
+import { $isHFLoginToastOpen } from '../../../hooks/useHFLoginToast';
+  
+  export const HFToken = () => {
+    const { t } = useTranslation();
+    const isHFTokenEnabled = useFeatureStatus('hfToken');
+    const [token, setToken] = useState('');
+    const [tokenStatus, setTokenStatus] = useState()
+    const { currentData } = useGetHFTokenStatusQuery(isHFTokenEnabled ? 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,
+            });
+            $isHFLoginToastOpen.set(false)
+          }
+        });
+    }, [t, toast, token, trigger]);
+  
+    const error = useMemo(() => {
+      if (!currentData || isLoading) {
+        return null;
+      }
+      if (token.length && currentData === 'invalid') {
+        return t('modelManager.hfTokenInvalidErrorMessage');
+      }
+      if (token.length && currentData === 'unknown') {
+        return t('modelManager.hfTokenUnableToVerifyErrorMessage');
+      }
+      return null;
+    }, [currentData, isLoading, t, token.length]);
+  
+    if (!currentData || currentData === 'valid') {
+      return null;
+    }
+  
+    return (
+      <Flex borderRadius="base" w="full">
+        <FormControl isInvalid={Boolean(error)} orientation="vertical">
+          <FormLabel>{t('modelManager.hfTokenLabel')}</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/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceForm.tsx

@@ -1,17 +1,22 @@
-import { Button, Flex, FormControl, FormErrorMessage, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library';
+import { Button, Divider, Flex, FormControl, FormErrorMessage, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library';
 import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
 import type { ChangeEventHandler } from 'react';
 import { memo, useCallback, useState } from 'react';
 import { useTranslation } from 'react-i18next';
-import { useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models';
+import { useGetHFTokenStatusQuery, useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models';
 
 import { HuggingFaceResults } from './HuggingFaceResults';
+import { HFToken } from './HFToken';
+import { skipToken } from '@reduxjs/toolkit/query';
+import { useFeatureStatus } from '../../../../system/hooks/useFeatureStatus';
 
 export const HuggingFaceForm = memo(() => {
   const [huggingFaceRepo, setHuggingFaceRepo] = useState('');
   const [displayResults, setDisplayResults] = useState(false);
   const [errorMessage, setErrorMessage] = useState('');
   const { t } = useTranslation();
+  const isHFTokenEnabled = useFeatureStatus('hfToken');
+  const { currentData } = useGetHFTokenStatusQuery(isHFTokenEnabled ? undefined : skipToken);
 
   const [_getHuggingFaceModels, { isLoading, data }] = useLazyGetHuggingFaceModelsQuery();
   const [installModel] = useInstallModel();
@@ -41,7 +46,7 @@ export const HuggingFaceForm = memo(() => {
   }, []);
 
   return (
-    <Flex flexDir="column" height="100%" gap={3}>
+    <Flex flexDir="column" height="100%" gap={4}>
       <FormControl isInvalid={!!errorMessage.length} w="full" orientation="vertical" flexShrink={0}>
         <FormLabel>{t('modelManager.huggingFaceRepoID')}</FormLabel>
         <Flex gap={3} alignItems="center" w="full">
@@ -63,6 +68,7 @@ export const HuggingFaceForm = memo(() => {
         <FormHelperText>{t('modelManager.huggingFaceHelper')}</FormHelperText>
         {!!errorMessage.length && <FormErrorMessage>{errorMessage}</FormErrorMessage>}
       </FormControl>
+      {currentData !== 'valid' && <HFToken />}
       {data && data.urls && displayResults && <HuggingFaceResults results={data.urls} />}
     </Flex>
   );

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManager.tsx

@@ -7,6 +7,8 @@ import { PiPlusBold } from 'react-icons/pi';
 
 import ModelList from './ModelManagerPanel/ModelList';
 import { ModelListNavigation } from './ModelManagerPanel/ModelListNavigation';
+import { useHFLoginToast } from '../hooks/useHFLoginToast';
+import { useHFForbiddenToast } from '../hooks/useHFForbiddenToast';
 
 export const ModelManager = memo(() => {
   const { t } = useTranslation();
@@ -16,6 +18,9 @@ export const ModelManager = memo(() => {
   }, [dispatch]);
   const selectedModelKey = useAppSelector(selectSelectedModelKey);
 
+  useHFLoginToast();
+  useHFForbiddenToast()
+
   return (
     <Flex flexDir="column" layerStyle="first" p={4} gap={4} borderRadius="base" w="50%" h="full">
       <Flex w="full" gap={4} justifyContent="space-between" alignItems="center">