feat: dynamic model fields in workflow editor

[psychedelicious]

Summary

This PR adds a new system for dynamically generating Workflow Editor model drop-downs. This pattern is has a few benefits:

• It is much simpler, conceptually and technically.
• It allows node authors to be more specific about what kinds of models a field can accept, reducing footguns in the Workflow Editor.
• It greatly reduces the boilerplate required to get a new model architecture working in the Workflow Editor (does not improve things for the linear tabs).

Context

Nodes often need a drop-down in the UI that lists a specific base, type, and/or variant of model. For example, the Main Model - SD1.5 (main_model_loader) node needs to display a list of SD1.5 main models only.

Informing the frontend that it should only show SD1.5 main models was a small technical challenge. Pydantic creates OpenAPI schemas for each node class, and we parse those into something the Workflow Editor can use to dynamically build node UIs. Python type annotations are what determine the "type" of the field.

For example, in the Main Model - SD1.5 node, we had a field definition like this:

model: ModelIdentifierField = InputField(description=FieldDescriptions.main_model # ...

ModelIdentifierField - a pydantic model - ends up in the OpenAPI schema, then the frontend sees this annotation and knows this field can accept a model. But, as described above, we don't need just any model - we need exclusively SD1.5 main models.

We could have created new python classes for every type of model, using them as type annotations. For example:

class SD1_5MainModelField(BaseModel):
    # rest of class definition

# Use it in the node:
 model: SD1_5MainModelField = InputField(description=FieldDescriptions.main_model # ...

This could get pretty hairy, with many different classes, each of which with their own pydantic model. In an effort to make things a bit simpler and more flexible, we added the ui_type arg. This doesn't affect pydantic or OpenAPI schema generation, but it gives the frontend an extra hint.

In the Main Model - SD1.5 node, we set the arg to UIType.MainModel:

model: ModelIdentifierField = InputField(description=FieldDescriptions.main_model, ui_type=UIType.MainModel)

This tells the frontend "This is a ModelIdentifierField, but it has this extra constraint of needing only SD1.5 main models.".

The frontend requires a lot of special handling and boilerplate to understand each of the UIType values:

• Zod schemas for each, plus inferred TS types and type guards
• Model-fetching hooks
• Zod schemas for the corresponding field types, field instances and field templates
• A component to render the drop-down with the appropriate models
• Update OpenAPI schema parsing

(We'd still need this special handling if we used different pydantic classes for each model, but we'd also have a crapload of special handling in the backend in that case.)

Solution

It's really simple. Instead of exposing a ui_type arg, we expose a few optional args that let the node describe the models it needs. Each of these args supports either a single or list of allowed types.

• ui_model_base: BaseModelType | list[BaseModelType] indicates which architectures are allowed (e.g. SD1.5, SD1.5 and SDXL, FLUX, etc).
• ui_model_type: ModelType | list[ModelType] indicates which types of models are allowed (e.g. main, ControlNet, LoRA, etc.
• ui_model_variant: ModelVariantType | ClipVariantType | list[ModelVariantType | ClipVariantType] indicates which variants are allowed. This is currently only used to differentiate between CLIP-L and CLIP-G.

So for that Main Model - SD1.5 node, we now write the model field like this:

model: ModelIdentifierField = InputField(
    description=FieldDescriptions.main_model,
    ui_model_base=BaseModelType.StableDiffusion1,
    ui_model_type=ModelType.Main,
)

In the Workflow Editor frontend codebase, we now have only 1 model field component, compared to like 20 before. This one component filters its downdown based on the args.

The change is fully backwards compatible with existing workflows, with no data migration or compatibility layer required. Workflow Editor field types are not stored in the user's workflows, so we can modify them freely without breaking anything.

Less Boilerplate

Let's say we want to add support for Qwen-Image. We'd do all the same backend changes, then create the Qwen-Image model loader. It would have a model field that looks like this:

model: ModelIdentifierField = InputField(
    description=FieldDescriptions.main_model,
    ui_model_base=BaseModelType.QwenImage,
    ui_model_type=ModelType.Main,
)

That's it! The frontend will see BaseModelType.QwenImage and automatically show Qwen-Image models in the node. Zero frontend code changes required to support it in the Workflow Editor.

Reduced Footguns

There are a bajillion combinations of base model and model types. For example, ControlNet works with SD1.5, SDXL, and FLUX (maybe others I'm forgetting). We have a node for SD1.5/SDXL and another for FLUX.

But our ControlNet model fields were all defined like this:

control_model: ModelIdentifierField = InputField( description=FieldDescriptions.controlnet_model, ui_type=UIType.ControlNetModel)

Note that there is nothing that tells the frontend if this is for SD1.5/SDXL, or FLUX. Regardless of the base model of the node, we showed a list of all ControlNet models. It was possible to select FLUX models on the SD1.5/SDXL node and vice-versa. If you chose the wrong model, you'd get a cryptic error during denoising.

With the new system, we can be more specific and only show the correct models:

# SD1.5/SDXL ControlNet model field
control_model: ModelIdentifierField = InputField(
    description=FieldDescriptions.controlnet_model,
    # Only allow SD1.5 and SDXL bases
    ui_model_base=[BaseModelType.StableDiffusion1, BaseModelType.StableDiffusionXL],
    ui_model_type=ModelType.ControlNet,
)

# FLUX ControlNet model field
control_model: ModelIdentifierField = InputField(
    description=FieldDescriptions.controlnet_model,
    # Only allow FLUX base
    ui_model_base=BaseModelType.Flux,
    ui_model_type=ModelType.ControlNet,
)

Simpler Implementation and Extension

The schema parsing required for the Workflow Editor is simpler thanks to this change. We don't need special handling and can simply filter models based on the args to narrow the list shown in a drop-down.

Extending the system is pretty easy, too. For example, if we needed to only show checkpoint-format models in a node for whatever reason, we would could ui_model_format. It would require editing only a handful of files:

• invokeai/app/invocations/fields.py to add the arg to the input field definitions
• invokeai/frontend/web/src/features/nodes/types/field.ts to add the arg to the workflow editor's input field template schema (requires a matching zod schema for the format enum)
• invokeai/frontend/web/src/features/nodes/util/schema/buildFieldInputTemplate.ts to parse the value from OpenAPI schema and add it to the input field template
• invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/ModelIdentifierFieldInputComponent.tsx to filter the model dropdown list based on the new arg

Backwards Compatible

As mentioned above, the change is fully backwards compatible with existing workflows. What about the existing ui_type: UIType in the InputField helper, and custom nodes?

Well, all of the model-related UITypes map directly to some combination of base, type and variant. In invokeai/app/invocations/fields.py you'll see a long (but very simple) conditional that sets the appropriate ui_model_[base|type|variant] for each UIType.

I already migrated all core nodes from UIType to the new args, but custom nodes should continue to work with no code changes required. The conditional logic just kinda migrates them automatically at runtime.

TODO

• Probably should add some deprecation messages when we detect users.
• We support UIType for output fields. I need to think more about it, but I think it's fine to just ignore UIType for outputs. I don't think it serves any functional or user-facing purpose. We do need UIType for output fields. It's used for Any and Scheduler fields.

Related Issues / Discussions

n/a

QA Instructions

Try out the workflow editor, especially custom nodes. Besides the reduced footguns for model fields (i.e. you may no longer select FLUX ControlNets in the SD1.5/SDXL ControlNet node), this should be a transparent change.

Merge Plan

n/a

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

[maryhipp]

That diff tho 🤩

[psychedelicious]

Though we have no current use for it, for the sake of future node authors' ease of experimentation, I added ui_model_format. This indicates the model's file format (e.g. diffusers, safetensors, bnb, etc.). I also added some deprecation warnings.