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
@@ -85,8 +87,7 @@ async def get_model_record(
 )
 async def update_model_record(
         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,8 +257,7 @@ _ControlNetConfig = Annotated[
 _VaeConfig = Annotated[Union[VaeDiffusersConfig, VaeCheckpointConfig], Field(discriminator="format")]
 _MainModelConfig = Annotated[Union[MainDiffusersConfig, MainCheckpointConfig], Field(discriminator="format")]
 
-AnyModelConfig = Annotated[
+AnyModelConfig = Union[
-    Union[
     _MainModelConfig,
     _ONNXConfig,
     _VaeConfig,
@@ -268,10 +267,24 @@ AnyModelConfig = Annotated[
     IPAdapterConfig,
     CLIPVisionDiffusersConfig,
     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,9 +308,9 @@ 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:

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",