From e282d2ee7ae3b37b3a5d87984b3e8ed1a98b82af Mon Sep 17 00:00:00 2001
From: Kent Keirsey <31807370+hipsterusername@users.noreply.github.com>
Date: Tue, 28 Mar 2023 10:39:20 -0400
Subject: [PATCH] Updated to fix Annotated pydantic errors on modelInfo

---
 invokeai/app/api/routers/models.py | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/invokeai/app/api/routers/models.py b/invokeai/app/api/routers/models.py
index e58accbdbd..d7cc6a0111 100644
--- a/invokeai/app/api/routers/models.py
+++ b/invokeai/app/api/routers/models.py
@@ -34,19 +34,17 @@ class DiffusersModelInfo(ModelInfo):
     repo_id: Optional[str] = Field(description="The repo ID to use for this model")
     path: Optional[str] = Field(description="The path to the model")
 
+class modelInfo(ModelInfo):
+    info: Annotated[Union[CkptModelInfo,DiffusersModelInfo], Field(discriminator="format")] 
+
 class CreateModelRequest (BaseModel):
     name: str = Field(description="The name of the model")
-    info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model details and configuration")
-
-class CreateModelResponse (BaseModel):
-    name: str = Field(description="The name of the new model")
-    info: Union[CkptModelInfo, DiffusersModelInfo] = Field(..., discriminator="format", description="The model details and configuration")
-    status: str = Field(description="The status of the API response")
-
-class CreateModelResponse (BaseModel):
-    name: str = Field(description="The name of the new model")
     info: Annotated[Union[(CkptModelInfo,DiffusersModelInfo)], Field(discriminator="format")] = Field(description="The model info")
 
+class CreateModelResponse (BaseModel):
+    name: str = Field(description="The name of the new model")
+    info: modelInfo = Field(description="The model details and configuration")
+    status: str = Field(description="The status of the API response")
 
 class ModelsList(BaseModel):
     models: dict[str, Annotated[Union[(CkptModelInfo,DiffusersModelInfo)], Field(discriminator="format")]]
@@ -95,7 +93,7 @@ async def update_model(
             model_attributes=model_request["info"],
             clobber=True,
         )
-        model_response = CreateModelResponse(status="success")
+        model_response = CreateModelResponse(name=model_request.name, info=model_request.info, status="success")
 
     except Exception as e:
         # Handle any exceptions thrown during the execution of the method