mirror/stable-diffusion-webui
2
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-11-27 06:40:10 +08:00
Code Issues 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
No known key found for this signature in database
GPG Key ID: 9D27E57D1F43E750
3 changed files with 6 additions and 3 deletions
Show Stats 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)

3
modules/textual_inversion/preprocess.py
View File

@ -12,12 +12,13 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
height = process_height
src = os.path.abspath(process_src)
dst = os.path.abspath(process_dst)
extns = [".jpg",".jpeg",".png"]
assert src != dst, 'same directory specified as source and destination'
os.makedirs(dst, exist_ok=True)
files = os.listdir(src)
files = [i for i in os.listdir(src) if os.path.splitext(i.casefold())[1] in extns]
shared.state.textinfo = "Preprocessing..."
shared.state.job_count = len(files)

3
modules/textual_inversion/textual_inversion.py
View File

@ -161,6 +161,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
shared.state.textinfo = "Initializing textual inversion training..."
shared.state.job_count = steps
extns = [".jpg",".jpeg",".png"]
filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')
@ -200,7 +201,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
if ititial_step > steps:
return embedding, filename
tr_img_len = len([os.path.join(data_root, file_path) for file_path in os.listdir(data_root)])
tr_img_len = len([os.path.join(data_root, file_path) for file_path in os.listdir(data_root) if os.path.splitext(file_path.casefold())[1] in extns])
epoch_len = (tr_img_len * num_repeats) + tr_img_len
pbar = tqdm.tqdm(enumerate(ds), total=steps-ititial_step)