change model store to object, update main model and vae dropdowns

invokeai/frontend/web/src/features/nodes/util/graphBuilders/addVAEToGraph.ts

@@ -16,10 +16,12 @@ export const addVAEToGraph = (
   graph: NonNullableGraph,
   state: RootState
 ): void => {
-  const { vae: vaeId } = state.generation;
-  const vae_model = modelIdToVAEModelField(vaeId);
+  const { vae } = state.generation;
+  const vae_model = modelIdToVAEModelField(vae?.id || '');
 
-  if (vaeId !== 'auto') {
+  const isAutoVae = vae?.id === 'auto';
+
+  if (!isAutoVae) {
     graph.nodes[VAE_LOADER] = {
       type: 'vae_loader',
       id: VAE_LOADER,
@@ -30,7 +32,7 @@ export const addVAEToGraph = (
   if (graph.id === TEXT_TO_IMAGE_GRAPH || graph.id === IMAGE_TO_IMAGE_GRAPH) {
     graph.edges.push({
       source: {
-        node_id: vaeId === 'auto' ? MAIN_MODEL_LOADER : VAE_LOADER,
+        node_id: isAutoVae ? MAIN_MODEL_LOADER : VAE_LOADER,
         field: 'vae',
       },
       destination: {
@@ -43,7 +45,7 @@ export const addVAEToGraph = (
   if (graph.id === IMAGE_TO_IMAGE_GRAPH) {
     graph.edges.push({
       source: {
-        node_id: vaeId === 'auto' ? MAIN_MODEL_LOADER : VAE_LOADER,
+        node_id: isAutoVae ? MAIN_MODEL_LOADER : VAE_LOADER,
         field: 'vae',
       },
       destination: {
@@ -56,7 +58,7 @@ export const addVAEToGraph = (
   if (graph.id === INPAINT_GRAPH) {
     graph.edges.push({
       source: {
-        node_id: vaeId === 'auto' ? MAIN_MODEL_LOADER : VAE_LOADER,
+        node_id: isAutoVae ? MAIN_MODEL_LOADER : VAE_LOADER,
         field: 'vae',
       },
       destination: {

invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasImageToImageGraph.ts

@@ -36,7 +36,7 @@ export const buildCanvasImageToImageGraph = (
   const {
     positivePrompt,
     negativePrompt,
-    model: modelId,
+    model: currentModel,
     cfgScale: cfg_scale,
     scheduler,
     steps,
@@ -50,7 +50,7 @@ export const buildCanvasImageToImageGraph = (
   // The bounding box determines width and height, not the width and height params
   const { width, height } = state.canvas.boundingBoxDimensions;
 
-  const model = modelIdToMainModelField(modelId);
+  const model = modelIdToMainModelField(currentModel?.id || '');
 
   /**
    * The easiest way to build linear graphs is to do it in the node editor, then copy and paste the

invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasInpaintGraph.ts

@@ -35,7 +35,7 @@ export const buildCanvasInpaintGraph = (
   const {
     positivePrompt,
     negativePrompt,
-    model: modelId,
+    model: currentModel,
     cfgScale: cfg_scale,
     scheduler,
     steps,
@@ -59,7 +59,7 @@ export const buildCanvasInpaintGraph = (
   // We may need to set the inpaint width and height to scale the image
   const { scaledBoundingBoxDimensions, boundingBoxScaleMethod } = state.canvas;
 
-  const model = modelIdToMainModelField(modelId);
+  const model = modelIdToMainModelField(currentModel?.id || '');
 
   const graph: NonNullableGraph = {
     id: INPAINT_GRAPH,

invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildCanvasTextToImageGraph.ts

@@ -25,7 +25,7 @@ export const buildCanvasTextToImageGraph = (
   const {
     positivePrompt,
     negativePrompt,
-    model: modelId,
+    model: currentModel,
     cfgScale: cfg_scale,
     scheduler,
     steps,
@@ -38,7 +38,7 @@ export const buildCanvasTextToImageGraph = (
   // The bounding box determines width and height, not the width and height params
   const { width, height } = state.canvas.boundingBoxDimensions;
 
-  const model = modelIdToMainModelField(modelId);
+  const model = modelIdToMainModelField(currentModel?.id || '');
 
   /**
    * The easiest way to build linear graphs is to do it in the node editor, then copy and paste the

invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearImageToImageGraph.ts

@@ -38,7 +38,7 @@ export const buildLinearImageToImageGraph = (
   const {
     positivePrompt,
     negativePrompt,
-    model: modelId,
+    model: currentModel,
     cfgScale: cfg_scale,
     scheduler,
     steps,
@@ -73,7 +73,7 @@ export const buildLinearImageToImageGraph = (
     throw new Error('No initial image found in state');
   }
 
-  const model = modelIdToMainModelField(modelId);
+  const model = modelIdToMainModelField(currentModel?.id || '');
 
   // copy-pasted graph from node editor, filled in with state values & friendly node ids
   const graph: NonNullableGraph = {

invokeai/frontend/web/src/features/nodes/util/graphBuilders/buildLinearTextToImageGraph.ts

@@ -22,7 +22,7 @@ export const buildLinearTextToImageGraph = (
   const {
     positivePrompt,
     negativePrompt,
-    model: modelId,
+    model: currentModel,
     cfgScale: cfg_scale,
     scheduler,
     steps,
@@ -31,7 +31,7 @@ export const buildLinearTextToImageGraph = (
     clipSkip,
   } = state.generation;
 
-  const model = modelIdToMainModelField(modelId);
+  const model = modelIdToMainModelField(currentModel?.id || '');
 
   /**
    * The easiest way to build linear graphs is to do it in the node editor, then copy and paste the

invokeai/frontend/web/src/features/parameters/store/generationSlice.ts

@@ -16,7 +16,6 @@ import {
   SeedParam,
   StepsParam,
   StrengthParam,
-  VAEParam,
   WidthParam,
 } from './parameterZodSchemas';
 
@@ -50,7 +49,7 @@ export interface GenerationState {
   horizontalSymmetrySteps: number;
   verticalSymmetrySteps: number;
   model: ModelParam;
-  vae: VAEParam;
+  vae: ModelParam;
   seamlessXAxis: boolean;
   seamlessYAxis: boolean;
   clipSkip: number;
@@ -84,8 +83,8 @@ export const initialGenerationState: GenerationState = {
   shouldUseSymmetry: false,
   horizontalSymmetrySteps: 0,
   verticalSymmetrySteps: 0,
-  model: '',
-  vae: '',
+  model: null,
+  vae: null,
   seamlessXAxis: false,
   seamlessYAxis: false,
   clipSkip: 0,
@@ -216,16 +215,17 @@ export const generationSlice = createSlice({
       state.initialImage = { imageName: image_name, width, height };
     },
     modelSelected: (state, action: PayloadAction<string>) => {
-      state.model = action.payload;
+      const [base_model, type, name] = action.payload.split('/');
 
       // Clamp ClipSkip Based On Selected Model
-      const clipSkipMax =
-        clipSkipMap[action.payload.split('/')[0] as keyof typeof clipSkipMap]
-          .maxClip;
-      state.clipSkip = clamp(state.clipSkip, 0, clipSkipMax);
+      const { maxClip } = clipSkipMap[base_model as keyof typeof clipSkipMap];
+      state.clipSkip = clamp(state.clipSkip, 0, maxClip);
+
+      state.model = { id: action.payload, base_model, name, type };
     },
     vaeSelected: (state, action: PayloadAction<string>) => {
-      state.vae = action.payload;
+      const [base_model, type, name] = action.payload.split('/');
+      state.vae = { id: action.payload, base_model, name, type };
     },
     setClipSkip: (state, action: PayloadAction<number>) => {
       state.clipSkip = action.payload;
@@ -235,7 +235,13 @@ export const generationSlice = createSlice({
     builder.addCase(configChanged, (state, action) => {
       const defaultModel = action.payload.sd?.defaultModel;
       if (defaultModel && !state.model) {
-        state.model = defaultModel;
+        const [base_model, model_type, model_name] = defaultModel.split('/');
+        state.model = {
+          id: defaultModel,
+          name: model_name,
+          type: model_type,
+          base_model: base_model,
+        };
       }
     });
     builder.addCase(setShouldShowAdvancedOptions, (state, action) => {

invokeai/frontend/web/src/features/parameters/store/parameterZodSchemas.ts

@@ -130,20 +130,21 @@ export const isValidHeight = (val: unknown): val is HeightParam =>
  * Zod schema for model parameter
  * TODO: Make this a dynamically generated enum?
  */
-export const zModel = z.string();
+const zModel = z.object({
+  id: z.string(),
+  name: z.string(),
+  type: z.string(),
+  base_model: z.string(),
+});
+
 /**
  * Type alias for model parameter, inferred from its zod schema
  */
-export type ModelParam = z.infer<typeof zModel>;
-/**
- * Zod schema for VAE parameter
- * TODO: Make this a dynamically generated enum?
- */
-export const zVAE = z.string();
+export type ModelParam = z.infer<typeof zModel> | null;
 /**
  * Type alias for model parameter, inferred from its zod schema
  */
-export type VAEParam = z.infer<typeof zVAE>;
+export type VAEParam = z.infer<typeof zModel> | null;
 /**
  * Validates/type-guards a value as a model parameter
  */

invokeai/frontend/web/src/features/system/components/ModelSelect.tsx

@@ -19,7 +19,7 @@ const ModelSelect = () => {
   const dispatch = useAppDispatch();
   const { t } = useTranslation();
 
-  const selectedModelId = useAppSelector(
+  const currentModel = useAppSelector(
     (state: RootState) => state.generation.model
   );
 
@@ -48,8 +48,8 @@ const ModelSelect = () => {
   }, [mainModels]);
 
   const selectedModel = useMemo(
-    () => mainModels?.entities[selectedModelId],
-    [mainModels?.entities, selectedModelId]
+    () => mainModels?.entities[currentModel?.id || ''],
+    [mainModels?.entities, currentModel]
   );
 
   const handleChangeModel = useCallback(
@@ -63,10 +63,6 @@ const ModelSelect = () => {
   );
 
   useEffect(() => {
-    if (selectedModelId && mainModels?.ids.includes(selectedModelId)) {
-      return;
-    }
-
     const firstModel = mainModels?.ids[0];
 
     if (!isString(firstModel)) {
@@ -74,7 +70,7 @@ const ModelSelect = () => {
     }
 
     handleChangeModel(firstModel);
-  }, [handleChangeModel, mainModels?.ids, selectedModelId]);
+  }, [handleChangeModel, mainModels?.ids]);
 
   return isLoading ? (
     <IAIMantineSelect
@@ -87,7 +83,7 @@ const ModelSelect = () => {
     <IAIMantineSelect
       tooltip={selectedModel?.description}
       label={t('modelManager.model')}
-      value={selectedModelId}
+      value={selectedModel?.id}
       placeholder={data.length > 0 ? 'Select a model' : 'No models detected!'}
       data={data}
       error={data.length === 0}

invokeai/frontend/web/src/features/system/components/VAESelect.tsx

@@ -18,7 +18,7 @@ const VAESelect = () => {
 
   const { data: vaeModels } = useGetVaeModelsQuery();
 
-  const selectedModelId = useAppSelector(
+  const currentModel = useAppSelector(
     (state: RootState) => state.generation.vae
   );
 
@@ -51,8 +51,8 @@ const VAESelect = () => {
   }, [vaeModels]);
 
   const selectedModel = useMemo(
-    () => vaeModels?.entities[selectedModelId],
-    [vaeModels?.entities, selectedModelId]
+    () => vaeModels?.entities[currentModel?.id || ''],
+    [vaeModels?.entities, currentModel]
   );
 
   const handleChangeModel = useCallback(
@@ -66,17 +66,17 @@ const VAESelect = () => {
   );
 
   useEffect(() => {
-    if (selectedModelId && vaeModels?.ids.includes(selectedModelId)) {
+    if (currentModel?.id && vaeModels?.ids.includes(currentModel?.id)) {
       return;
     }
     handleChangeModel('auto');
-  }, [handleChangeModel, vaeModels?.ids, selectedModelId]);
+  }, [handleChangeModel, vaeModels?.ids, currentModel?.id]);
 
   return (
     <IAIMantineSelect
       tooltip={selectedModel?.description}
       label={t('modelManager.vae')}
-      value={selectedModelId}
+      value={currentModel?.id}
       placeholder="Pick one"
       data={data}
       onChange={handleChangeModel}