1580158
/
stable-diffusion-webui-dataset-tag-editor
mirror of https://github.com/toshiaki1729/stable-diffusion-webui-dataset-tag-editor.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix #97 

and a bug: caption file extension not work on save
bugfix/interrogate-if-empty
toshiaki1729 2024-05-20 19:07:06 +09:00
parent 2f99cc54d8
commit 77388e481d
2 changed files with 14 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
scripts/dataset_tag_editor/dte_logic.py
View File

@ -768,7 +768,13 @@ class DatasetTagEditor(Singleton):
taglists = load_captions(imgpaths)
interrogate_tags = {img_path : [] for img_path in imgpaths}
if interrogate_method != self.InterrogateMethod.NONE:
img_to_interrogate = [
img_path for i, img_path in enumerate(imgpaths)
if (not taglists[i] or interrogate_method != self.InterrogateMethod.PREFILL)
]
if interrogate_method != self.InterrogateMethod.NONE and img_to_interrogate:
logger.write("Preprocess images...")
max_workers = shared.opts.dataset_editor_num_cpu_workers
if max_workers < 0:
@ -781,7 +787,7 @@ class DatasetTagEditor(Singleton):
from concurrent.futures import ThreadPoolExecutor
with ThreadPoolExecutor(max_workers=max_workers) as executor:
result = list(executor.map(convert_rgb, gen_data(imgpaths, self.images)))
result = list(executor.map(convert_rgb, gen_data(img_to_interrogate, self.images)))
logger.write("Preprocess completed")
for tg, th in tqdm(tagger_thresholds):
@ -798,10 +804,10 @@ class DatasetTagEditor(Singleton):
continue
try:
if use_pipe:
for img_path, tags in tqdm(zip(imgpaths, tg.predict_pipe(result, th)), desc=tg.name(), total=len(imgpaths)):
for img_path, tags in tqdm(zip(img_to_interrogate, tg.predict_pipe(result, th)), desc=tg.name(), total=len(img_to_interrogate)):
interrogate_tags[img_path] += tags
else:
for img_path, data in tqdm(zip(imgpaths, result), desc=tg.name(), total=len(imgpaths)):
for img_path, data in tqdm(zip(img_to_interrogate, result), desc=tg.name(), total=len(img_to_interrogate)):
interrogate_tags[img_path] += tg.predict(data, th)
except Exception as e:
tb = sys.exc_info()[2]
@ -814,7 +820,7 @@ class DatasetTagEditor(Singleton):
tags = interrogate_tags[img_path]
elif interrogate_method == self.InterrogateMethod.PREPEND:
tags = interrogate_tags[img_path] + tags
else:
elif interrogate_method != self.InterrogateMethod.PREFILL:
tags = tags + interrogate_tags[img_path]
self.set_tags_by_image_path(img_path, tags)

6
scripts/tag_editor_ui/block_toprow.py
View File

@ -33,14 +33,14 @@ class ToprowUI(UIBase):
def set_callbacks(self, load_dataset:LoadDatasetUI):
def save_all_changes(backup: bool, save_kohya_metadata:bool, metadata_output:str, metadata_input:str, metadata_overwrite:bool, metadata_as_caption:bool, metadata_use_fullpath:bool):
def save_all_changes(backup: bool, save_kohya_metadata:bool, metadata_output:str, metadata_input:str, metadata_overwrite:bool, metadata_as_caption:bool, metadata_use_fullpath:bool, caption_file_ext:str):
if not metadata_input:
metadata_input = None
dte_instance.save_dataset(backup, load_dataset.caption_file_ext, save_kohya_metadata, metadata_output, metadata_input, metadata_overwrite, metadata_as_caption, metadata_use_fullpath)
dte_instance.save_dataset(backup, caption_file_ext, save_kohya_metadata, metadata_output, metadata_input, metadata_overwrite, metadata_as_caption, metadata_use_fullpath)
self.btn_save_all_changes.click(
fn=save_all_changes,
inputs=[self.cb_backup, self.cb_save_kohya_metadata, self.tb_metadata_output, self.tb_metadata_input, self.cb_metadata_overwrite, self.cb_metadata_as_caption, self.cb_metadata_use_fullpath]
inputs=[self.cb_backup, self.cb_save_kohya_metadata, self.tb_metadata_output, self.tb_metadata_input, self.cb_metadata_overwrite, self.cb_metadata_as_caption, self.cb_metadata_use_fullpath, load_dataset.tb_caption_file_ext]
)
self.cb_save_kohya_metadata.change(