ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: print out the path where the image encode failed (#107)

Browse Source
This commit is contained in:
liaoliaojun
2024-08-23 11:34:35 +08:00
committed by GitHub
parent 34db804c76
commit e127c079da
1 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
toolkit/dataloader_mixins.py
View File

@@ -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([