Skip to content

Commit ee40026

Browse files
psychedeliciouspsychedelicious
committed
feat(ui): add UI to reset hf token
1 parent fd20582 commit ee40026

File tree

5 files changed

+111
-51
lines changed

5 files changed

+111
-51
lines changed

‎invokeai/frontend/web/public/locales/en.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -798,6 +798,7 @@
798798
"hfTokenUnableToVerify": "Unable to Verify HF Token",
799799
"hfTokenUnableToVerifyErrorMessage": "Unable to verify HuggingFace token. This is likely due to a network error. Please try again later.",
800800
"hfTokenSaved": "HF Token Saved",
801+
"hfTokenReset": "HF Token Reset",
801802
"urlUnauthorizedErrorMessage": "You may need to configure an API token to access this model.",
802803
"urlUnauthorizedErrorMessage2": "Learn how here.",
803804
"imageEncoderModelId": "Image Encoder Model ID",

‎invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HFToken.tsx

Lines changed: 93 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -12,63 +12,45 @@ import { skipToken } from '@reduxjs/toolkit/query';
1212
import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
1313
import { toast } from 'features/toast/toast';
1414
import type { ChangeEvent } from 'react';
15-
import { useCallback, useMemo, useState } from 'react';
15+
import { memo, useCallback, useMemo, useState } from 'react';
1616
import { useTranslation } from 'react-i18next';
17-
import { useGetHFTokenStatusQuery, useSetHFTokenMutation } from 'services/api/endpoints/models';
18-
import { UNAUTHORIZED_TOAST_ID } from 'services/events/onModelInstallError';
17+
import {
18+
useGetHFTokenStatusQuery,
19+
useResetHFTokenMutation,
20+
useSetHFTokenMutation,
21+
} from 'services/api/endpoints/models';
22+
import type { Equals } from 'tsafe';
23+
import { assert } from 'tsafe';
1924

2025
export const HFToken = () => {
2126
const { t } = useTranslation();
2227
const isHFTokenEnabled = useFeatureStatus('hfToken');
23-
const [token, setToken] = useState('');
2428
const { currentData } = useGetHFTokenStatusQuery(isHFTokenEnabled ? undefined : skipToken);
25-
const [trigger, { isLoading, isUninitialized }] = useSetHFTokenMutation();
26-
const onChange = useCallback((e: ChangeEvent<HTMLInputElement>) => {
27-
setToken(e.target.value);
28-
}, []);
29-
const onClick = useCallback(() => {
30-
trigger({ token })
31-
.unwrap()
32-
.then((res) => {
33-
if (res === 'valid') {
34-
setToken('');
35-
toast({
36-
id: UNAUTHORIZED_TOAST_ID,
37-
title: t('modelManager.hfTokenSaved'),
38-
status: 'success',
39-
duration: 3000,
40-
});
41-
}
42-
});
43-
}, [t, token, trigger]);
4429

4530
const error = useMemo(() => {
46-
if (!currentData || isUninitialized || isLoading) {
47-
return null;
48-
}
49-
if (currentData === 'invalid') {
50-
return t('modelManager.hfTokenInvalidErrorMessage');
51-
}
52-
if (currentData === 'unknown') {
53-
return t('modelManager.hfTokenUnableToVerifyErrorMessage');
31+
switch (currentData) {
32+
case 'invalid':
33+
return t('modelManager.hfTokenInvalidErrorMessage');
34+
case 'unknown':
35+
return t('modelManager.hfTokenUnableToVerifyErrorMessage');
36+
case 'valid':
37+
case undefined:
38+
return null;
39+
default:
40+
assert<Equals<never, typeof currentData>>(false, 'Unexpected HF token status');
5441
}
55-
return null;
56-
}, [currentData, isLoading, isUninitialized, t]);
42+
}, [currentData, t]);
5743

