feat(ui): update invocation parsing to handle new ui_model_[base|type|variant] attrs

Author: psychedelicious

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/ModelIdentifierFieldInputComponent.tsx

@@ -12,7 +12,7 @@ import type { FieldComponentProps } from './types';
 type Props = FieldComponentProps<ModelIdentifierFieldInputInstance, ModelIdentifierFieldInputTemplate>;
 
 const ModelIdentifierFieldInputComponent = (props: Props) => {
-  const { nodeId, field } = props;
+  const { nodeId, field, fieldTemplate } = props;
   const dispatch = useAppDispatch();
   const { data, isLoading } = useGetModelConfigsQuery();
   const onChange = useCallback(
@@ -36,8 +36,28 @@ const ModelIdentifierFieldInputComponent = (props: Props) => {
       return EMPTY_ARRAY;
     }
 
-    return modelConfigsAdapterSelectors.selectAll(data);
-  }, [data]);
+    if (!fieldTemplate.ui_model_base && !fieldTemplate.ui_model_type) {
+      return modelConfigsAdapterSelectors.selectAll(data);
+    }
+
+    return modelConfigsAdapterSelectors.selectAll(data).filter((config) => {
+      if (fieldTemplate.ui_model_base && !fieldTemplate.ui_model_base.includes(config.base)) {
+        return false;
+      }
+      if (fieldTemplate.ui_model_type && !fieldTemplate.ui_model_type.includes(config.type)) {
+        return false;
+      }
+      if (
+        fieldTemplate.ui_model_variant &&
+        'variant' in config &&
+        config.variant &&
+        !fieldTemplate.ui_model_variant.includes(config.variant)
+      ) {
+        return false;
+      }
+      return true;
+    });
+  }, [data, fieldTemplate.ui_model_base, fieldTemplate.ui_model_type, fieldTemplate.ui_model_variant]);
 
   return (
     <ModelFieldCombobox

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

@@ -12,11 +12,14 @@ import type {
   SchedulerField,
   SubModelType,
   T2IAdapterField,
+  zClipVariantType,
+  zModelVariantType,
 } from 'features/nodes/types/common';
-import type { Invocation, ModelType, S } from 'services/api/types';
+import type { Invocation, S } from 'services/api/types';
 import type { Equals, Extends } from 'tsafe';
 import { assert } from 'tsafe';
 import { describe, test } from 'vitest';
+import type z from 'zod';
 
 /**
  * These types originate from the server and are recreated as zod schemas manually, for use at runtime.
@@ -38,7 +41,8 @@ describe('Common 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<ModelType, S['ModelType']>>());
+  test('ClipVariantType', () => assert<Equals<z.infer<typeof zClipVariantType>, S['ClipVariantType']>>());
+  test('ModelVariantType', () => assert<Equals<z.infer<typeof zModelVariantType>, S['ModelVariantType']>>());
 
   // Misc types
   test('ProgressImage', () => assert<Equals<ProgressImage, S['ProgressImage']>>());

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

@@ -73,7 +73,7 @@ export type SchedulerField = z.infer<typeof zSchedulerField>;
 // #endregion
 
 // #region Model-related schemas
-const zBaseModel = z.enum([
+export const zBaseModelType = z.enum([
   'any',
   'sd-1',
   'sd-2',
@@ -90,7 +90,7 @@ const zBaseModel = z.enum([
   'veo3',
   'runway',
 ]);
-export type BaseModelType = z.infer<typeof zBaseModel>;
+export type BaseModelType = z.infer<typeof zBaseModelType>;
 export const zMainModelBase = z.enum([
   'sd-1',
   'sd-2',
@@ -143,11 +143,15 @@ const zSubModelType = z.enum([
   'safety_checker',
 ]);
 export type SubModelType = z.infer<typeof zSubModelType>;
+
+export const zClipVariantType = z.enum(['large', 'gigantic']);
+export const zModelVariantType = z.enum(['normal', 'inpaint', 'depth']);
+
 export const zModelIdentifierField = z.object({
   key: z.string().min(1),
   hash: z.string().min(1),
   name: z.string().min(1),
-  base: zBaseModel,
+  base: zBaseModelType,
   type: zModelType,
   submodel_type: zSubModelType.nullish(),
 });

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

@@ -9,7 +9,17 @@ import { assert } from 'tsafe';
 import { z } from 'zod';
 
 import type { ImageField } from './common';
-import { zBoardField, zColorField, zImageField, zModelIdentifierField, zSchedulerField } from './common';
+import {
+  zBaseModelType,
+  zBoardField,
+  zClipVariantType,
+  zColorField,
+  zImageField,
+  zModelIdentifierField,
+  zModelType,
+  zModelVariantType,
+  zSchedulerField,
+} from './common';
 
 /**
  * zod schemas & inferred types for fields.
@@ -60,6 +70,9 @@ const zFieldInputTemplateBase = zFieldTemplateBase.extend({
   default: z.undefined(),
   ui_component: zFieldUIComponent.nullish(),
   ui_choice_labels: z.record(z.string(), z.string()).nullish(),
+  ui_model_base: z.array(zBaseModelType).nullish(),
+  ui_model_type: z.array(zModelType).nullish(),
+  ui_model_variant: z.array(zModelVariantType.or(zClipVariantType)).nullish(),
 });
 const zFieldOutputTemplateBase = zFieldTemplateBase.extend({
   fieldKind: z.literal('output'),

invokeai/frontend/web/src/features/nodes/util/schema/buildFieldInputTemplate.ts

@@ -892,7 +892,18 @@ export const buildFieldInputTemplate = (
   fieldName: string,
   fieldType: FieldType
 ): FieldInputTemplate => {
-  const { input, ui_hidden, ui_component, ui_type, ui_order, ui_choice_labels, orig_required: required } = fieldSchema;
+  const {
+    input,
+    ui_hidden,
+    ui_component,
+    ui_type,
+    ui_order,
+    ui_choice_labels,
+    orig_required: required,
+    ui_model_base,
+    ui_model_type,
+    ui_model_variant,
+  } = fieldSchema;
 
   // This is the base field template that is common to all fields. The builder function will add all other
   // properties to this template.
@@ -908,6 +919,9 @@ export const buildFieldInputTemplate = (
     ui_type,
     ui_order,
     ui_choice_labels,
+    ui_model_base,
+    ui_model_type,
+    ui_model_variant,
   };
 
   if (isStatefulFieldType(fieldType)) {