During textual inversion training, skip over non-image files (#2747)

- The TI script was looping over all files in the training image
directory, regardless of whether they were image files or not. This PR
adds a check for image file extensions.
- 
- Closes #2715
This commit is contained in:
Lincoln Stein 2023-02-20 16:17:32 -05:00 committed by GitHub
commit bac6b50dd1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

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