Merge pull request #3197 from AUTOMATIC1111/training-help-text

Training UI Changes

1 files changed, 28 insertions, 12 deletions

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
index 886cf0c3..6bba3852 100644
--- a/modules/textual_inversion/preprocess.py
+++ b/modules/textual_inversion/preprocess.py
@@ -11,7 +11,7 @@ if cmd_opts.deepdanbooru:
     import modules.deepbooru as deepbooru

 

 

-def preprocess(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):

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

     try:

         if process_caption:

             shared.interrogator.load()

@@ -21,7 +21,7 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
             db_opts[deepbooru.OPT_INCLUDE_RANKS] = False

             deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, db_opts)

 

-        preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru)

+        preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_flip, process_split, process_caption, process_caption_deepbooru)

 

     finally:

 

@@ -33,7 +33,7 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
 

 

 

-def preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):

+def preprocess_work(process_src, process_dst, process_width, process_height, preprocess_txt_action, process_flip, process_split, process_caption, process_caption_deepbooru=False):

     width = process_width

     height = process_height

     src = os.path.abspath(process_src)

@@ -48,7 +48,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
     shared.state.textinfo = "Preprocessing..."

     shared.state.job_count = len(files)

 

-    def save_pic_with_caption(image, index):

+    def save_pic_with_caption(image, index, existing_caption=None):

         caption = ""

 

         if process_caption:

@@ -66,17 +66,26 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
         basename = f"{index:05}-{subindex[0]}-{filename_part}"

         image.save(os.path.join(dst, f"{basename}.png"))

 

+        if preprocess_txt_action == 'prepend' and existing_caption:

+            caption = existing_caption + ' ' + caption

+        elif preprocess_txt_action == 'append' and existing_caption:

+            caption = caption + ' ' + existing_caption

+        elif preprocess_txt_action == 'copy' and existing_caption:

+            caption = existing_caption

+

+        caption = caption.strip()

+        

         if len(caption) > 0:

             with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file:

                 file.write(caption)

 

         subindex[0] += 1

 

-    def save_pic(image, index):

-        save_pic_with_caption(image, index)

+    def save_pic(image, index, existing_caption=None):

+        save_pic_with_caption(image, index, existing_caption=existing_caption)

 

         if process_flip:

-            save_pic_with_caption(ImageOps.mirror(image), index)

+            save_pic_with_caption(ImageOps.mirror(image), index, existing_caption=existing_caption)

 

     for index, imagefile in enumerate(tqdm.tqdm(files)):

         subindex = [0]

@@ -86,6 +95,13 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
         except Exception:

             continue

 

+        existing_caption = None

+

+        try:

+            existing_caption = open(os.path.splitext(filename)[0] + '.txt', 'r').read()

+        except Exception as e:

+            print(e)

+

         if shared.state.interrupted:

             break

 

@@ -97,20 +113,20 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
             img = img.resize((width, height * img.height // img.width))

 

             top = img.crop((0, 0, width, height))

-            save_pic(top, index)

+            save_pic(top, index, existing_caption=existing_caption)

 

             bot = img.crop((0, img.height - height, width, img.height))

-            save_pic(bot, index)

+            save_pic(bot, index, existing_caption=existing_caption)

         elif process_split and is_wide:

             img = img.resize((width * img.width // img.height, height))

 

             left = img.crop((0, 0, width, height))

-            save_pic(left, index)

+            save_pic(left, index, existing_caption=existing_caption)

 

             right = img.crop((img.width - width, 0, img.width, height))

-            save_pic(right, index)

+            save_pic(right, index, existing_caption=existing_caption)

         else:

             img = images.resize_image(1, img, width, height)

-            save_pic(img, index)

+            save_pic(img, index, existing_caption=existing_caption)

 

         shared.state.nextjob()