From a7e6279cdb34791a990b39d6ce5e48db0a23dd65 Mon Sep 17 00:00:00 2001 From: kandy Date: Wed, 22 May 2024 00:16:41 -0500 Subject: [PATCH 1/2] Add cache --- wd14tagger.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/wd14tagger.py b/wd14tagger.py index c88f583..1c6efc9 100644 --- a/wd14tagger.py +++ b/wd14tagger.py @@ -7,6 +7,7 @@ import csv import os import sys +import hashlib import onnxruntime as ort from onnxruntime import InferenceSession from PIL import Image @@ -179,6 +180,7 @@ def INPUT_TYPES(s): "character_threshold": ("FLOAT", {"default": defaults["character_threshold"], "min": 0.0, "max": 1, "step": 0.05}), "replace_underscore": ("BOOLEAN", {"default": defaults["replace_underscore"]}), "trailing_comma": ("BOOLEAN", {"default": defaults["trailing_comma"]}), + "cache": ("BOOLEAN", {"default": False}), "exclude_tags": ("STRING", {"default": defaults["exclude_tags"]}), }} @@ -189,7 +191,16 @@ def INPUT_TYPES(s): CATEGORY = "image" - def tag(self, image, model, threshold, character_threshold, exclude_tags="", replace_underscore=False, trailing_comma=False): + def tag(self, image, model, threshold, character_threshold, exclude_tags="", replace_underscore=False, trailing_comma=False, cache=False): + + hash = image.__hash__() + + # use tags hasing + if cache and hasattr(self, "imghash") and self.imghash == hash: + print(hash) + return {"ui": {"tags": self.tags}, "result": (self.tags,)} + + tensor = image*255 tensor = np.array(tensor, dtype=np.uint8) @@ -199,6 +210,8 @@ def tag(self, image, model, threshold, character_threshold, exclude_tags="", rep image = Image.fromarray(tensor[i]) tags.append(wait_for_async(lambda: tag(image, model, threshold, character_threshold, exclude_tags, replace_underscore, trailing_comma))) pbar.update(1) + self.imghash = hash + self.tags = tags return {"ui": {"tags": tags}, "result": (tags,)} From 7c66893db4ea68561c878d1ac5a4cfc7ea04b36e Mon Sep 17 00:00:00 2001 From: Andrii Kasian Date: Sat, 25 May 2024 13:12:07 -0500 Subject: [PATCH 2/2] Update wd14tagger.py --- wd14tagger.py | 1 - 1 file changed, 1 deletion(-) diff --git a/wd14tagger.py b/wd14tagger.py index 1c6efc9..4418f82 100644 --- a/wd14tagger.py +++ b/wd14tagger.py @@ -7,7 +7,6 @@ import csv import os import sys -import hashlib import onnxruntime as ort from onnxruntime import InferenceSession from PIL import Image