refactor(ui): move model categorisation-ish logic to central location, simplify model manager models list

Author: psychedelicious

invokeai/backend/model_manager/config.py

@@ -137,7 +137,7 @@ class ModelConfigBase(ABC, BaseModel):
 
     @staticmethod
     def json_schema_extra(schema: dict[str, Any]) -> None:
-        schema["required"].extend(["key", "type", "format"])
+        schema["required"].extend(["key", "base", "type", "format"])
 
     model_config = ConfigDict(validate_assignment=True, json_schema_extra=json_schema_extra)
 
@@ -172,7 +172,8 @@ def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         if issubclass(cls, LegacyProbeMixin):
             ModelConfigBase.USING_LEGACY_PROBE.add(cls)
-        elif cls is not UnknownModelConfig:
+        # Cannot use `elif isinstance(cls, UnknownModelConfig)` because UnknownModelConfig is not defined yet
+        elif cls.__name__ != "UnknownModelConfig":
             ModelConfigBase.USING_CLASSIFY_API.add(cls)
 
     @staticmethod
@@ -274,16 +275,17 @@ def from_model_on_disk(cls, mod: ModelOnDisk, **overrides):
 
 
 class UnknownModelConfig(ModelConfigBase):
+    base: Literal[BaseModelType.Any] = BaseModelType.Any
     type: Literal[ModelType.Unknown] = ModelType.Unknown
     format: Literal[ModelFormat.Unknown] = ModelFormat.Unknown
 
     @classmethod
-    def matches(cls, *args, **kwargs) -> bool:
-        raise NotImplementedError("UnknownModelConfig cannot match anything")
+    def matches(cls, mod: ModelOnDisk) -> bool:
+        return False
 
     @classmethod
-    def parse(cls, *args, **kwargs) -> dict[str, Any]:
-        raise NotImplementedError("UnknownModelConfig cannot parse anything")
+    def parse(cls, mod: ModelOnDisk) -> dict[str, Any]:
+        return {}
 
 
 class CheckpointConfigBase(ABC, BaseModel):

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts

@@ -11,8 +11,8 @@ import {
   selectCanvasSlice,
 } from 'features/controlLayers/store/selectors';
 import { getEntityIdentifier } from 'features/controlLayers/store/types';
+import { SUPPORTS_REF_IMAGES_BASE_MODELS } from 'features/modelManagerV2/models';
 import { modelSelected } from 'features/parameters/store/actions';
-import { SUPPORTS_REF_IMAGES_BASE_MODELS } from 'features/parameters/types/constants';
 import { zParameterModel } from 'features/parameters/types/parameterSchemas';
 import { toast } from 'features/toast/toast';
 import { t } from 'i18next';

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelsLoaded.ts

