diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py index b21d1f425b..ccb1443120 100644 --- a/invokeai/app/invocations/fields.py +++ b/invokeai/app/invocations/fields.py @@ -7,6 +7,7 @@ from pydantic_core import PydanticUndefined 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