58-
if (!currentData || currentData === 'valid') {
44+
if (!currentData) {
5945
return null;
6046
}
6147

6248
return (
6349
<Flex borderRadius="base" w="full">
64-
<FormControl isInvalid={!isUninitialized && Boolean(error)} orientation="vertical">
50+
<FormControl isInvalid={Boolean(error)} orientation="vertical">
6551
<FormLabel>{t('modelManager.hfTokenLabel')}</FormLabel>
66-
<Flex gap={3} alignItems="center" w="full">
67-
<Input type="password" value={token} onChange={onChange} />
68-
<Button onClick={onClick} size="sm" isDisabled={token.trim().length === 0} isLoading={isLoading}>
69-
{t('common.save')}
70-
</Button>
71-
</Flex>
52+
{error && <SetHFTokenInput />}
53+
{!error && <ResetHFTokenButton />}
7254
<FormHelperText>
7355
<ExternalLink label={t('modelManager.hfTokenHelperText')} href="https://huggingface.co/settings/tokens" />
7456
</FormHelperText>
@@ -77,3 +59,73 @@ export const HFToken = () => {
7759
</Flex>
7860
);
7961
};
62+
63+
const PLACEHOLDER_TOKEN = Array.from({ length: 37 }, () => 'a').join('');
64+
65+
const ResetHFTokenButton = memo(() => {
66+
const { t } = useTranslation();
67+
const [resetHFToken, { isLoading }] = useResetHFTokenMutation();
68+
69+
const onClick = useCallback(() => {
70+
resetHFToken()
71+
.unwrap()
72+
.then(() => {
73+
toast({
74+
title: t('modelManager.hfTokenReset'),
75+
status: 'info',
76+
});
77+
});
78+
}, [resetHFToken, t]);
79+
80+
return (
81+
<Flex gap={3} alignItems="center" w="full">
82+
<Input type="password" value={PLACEHOLDER_TOKEN} isDisabled />
83+
<Button onClick={onClick} size="sm" isLoading={isLoading}>
84+
{t('common.reset')}
85+
</Button>
86+
</Flex>
87+
);
88+
});
89+
ResetHFTokenButton.displayName = 'ResetHFTokenButton';
90+
91+
const SetHFTokenInput = memo(() => {
92+
const { t } = useTranslation();
93+
const [token, setToken] = useState('');
94+
const [trigger, { isLoading }] = useSetHFTokenMutation();
95+
const onChange = useCallback((e: ChangeEvent<HTMLInputElement>) => {
96+
setToken(e.target.value);
97+
}, []);
98+
const onClick = useCallback(() => {
99+
trigger({ token })
100+
.unwrap()
101+
.then((res) => {
102+
switch (res) {
103+
case 'valid':
104+
setToken('');
105+
toast({
106+
title: t('modelManager.hfTokenSaved'),
107+
status: 'success',
108+
});
109+
break;
110+
case 'invalid':
111+
case 'unknown':
112+
default:
113+
toast({
114+
title: t('modelManager.hfTokenUnableToVerify'),
115+
status: 'error',
116+
});
117+
break;
118+
}
119+
});
120+
}, [t, token, trigger]);
121+
122+
return (
123+
<Flex gap={3} alignItems="center" w="full">
124+
<Input type="password" value={token} onChange={onChange} />
125+
<Button onClick={onClick} size="sm" isDisabled={token.trim().length === 0} isLoading={isLoading}>
126+
{t('common.save')}
127+
</Button>
128+
</Flex>
129+
);
130+
});
131+
SetHFTokenInput.displayName = 'SetHFTokenInput';

‎invokeai/frontend/web/src/features/modelManagerV2/subpanels/AddModelPanel/HuggingFaceFolder/HuggingFaceForm.tsx

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11
import { Button, Flex, FormControl, FormErrorMessage, FormHelperText, FormLabel, Input } from '@invoke-ai/ui-library';
2-
import { skipToken } from '@reduxjs/toolkit/query';
32
import { useInstallModel } from 'features/modelManagerV2/hooks/useInstallModel';
43
import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
54
import type { ChangeEventHandler } from 'react';
65
import { memo, useCallback, useState } from 'react';
76
import { useTranslation } from 'react-i18next';
8-
import { useGetHFTokenStatusQuery, useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models';
7+
import { useLazyGetHuggingFaceModelsQuery } from 'services/api/endpoints/models';
98

109
import { HFToken } from './HFToken';
1110
import { HuggingFaceResults } from './HuggingFaceResults';
@@ -16,7 +15,6 @@ export const HuggingFaceForm = memo(() => {
1615
const [errorMessage, setErrorMessage] = useState('');
1716
const { t } = useTranslation();
1817
const isHFTokenEnabled = useFeatureStatus('hfToken');
19-
const { currentData } = useGetHFTokenStatusQuery(isHFTokenEnabled ? undefined : skipToken);
2018

2119
const [_getHuggingFaceModels, { isLoading, data }] = useLazyGetHuggingFaceModelsQuery();
2220
const [installModel] = useInstallModel();
@@ -68,7 +66,7 @@ export const HuggingFaceForm = memo(() => {
6866
<FormHelperText>{t('modelManager.huggingFaceHelper')}</FormHelperText>
6967
{!!errorMessage.length && <FormErrorMessage>{errorMessage}</FormErrorMessage>}
7068
</FormControl>
71-
{currentData !== 'valid' && <HFToken />}
69+
{isHFTokenEnabled && <HFToken />}
7270
{data && data.urls && displayResults && <HuggingFaceResults results={data.urls} />}
7371
</Flex>
7472
);

‎invokeai/frontend/web/src/services/api/endpoints/models.ts

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,14 @@ import { getSelectorsOptions } from 'app/store/createMemoizedSelector';
44
import { selectParamsSlice } from 'features/controlLayers/store/paramsSlice';
55
import queryString from 'query-string';
66
import type { operations, paths } from 'services/api/schema';
7-
import {
8-
type AnyModelConfig,
9-
type GetHFTokenStatusResponse,
10-
isNonRefinerMainModelConfig,
11-
type SetHFTokenArg,
12-
type SetHFTokenResponse,
7+
import type {
8+
AnyModelConfig,
9+
GetHFTokenStatusResponse,
10+
ResetHFTokenResponse,
11+
SetHFTokenArg,
12+
SetHFTokenResponse,
1313
} from 'services/api/types';
14+
import { isNonRefinerMainModelConfig } from 'services/api/types';
1415
import type { Param0 } from 'tsafe';
1516

1617
import type { ApiTagDescription } from '..';
@@ -293,6 +294,10 @@ export const modelsApi = api.injectEndpoints({
293294
}
294295
},
295296
}),
297+
resetHFToken: build.mutation<ResetHFTokenResponse, void>({
298+
query: () => ({ url: buildModelsUrl('hf_login'), method: 'DELETE' }),
299+
invalidatesTags: ['HFTokenStatus'],
300+
}),
296301
emptyModelCache: build.mutation<void, void>({
297302
query: () => ({ url: buildModelsUrl('empty_model_cache'), method: 'POST' }),
298303
}),
@@ -316,6 +321,7 @@ export const {
316321
useGetStarterModelsQuery,
317322
useGetHFTokenStatusQuery,
318323
useSetHFTokenMutation,
324+
useResetHFTokenMutation,
319325
useEmptyModelCacheMutation,
320326
} = modelsApi;
321327

‎invokeai/frontend/web/src/services/api/types.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,9 @@ export type GetHFTokenStatusResponse =
321321
export type SetHFTokenResponse = NonNullable<
322322
paths['/api/v2/models/hf_login']['post']['responses']['200']['content']['application/json']
323323
>;
324+
export type ResetHFTokenResponse = NonNullable<
325+
paths['/api/v2/models/hf_login']['delete']['responses']['200']['content']['application/json']
326+
>;
324327
export type SetHFTokenArg = NonNullable<
325328
paths['/api/v2/models/hf_login']['post']['requestBody']['content']['application/json']
326329
>;

0 commit comments

Comments
 (0)