michael/stable-diffusion-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

add existing caption file handling

Browse source
This commit is contained in:
DepFA 2022-10-20 00:46:54 +01:00 committed by GitHub
parent 4d6b9f76a5
commit fbcce66601
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 24 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
modules/textual_inversion/preprocess.py
View file

@ -48,7 +48,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
shared.state.textinfo = "Preprocessing..." shared.state.textinfo = "Preprocessing..."
shared.state.job_count = len(files) shared.state.job_count = len(files)
def save_pic_with_caption(image, index): def save_pic_with_caption(image, index, existing_caption=None):
caption = "" caption = ""
if process_caption: if process_caption:
@ -66,17 +66,26 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
basename = f"{index:05}-{subindex[0]}-{filename_part}" basename = f"{index:05}-{subindex[0]}-{filename_part}"
image.save(os.path.join(dst, f"{basename}.png")) image.save(os.path.join(dst, f"{basename}.png"))
if preprocess_txt_action == 'prepend' and existing_caption:
caption = existing_caption + ' ' + caption
elif preprocess_txt_action == 'append' and existing_caption:
caption = caption + ' ' + existing_caption
elif preprocess_txt_action == 'copy' and existing_caption:
caption = existing_caption
caption = caption.strip()
if len(caption) > 0: if len(caption) > 0:
with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file: with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file:
file.write(caption) file.write(caption)
subindex[0] += 1 subindex[0] += 1
def save_pic(image, index): def save_pic(image, index, existing_caption=None):
save_pic_with_caption(image, index) save_pic_with_caption(image, index)
if process_flip: if process_flip:
save_pic_with_caption(ImageOps.mirror(image), index) save_pic_with_caption(ImageOps.mirror(image), index, existing_caption=existing_caption)
for index, imagefile in enumerate(tqdm.tqdm(files)): for index, imagefile in enumerate(tqdm.tqdm(files)):
subindex = [0] subindex = [0]
@ -86,6 +95,13 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
except Exception: except Exception:
continue continue
existing_caption = None
try:
existing_caption = open(os.path.splitext(filename)[0] + '.txt', 'r').read()
except Exception as e:
print(e)
if shared.state.interrupted: if shared.state.interrupted:
break break
@ -97,20 +113,20 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pro
img = img.resize((width, height * img.height // img.width)) img = img.resize((width, height * img.height // img.width))
top = img.crop((0, 0, width, height)) top = img.crop((0, 0, width, height))
save_pic(top, index) save_pic(top, index, existing_caption=existing_caption)
bot = img.crop((0, img.height - height, width, img.height)) bot = img.crop((0, img.height - height, width, img.height))
save_pic(bot, index) save_pic(bot, index, existing_caption=existing_caption)
elif process_split and is_wide: elif process_split and is_wide:
img = img.resize((width * img.width // img.height, height)) img = img.resize((width * img.width // img.height, height))
left = img.crop((0, 0, width, height)) left = img.crop((0, 0, width, height))
save_pic(left, index) save_pic(left, index, existing_caption=existing_caption)
right = img.crop((img.width - width, 0, img.width, height)) right = img.crop((img.width - width, 0, img.width, height))
save_pic(right, index) save_pic(right, index, existing_caption=existing_caption)
else: else:
img = images.resize_image(1, img, width, height) img = images.resize_image(1, img, width, height)
save_pic(img, index) save_pic(img, index, existing_caption=existing_caption)
shared.state.nextjob() shared.state.nextjob()