refactored the deepbooru module to improve speed on running multiple interogations in a row. Added the option to generate deepbooru tags for textual inversion preproccessing.

1 files changed, 20 insertions, 2 deletions

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index f1c002a2..9f63c9a4 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -3,11 +3,14 @@ from PIL import Image, ImageOps
 import platform

 import sys

 import tqdm

+import time

 

 from modules import shared, images

+from modules.shared import opts, cmd_opts

+if cmd_opts.deepdanbooru:

+    import modules.deepbooru as deepbooru

 

-

-def preprocess(process_src, process_dst, process_flip, process_split, process_caption):

+def preprocess(process_src, process_dst, process_flip, process_split, process_caption, process_caption_deepbooru=False):

     size = 512

     src = os.path.abspath(process_src)

     dst = os.path.abspath(process_dst)

@@ -24,10 +27,21 @@ def preprocess(process_src, process_dst, process_flip, process_split, process_ca
     if process_caption:

         shared.interrogator.load()

 

+    if process_caption_deepbooru:

+        deepbooru.create_deepbooru_process()

+

     def save_pic_with_caption(image, index):

         if process_caption:

             caption = "-" + shared.interrogator.generate_caption(image)

             caption = sanitize_caption(os.path.join(dst, f"{index:05}-{subindex[0]}"), caption, ".png")

+        elif process_caption_deepbooru:

+            shared.deepbooru_process_return["value"] = -1

+            shared.deepbooru_process_queue.put(image)

+            while shared.deepbooru_process_return["value"] == -1:

+                time.sleep(0.2)

+            caption = "-" + shared.deepbooru_process_return["value"]

+            caption = sanitize_caption(os.path.join(dst, f"{index:05}-{subindex[0]}"), caption, ".png")

+            shared.deepbooru_process_return["value"] = -1

         else:

             caption = filename

             caption = os.path.splitext(caption)[0]

@@ -79,6 +93,10 @@ def preprocess(process_src, process_dst, process_flip, process_split, process_ca
     if process_caption:

         shared.interrogator.send_blip_to_ram()

 

+    if process_caption_deepbooru:

+        deepbooru.release_process()

+

+

 def sanitize_caption(base_path, original_caption, suffix):

     operating_system = platform.system().lower()

     if (operating_system == "windows"):