tidy(ui): prefer types from zod schemas for model attrs

Author: psychedelicious

invokeai/frontend/web/src/features/metadata/parsing.tsx

@@ -49,7 +49,7 @@ import {
   zVideoDuration,
   zVideoResolution,
 } from 'features/controlLayers/store/types';
-import type { ModelIdentifierField } from 'features/nodes/types/common';
+import type { ModelIdentifierField, ModelType } from 'features/nodes/types/common';
 import { zModelIdentifierField } from 'features/nodes/types/common';
 import { zModelIdentifier } from 'features/nodes/types/v2/common';
 import { modelSelected } from 'features/parameters/store/actions';
@@ -108,7 +108,7 @@ import { useCallback, useEffect, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { imagesApi } from 'services/api/endpoints/images';
 import { modelsApi } from 'services/api/endpoints/models';
-import type { AnyModelConfig, ModelType } from 'services/api/types';
+import type { AnyModelConfig } from 'services/api/types';
 import { assert } from 'tsafe';
 import z from 'zod';
 

invokeai/frontend/web/src/features/modelManagerV2/subpanels/ModelManagerPanel/ModelBaseBadge.tsx

@@ -1,7 +1,7 @@
 import { Badge } from '@invoke-ai/ui-library';
 import { MODEL_BASE_TO_COLOR, MODEL_BASE_TO_SHORT_NAME } from 'features/modelManagerV2/models';
+import type { BaseModelType } from 'features/nodes/types/common';
 import { memo } from 'react';
-import type { BaseModelType } from 'services/api/types';
 
 type Props = {
   base: BaseModelType;

invokeai/frontend/web/src/features/nodes/types/common.test-d.ts

@@ -1,5 +1,4 @@
 import type {
-  BaseModelType,
   BoardField,
   Classification,
   ColorField,
@@ -10,11 +9,12 @@ import type {
   ModelIdentifierField,
   ProgressImage,
   SchedulerField,
-  SubModelType,
   T2IAdapterField,
+  zBaseModelType,
   zClipVariantType,
   zModelFormat,
   zModelVariantType,
+  zSubModelType,
 } from 'features/nodes/types/common';
 import type { Invocation, S } from 'services/api/types';
 import type { Equals, Extends } from 'tsafe';
@@ -24,7 +24,8 @@ import type z from 'zod';
 
 /**
  * These types originate from the server and are recreated as zod schemas manually, for use at runtime.
- * The tests ensure that the types are correctly recreated.
+ * The tests ensure that the types are correctly recreated. If one of these tests fails, it means the zod
+ * schema and the type have diverged and need to be reconciled - update the zod schema.
  */
 
 describe('Common types', () => {
@@ -40,8 +41,8 @@ describe('Common types', () => {
 
   // Model component types
   test('ModelIdentifier', () => assert<Equals<ModelIdentifierField, S['ModelIdentifierField']>>());
-  test('ModelIdentifier', () => assert<Equals<BaseModelType, S['BaseModelType']>>());
-  test('ModelIdentifier', () => assert<Equals<SubModelType, S['SubModelType']>>());
+  test('ModelIdentifier', () => assert<Equals<z.infer<typeof zBaseModelType>, S['BaseModelType']>>());
+  test('ModelIdentifier', () => assert<Equals<z.infer<typeof zSubModelType>, S['SubModelType']>>());
   test('ClipVariantType', () => assert<Equals<z.infer<typeof zClipVariantType>, S['ClipVariantType']>>());
   test('ModelVariantType', () => assert<Equals<z.infer<typeof zModelVariantType>, S['ModelVariantType']>>());
   test('ModelFormat', () => assert<Equals<z.infer<typeof zModelFormat>, S['ModelFormat']>>());

invokeai/frontend/web/src/features/nodes/types/common.ts

@@ -130,7 +130,7 @@ export const zModelType = z.enum([
   'unknown',
 ]);
 export type ModelType = z.infer<typeof zModelType>;
-const zSubModelType = z.enum([
+export const zSubModelType = z.enum([
   'unet',
   'transformer',
   'text_encoder',

invokeai/frontend/web/src/features/nodes/util/graph/buildLinearBatchConfig.ts

@@ -3,9 +3,10 @@ import { generateSeeds } from 'common/util/generateSeeds';
 import { range } from 'es-toolkit/compat';
 import type { SeedBehaviour } from 'features/dynamicPrompts/store/dynamicPromptsSlice';
 import { API_BASE_MODELS, VIDEO_BASE_MODELS } from 'features/modelManagerV2/models';
+import type { BaseModelType } from 'features/nodes/types/common';
 import type { Graph } from 'features/nodes/util/graph/generation/Graph';
 import type { components } from 'services/api/schema';
-import type { BaseModelType, Batch, EnqueueBatchArg, Invocation } from 'services/api/types';
+import type { Batch, EnqueueBatchArg, Invocation } from 'services/api/types';
 
 const getExtendedPrompts = (arg: {
   seedBehaviour: SeedBehaviour;

invokeai/frontend/web/src/features/parameters/components/ModelPicker.tsx

@@ -32,6 +32,7 @@ import {
 } from 'features/modelManagerV2/models';
 import { setInstallModelsTabByName } from 'features/modelManagerV2/store/installModelsStore';
 import ModelImage from 'features/modelManagerV2/subpanels/ModelManagerPanel/ModelImage';
+import type { BaseModelType } from 'features/nodes/types/common';
 import { NavigateToModelManagerButton } from 'features/parameters/components/MainModel/NavigateToModelManagerButton';
 import { useFeatureStatus } from 'features/system/hooks/useFeatureStatus';
 import { selectIsModelsTabDisabled } from 'features/system/store/configSlice';
@@ -41,7 +42,7 @@ import { memo, useCallback, useMemo, useRef } from 'react';
 import { Trans, useTranslation } from 'react-i18next';
 import { PiCaretDownBold, PiLinkSimple } from 'react-icons/pi';
 import { useGetRelatedModelIdsBatchQuery } from 'services/api/endpoints/modelRelationships';
-import type { AnyModelConfig, BaseModelType } from 'services/api/types';
+import type { AnyModelConfig } from 'services/api/types';
 
 const selectSelectedModelKeys = createMemoizedSelector(selectParamsSlice, selectLoRAsSlice, (params, loras) => {
   const keys: string[] = [];

invokeai/frontend/web/src/features/parameters/types/constants.ts

@@ -1,5 +1,5 @@
 import type { ComboboxOption } from '@invoke-ai/ui-library';
-import type { BaseModelType } from 'services/api/types';
+import type { BaseModelType } from 'features/nodes/types/common';
 
 /**
  * Mapping of base model to CLIP skip parameter constraints

invokeai/frontend/web/src/features/parameters/util/optimalDimension.ts

@@ -1,4 +1,4 @@
-import type { BaseModelType } from 'services/api/types';
+import type { BaseModelType } from 'features/nodes/types/common';
 
 /**
  * Gets the optimal dimension for a given base model:

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

@@ -105,16 +105,9 @@ export const isVideoDTO = (dto: ImageDTO | VideoDTO): dto is VideoDTO => {
   return 'video_id' in dto;
 };
 
-// Models
-export type ModelType = S['ModelType'];
-export type BaseModelType = S['BaseModelType'];
-
 // Model Configs
-
 export type ControlLoRAModelConfig = S['ControlLoRALyCORISConfig'] | S['ControlLoRADiffusersConfig'];
-// TODO(MM2): Can we make key required in the pydantic model?
 export type LoRAModelConfig = S['LoRADiffusersConfig'] | S['LoRALyCORISConfig'] | S['LoRAOmiConfig'];
-// TODO(MM2): Can we rename this from Vae -> VAE
 export type VAEModelConfig = S['VAECheckpointConfig'] | S['VAEDiffusersConfig'];
 export type ControlNetModelConfig = S['ControlNetDiffusersConfig'] | S['ControlNetCheckpointConfig'];
 export type IPAdapterModelConfig = S['IPAdapterInvokeAIConfig'] | S['IPAdapterCheckpointConfig'];