feat(nodes): add ui_model_[base|type|variant] to InputField args for dynamic UI generation

Author: psychedelicious

invokeai/app/invocations/fields.py

@@ -7,6 +7,7 @@
 
 from invokeai.app.util.metaenum import MetaEnum
 from invokeai.backend.image_util.segment_anything.shared import BoundingBox
+from invokeai.backend.model_manager.taxonomy import BaseModelType, ClipVariantType, ModelType, ModelVariantType
 from invokeai.backend.util.logging import InvokeAILogger
 
 logger = InvokeAILogger.get_logger()
@@ -39,7 +40,7 @@ class UIType(str, Enum, metaclass=MetaEnum):
     used, and the type will be ignored. They are included here for backwards compatibility.
     """
 
-    # region Model Field Types
+    # region Quasi-deprecated Model Field Types - use ui_model_[base|type|variant] instead
     MainModel = "MainModelField"
     CogView4MainModel = "CogView4MainModelField"
     FluxMainModel = "FluxMainModelField"
@@ -409,6 +410,9 @@ class InputFieldJSONSchemaExtra(BaseModel):
     ui_component: Optional[UIComponent] = None
     ui_order: Optional[int] = None
     ui_choice_labels: Optional[dict[str, str]] = None
+    ui_model_base: Optional[list[BaseModelType]] = None
+    ui_model_type: Optional[list[ModelType]] = None
+    ui_model_variant: Optional[list[ClipVariantType | ModelVariantType]] = None
 
     model_config = ConfigDict(
         validate_assignment=True,
@@ -501,6 +505,9 @@ def InputField(
     ui_hidden: Optional[bool] = None,
     ui_order: Optional[int] = None,
     ui_choice_labels: Optional[dict[str, str]] = None,
+    ui_model_base: Optional[BaseModelType | list[BaseModelType]] = None,
+    ui_model_type: Optional[ModelType | list[ModelType]] = None,
+    ui_model_variant: Optional[ClipVariantType | ModelVariantType | list[ClipVariantType | ModelVariantType]] = None,
 ) -> Any:
     """
     Creates an input field for an invocation.
@@ -538,7 +545,87 @@ def InputField(
     )
 
     if ui_type is not None:
-        json_schema_extra_.ui_type = ui_type
+        if ui_model_base is not None or ui_model_type is not None or ui_model_variant is not None:
+            logger.warning("Use either ui_type or ui_model_[base|type|variant]. Ignoring ui_type.")
+        # Map old-style UIType to new-style ui_model_[base|type|variant]
+        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:
@@ -547,6 +634,21 @@ def InputField(
         json_schema_extra_.ui_order = ui_order
     if ui_choice_labels is not None:
         json_schema_extra_.ui_choice_labels = ui_choice_labels
+    if ui_model_base is not None:
+        if isinstance(ui_model_base, list):
+            json_schema_extra_.ui_model_base = ui_model_base
+        else:
+            json_schema_extra_.ui_model_base = [ui_model_base]
+    if ui_model_type is not None:
+        if isinstance(ui_model_type, list):
+            json_schema_extra_.ui_model_type = ui_model_type
+        else:
+            json_schema_extra_.ui_model_type = [ui_model_type]
+    if ui_model_variant is not None:
+        if isinstance(ui_model_variant, list):
+            json_schema_extra_.ui_model_variant = ui_model_variant
+        else:
+            json_schema_extra_.ui_model_variant = [ui_model_variant]
 
     """
     There is a conflict between the typing of invocation definitions and the typing of an invocation's