add interrogateapi

AIRI-AUBE · May 22, 2023 · c74d4fe · c74d4fe
1 parent 17ea042
commit c74d4fe
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/modules/api/api.py b/modules/api/api.py
@@ -234,7 +234,7 @@ def __init__(self, app: FastAPI, queue_lock: Lock):
         self.add_api_route("/sdapi/v1/unload-checkpoint", self.unloadapi, methods=["POST"])
         self.add_api_route("/sdapi/v1/reload-checkpoint", self.reloadapi, methods=["POST"])
         self.add_api_route("/sdapi/v1/scripts", self.get_scripts_list, methods=["GET"], response_model=ScriptsList)
-        self.add_api_route("/invocations", self.invocations, methods=["POST"], response_model=Union[TextToImageResponse, ImageToImageResponse, ExtrasSingleImageResponse, ExtrasBatchImagesResponse, InvocationsErrorResponse])
+        self.add_api_route("/invocations", self.invocations, methods=["POST"], response_model=Union[TextToImageResponse, ImageToImageResponse, ExtrasSingleImageResponse, ExtrasBatchImagesResponse, InvocationsErrorResponse, InterrogateResponse])
         self.add_api_route("/ping", self.ping, methods=["GET"], response_model=PingResponse)
 
         self.default_script_arg_txt2img = []
@@ -806,6 +806,9 @@ def invocations(self, req: InvocationsRequest):
                 response = self.extras_batch_images_api(req.extras_batch_payload)
                 response.images = self.post_invocations(response.images, quality)
                 return response
+            elif req.task == 'interrogate':
+                response = self.interrogateapi(req.interrogate_payload)
+                return response
             else:
                 return InvocationsErrorResponse(error = f'Invalid task - {req.task}')
 

diff --git a/modules/api/models.py b/modules/api/models.py
@@ -302,6 +302,7 @@ class InvocationsRequest(BaseModel):
     img2img_payload: Optional[StableDiffusionImg2ImgProcessingAPI]
     extras_single_payload: Optional[ExtrasSingleImageRequest]
     extras_batch_payload: Optional[ExtrasBatchImagesRequest]
+    interrogate_payload: Optional[InterrogateRequest]
 
 class InvocationsErrorResponse(BaseModel):
     error: str = Field(title="Invocation error", description="Error response from invocation.")