From e127c079da2fc6f8cafc03816e030d705a1ae43d Mon Sep 17 00:00:00 2001
From: liaoliaojun <1030219089a@gmail.com>
Date: Fri, 23 Aug 2024 11:34:35 +0800
Subject: [PATCH] fix: print out the path where the image encode failed (#107)

---
 toolkit/dataloader_mixins.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py
index df1784e1..80f00cf0 100644
--- a/toolkit/dataloader_mixins.py
+++ b/toolkit/dataloader_mixins.py
@@ -1385,8 +1385,13 @@ class LatentCachingMixin:
                 dtype = self.sd.torch_dtype
                 device = self.sd.device_torch
                 # add batch dimension
-                imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
-                latent = self.sd.encode_images(imgs).squeeze(0)
+                try:
+                    imgs = file_item.tensor.unsqueeze(0).to(device, dtype=dtype)
+                    latent = self.sd.encode_images(imgs).squeeze(0)
+                except Exception as e:
+                    print(f"Error processing image: {file_item.path}")
+                    print(f"Error: {str(e)}")
+                    raise e
                 # save_latent
                 if to_disk:
                     state_dict = OrderedDict([