Added some split prompting started code, adamw8bit, replacements improving, learnable snr gos. A lot of good stuff.

extensions_built_in/dataset_tools/SuperTagger.py

@@ -12,7 +12,7 @@ from .tools.dataset_tools_config_modules import RAW_DIR, TRAIN_DIR, Step, ImgInf
 from .tools.fuyu_utils import FuyuImageProcessor
 from .tools.image_tools import load_image, ImageProcessor, resize_to_max
 from .tools.llava_utils import LLaVAImageProcessor
-from .tools.caption import default_long_prompt, default_short_prompt
+from .tools.caption import default_long_prompt, default_short_prompt, default_replacements
 from jobs.process import BaseExtensionProcess
 from .tools.sync_tools import get_img_paths
 
@@ -39,6 +39,8 @@ class SuperTagger(BaseExtensionProcess):
         self.caption_prompt = config.get('caption_prompt', default_long_prompt)
         self.caption_short_prompt = config.get('caption_short_prompt', default_short_prompt)
         self.force_reprocess_img = config.get('force_reprocess_img', False)
+        self.caption_replacements = config.get('caption_replacements', default_replacements)
+        self.caption_short_replacements = config.get('caption_short_replacements', default_replacements)
         self.master_dataset_dict = OrderedDict()
         self.dataset_master_config_file = config.get('dataset_master_config_file', None)
         if parent_dir is not None and len(self.dataset_paths) == 0:
@@ -118,7 +120,8 @@ class SuperTagger(BaseExtensionProcess):
 
                 img_info.caption = self.image_processor.generate_caption(
                     image=caption_image,
-                    prompt=self.caption_prompt
+                    prompt=self.caption_prompt,
+                    replacements=self.caption_replacements
                 )
                 img_info.mark_step_complete(step)
             elif step == 'caption_short':
@@ -134,7 +137,8 @@ class SuperTagger(BaseExtensionProcess):
                     self.image_processor.load_model()
                 img_info.caption_short = self.image_processor.generate_caption(
                     image=caption_image,
-                    prompt=self.caption_short_prompt
+                    prompt=self.caption_short_prompt,
+                    replacements=self.caption_short_replacements
                 )
                 img_info.mark_step_complete(step)
             elif step == 'contrast_stretch':

extensions_built_in/dataset_tools/tools/caption.py

@@ -33,7 +33,13 @@ def clean_caption(cap, replacements=None):
     cap = " ".join(cap.split())
 
     for replacement in replacements:
-        cap = cap.replace(replacement[0], replacement[1])
+        if replacement[0].startswith('*'):
+            # we are removing all text if it starts with this and the rest matches
+            search_text = replacement[0][1:]
+            if cap.startswith(search_text):
+                cap = ""
+        else:
+            cap = cap.replace(replacement[0].lower(), replacement[1].lower())
 
     cap_list = cap.split(",")
     # trim whitespace

extensions_built_in/dataset_tools/tools/llava_utils.py

@@ -77,7 +77,7 @@ class LLaVAImageProcessor:
             output_ids = self.model.generate(
                 input_ids, images=image_tensor, do_sample=True, temperature=0.1,
                 max_new_tokens=max_new_tokens, use_cache=True, stopping_criteria=[stopping_criteria],
-                top_p=0.9
+                top_p=0.8
             )
         outputs = self.tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
         conv.messages[-1][-1] = outputs