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

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]
                 )