Add a barebones interrogate API

2 files changed, 36 insertions, 2 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index 6e9d6097..eabdb7b8 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -1,4 +1,4 @@
-from modules.api.models import StableDiffusionTxt2ImgProcessingAPI, StableDiffusionImg2ImgProcessingAPI
+from modules.api.models import StableDiffusionTxt2ImgProcessingAPI, StableDiffusionImg2ImgProcessingAPI, InterrogateAPI
 from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
 from modules.sd_samplers import all_samplers
 from modules.extras import run_pnginfo
@@ -25,6 +25,11 @@ class ImageToImageResponse(BaseModel):
     parameters: Json
     info: Json
 
+class InterrogateResponse(BaseModel):
+    caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")
+    parameters: Json
+    info: Json
+
 
 class Api:
     def __init__(self, app, queue_lock):
@@ -33,6 +38,7 @@ class Api:
         self.queue_lock = queue_lock
         self.app.add_api_route("/sdapi/v1/txt2img", self.text2imgapi, methods=["POST"])
         self.app.add_api_route("/sdapi/v1/img2img", self.img2imgapi, methods=["POST"])
+        self.app.add_api_route("/sdapi/v1/interrogate", self.interrogateapi, methods=["POST"])
 
     def __base64_to_image(self, base64_string):
         # if has a comma, deal with prefix
@@ -118,6 +124,23 @@ class Api:
 
         return ImageToImageResponse(images=b64images, parameters=json.dumps(vars(img2imgreq)), info=processed.js())
 
+    def interrogateapi(self, interrogatereq: InterrogateAPI):
+        image_b64 = interrogatereq.image
+        if image_b64 is None:
+            raise HTTPException(status_code=404, detail="Image not found") 
+
+        populate = interrogatereq.copy(update={ # Override __init__ params
+            }
+        )
+
+        img = self.__base64_to_image(image_b64)
+
+        # Override object param
+        with self.queue_lock:
+            processed = shared.interrogator.interrogate(img)
+        
+        return InterrogateResponse(caption=processed, parameters=json.dumps(vars(interrogatereq)), info=None)
+
     def extrasapi(self):
         raise NotImplementedError
 
diff --git a/modules/api/models.py b/modules/api/models.py
index 079e33d9..8be64749 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -63,7 +63,12 @@ class PydanticModelGenerator:
             
                 
         self._model_name = model_name
-        self._class_data = merge_class_params(class_instance)
+
+        if class_instance is not None:
+            self._class_data = merge_class_params(class_instance)
+        else:
+            self._class_data = {}            
+
         self._model_def = [
             ModelDef(
                 field=underscore(k),
@@ -105,4 +110,10 @@ StableDiffusionImg2ImgProcessingAPI = PydanticModelGenerator(
     "StableDiffusionProcessingImg2Img", 
     StableDiffusionProcessingImg2Img,
     [{"key": "sampler_index", "type": str, "default": "Euler"}, {"key": "init_images", "type": list, "default": None}, {"key": "denoising_strength", "type": float, "default": 0.75}, {"key": "mask", "type": str, "default": None}, {"key": "include_init_images", "type": bool, "default": False, "exclude" : True}]
+).generate_model()
+
+InterrogateAPI = PydanticModelGenerator(
+    "Interrogate", 
+    None,
+    [{"key": "image", "type": str, "default": None}]
 ).generate_model()
\ No newline at end of file