mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
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:
commit
bac6b50dd1
@ -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)
|
||||
|
Loading…
Reference in New Issue
Block a user