Fix caption dropout so it now works identically when caching text embeddings.

toolkit/dataloader_mixins.py

@@ -2090,6 +2090,9 @@ class TextEmbeddingFileItemDTOMixin:
         # diff output preservation embeds (caption with trigger word replaced by class)
         self.dop_prompt_embeds: Union[PromptEmbeds, None] = None
         self._dop_text_embedding_path: Union[str, None] = None
+        # blank caption embeds used for caption dropout when caching text embeddings
+        self._blank_text_embedding_path: Union[str, None] = None
+        self._loaded_text_embedding_path: Union[str, None] = None
         self.is_text_embedding_cached = False
         self.text_embedding_load_device = 'cpu'
         self.text_embedding_version = 1
@@ -2151,6 +2154,25 @@ class TextEmbeddingFileItemDTOMixin:
 
         return self._dop_text_embedding_path
 
+    def get_dropout_caption(self: 'FileItemDTO'):
+        # when encoding live, dropped captions still get the trigger word injected
+        # downstream (add_if_not_present when not a reg image), so match that here
+        if self.trigger_word is not None and not self.is_reg:
+            return inject_trigger_into_prompt('', trigger=self.trigger_word, add_if_not_present=True)
+        return ''
+
+    def get_blank_text_embedding_path(self: 'FileItemDTO', recalculate=False):
+        if self._blank_text_embedding_path is not None and not recalculate:
+            return self._blank_text_embedding_path
+        else:
+            # if the dropout caption matches the normal caption, this hashes to the
+            # same path as the normal embedding and the cache file is shared
+            self._blank_text_embedding_path = self._build_text_embedding_path(
+                caption_override=self.get_dropout_caption()
+            )
+
+        return self._blank_text_embedding_path
+
     def cleanup_text_embedding(self):
         if self.prompt_embeds is not None:
             # we are caching on disk, don't save in memory
@@ -2162,11 +2184,19 @@ class TextEmbeddingFileItemDTOMixin:
         if not self.is_text_embedding_cached:
             return
         if self.prompt_embeds is None:
+            te_path = self.get_text_embedding_path()
+            if self.dataset_config.caption_dropout_rate > 0:
+                # get a random float form 0 to 1
+                rand = random.random()
+                if rand < self.dataset_config.caption_dropout_rate:
+                    # drop the caption by using the cached blank embedding
+                    te_path = self.get_blank_text_embedding_path()
             # load it from disk
-            self.prompt_embeds = PromptEmbeds.load(self.get_text_embedding_path())
+            self.prompt_embeds = PromptEmbeds.load(te_path)
+            self._loaded_text_embedding_path = te_path
         if self.dataset_config.diff_output_preservation and self.dop_prompt_embeds is None:
             dop_path = self.get_dop_text_embedding_path()
-            if dop_path == self.get_text_embedding_path():
+            if dop_path == self._loaded_text_embedding_path:
                 # no trigger word in caption, same embedding
                 self.dop_prompt_embeds = self.prompt_embeds
             else:
@@ -2199,6 +2229,11 @@ class TextEmbeddingCachingMixin:
                     if dop_path != text_embedding_path:
                         # trigger word was in the caption, cache the DOP version too
                         encode_targets.append((dop_path, file_item.caption_dop))
+                if self.dataset_config.caption_dropout_rate > 0:
+                    blank_path = file_item.get_blank_text_embedding_path(recalculate=True)
+                    if blank_path != text_embedding_path:
+                        # cache the dropout caption embedding (blank, or trigger word only)
+                        encode_targets.append((blank_path, file_item.get_dropout_caption()))
                 # only process if not saved to disk
                 encode_targets = [t for t in encode_targets if not os.path.exists(t[0])]
                 if len(encode_targets) > 0:

ui/src/app/jobs/new/SimpleJob.tsx

@@ -1233,6 +1233,7 @@ export default function SimpleJob({
                       <NumberInput
                         label="Caption Dropout Rate"
                         className="pt-2"
+                        docKey="datasets.caption_dropout_rate"
                         value={dataset.caption_dropout_rate}
                         onChange={value => setJobConfig(value, `config.process[0].datasets[${i}].caption_dropout_rate`)}
                         placeholder="eg. 0.05"

ui/src/docs.tsx

@@ -359,6 +359,24 @@ const docs: { [key: string]: ConfigDoc } = {
       </>
     ),
   },
+  'datasets.caption_dropout_rate': {
+    title: 'Caption Dropout Rate',
+    description: (
+      <>
+        Caption dropout rate is the probability that the caption for an image will be dropped (replaced with a blank
+        caption) for any given training step. For example, a value of 0.05 will drop the caption around 5% of the time.
+        Dropping captions helps the model learn the concept being trained without relying entirely on the caption,
+        and helps preserve the model&apos;s ability to generate without a prompt. If a trigger word is set, the trigger
+        word is still used when the caption is dropped, so the model still associates the dropped samples with your
+        trigger word. Regularization images, or images without a trigger word, drop to a fully blank caption.
+        <br />
+        <br />
+        Caption dropout also works when caching text embeddings. An additional embedding for the dropout caption
+        (blank, or the trigger word alone) is cached to disk alongside the normal one, and it is randomly swapped in
+        at train time at this rate.
+      </>
+    ),
+  },
 };
 
 export const getDoc = (key: string | null | undefined): ConfigDoc | null => {