From 689a60c4b9ee06ee0fe0e4ff5b66ef79e6d8407d Mon Sep 17 00:00:00 2001
From: unknown <530541684@qq.com>
Date: Thu, 13 Jun 2024 15:14:01 +0800
Subject: [PATCH] Add WD14 image reverse inference prompt words

---
 webuiapi/webuiapi.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/webuiapi/webuiapi.py b/webuiapi/webuiapi.py
index 576fa51..9e65447 100644
--- a/webuiapi/webuiapi.py
+++ b/webuiapi/webuiapi.py
@@ -401,6 +401,7 @@ def __init__(
         host="127.0.0.1",
         port=7860,
         baseurl=None,
+        taggerurl=None,
         sampler="Euler a",
         scheduler="automatic",
         steps=20,
@@ -413,8 +414,15 @@ def __init__(
                 baseurl = f"https://{host}:{port}/sdapi/v1"
             else:
                 baseurl = f"http://{host}:{port}/sdapi/v1"
+                
+        if taggerurl is None:
+            if use_https:
+                taggerurl = f"https://{host}:{port}/tagger/v1"
+            else:
+                taggerurl = f"http://{host}:{port}/tagger/v1"
 
         self.baseurl = baseurl
+        self.taggerurl = taggerurl
         self.default_sampler = sampler
         self.default_scheduler = scheduler
         self.default_steps = steps
@@ -948,6 +956,20 @@ def interrogate(self, image, model="clip"):
 
         response = self.session.post(url=f"{self.baseurl}/interrogate", json=payload)
         return self._to_api_result(response)
+    
+    def WD14_interrogate(self, image, model="wd14-vit-v2-git",threshold=0.35):
+        payload = {
+            "image": b64_img(image) if isinstance(image, Image.Image) else image,
+            "model": model,
+            'threshold':threshold,
+        }
+
+        response = self.session.post(url=f"{self.taggerurl}/interrogate", json=payload)
+        return self._to_api_result(response)
+    
+    def WD14_interrogators(self):
+        response = self.session.get(url=f"{self.taggerurl}/interrogators")
+        return response.json()
         
     def list_prompt_gen_models(self):
         r = self.custom_get("promptgen/list_models")