Adding DeepDanbooru to the interrogation API

2 files changed, 15 insertions, 2 deletions

diff --git a/modules/api/api.py b/modules/api/api.py
index 688469ad..596a6616 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -15,6 +15,9 @@ from modules.sd_models import checkpoints_list
 from modules.realesrgan_model import get_realesrgan_models
 from typing import List
 
+if shared.cmd_opts.deepdanbooru:
+    from modules.deepbooru import get_deepbooru_tags
+
 def upscaler_to_index(name: str):
     try:
         return [x.name.lower() for x in shared.sd_upscalers].index(name.lower())
@@ -220,11 +223,20 @@ class Api:
         if image_b64 is None:
             raise HTTPException(status_code=404, detail="Image not found") 
 
-        img = self.__base64_to_image(image_b64)
+        img = decode_base64_to_image(image_b64)
+        img = img.convert('RGB')
 
         # Override object param
         with self.queue_lock:
-            processed = shared.interrogator.interrogate(img)
+            if interrogatereq.model == "clip":
+                processed = shared.interrogator.interrogate(img)
+            elif interrogatereq.model == "deepdanbooru":
+                if shared.cmd_opts.deepdanbooru:
+                    processed = get_deepbooru_tags(img)
+                else:
+                    raise HTTPException(status_code=404, detail="Model not found. Add --deepdanbooru when launching for using the model.")
+            else:
+                raise HTTPException(status_code=404, detail="Model not found")
         
         return InterrogateResponse(caption=processed)
 
diff --git a/modules/api/models.py b/modules/api/models.py
index 34dbfa16..f9cd929e 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -170,6 +170,7 @@ class ProgressResponse(BaseModel):
 
 class InterrogateRequest(BaseModel):
     image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
+    model: str = Field(default="clip", title="Model", description="The interrogate model used.")
 
 class InterrogateResponse(BaseModel):
     caption: str = Field(default=None, title="Caption", description="The generated caption for the image.")