feat(ui): add button to reidentify model to mm

Author: psychedelicious

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

@@ -844,6 +844,11 @@
         "clipLEmbed": "CLIP-L Embed",
         "clipGEmbed": "CLIP-G Embed",
         "config": "Config",
+        "reidentify": "Reidentify",
+        "reidentifyTooltip": "If a model didn't install correctly (e.g. it has the wrong type or doesn't work), you can try reidentifying it. This will reset any custom settings you may have applied.",
+        "reidentifySuccess": "Model reidentified successfully",
+        "reidentifyUnknown": "Unable to identify model",
+        "reidentifyError": "Error reidentifying model",
         "convert": "Convert",
         "convertingModelBegin": "Converting Model. Please wait.",
         "convertToDiffusers": "Convert To Diffusers",

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelReidentifyButton.tsx

@@ -0,0 +1,58 @@
+import { Button } from '@invoke-ai/ui-library';
+import { toast } from 'features/toast/toast';
+import { memo, useCallback } from 'react';
+import { useTranslation } from 'react-i18next';
+import { PiSparkleFill } from 'react-icons/pi';
+import { useReidentifyModelMutation } from 'services/api/endpoints/models';
+import type { AnyModelConfig } from 'services/api/types';
+
+interface Props {
+  modelConfig: AnyModelConfig;
+}
+
+export const ModelReidentifyButton = memo(({ modelConfig }: Props) => {
+  const { t } = useTranslation();
+  const [reidentifyModel, { isLoading }] = useReidentifyModelMutation();
+
+  const onClick = useCallback(() => {
+    reidentifyModel({ key: modelConfig.key })
+      .unwrap()
+      .then(({ type }) => {
+        if (type === 'unknown') {
+          toast({
+            id: 'MODEL_REIDENTIFY_UNKNOWN',
+            title: t('modelManager.reidentifyUnknown'),
+            status: 'warning',
+          });
+        }
+        toast({
+          id: 'MODEL_REIDENTIFY_SUCCESS',
+          title: t('modelManager.reidentifySuccess'),
+          status: 'success',
+        });
+      })
+      .catch((_) => {
+        toast({
+          id: 'MODEL_REIDENTIFY_ERROR',
+          title: t('modelManager.reidentifyError'),
+          status: 'error',
+        });
+      });
+  }, [modelConfig.key, reidentifyModel, t]);
+
+  return (
+    <Button
+      onClick={onClick}
+      size="sm"
+      aria-label={t('modelManager.reidentifyTooltip')}
+      tooltip={t('modelManager.reidentifyTooltip')}
+      isLoading={isLoading}
+      flexShrink={0}
+      leftIcon={<PiSparkleFill />}
+    >
+      {t('modelManager.reidentify')}
+    </Button>
+  );
+});
+
+ModelReidentifyButton.displayName = 'ModelReidentifyButton';

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelPanel/ModelView.tsx

@@ -13,6 +13,7 @@ import type { AnyModelConfig } from 'services/api/types';
 import { MainModelDefaultSettings } from './MainModelDefaultSettings/MainModelDefaultSettings';
 import { ModelAttrView } from './ModelAttrView';
 import { ModelFooter } from './ModelFooter';
+import { ModelReidentifyButton } from './ModelReidentifyButton';
 import { RelatedModels } from './RelatedModels';
 
 type Props = {
@@ -21,6 +22,7 @@ type Props = {
 
 export const ModelView = memo(({ modelConfig }: Props) => {
   const { t } = useTranslation();
+
   const withSettings = useMemo(() => {
     if (modelConfig.type === 'main' && modelConfig.base !== 'sdxl-refiner') {
       return true;
@@ -46,6 +48,7 @@ export const ModelView = memo(({ modelConfig }: Props) => {
           <ModelConvertButton modelConfig={modelConfig} />
         )}
         <ModelEditButton />
+        <ModelReidentifyButton modelConfig={modelConfig} />
       </ModelHeader>
       <Divider />
       <Flex flexDir="column" gap={4}>

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

@@ -299,6 +299,40 @@ export const modelsApi = api.injectEndpoints({
     emptyModelCache: build.mutation<void, void>({
       query: () => ({ url: buildModelsUrl('empty_model_cache'), method: 'POST' }),
     }),
+    reidentifyModel: build.mutation<
+      paths['/api/v2/models/i/{key}/reidentify']['post']['responses']['200']['content']['application/json'],
+      { key: string }
+    >({
+      query: ({ key }) => {
+        return {
+          url: buildModelsUrl(`i/${key}/reidentify`),
+          method: 'POST',
+        };
+      },
+      onQueryStarted: async (_, { dispatch, queryFulfilled }) => {
+        try {
+          const { data } = await queryFulfilled;
+
+          // Update the individual model query caches
+          dispatch(modelsApi.util.upsertQueryData('getModelConfig', data.key, data));
+
+          const { base, name, type } = data;
+          dispatch(modelsApi.util.upsertQueryData('getModelConfigByAttrs', { base, name, type }, data));
+
+          // Update the list query cache
+          dispatch(
+            modelsApi.util.updateQueryData('getModelConfigs', undefined, (draft) => {
+              modelConfigsAdapter.updateOne(draft, {
+                id: data.key,
+                changes: data,
+              });
+            })
+          );
+        } catch {
+          // no-op
+        }
+      },
+    }),
   }),
 });
 
@@ -321,6 +355,7 @@ export const {
   useSetHFTokenMutation,
   useResetHFTokenMutation,
   useEmptyModelCacheMutation,
+  useReidentifyModelMutation,
 } = modelsApi;
 
 export const selectModelConfigsQuery = modelsApi.endpoints.getModelConfigs.select();