in textual inversion training, skip over non-image files

- Closes #2715
2024-08-30 20:32:17 +00:00 · 2023-02-20 14:44:10 -05:00 · 2023-02-20 14:44:10 -05:00 · 47ddc00c6a
commit 47ddc00c6a
parent 63e790b79b
1 changed files with 1 additions and 0 deletions
--- a/ldm/invoke/training/textual_inversion_training.py
+++ b/ldm/invoke/training/textual_inversion_training.py
@ -441,6 +441,7 @@ class TextualInversionDataset(Dataset):
        self.image_paths = [
            os.path.join(self.data_root, file_path)
            for file_path in os.listdir(self.data_root)
+            if os.path.isfile(file_path) and file_path.endswith(('.png','.PNG','.jpg','.JPG','.jpeg','.JPEG','.gif','.GIF'))
        ]

        self.num_images = len(self.image_paths)