Added Differential Output Preservation Loss to trainer and ui

toolkit/config_modules.py

@@ -341,6 +341,12 @@ class TrainConfig:
         # unmasked reign. It is unmasked regularization basically
         self.inverted_mask_prior = kwargs.get('inverted_mask_prior', False)
         self.inverted_mask_prior_multiplier = kwargs.get('inverted_mask_prior_multiplier', 0.5)
+        
+        # DOP will will run the same image and prompt through the network without the trigger word blank and use it as a target
+        self.diff_output_preservation = kwargs.get('diff_output_preservation', False)
+        self.diff_output_preservation_multiplier = kwargs.get('diff_output_preservation_multiplier', 1.0)
+        # If the trigger word is in the prompt, we will use this class name to replace it eg. "sks woman" -> "woman"
+        self.diff_output_preservation_class = kwargs.get('diff_output_preservation_class', '')
 
         # legacy
         if match_adapter_assist and self.match_adapter_chance == 0.0:

toolkit/prompt_utils.py

@@ -62,6 +62,20 @@ class PromptEmbeds:
             prompt_embeds.attention_mask = self.attention_mask.clone()
         return prompt_embeds
 
+    def expand_to_batch(self, batch_size):
+        pe = self.clone()
+        current_batch_size = pe.text_embeds.shape[0]
+        if current_batch_size == batch_size:
+            return pe
+        if current_batch_size != 1:
+            raise Exception("Can only expand batch size for batch size 1")
+        pe.text_embeds = pe.text_embeds.expand(batch_size, -1)
+        if pe.pooled_embeds is not None:
+            pe.pooled_embeds = pe.pooled_embeds.expand(batch_size, -1)
+        if pe.attention_mask is not None:
+            pe.attention_mask = pe.attention_mask.expand(batch_size, -1)
+        return pe
+
 
 class EncodedPromptPair:
     def __init__(