aboutsummaryrefslogtreecommitdiffstats
path: root/modules/interrogate.py
diff options
context:
space:
mode:
authorGreg Fuller <gfuller23@gmail.com>2022-10-12 01:02:41 +0000
committerGreg Fuller <gfuller23@gmail.com>2022-10-12 01:02:41 +0000
commitd717eb079cd6b7fa7a4f97c0a10d400bdec753fb (patch)
tree5e63e2a59b4f0de00bf3b7cb960e4f86c23b3be0 /modules/interrogate.py
parent6be32b31d181e42c639dad3451229aa7b9cfd1cf (diff)
downloadstable-diffusion-webui-gfx803-d717eb079cd6b7fa7a4f97c0a10d400bdec753fb.tar.gz
stable-diffusion-webui-gfx803-d717eb079cd6b7fa7a4f97c0a10d400bdec753fb.tar.bz2
stable-diffusion-webui-gfx803-d717eb079cd6b7fa7a4f97c0a10d400bdec753fb.zip
Interrogate: add option to include ranks in output
Since the UI also allows users to specify ranks, it can be useful to show people what ranks are being returned by interrogate This can also give much better results when feeding the interrogate results back into either img2img or txt2img, especially when trying to generate a specific character or scene for which you have a similar concept image Testing Steps: Launch Webui with command line arg: --deepdanbooru Navigate to img2img tab, use interrogate DeepBooru, verify tags appears as before. Use "Interrogate CLIP", verify prompt appears as before Navigate to Settings tab, enable new option, click "apply settings" Navigate to img2img, Interrogate DeepBooru again, verify that weights appear and are properly formatted. Note that "Interrogate CLIP" prompt is still unchanged In my testing, this change has no effect to "Interrogate CLIP", as it seems to generate a sentence-structured caption, and not a set of tags. (reproduce changes from https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/2149/commits/6ed4faac46c45ca7353f228aca9b436bbaba7bc7)
Diffstat (limited to 'modules/interrogate.py')
-rw-r--r--modules/interrogate.py7
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)