aboutsummaryrefslogtreecommitdiffstats
path: root/modules/interrogate.py
diff options
context:
space:
mode:
authormissionfloyd <missionfloyd@users.noreply.github.com>2023-05-26 00:53:33 +0000
committerGitHub <noreply@github.com>2023-05-26 00:53:33 +0000
commit6645f23c4c715b1bc704c88a499b2f4224d7f1e6 (patch)
tree6aeb51e366254fe8993856a3db341690bb39dca5 /modules/interrogate.py
parent43bdaa2f0eda79c685792b06a2bd84c65806a48f (diff)
parenta6e653be26cc05f4438145fa0082816e9fbbf5fc (diff)
downloadstable-diffusion-webui-gfx803-6645f23c4c715b1bc704c88a499b2f4224d7f1e6.tar.gz
stable-diffusion-webui-gfx803-6645f23c4c715b1bc704c88a499b2f4224d7f1e6.tar.bz2
stable-diffusion-webui-gfx803-6645f23c4c715b1bc704c88a499b2f4224d7f1e6.zip
Merge branch 'dev' into reorder-hotkeys
Diffstat (limited to 'modules/interrogate.py')
-rw-r--r--modules/interrogate.py7
1 files changed, 3 insertions, 4 deletions
diff --git a/modules/interrogate.py b/modules/interrogate.py
index 9f7d657f..111b1322 100644
--- a/modules/interrogate.py
+++ b/modules/interrogate.py
@@ -11,7 +11,6 @@ import torch.hub
from torchvision import transforms
from torchvision.transforms.functional import InterpolationMode
-import modules.shared as shared
from modules import devices, paths, shared, lowvram, modelloader, errors
blip_image_eval_size = 384
@@ -160,7 +159,7 @@ class InterrogateModels:
text_array = text_array[0:int(shared.opts.interrogate_clip_dict_limit)]
top_count = min(top_count, len(text_array))
- text_tokens = clip.tokenize([text for text in text_array], truncate=True).to(devices.device_interrogate)
+ text_tokens = clip.tokenize(list(text_array), truncate=True).to(devices.device_interrogate)
text_features = self.clip_model.encode_text(text_tokens).type(self.dtype)
text_features /= text_features.norm(dim=-1, keepdim=True)
@@ -208,8 +207,8 @@ class InterrogateModels:
image_features /= image_features.norm(dim=-1, keepdim=True)
- for name, topn, items in self.categories():
- matches = self.rank(image_features, items, top_count=topn)
+ for cat in self.categories():
+ matches = self.rank(image_features, cat.items, top_count=cat.topn)
for match, score in matches:
if shared.opts.interrogate_return_ranks:
res += f", ({match}:{score/100:.3f})"