@@ -37,7 +37,7 @@ import type { Logger } from 'roarr';
 import { modelConfigsAdapterSelectors, modelsApi } from 'services/api/endpoints/models';
 import type { AnyModelConfig } from 'services/api/types';
 import {
-  isCLIPEmbedModelConfig,
+  isCLIPEmbedModelConfigOrSubmodel,
   isControlLayerModelConfig,
   isControlNetModelConfig,
   isFluxReduxModelConfig,
@@ -48,7 +48,7 @@ import {
   isNonRefinerMainModelConfig,
   isRefinerMainModelModelConfig,
   isSpandrelImageToImageModelConfig,
-  isT5EncoderModelConfig,
+  isT5EncoderModelConfigOrSubmodel,
   isVideoModelConfig,
 } from 'services/api/types';
 import type { JsonObject } from 'type-fest';
@@ -418,7 +418,7 @@ const handleTileControlNetModel: ModelHandler = (models, state, dispatch, log) =
 
 const handleT5EncoderModels: ModelHandler = (models, state, dispatch, log) => {
   const selectedT5EncoderModel = state.params.t5EncoderModel;
-  const t5EncoderModels = models.filter((m) => isT5EncoderModelConfig(m));
+  const t5EncoderModels = models.filter((m) => isT5EncoderModelConfigOrSubmodel(m));
 
   // If the currently selected model is available, we don't need to do anything
   if (selectedT5EncoderModel && t5EncoderModels.some((m) => m.key === selectedT5EncoderModel.key)) {
@@ -446,7 +446,7 @@ const handleT5EncoderModels: ModelHandler = (models, state, dispatch, log) => {
 
 const handleCLIPEmbedModels: ModelHandler = (models, state, dispatch, log) => {
   const selectedCLIPEmbedModel = state.params.clipEmbedModel;
-  const CLIPEmbedModels = models.filter((m) => isCLIPEmbedModelConfig(m));
+  const CLIPEmbedModels = models.filter((m) => isCLIPEmbedModelConfigOrSubmodel(m));
 
   // If the currently selected model is available, we don't need to do anything
   if (selectedCLIPEmbedModel && CLIPEmbedModels.some((m) => m.key === selectedCLIPEmbedModel.key)) {

invokeai/frontend/web/src/features/controlLayers/konva/CanvasEntity/CanvasEntityFilterer.ts

@@ -17,7 +17,7 @@ import Konva from 'konva';
 import { atom, computed } from 'nanostores';
 import type { Logger } from 'roarr';
 import { serializeError } from 'serialize-error';
-import { buildSelectModelConfig } from 'services/api/hooks/modelsByType';
+import { modelConfigsAdapterSelectors, selectModelConfigsQuery } from 'services/api/endpoints/models';
 import { isControlLayerModelConfig } from 'services/api/types';
 import stableHash from 'stable-hash';
 import type { Equals } from 'tsafe';
@@ -202,11 +202,19 @@ export class CanvasEntityFilterer extends CanvasModuleBase {
   createInitialFilterConfig = (): FilterConfig => {
     if (this.parent.type === 'control_layer_adapter' && this.parent.state.controlAdapter.model) {
       // If the parent is a control layer adapter, we should check if the model has a default filter and set it if so
-      const selectModelConfig = buildSelectModelConfig(
-        this.parent.state.controlAdapter.model.key,
-        isControlLayerModelConfig
-      );
-      const modelConfig = this.manager.stateApi.runSelector(selectModelConfig);
+      const key = this.parent.state.controlAdapter.model.key;
+      const modelConfig = this.manager.stateApi.runSelector((state) => {
+        const { data } = selectModelConfigsQuery(state);
+        if (!data) {
+          return null;
+        }
+        return (
+          modelConfigsAdapterSelectors
+            .selectAll(data)
+            .filter(isControlLayerModelConfig)
+            .find((m) => m.key === key) ?? null
+        );
+      });
       // This always returns a filter
       const filter = getFilterForModel(modelConfig) ?? IMAGE_FILTERS.canny_edge_detection;
       return filter.buildDefaults();

invokeai/frontend/web/src/features/controlLayers/konva/CanvasTool/CanvasBboxToolModule.ts

@@ -13,8 +13,8 @@ import { selectBboxOverlay } from 'features/controlLayers/store/canvasSettingsSl
 import { selectModel } from 'features/controlLayers/store/paramsSlice';
 import { selectBbox } from 'features/controlLayers/store/selectors';
 import type { Coordinate, Rect, Tool } from 'features/controlLayers/store/types';
+import { API_BASE_MODELS } from 'features/modelManagerV2/models';
 import type { ModelIdentifierField } from 'features/nodes/types/common';
-import { API_BASE_MODELS } from 'features/parameters/types/constants';
 import Konva from 'konva';
 import { atom } from 'nanostores';
 import type { Logger } from 'roarr';

invokeai/frontend/web/src/features/controlLayers/store/canvasSlice.ts

@@ -35,8 +35,8 @@ import {
   getScaledBoundingBoxDimensions,
 } from 'features/controlLayers/util/getScaledBoundingBoxDimensions';
 import { simplifyFlatNumbersArray } from 'features/controlLayers/util/simplify';
+import { API_BASE_MODELS } from 'features/modelManagerV2/models';
 import { isMainModelBase, zModelIdentifierField } from 'features/nodes/types/common';
-import { API_BASE_MODELS } from 'features/parameters/types/constants';
 import { getGridSize, getIsSizeOptimal, getOptimalDimension } from 'features/parameters/util/optimalDimension';
 import type { IRect } from 'konva/lib/types';
 import type { UndoableOptions } from 'redux-undo';

invokeai/frontend/web/src/features/controlLayers/store/paramsSlice.ts

@@ -25,14 +25,14 @@ import {
 import { calculateNewSize } from 'features/controlLayers/util/getScaledBoundingBoxDimensions';
 import {
   API_BASE_MODELS,
-  CLIP_SKIP_MAP,
   SUPPORTS_ASPECT_RATIO_BASE_MODELS,
   SUPPORTS_NEGATIVE_PROMPT_BASE_MODELS,
   SUPPORTS_OPTIMIZED_DENOISING_BASE_MODELS,
   SUPPORTS_PIXEL_DIMENSIONS_BASE_MODELS,
   SUPPORTS_REF_IMAGES_BASE_MODELS,
   SUPPORTS_SEED_BASE_MODELS,
-} from 'features/parameters/types/constants';
+} from 'features/modelManagerV2/models';
+import { CLIP_SKIP_MAP } from 'features/parameters/types/constants';
 import type {
   ParameterCanvasCoherenceMode,
   ParameterCFGRescaleMultiplier,

invokeai/frontend/web/src/features/lora/components/LoRASelect.tsx

@@ -6,8 +6,8 @@ import { InformationalPopover } from 'common/components/InformationalPopover/Inf
 import type { GroupStatusMap } from 'common/components/Picker/Picker';
 import { loraAdded, selectLoRAsSlice } from 'features/controlLayers/store/lorasSlice';
 import { selectBase } from 'features/controlLayers/store/paramsSlice';
+import { API_BASE_MODELS } from 'features/modelManagerV2/models';
 import { ModelPicker } from 'features/parameters/components/ModelPicker';
-import { API_BASE_MODELS } from 'features/parameters/types/constants';
 import { memo, useCallback, useMemo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useLoRAModels } from 'services/api/hooks/modelsByType';

invokeai/frontend/web/src/features/modelManagerV2/models.ts

@@ -0,0 +1,214 @@
+import type { BaseModelType } from 'features/nodes/types/common';
+import type { AnyModelConfig } from 'services/api/types';
+import {
+  isCLIPEmbedModelConfig,
+  isCLIPVisionModelConfig,
+  isControlLoRAModelConfig,
+  isControlNetModelConfig,
+  isFluxReduxModelConfig,
+  isIPAdapterModelConfig,
+  isLLaVAModelConfig,
+  isLoRAModelConfig,
+  isNonRefinerMainModelConfig,
+  isRefinerMainModelModelConfig,
+  isSigLipModelConfig,
+  isSpandrelImageToImageModelConfig,
+  isT2IAdapterModelConfig,
+  isT5EncoderModelConfig,
+  isTIModelConfig,
+  isUnknownModelConfig,
+  isVAEModelConfig,
+} from 'services/api/types';
+
+type ModelCategoryData = {
+  i18nKey: string;
+  filter: (config: AnyModelConfig) => boolean;
+};
+
+export const MODEL_CATEGORIES: Record<string, ModelCategoryData> = {
+  main: {
+    i18nKey: 'model_manager.category.main_models',
+    filter: isNonRefinerMainModelConfig,
+  },
+  refiner: {
+    i18nKey: 'model_manager.category.refiner_models',
+    filter: isRefinerMainModelModelConfig,
+  },
+  lora: {
+    i18nKey: 'model_manager.category.lora_models',
+    filter: isLoRAModelConfig,
+  },
+  embedding: {
+    i18nKey: 'model_manager.category.embedding_models',
+    filter: isTIModelConfig,
+  },
+  controlnet: {
+    i18nKey: 'model_manager.category.controlnet_models',
+    filter: isControlNetModelConfig,
+  },
+  t2i_adapter: {
+    i18nKey: 'model_manager.category.t2i_adapter_models',
+    filter: isT2IAdapterModelConfig,
+  },
+  t5_encoder: {
+    i18nKey: 'model_manager.category.t5_encoder_models',
+    filter: isT5EncoderModelConfig,
+  },
+  control_lora: {
+    i18nKey: 'model_manager.category.control_lora_models',
+    filter: isControlLoRAModelConfig,
+  },
+  clip_embed: {
+    i18nKey: 'model_manager.category.clip_embed_models',
+    filter: isCLIPEmbedModelConfig,
+  },
+  spandrel: {
+    i18nKey: 'model_manager.category.spandrel_image_to_image_models',
+    filter: isSpandrelImageToImageModelConfig,
+  },
+  ip_adapter: {
+    i18nKey: 'model_manager.category.ip_adapter_models',
+    filter: isIPAdapterModelConfig,
+  },
+  vae: {
+    i18nKey: 'model_manager.category.vae_models',
+    filter: isVAEModelConfig,
+  },
+  clip_vision: {
+    i18nKey: 'model_manager.category.clip_vision_models',
+    filter: isCLIPVisionModelConfig,
+  },
+  siglip: {
+    i18nKey: 'model_manager.category.siglip_models',
+    filter: isSigLipModelConfig,
+  },
+  flux_redux: {
+    i18nKey: 'model_manager.category.flux_redux_models',
+    filter: isFluxReduxModelConfig,
+  },
+  llava_one_vision: {
+    i18nKey: 'model_manager.category.llava_one_vision_models',
+    filter: isLLaVAModelConfig,
+  },
+  unknown: {
+    i18nKey: 'model_manager.category.unknown_models',
+    filter: isUnknownModelConfig,
+  },
+};
+
+/**
+ * Mapping of model base to its color
+ */
+export const MODEL_BASE_TO_COLOR: Record<BaseModelType, string> = {
+  any: 'base',
+  'sd-1': 'green',
+  'sd-2': 'teal',
+  'sd-3': 'purple',
+  sdxl: 'invokeBlue',
+  'sdxl-refiner': 'invokeBlue',
+  flux: 'gold',
+  cogview4: 'red',
+  imagen3: 'pink',
+  imagen4: 'pink',
+  'chatgpt-4o': 'pink',
+  'flux-kontext': 'pink',
+  'gemini-2.5': 'pink',
+  veo3: 'purple',
+  runway: 'green',
+};
+
+/**
+ * Mapping of model base to human readable name
+ */
+export const MODEL_BASE_TO_LONG_NAME: Record<BaseModelType, string> = {
+  any: 'Any',
+  'sd-1': 'Stable Diffusion 1.x',
+  'sd-2': 'Stable Diffusion 2.x',
+  'sd-3': 'Stable Diffusion 3.x',
+  sdxl: 'Stable Diffusion XL',
+  'sdxl-refiner': 'Stable Diffusion XL Refiner',
+  flux: 'FLUX',
+  cogview4: 'CogView4',
+  imagen3: 'Imagen3',
+  imagen4: 'Imagen4',
+  'chatgpt-4o': 'ChatGPT 4o',
+  'flux-kontext': 'Flux Kontext',
+  'gemini-2.5': 'Gemini 2.5',
+  veo3: 'Veo3',
+  runway: 'Runway',
+};
+
+/**
+ * Mapping of model base to short human readable name
+ */
+export const MODEL_BASE_TO_SHORT_NAME: Record<BaseModelType, string> = {
+  any: 'Any',
+  'sd-1': 'SD1.X',
+  'sd-2': 'SD2.X',
+  'sd-3': 'SD3.X',
+  sdxl: 'SDXL',
+  'sdxl-refiner': 'SDXLR',
+  flux: 'FLUX',
+  cogview4: 'CogView4',
+  imagen3: 'Imagen3',
+  imagen4: 'Imagen4',
+  'chatgpt-4o': 'ChatGPT 4o',
+  'flux-kontext': 'Flux Kontext',
+  'gemini-2.5': 'Gemini 2.5',
+  veo3: 'Veo3',
+  runway: 'Runway',
+};
+
+/**
+ * List of base models that make API requests
+ */
+export const API_BASE_MODELS: BaseModelType[] = ['imagen3', 'imagen4', 'chatgpt-4o', 'flux-kontext', 'gemini-2.5'];
+
+export const SUPPORTS_SEED_BASE_MODELS: BaseModelType[] = ['sd-1', 'sd-2', 'sd-3', 'sdxl', 'flux', 'cogview4'];
+
+export const SUPPORTS_OPTIMIZED_DENOISING_BASE_MODELS: BaseModelType[] = ['flux', 'sd-3'];
+
+export const SUPPORTS_REF_IMAGES_BASE_MODELS: BaseModelType[] = [
+  'sd-1',
+  'sdxl',
+  'flux',
+  'flux-kontext',
+  'chatgpt-4o',
+  'gemini-2.5',
+];
+
+export const SUPPORTS_NEGATIVE_PROMPT_BASE_MODELS: BaseModelType[] = [
+  'sd-1',
+  'sd-2',
+  'sdxl',
+  'cogview4',
+  'sd-3',
+  'imagen3',
+  'imagen4',
+];
+
+export const SUPPORTS_PIXEL_DIMENSIONS_BASE_MODELS: BaseModelType[] = [
+  'sd-1',
+  'sd-2',
+  'sd-3',
+  'sdxl',
+  'flux',
+  'cogview4',
+];
+
+export const SUPPORTS_ASPECT_RATIO_BASE_MODELS: BaseModelType[] = [
+  'sd-1',
+  'sd-2',
+  'sd-3',
+  'sdxl',
+  'flux',
+  'cogview4',
+  'imagen3',
+  'imagen4',
+  'flux-kontext',
+  'chatgpt-4o',
+];
+
+export const VIDEO_BASE_MODELS = ['veo3', 'runway'];
+
+export const REQUIRES_STARTING_FRAME_BASE_MODELS = ['runway'];