add vqa api

modules/api/api.py

@@ -78,6 +78,7 @@ class Api:
         # functional api
         self.add_api_route("/sdapi/v1/png-info", endpoints.post_pnginfo, methods=["POST"], response_model=models.ResImageInfo)
         self.add_api_route("/sdapi/v1/interrogate", endpoints.post_interrogate, methods=["POST"])
+        self.add_api_route("/sdapi/v1/vqa", endpoints.post_vqa, methods=["POST"])
         self.add_api_route("/sdapi/v1/refresh-checkpoints", endpoints.post_refresh_checkpoints, methods=["POST"])
         self.add_api_route("/sdapi/v1/unload-checkpoint", endpoints.post_unload_checkpoint, methods=["POST"])
         self.add_api_route("/sdapi/v1/reload-checkpoint", endpoints.post_reload_checkpoint, methods=["POST"])

modules/api/endpoints.py

@@ -100,6 +100,16 @@ def post_interrogate(req: models.ReqInterrogate):
             medium, artist, movement, trending, flavor = analyze_image(image, model=req.model)
             return models.ResInterrogate(caption=caption, medium=medium, artist=artist, movement=movement, trending=trending, flavor=flavor)
 
+def post_vqa(req: models.ReqVQA):
+    if req.image is None or len(req.image) < 64:
+        raise HTTPException(status_code=404, detail="Image not found")
+    image = helpers.decode_base64_to_image(req.image)
+    image = image.convert('RGB')
+    from modules import vqa
+    print('HERE', req.question, req.model)
+    answer = vqa.interrogate(req.question, image, req.model)
+    return models.ResVQA(answer=answer)
+
 def post_unload_checkpoint():
     from modules import sd_models
     sd_models.unload_model_weights(op='model')

modules/api/models.py

@@ -305,6 +305,14 @@ class ResInterrogate(BaseModel):
     trending: Optional[str] = Field(default=None, title="Medium", description="Image trending.")
     flavor: Optional[str] = Field(default=None, title="Medium", description="Image flavor.")
 
+class ReqVQA(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="Moondream 2", title="Model", description="The interrogate model used.")
+    question: str = Field(default="describe the image", title="Question", description="Question to ask the model.")
+
+class ResVQA(BaseModel):
+    answer: Optional[str] = Field(default=None, title="Answer", description="The generated answer for the image.")
+
 class ResTrain(BaseModel):
     info: str = Field(title="Train info", description="Response string from train embedding or hypernetwork task.")
 

modules/vqa.py

@@ -124,13 +124,18 @@ def moondream(question: str, image: Image.Image, repo: str = None):
     return response
 
 
-def interrogate(vqa_question, vqa_image, vqa_model):
-    vqa_model = MODELS.get(vqa_model, None)
+def interrogate(vqa_question, vqa_image, vqa_model_req):
+    vqa_model = MODELS.get(vqa_model_req, None)
     shared.log.debug(f'VQA: model="{vqa_model}" question="{vqa_question}" image={vqa_image}')
     if vqa_image is None:
         answer = 'no image provided'
-    if vqa_model is None:
+        return answer
+    if vqa_model_req is None:
         answer = 'no model selected'
+        return answer
+    if vqa_model is None:
+        answer = f'unknown: model={vqa_model_req} available={MODELS.keys()}'
+        return answer
     if 'git' in vqa_model.lower():
         answer = git(vqa_question, vqa_image, vqa_model)
     if 'vilt' in vqa_model.lower():