implement workaround for FastAPI and discriminated unions in Body parameter

invokeai/app/api/routers/model_records.py

@@ -3,6 +3,8 @@
 
 
 from typing import List, Optional
+from hashlib import sha1
+from random import randbytes
 
 from fastapi import Body, Path, Query, Response
 from fastapi.routing import APIRouter
@@ -34,8 +36,8 @@ ModelsListValidator = TypeAdapter(ModelsList)
     operation_id="list_model_records",
 )
 async def list_model_records(
-    base_models: Optional[List[BaseModelType]] = Query(default=None, description="Base models to include"),
+        base_models: Optional[List[BaseModelType]] = Query(default=None, description="Base models to include"),
-    model_type: Optional[ModelType] = Query(default=None, description="The type of model to get"),
+        model_type: Optional[ModelType] = Query(default=None, description="The type of model to get"),
 ) -> ModelsList:
     """Get a list of models."""
     record_store = ApiDependencies.invoker.services.model_records
@@ -61,7 +63,7 @@ async def list_model_records(
     },
 )
 async def get_model_record(
-    key: str = Path(description="Key of the model record to fetch."),
+        key: str = Path(description="Key of the model record to fetch."),
 ) -> AnyModelConfig:
     """Get a model record"""
     record_store = ApiDependencies.invoker.services.model_records
@@ -84,9 +86,8 @@ async def get_model_record(
     response_model=AnyModelConfig,
 )
 async def update_model_record(
-    key: Annotated[str, Path(description="Unique key of model")],
+        key: Annotated[str, Path(description="Unique key of model")],
-    # info: Annotated[AnyModelConfig, Body(description="Model configuration")],
+        info: Annotated[AnyModelConfig, Body(description="Model config", discriminator="type")]
-    info: AnyModelConfig,
 ) -> AnyModelConfig:
     """Update model contents with a new config. If the model name or base fields are changed, then the model is renamed."""
     logger = ApiDependencies.invoker.services.logger
@@ -134,7 +135,7 @@ async def del_model_record(
     status_code=201,
 )
 async def add_model_record(
-    config: AnyModelConfig,
+    config: Annotated[AnyModelConfig,  Body(description="Model config", discriminator="type")]
 ) -> AnyModelConfig:
     """
     Add a model using the configuration information appropriate for its type.

invokeai/backend/model_manager/config.py

@@ -257,21 +257,34 @@ _ControlNetConfig = Annotated[
 _VaeConfig = Annotated[Union[VaeDiffusersConfig, VaeCheckpointConfig], Field(discriminator="format")]
 _MainModelConfig = Annotated[Union[MainDiffusersConfig, MainCheckpointConfig], Field(discriminator="format")]
 
-AnyModelConfig = Annotated[
+AnyModelConfig = Union[
-    Union[
+    _MainModelConfig,
-        _MainModelConfig,
+    _ONNXConfig,
-        _ONNXConfig,
+    _VaeConfig,
-        _VaeConfig,
+    _ControlNetConfig,
-        _ControlNetConfig,
+    LoRAConfig,
-        LoRAConfig,
+    TextualInversionConfig,
-        TextualInversionConfig,
+    IPAdapterConfig,
-        IPAdapterConfig,
+    CLIPVisionDiffusersConfig,
-        CLIPVisionDiffusersConfig,
+    T2IConfig,
-        T2IConfig,
-    ],
-    Body(discriminator="type"),
 ]
 
+# Preferred alternative is a discriminated Union, but it breaks FastAPI when applied to a route.
+# AnyModelConfig = Annotated[
+#     Union[
+#         _MainModelConfig,
+#         _ONNXConfig,
+#         _VaeConfig,
+#         _ControlNetConfig,
+#         LoRAConfig,
+#         TextualInversionConfig,
+#         IPAdapterConfig,
+#         CLIPVisionDiffusersConfig,
+#         T2IConfig,
+#     ],
+#     Field(discriminator="type"),
+# ]
+
 AnyModelConfigValidator = TypeAdapter(AnyModelConfig)
 
 
@@ -295,11 +308,11 @@ class ModelConfigFactory(object):
         be selected automatically.
         """
         if isinstance(model_data, ModelConfigBase):
-            if key:
+            model = model_data
-                model_data.key = key
+        elif dest_class:
-            return model_data
+            model = dest_class.validate_python(model_data)
         else:
             model = AnyModelConfigValidator.validate_python(model_data)
-            if key:
+        if key:
-                model.key = key
+            model.key = key
-            return model
+        return model

tests/app/services/model_records/test_model_records_sql.py

@@ -44,6 +44,12 @@ def example_config() -> TextualInversionConfig:
     )
 
 
+def test_type(store: ModelRecordServiceBase):
+    config = example_config()
+    store.add_model("key1", config)
+    config1 = store.get_model("key1")
+    assert type(config1) == TextualInversionConfig
+    
 def test_add(store: ModelRecordServiceBase):
     raw = dict(
         path="/tmp/foo.ckpt",