Add ability to recover from a truncated image file in the dataset

toolkit/data_loader.py

@@ -611,9 +611,29 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
             return len(self.batch_indices)
         return len(self.file_list)
 
-    def _get_single_item(self, index) -> 'FileItemDTO':
+    def _get_replacement_index(self, index) -> int:
+        # when an image fails to load we have to swap in a different one. With buckets the
+        # replacement must come from the same bucket so the collated shapes still match.
+        if self.dataset_config.buckets:
+            for bucket in self.buckets.values():
+                if index in bucket.file_list_idx:
+                    candidates = [i for i in bucket.file_list_idx if i != index]
+                    if candidates:
+                        return random.choice(candidates)
+                    break
+        return random.randint(0, len(self.file_list) - 1)
+
+    def _get_single_item(self, index, _attempts=0) -> 'FileItemDTO':
         file_item: 'FileItemDTO' = copy.deepcopy(self.file_list[index])
-        file_item.load_and_process_image(self.transform)
+        try:
+            file_item.load_and_process_image(self.transform)
+        except Exception as e:
+            print(f"Error loading image, skipping and loading a different one: {file_item.path} ({e})")
+            if _attempts >= 10:
+                # avoid infinite recursion if many files are corrupt
+                raise
+            new_index = self._get_replacement_index(index)
+            return self._get_single_item(new_index, _attempts=_attempts + 1)
         file_item.load_caption(self.caption_dict)
         return file_item