hoelee/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-08-05 03:32:37 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Textual Inversion: Preprocess and Training will only pick-up image files

Browse Source
This commit is contained in:
alg-wiki
2022-10-11 04:30:13 +09:00
parent 04c745ea4f
commit bc3e183b73
3 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
modules/textual_inversion/dataset.py
View File

@@ -22,6 +22,7 @@ class PersonalizedBase(Dataset):
self.width = width
self.height = height
self.flip = transforms.RandomHorizontalFlip(p=flip_p)
self.extns = [".jpg",".jpeg",".png"]
self.dataset = []
@@ -32,7 +33,7 @@ class PersonalizedBase(Dataset):
assert data_root, 'dataset directory not specified'
self.image_paths = [os.path.join(data_root, file_path) for file_path in os.listdir(data_root)]
self.image_paths = [os.path.join(data_root, file_path) for file_path in os.listdir(data_root) if os.path.splitext(file_path.casefold())[1] in self.extns]
print("Preparing dataset...")
for path in tqdm.tqdm(self.image_paths):
image = Image.open(path)