Added a flag to keep loading the image when latents are cached. Useful for DFE and other methods that target pixelspace losses.

toolkit/config_modules.py

@@ -962,6 +962,7 @@ class DatasetConfig:
         self.cache_latents_to_disk: bool = kwargs.get('cache_latents_to_disk', False)
         self.cache_clip_vision_to_disk: bool = kwargs.get('cache_clip_vision_to_disk', False)
         self.cache_text_embeddings: bool = kwargs.get('cache_text_embeddings', False)
+        self.load_image_when_caching_latents: bool = kwargs.get('load_image_when_caching_latents', False)
 
         self.standardize_images: bool = kwargs.get('standardize_images', False)
 

toolkit/data_transfer_object/data_loader.py

@@ -213,8 +213,9 @@ class DataLoaderBatchDTO:
             
             self.num_frames: int = self.file_items[0].num_frames
 
-            if not is_latents_cached:
-                # only return a tensor if latents are not cached
+            if not is_latents_cached or self.file_items[0].dataset_config.load_image_when_caching_latents:
+                # only return a tensor if latents are not cached, or if we are explicitly
+                # loading the raw image alongside the cached latents
                 self.tensor: torch.Tensor = torch.cat(
                     [x.tensor.unsqueeze(0) for x in self.file_items]
                 )

toolkit/dataloader_mixins.py

@@ -754,17 +754,20 @@ class ImageProcessingDTOMixin:
         # if we are caching latents, just do that
         if self.is_latent_cached:
             self.get_latent()
-            if self.has_control_image:
-                self.load_control_image()
-            if self.has_inpaint_image:
-                self.load_inpaint_image()
-            if self.has_clip_image:
-                self.load_clip_image()
-            if self.has_mask_image:
-                self.load_mask_image()
-            if self.has_unconditional:
-                self.load_unconditional_image()
-            return
+            # if load_image_when_caching_latents is set, we still need the raw image
+            # tensor in addition to the cached latent, so fall through to load it below
+            if not self.dataset_config.load_image_when_caching_latents:
+                if self.has_control_image:
+                    self.load_control_image()
+                if self.has_inpaint_image:
+                    self.load_inpaint_image()
+                if self.has_clip_image:
+                    self.load_clip_image()
+                if self.has_mask_image:
+                    self.load_mask_image()
+                if self.has_unconditional:
+                    self.load_unconditional_image()
+                return
         if self.is_audio_model:
             self.load_and_process_audio()
             return