Merge branch 'master' of https://github.com/yfszzx/stable-diffusion-webui-plus

1 files changed, 5 insertions, 2 deletions

diff --git a/modules/interrogate.py b/modules/interrogate.py
index 635e266e..af858cc0 100644
--- a/modules/interrogate.py
+++ b/modules/interrogate.py
@@ -123,7 +123,7 @@ class InterrogateModels:
 

         return caption[0]

 

-    def interrogate(self, pil_image):

+    def interrogate(self, pil_image, include_ranks=False):

         res = None

 

         try:

@@ -156,7 +156,10 @@ class InterrogateModels:
                 for name, topn, items in self.categories:

                     matches = self.rank(image_features, items, top_count=topn)

                     for match, score in matches:

-                        res += ", " + match

+                        if include_ranks:

+                            res += ", " + match

+                        else:

+                            res += f", ({match}:{score})"

 

         except Exception:

             print(f"Error interrogating", file=sys.stderr)