feat(nodes): better deprecation handling for ui_type

psychedelicious · psychedelicious · commit 28b5f7a1c5f1 · 2025-09-19T22:13:12.000+10:00
- Move migration of model-specific ui_types into BaseInvocation. This
gives us access to the node and field names, so the warnings are more
useful to the end user.
- Ensure we serialize the fields' json_schema_extra with enum values.
This wasn't a problem until now, when it interferes with migrating
ui_type cleanly. It's a transparent change.
- Improve warnings when validating fields (which includes the ui_type
migration logic)
diff --git a/invokeai/app/invocations/baseinvocation.py b/invokeai/app/invocations/baseinvocation.py
@@ -36,6 +36,9 @@
 from invokeai.app.invocations.fields import (
     FieldKind,
     Input,
+    InputFieldJSONSchemaExtra,
+    UIType,
+    migrate_model_ui_type,
 )
 from invokeai.app.services.config.config_default import get_config
 from invokeai.app.services.shared.invocation_context import InvocationContext
@@ -256,7 +259,9 @@ def invoke_internal(self, context: InvocationContext, services: "InvocationServi
     is_intermediate: bool = Field(
         default=False,
         description="Whether or not this is an intermediate invocation.",
-        json_schema_extra={"ui_type": "IsIntermediate", "field_kind": FieldKind.NodeAttribute},
+        json_schema_extra=InputFieldJSONSchemaExtra(
+            input=Input.Direct, field_kind=FieldKind.NodeAttribute, ui_type=UIType._IsIntermediate
+        ).model_dump(exclude_none=True),
     )
     use_cache: bool = Field(
         default=True,
@@ -445,6 +450,15 @@ class _Model(BaseModel):
     RESERVED_PYDANTIC_FIELD_NAMES = {m[0] for m in inspect.getmembers(_Model())}
 
 
+def is_enum_member(value: Any, enum_class: type[Enum]) -> bool:
+    """Checks if a value is a member of an enum class."""
+    try:
+        enum_class(value)
+        return True
+    except ValueError:
+        return False
+
+
 def validate_fields(model_fields: dict[str, FieldInfo], model_type: str) -> None:
     """
     Validates the fields of an invocation or invocation output:
@@ -456,51 +470,99 @@ def validate_fields(model_fields: dict[str, FieldInfo], model_type: str) -> None
     """
     for name, field in model_fields.items():
         if name in RESERVED_PYDANTIC_FIELD_NAMES:
-            raise InvalidFieldError(f'Invalid field name "{name}" on "{model_type}" (reserved by pydantic)')
+            raise InvalidFieldError(f"{model_type}.{name}: Invalid field name (reserved by pydantic)")
 
         if not field.annotation:
-            raise InvalidFieldError(f'Invalid field type "{name}" on "{model_type}" (missing annotation)')
+            raise InvalidFieldError(f"{model_type}.{name}: Invalid field type (missing annotation)")
 
         if not isinstance(field.json_schema_extra, dict):
-            raise InvalidFieldError(
-                f'Invalid field definition for "{name}" on "{model_type}" (missing json_schema_extra dict)'
-            )
+            raise InvalidFieldError(f"{model_type}.{name}: Invalid field definition (missing json_schema_extra dict)")
 
         field_kind = field.json_schema_extra.get("field_kind", None)
 
         # must have a field_kind
-        if not isinstance(field_kind, FieldKind):
+        if not is_enum_member(field_kind, FieldKind):
             raise InvalidFieldError(
-                f'Invalid field definition for "{name}" on "{model_type}" (maybe it\'s not an InputField or OutputField?)'
+                f"{model_type}.{name}: Invalid field definition for (maybe it's not an InputField or OutputField?)"
             )
 
-        if field_kind is FieldKind.Input and (
+        if field_kind == FieldKind.Input.value and (
             name in RESERVED_NODE_ATTRIBUTE_FIELD_NAMES or name in RESERVED_INPUT_FIELD_NAMES
         ):
-            raise InvalidFieldError(f'Invalid field name "{name}" on "{model_type}" (reserved input field name)')
+            raise InvalidFieldError(f"{model_type}.{name}: Invalid field name (reserved input field name)")
 
-        if field_kind is FieldKind.Output and name in RESERVED_OUTPUT_FIELD_NAMES:
-            raise InvalidFieldError(f'Invalid field name "{name}" on "{model_type}" (reserved output field name)')
+        if field_kind == FieldKind.Output.value and name in RESERVED_OUTPUT_FIELD_NAMES:
+            raise InvalidFieldError(f"{model_type}.{name}: Invalid field name (reserved output field name)")
 
-        if (field_kind is FieldKind.Internal) and name not in RESERVED_INPUT_FIELD_NAMES:
-            raise InvalidFieldError(
-                f'Invalid field name "{name}" on "{model_type}" (internal field without reserved name)'
-            )
+        if field_kind == FieldKind.Internal.value and name not in RESERVED_INPUT_FIELD_NAMES:
+            raise InvalidFieldError(f"{model_type}.{name}: Invalid field name (internal field without reserved name)")
 
         # node attribute fields *must* be in the reserved list
         if (
-            field_kind is FieldKind.NodeAttribute
+            field_kind == FieldKind.NodeAttribute.value
             and name not in RESERVED_NODE_ATTRIBUTE_FIELD_NAMES
             and name not in RESERVED_OUTPUT_FIELD_NAMES
         ):
             raise InvalidFieldError(
-                f'Invalid field name "{name}" on "{model_type}" (node attribute field without reserved name)'
+                f"{model_type}.{name}: Invalid field name (node attribute field without reserved name)"
             )
 
         ui_type = field.json_schema_extra.get("ui_type", None)
-        if isinstance(ui_type, str) and ui_type.startswith("DEPRECATED_"):
-            logger.warning(f'"UIType.{ui_type.split("_")[-1]}" is deprecated, ignoring')
-            field.json_schema_extra.pop("ui_type")
+        ui_model_base = field.json_schema_extra.get("ui_model_base", None)
+        ui_model_type = field.json_schema_extra.get("ui_model_type", None)
+        ui_model_variant = field.json_schema_extra.get("ui_model_variant", None)
+        ui_model_format = field.json_schema_extra.get("ui_model_format", None)
+
+        if ui_type is not None:
+            # There are 3 cases where we may need to take action:
+            #
+            # 1. The ui_type is a migratable, deprecated value. For example, ui_type=UIType.MainModel value is
+            #    deprecated and should be migrated to:
+            #       - ui_model_base=[BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]
+            #       - ui_model_type=[ModelType.Main]
+            #
+            # 2. ui_type was set in conjunction with any of the new ui_model_[base|type|variant|format] fields, which
+            #    is not allowed (they are mutually exclusive). In this case, we ignore ui_type and log a warning.
+            #
+            # 3. ui_type is a deprecated value that is not migratable. For example, ui_type=UIType.Image is deprecated;
+            #    Image fields are now automatically detected based on the field's type annotation. In this case, we
+            #    ignore ui_type and log a warning.
+            #
+            # The cases must be checked in this order to ensure proper handling.
+
+            # Easier to work with as an enum
+            ui_type = UIType(ui_type)
+
+            # The enum member values are not always the same as their names - we want to log the name so the user can
+            # easily review their code and see where the deprecated enum member is used.
+            human_readable_name = f"UIType.{ui_type.name}"
+
+            # Case 1: migratable deprecated value
+            did_migrate = migrate_model_ui_type(ui_type, field.json_schema_extra)
+
+            if did_migrate:
+                logger.warning(
+                    f'{model_type}.{name}: Migrated deprecated "ui_type" "{human_readable_name}" to new ui_model_[base|type|variant|format] fields'
+                )
+                field.json_schema_extra.pop("ui_type")
+
+            # Case 2: mutually exclusive with new fields
+            elif (
+                ui_model_base is not None
+                or ui_model_type is not None
+                or ui_model_variant is not None
+                or ui_model_format is not None
+            ):
+                logger.warning(
+                    f'{model_type}.{name}: "ui_type" is mutually exclusive with "ui_model_[base|type|format|variant]", ignoring "ui_type"'
+                )
+                field.json_schema_extra.pop("ui_type")
+
+            # Case 3: deprecated value that is not migratable
+            elif ui_type.startswith("DEPRECATED_"):
+                logger.warning(f'{model_type}.{name}: Deprecated "ui_type" "{human_readable_name}", ignoring')
+                field.json_schema_extra.pop("ui_type")
+
     return None
 
 
diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py
@@ -54,6 +54,7 @@ class UIType(str, Enum, metaclass=MetaEnum):
     # region Internal Field Types
     _Collection = "CollectionField"
     _CollectionItem = "CollectionItemField"
+    _IsIntermediate = "IsIntermediate"
     # endregion
 
     # region DEPRECATED
@@ -91,7 +92,6 @@ class UIType(str, Enum, metaclass=MetaEnum):
     CollectionItem = "DEPRECATED_CollectionItem"
     Enum = "DEPRECATED_Enum"
     WorkflowField = "DEPRECATED_WorkflowField"
-    IsIntermediate = "DEPRECATED_IsIntermediate"
     BoardField = "DEPRECATED_BoardField"
     MetadataItem = "DEPRECATED_MetadataItem"
     MetadataItemCollection = "DEPRECATED_MetadataItemCollection"
@@ -423,6 +423,7 @@ class InputFieldJSONSchemaExtra(BaseModel):
     model_config = ConfigDict(
         validate_assignment=True,
         json_schema_serialization_defaults_required=True,
+        use_enum_values=True,
     )
 
 
@@ -482,9 +483,114 @@ class OutputFieldJSONSchemaExtra(BaseModel):
     model_config = ConfigDict(
         validate_assignment=True,
         json_schema_serialization_defaults_required=True,
+        use_enum_values=True,
     )
 
 
+def migrate_model_ui_type(ui_type: UIType | str, json_schema_extra: dict[str, Any]) -> bool:
+    """Migrate deprecated model-specifier ui_type values to new-style ui_model_[base|type|variant|format] in json_schema_extra."""
+    if not isinstance(ui_type, UIType):
+        ui_type = UIType(ui_type)
+
+    ui_model_type: list[ModelType] | None = None
+    ui_model_base: list[BaseModelType] | None = None
+    ui_model_format: list[ModelFormat] | None = None
+    ui_model_variant: list[ClipVariantType | ModelVariantType] | None = None
+
+    match ui_type:
+        case UIType.MainModel:
+            ui_model_base = [BaseModelType.StableDiffusion1, BaseModelType.StableDiffusion2]
+            ui_model_type = [ModelType.Main]
+        case UIType.CogView4MainModel:
+            ui_model_base = [BaseModelType.CogView4]
+            ui_model_type = [ModelType.Main]
+        case UIType.FluxMainModel:
+            ui_model_base = [BaseModelType.Flux]
+            ui_model_type = [ModelType.Main]
+        case UIType.SD3MainModel:
+            ui_model_base = [BaseModelType.StableDiffusion3]
+            ui_model_type = [ModelType.Main]
+        case UIType.SDXLMainModel:
+            ui_model_base = [BaseModelType.StableDiffusionXL]
+            ui_model_type = [ModelType.Main]
+        case UIType.SDXLRefinerModel:
+            ui_model_base = [BaseModelType.StableDiffusionXLRefiner]
+            ui_model_type = [ModelType.Main]
+        case UIType.VAEModel:
+            ui_model_type = [ModelType.VAE]
+        case UIType.FluxVAEModel:
+            ui_model_base = [BaseModelType.Flux]
+            ui_model_type = [ModelType.VAE]
+        case UIType.LoRAModel:
+            ui_model_type = [ModelType.LoRA]
+        case UIType.ControlNetModel:
+            ui_model_type = [ModelType.ControlNet]
+        case UIType.IPAdapterModel:
+            ui_model_type = [ModelType.IPAdapter]
+        case UIType.T2IAdapterModel:
+            ui_model_type = [ModelType.T2IAdapter]
+        case UIType.T5EncoderModel:
+            ui_model_type = [ModelType.T5Encoder]
+        case UIType.CLIPEmbedModel:
+            ui_model_type = [ModelType.CLIPEmbed]
+        case UIType.CLIPLEmbedModel:
+            ui_model_type = [ModelType.CLIPEmbed]
+            ui_model_variant = [ClipVariantType.L]
+        case UIType.CLIPGEmbedModel:
+            ui_model_type = [ModelType.CLIPEmbed]
+            ui_model_variant = [ClipVariantType.G]
+        case UIType.SpandrelImageToImageModel:
+            ui_model_type = [ModelType.SpandrelImageToImage]
+        case UIType.ControlLoRAModel:
+            ui_model_type = [ModelType.ControlLoRa]
+        case UIType.SigLipModel:
+            ui_model_type = [ModelType.SigLIP]
+        case UIType.FluxReduxModel:
+            ui_model_type = [ModelType.FluxRedux]
+        case UIType.LlavaOnevisionModel:
+            ui_model_type = [ModelType.LlavaOnevision]
+        case UIType.Imagen3Model:
+            ui_model_base = [BaseModelType.Imagen3]
+            ui_model_type = [ModelType.Main]
+        case UIType.Imagen4Model:
+            ui_model_base = [BaseModelType.Imagen4]
+            ui_model_type = [ModelType.Main]
+        case UIType.ChatGPT4oModel:
+            ui_model_base = [BaseModelType.ChatGPT4o]
+            ui_model_type = [ModelType.Main]
+        case UIType.Gemini2_5Model:
+            ui_model_base = [BaseModelType.Gemini2_5]
+            ui_model_type = [ModelType.Main]
+        case UIType.FluxKontextModel:
+            ui_model_base = [BaseModelType.FluxKontext]
+            ui_model_type = [ModelType.Main]
+        case UIType.Veo3Model:
+            ui_model_base = [BaseModelType.Veo3]
+            ui_model_type = [ModelType.Video]
+        case UIType.RunwayModel:
+            ui_model_base = [BaseModelType.Runway]
+            ui_model_type = [ModelType.Video]
+        case _:
+            pass
+
+    did_migrate = False
+
+    if ui_model_type is not None:
+        json_schema_extra["ui_model_type"] = [m.value for m in ui_model_type]
+        did_migrate = True
+    if ui_model_base is not None:
+        json_schema_extra["ui_model_base"] = [m.value for m in ui_model_base]
+        did_migrate = True
+    if ui_model_format is not None:
+        json_schema_extra["ui_model_format"] = [m.value for m in ui_model_format]
+        did_migrate = True
+    if ui_model_variant is not None:
+        json_schema_extra["ui_model_variant"] = [m.value for m in ui_model_variant]
+        did_migrate = True
+
+    return did_migrate
+
+
 def InputField(
     # copied from pydantic's Field
     # TODO: Can we support default_factory?
@@ -575,93 +681,6 @@ def InputField(
         field_kind=FieldKind.Input,
     )
 
-    if ui_type is not None:
-        if (
-            ui_model_base is not None
-            or ui_model_type is not None
-            or ui_model_variant is not None
-            or ui_model_format is not None
-        ):
-            logger.warning("InputField: Use either ui_type or ui_model_[base|type|variant|format]. Ignoring ui_type.")
-        # Map old-style UIType to new-style ui_model_[base|type|variant|format]
-        elif ui_type is UIType.MainModel:
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.CogView4MainModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.CogView4]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.FluxMainModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.Flux]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.SD3MainModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.StableDiffusion3]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.SDXLMainModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.StableDiffusionXL]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.SDXLRefinerModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.StableDiffusionXLRefiner]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        # Think this UIType is unused...?
-        # elif ui_type is UIType.ONNXModel:
-        #     json_schema_extra_.ui_model_base =
-        #     json_schema_extra_.ui_model_type =
-        elif ui_type is UIType.VAEModel:
-            json_schema_extra_.ui_model_type = [ModelType.VAE]
-        elif ui_type is UIType.FluxVAEModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.Flux]
-            json_schema_extra_.ui_model_type = [ModelType.VAE]
-        elif ui_type is UIType.LoRAModel:
-            json_schema_extra_.ui_model_type = [ModelType.LoRA]
-        elif ui_type is UIType.ControlNetModel:
-            json_schema_extra_.ui_model_type = [ModelType.ControlNet]
-        elif ui_type is UIType.IPAdapterModel:
-            json_schema_extra_.ui_model_type = [ModelType.IPAdapter]
-        elif ui_type is UIType.T2IAdapterModel:
-            json_schema_extra_.ui_model_type = [ModelType.T2IAdapter]
-        elif ui_type is UIType.T5EncoderModel:
-            json_schema_extra_.ui_model_type = [ModelType.T5Encoder]
-        elif ui_type is UIType.CLIPEmbedModel:
-            json_schema_extra_.ui_model_type = [ModelType.CLIPEmbed]
-        elif ui_type is UIType.CLIPLEmbedModel:
-            json_schema_extra_.ui_model_type = [ModelType.CLIPEmbed]
-            json_schema_extra_.ui_model_variant = [ClipVariantType.L]
-        elif ui_type is UIType.CLIPGEmbedModel:
-            json_schema_extra_.ui_model_type = [ModelType.CLIPEmbed]
-            json_schema_extra_.ui_model_variant = [ClipVariantType.G]
-        elif ui_type is UIType.SpandrelImageToImageModel:
-            json_schema_extra_.ui_model_type = [ModelType.SpandrelImageToImage]
-        elif ui_type is UIType.ControlLoRAModel:
-            json_schema_extra_.ui_model_type = [ModelType.ControlLoRa]
-        elif ui_type is UIType.SigLipModel:
-            json_schema_extra_.ui_model_type = [ModelType.SigLIP]
-        elif ui_type is UIType.FluxReduxModel:
-            json_schema_extra_.ui_model_type = [ModelType.FluxRedux]
-        elif ui_type is UIType.LlavaOnevisionModel:
-            json_schema_extra_.ui_model_type = [ModelType.LlavaOnevision]
-        elif ui_type is UIType.Imagen3Model:
-            json_schema_extra_.ui_model_base = [BaseModelType.Imagen3]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.Imagen4Model:
-            json_schema_extra_.ui_model_base = [BaseModelType.Imagen4]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.ChatGPT4oModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.ChatGPT4o]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.Gemini2_5Model:
-            json_schema_extra_.ui_model_base = [BaseModelType.Gemini2_5]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.FluxKontextModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.FluxKontext]
-            json_schema_extra_.ui_model_type = [ModelType.Main]
-        elif ui_type is UIType.Veo3Model:
-            json_schema_extra_.ui_model_base = [BaseModelType.Veo3]
-            json_schema_extra_.ui_model_type = [ModelType.Video]
-        elif ui_type is UIType.RunwayModel:
-            json_schema_extra_.ui_model_base = [BaseModelType.Runway]
-            json_schema_extra_.ui_model_type = [ModelType.Video]
-        else:
-            json_schema_extra_.ui_type = ui_type
-
     if ui_component is not None:
         json_schema_extra_.ui_component = ui_component
     if ui_hidden is not None:
@@ -690,6 +709,8 @@ def InputField(
             json_schema_extra_.ui_model_format = ui_model_format
         else:
             json_schema_extra_.ui_model_format = [ui_model_format]
+    if ui_type is not None:
+        json_schema_extra_.ui_type = ui_type
 
     """
     There is a conflict between the typing of invocation definitions and the typing of an invocation's
