Add model loader node; unet, clip, vae fields; change compel node to clip field

invokeai/frontend/web/src/common/util/parseMetadata.ts

@@ -1,5 +1,12 @@
 import { forEach, size } from 'lodash-es';
-import { ImageField, LatentsField, ConditioningField } from 'services/api';
+import {
+  ImageField,
+  LatentsField,
+  ConditioningField,
+  UNetField,
+  ClipField,
+  VaeField,
+} from 'services/api';
 
 const OBJECT_TYPESTRING = '[object Object]';
 const STRING_TYPESTRING = '[object String]';
@@ -98,6 +105,101 @@ const parseConditioningField = (
   };
 };
 
+const _parseModelInfo = (modelInfo: unknown): ModelInfo | undefined => {
+  // Must be an object
+  if (!isObject(modelInfo)) {
+    return;
+  }
+
+  if (!('model_name' in modelInfo && typeof modelInfo.model_name == 'string')) {
+    return;
+  }
+
+  if (!('model_type' in modelInfo && typeof modelInfo.model_type == 'string')) {
+    return;
+  }
+
+  if (!('submodel' in modelInfo && typeof modelInfo.submodel == 'string')) {
+    return;
+  }
+
+  return {
+    model_name: modelInfo.model_name,
+    model_type: modelInfo.model_type,
+    submodel: modelInfo.submodel,
+  };
+};
+
+const parseUNetField = (unetField: unknown): UNetField | undefined => {
+  // Must be an object
+  if (!isObject(unetField)) {
+    return;
+  }
+
+  if (!('unet' in unetField && 'scheduler' in unetField)) {
+    return;
+  }
+
+  const unet = _parseModelInfo(unetField.unet);
+  const scheduler = _parseModelInfo(unetField.scheduler);
+
+  if (!(unet && scheduler)) {
+    return;
+  }
+
+  // Build a valid UNetField
+  return {
+    unet: unet,
+    scheduler: scheduler,
+  };
+};
+
+const parseClipField = (clipField: unknown): ClipField | undefined => {
+  // Must be an object
+  if (!isObject(clipField)) {
+    return;
+  }
+
+  if (!('tokenizer' in clipField && 'text_encoder' in clipField)) {
+    return;
+  }
+
+  const tokenizer = _parseModelInfo(clipField.tokenizer);
+  const text_encoder = _parseModelInfo(clipField.text_encoder);
+
+  if (!(tokenizer && text_encoder)) {
+    return;
+  }
+
+  // Build a valid ClipField
+  return {
+    tokenizer: tokenizer,
+    text_encoder: text_encoder,
+  };
+};
+
+const parseVaeField = (vaeField: unknown): VaeField | undefined => {
+  // Must be an object
+  if (!isObject(vaeField)) {
+    return;
+  }
+
+  if (!('vae' in vaeField)) {
+    return;
+  }
+
+  const vae = _parseModelInfo(vaeField.vae);
+
+  if (!vae) {
+    return;
+  }
+
+  // Build a valid VaeField
+  return {
+    vae: vae,
+  };
+};
+
 type NodeMetadata = {
   [key: string]:
     | string
@@ -105,7 +207,10 @@ type NodeMetadata = {
     | boolean
     | ImageField
     | LatentsField
-    | ConditioningField;
+    | ConditioningField
+    | UNetField
+    | ClipField
+    | VaeField;
 };
 
 type InvokeAIMetadata = {
@@ -131,7 +236,8 @@ export const parseNodeMetadata = (
       return;
     }
 
-    // the only valid object types are ImageField, LatentsField and ConditioningField
+    // valid object types are:
+    // ImageField, LatentsField ConditioningField, UNetField, ClipField, VaeField
     if (isObject(nodeItem)) {
       if ('image_name' in nodeItem || 'image_type' in nodeItem) {
         const imageField = parseImageField(nodeItem);
@@ -156,6 +262,27 @@ export const parseNodeMetadata = (
         }
         return;
       }
+
+      if ('unet' in nodeItem && 'tokenizer' in nodeItem) {
+        const unetField = parseUNetField(nodeItem);
+        if (unetField) {
+          parsed[nodeKey] = unetField;
+        }
+      }
+
+      if ('tokenizer' in nodeItem && 'text_encoder' in nodeItem) {
+        const clipField = parseClipField(nodeItem);
+        if (clipField) {
+          parsed[nodeKey] = clipField;
+        }
+      }
+
+      if ('vae' in nodeItem) {
+        const vaeField = parseVaeField(nodeItem);
+        if (vaeField) {
+          parsed[nodeKey] = vaeField;
+        }
+      }
     }
 
     // otherwise we accept any string, number or boolean

invokeai/frontend/web/src/features/nodes/components/InputFieldComponent.tsx

@@ -7,6 +7,9 @@ import EnumInputFieldComponent from './fields/EnumInputFieldComponent';
 import ImageInputFieldComponent from './fields/ImageInputFieldComponent';
 import LatentsInputFieldComponent from './fields/LatentsInputFieldComponent';
 import ConditioningInputFieldComponent from './fields/ConditioningInputFieldComponent';
+import UNetInputFieldComponent from './fields/UNetInputFieldComponent';
+import ClipInputFieldComponent from './fields/ClipInputFieldComponent';
+import VaeInputFieldComponent from './fields/VaeInputFieldComponent';
 import ModelInputFieldComponent from './fields/ModelInputFieldComponent';
 import NumberInputFieldComponent from './fields/NumberInputFieldComponent';
 import StringInputFieldComponent from './fields/StringInputFieldComponent';
@@ -97,6 +100,36 @@ const InputFieldComponent = (props: InputFieldComponentProps) => {
     );
   }
 
+  if (type === 'unet' && template.type === 'unet') {
+    return (
+      <UNetInputFieldComponent
+        nodeId={nodeId}
+        field={field}
+        template={template}
+      />
+    );
+  }
+
+  if (type === 'clip' && template.type === 'clip') {
+    return (
+      <ClipInputFieldComponent
+        nodeId={nodeId}
+        field={field}
+        template={template}
+      />
+    );
+  }
+
+  if (type === 'vae' && template.type === 'vae') {
+    return (
+      <VaeInputFieldComponent
+        nodeId={nodeId}
+        field={field}
+        template={template}
+      />
+    );
+  }
+
   if (type === 'model' && template.type === 'model') {
     return (
       <ModelInputFieldComponent

invokeai/frontend/web/src/features/nodes/components/fields/ClipInputFieldComponent.tsx

@@ -0,0 +1,16 @@
+import {
+  ClipInputFieldTemplate,
+  ClipInputFieldValue,
+} from 'features/nodes/types/types';
+import { memo } from 'react';
+import { FieldComponentProps } from './types';
+
+const ClipInputFieldComponent = (
+  props: FieldComponentProps<ClipInputFieldValue, ClipInputFieldTemplate>
+) => {
+  const { nodeId, field } = props;
+
+  return null;
+};
+
+export default memo(ClipInputFieldComponent);

invokeai/frontend/web/src/features/nodes/components/fields/UNetInputFieldComponent.tsx

@@ -0,0 +1,16 @@
+import {
+  UNetInputFieldTemplate,
+  UNetInputFieldValue,
+} from 'features/nodes/types/types';
+import { memo } from 'react';
+import { FieldComponentProps } from './types';
+
+const UNetInputFieldComponent = (
+  props: FieldComponentProps<UNetInputFieldValue, UNetInputFieldTemplate>
+) => {
+  const { nodeId, field } = props;
+
+  return null;
+};
+
+export default memo(UNetInputFieldComponent);

invokeai/frontend/web/src/features/nodes/components/fields/VaeInputFieldComponent.tsx

@@ -0,0 +1,16 @@
+import {
+  VaeInputFieldTemplate,
+  VaeInputFieldValue,
+} from 'features/nodes/types/types';
+import { memo } from 'react';
+import { FieldComponentProps } from './types';
+
+const VaeInputFieldComponent = (
+  props: FieldComponentProps<VaeInputFieldValue, VaeInputFieldTemplate>
+) => {
+  const { nodeId, field } = props;
+
+  return null;
+};
+
+export default memo(VaeInputFieldComponent);

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

@@ -11,6 +11,9 @@ export const FIELD_TYPE_MAP: Record<string, FieldType> = {
   ImageField: 'image',
   LatentsField: 'latents',
   ConditioningField: 'conditioning',
+  UNetField: 'unet',
+  ClipField: 'clip',
+  VaeField: 'vae',
   model: 'model',
   array: 'array',
   item: 'item',
@@ -71,6 +74,24 @@ export const FIELDS: Record<FieldType, FieldUIConfig> = {
     title: 'Conditioning',
     description: 'Conditioning may be passed between nodes.',
   },
+  unet: {
+    color: 'red',
+    colorCssVar: getColorTokenCssVariable('red'),
+    title: 'UNet',
+    description: 'UNet submodel.',
+  },
+  clip: {
+    color: 'green',
+    colorCssVar: getColorTokenCssVariable('green'),
+    title: 'Clip',
+    description: 'Tokenizer and text_encoder submodels.',
+  },
+  vae: {
+    color: 'blue',
+    colorCssVar: getColorTokenCssVariable('blue'),
+    title: 'Vae',
+    description: 'Vae submodel.',
+  },
   model: {
     color: 'teal',
     colorCssVar: getColorTokenCssVariable('teal'),

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

@@ -58,6 +58,9 @@ export type FieldType =
   | 'image'
   | 'latents'
   | 'conditioning'
+  | 'unet'
+  | 'clip'
+  | 'vae'
   | 'model'
   | 'array'
   | 'item'
@@ -79,6 +82,9 @@ export type InputFieldValue =
   | ImageInputFieldValue
   | LatentsInputFieldValue
   | ConditioningInputFieldValue
+  | UNetInputFieldValue
+  | ClipInputFieldValue
+  | VaeInputFieldValue
   | EnumInputFieldValue
   | ModelInputFieldValue
   | ArrayInputFieldValue
@@ -99,6 +105,9 @@ export type InputFieldTemplate =
   | ImageInputFieldTemplate
   | LatentsInputFieldTemplate
   | ConditioningInputFieldTemplate
+  | UNetInputFieldTemplate
+  | ClipInputFieldTemplate
+  | VaeInputFieldTemplate
   | EnumInputFieldTemplate
   | ModelInputFieldTemplate
   | ArrayInputFieldTemplate
@@ -177,6 +186,21 @@ export type ConditioningInputFieldValue = FieldValueBase & {
   value?: undefined;
 };
 
+export type UNetInputFieldValue = FieldValueBase & {
+  type: 'unet';
+  value?: undefined;
+};
+
+export type ClipInputFieldValue = FieldValueBase & {
+  type: 'clip';
+  value?: undefined;
+};
+
+export type VaeInputFieldValue = FieldValueBase & {
+  type: 'vae';
+  value?: undefined;
+};
+
 export type ImageInputFieldValue = FieldValueBase & {
   type: 'image';
   value?: Pick<ImageField, 'image_name' | 'image_type'>;

invokeai/frontend/web/src/features/nodes/util/fieldTemplateBuilders.ts

@@ -10,6 +10,9 @@ import {
   IntegerInputFieldTemplate,
   LatentsInputFieldTemplate,
   ConditioningInputFieldTemplate,
+  UNetInputFieldTemplate,
+  ClipInputFieldTemplate,
+  VaeInputFieldTemplate,
   StringInputFieldTemplate,
   ModelInputFieldTemplate,
   ArrayInputFieldTemplate,
@@ -215,6 +218,51 @@ const buildConditioningInputFieldTemplate = ({
   return template;
 };
 
+const buildUNetInputFieldTemplate = ({
+  schemaObject,
+  baseField,
+}: BuildInputFieldArg): UNetInputFieldTemplate => {
+  const template: UNetInputFieldTemplate = {
+    ...baseField,
+    type: 'unet',
+    inputRequirement: 'always',
+    inputKind: 'connection',
+    default: schemaObject.default ?? undefined,
+  };
+
+  return template;
+};
+
+const buildClipInputFieldTemplate = ({
+  schemaObject,
+  baseField,
+}: BuildInputFieldArg): ClipInputFieldTemplate => {
+  const template: ClipInputFieldTemplate = {
+    ...baseField,
+    type: 'clip',
+    inputRequirement: 'always',
+    inputKind: 'connection',
+    default: schemaObject.default ?? undefined,
+  };
+
+  return template;
+};
+
+const buildVaeInputFieldTemplate = ({
+  schemaObject,
+  baseField,
+}: BuildInputFieldArg): VaeInputFieldTemplate => {
+  const template: VaeInputFieldTemplate = {
+    ...baseField,
+    type: 'vae',
+    inputRequirement: 'always',
+    inputKind: 'connection',
+    default: schemaObject.default ?? undefined,
+  };
+
+  return template;
+};
+
 const buildEnumInputFieldTemplate = ({
   schemaObject,
   baseField,
@@ -331,6 +379,15 @@ export const buildInputFieldTemplate = (
   if (['conditioning'].includes(fieldType)) {
     return buildConditioningInputFieldTemplate({ schemaObject, baseField });
   }
+  if (['unet'].includes(fieldType)) {
+    return buildUNetInputFieldTemplate({ schemaObject, baseField });
+  }
+  if (['clip'].includes(fieldType)) {
+    return buildClipInputFieldTemplate({ schemaObject, baseField });
+  }
+  if (['vae'].includes(fieldType)) {
+    return buildVaeInputFieldTemplate({ schemaObject, baseField });
+  }
   if (['model'].includes(fieldType)) {
     return buildModelInputFieldTemplate({ schemaObject, baseField });
   }

invokeai/frontend/web/src/features/nodes/util/fieldValueBuilders.ts

@@ -52,6 +52,18 @@ export const buildInputFieldValue = (
       fieldValue.value = undefined;
     }
 
+    if (template.type === 'unet') {
+      fieldValue.value = undefined;
+    }
+
+    if (template.type === 'clip') {
+      fieldValue.value = undefined;
+    }
+
+    if (template.type === 'vae') {
+      fieldValue.value = undefined;
+    }
+
     if (template.type === 'model') {
       fieldValue.value = undefined;
     }