505 lines
17 KiB
Python
505 lines
17 KiB
Python
from typing import NamedTuple, Type, Dict, Any
|
|
from modules import shared, script_callbacks
|
|
from modules.shared import opts
|
|
import gradio as gr
|
|
from scripts.config import *
|
|
|
|
import scripts.tag_editor_ui as ui
|
|
|
|
# ================================================================
|
|
# General Callbacks
|
|
# ================================================================
|
|
|
|
config = Config()
|
|
|
|
|
|
def write_general_config(*args):
|
|
cfg = GeneralConfig(*args)
|
|
config.write(cfg._asdict(), "general")
|
|
|
|
|
|
def write_filter_config(*args):
|
|
hlen = len(args) // 2
|
|
cfg_p = FilterConfig(*args[:hlen])
|
|
cfg_n = FilterConfig(*args[hlen:])
|
|
config.write({"positive": cfg_p._asdict(), "negative": cfg_n._asdict()}, "filter")
|
|
|
|
|
|
def write_batch_edit_config(*args):
|
|
cfg = BatchEditConfig(*args)
|
|
config.write(cfg._asdict(), "batch_edit")
|
|
|
|
|
|
def write_edit_selected_config(*args):
|
|
cfg = EditSelectedConfig(*args)
|
|
config.write(cfg._asdict(), "edit_selected")
|
|
|
|
|
|
def write_move_delete_config(*args):
|
|
cfg = MoveDeleteConfig(*args)
|
|
config.write(cfg._asdict(), "file_move_delete")
|
|
|
|
|
|
def read_config(name: str, config_type: Type, default: NamedTuple, compat_func=None):
|
|
d = config.read(name)
|
|
cfg = default
|
|
if d:
|
|
if compat_func:
|
|
d = compat_func(d)
|
|
d = cfg._asdict() | d
|
|
d = {k: v for k, v in d.items() if k in cfg._asdict().keys()}
|
|
cfg = config_type(**d)
|
|
return cfg
|
|
|
|
|
|
def read_general_config():
|
|
# for compatibility
|
|
generalcfg_intterogator_names = [
|
|
("use_blip_to_prefill", "BLIP"),
|
|
("use_git_to_prefill", "GIT-large-COCO"),
|
|
("use_booru_to_prefill", "DeepDanbooru"),
|
|
("use_waifu_to_prefill", "wd-v1-4-vit-tagger"),
|
|
]
|
|
use_interrogator_names = []
|
|
|
|
def compat_func(d: Dict[str, Any]):
|
|
if "use_interrogator_names" in d.keys():
|
|
return d
|
|
for cfg in generalcfg_intterogator_names:
|
|
if d.get(cfg[0]):
|
|
use_interrogator_names.append(cfg[1])
|
|
d["use_interrogator_names"] = use_interrogator_names
|
|
return d
|
|
|
|
return read_config("general", GeneralConfig, CFG_GENERAL_DEFAULT, compat_func)
|
|
|
|
|
|
def read_filter_config():
|
|
d = config.read("filter")
|
|
d_p = d.get("positive") if d else None
|
|
d_n = d.get("negative") if d else None
|
|
cfg_p = CFG_FILTER_P_DEFAULT
|
|
cfg_n = CFG_FILTER_N_DEFAULT
|
|
if d_p:
|
|
d_p = cfg_p._asdict() | d_p
|
|
d_p = {k: v for k, v in d_p.items() if k in cfg_p._asdict().keys()}
|
|
cfg_p = FilterConfig(**d_p)
|
|
if d_n:
|
|
d_n = cfg_n._asdict() | d_n
|
|
d_n = {k: v for k, v in d_n.items() if k in cfg_n._asdict().keys()}
|
|
cfg_n = FilterConfig(**d_n)
|
|
return cfg_p, cfg_n
|
|
|
|
|
|
def read_batch_edit_config():
|
|
return read_config("batch_edit", BatchEditConfig, CFG_BATCH_EDIT_DEFAULT)
|
|
|
|
|
|
def read_edit_selected_config():
|
|
return read_config("edit_selected", EditSelectedConfig, CFG_EDIT_SELECTED_DEFAULT)
|
|
|
|
|
|
def read_move_delete_config():
|
|
return read_config("file_move_delete", MoveDeleteConfig, CFG_MOVE_DELETE_DEFAULT)
|
|
|
|
|
|
# ================================================================
|
|
# General Callbacks for Updating UIs
|
|
# ================================================================
|
|
|
|
|
|
def get_filters():
|
|
filters = [
|
|
ui.filter_by_tags.tag_filter_ui.get_filter(),
|
|
ui.filter_by_tags.tag_filter_ui_neg.get_filter(),
|
|
] + [ui.filter_by_selection.path_filter]
|
|
return filters
|
|
|
|
|
|
def update_gallery():
|
|
img_indices = ui.dte_instance.get_filtered_imgindices(filters=get_filters())
|
|
total_image_num = len(ui.dte_instance.dataset)
|
|
displayed_image_num = len(img_indices)
|
|
ui.gallery_state.register_value(
|
|
"Displayed Images", f"{displayed_image_num} / {total_image_num} total"
|
|
)
|
|
ui.gallery_state.register_value(
|
|
"Current Tag Filter",
|
|
f"{ui.filter_by_tags.tag_filter_ui.get_filter()} {' AND ' if ui.filter_by_tags.tag_filter_ui.get_filter().tags and ui.filter_by_tags.tag_filter_ui_neg.get_filter().tags else ''} {ui.filter_by_tags.tag_filter_ui_neg.get_filter()}",
|
|
)
|
|
ui.gallery_state.register_value(
|
|
"Current Selection Filter",
|
|
f"{len(ui.filter_by_selection.path_filter.paths)} images",
|
|
)
|
|
return [
|
|
[str(i) for i in img_indices],
|
|
1,
|
|
-1,
|
|
-1,
|
|
-1,
|
|
ui.gallery_state.get_current_gallery_txt(),
|
|
]
|
|
|
|
|
|
def update_filter_and_gallery():
|
|
return (
|
|
[
|
|
ui.filter_by_tags.tag_filter_ui.cbg_tags_update(),
|
|
ui.filter_by_tags.tag_filter_ui_neg.cbg_tags_update(),
|
|
]
|
|
+ update_gallery()
|
|
+ ui.batch_edit_captions.get_common_tags(get_filters, ui.filter_by_tags)
|
|
+ [", ".join(ui.filter_by_tags.tag_filter_ui.filter.tags)]
|
|
+ [ui.batch_edit_captions.tag_select_ui_remove.cbg_tags_update()]
|
|
+ ["", ""]
|
|
)
|
|
|
|
|
|
# ================================================================
|
|
# Script Callbacks
|
|
# ================================================================
|
|
|
|
|
|
def on_ui_tabs():
|
|
config.load()
|
|
|
|
cfg_general = read_general_config()
|
|
cfg_filter_p, cfg_filter_n = read_filter_config()
|
|
cfg_batch_edit = read_batch_edit_config()
|
|
cfg_edit_selected = read_edit_selected_config()
|
|
cfg_file_move_delete = read_move_delete_config()
|
|
|
|
ui.dte_instance.load_interrogators()
|
|
|
|
with gr.Blocks(analytics_enabled=False) as dataset_tag_editor_interface:
|
|
|
|
gr.HTML(
|
|
value="""
|
|
This extension works well with text captions in comma-separated style (such as the tags generated by DeepBooru interrogator).
|
|
"""
|
|
)
|
|
|
|
ui.toprow.create_ui(cfg_general)
|
|
|
|
with gr.Accordion(label="Reload/Save Settings (config.json)", open=False):
|
|
with gr.Row():
|
|
btn_reload_config_file = gr.Button(value="Reload settings")
|
|
btn_save_setting_as_default = gr.Button(value="Save current settings")
|
|
btn_restore_default = gr.Button(value="Restore settings to default")
|
|
|
|
with gr.Row(equal_height=False):
|
|
with gr.Column():
|
|
ui.load_dataset.create_ui(cfg_general)
|
|
ui.dataset_gallery.create_ui(opts.dataset_editor_image_columns)
|
|
ui.gallery_state.create_ui()
|
|
|
|
with gr.Tab(label="Filter by Tags"):
|
|
ui.filter_by_tags.create_ui(cfg_filter_p, cfg_filter_n, get_filters)
|
|
|
|
with gr.Tab(label="Filter by Selection"):
|
|
ui.filter_by_selection.create_ui(opts.dataset_editor_image_columns)
|
|
|
|
with gr.Tab(label="Batch Edit Captions"):
|
|
ui.batch_edit_captions.create_ui(cfg_batch_edit, get_filters)
|
|
|
|
with gr.Tab(label="Edit Caption of Selected Image"):
|
|
ui.edit_caption_of_selected_image.create_ui(cfg_edit_selected)
|
|
|
|
with gr.Tab(label="Move or Delete Files"):
|
|
ui.move_or_delete_files.create_ui(cfg_file_move_delete)
|
|
|
|
# ----------------------------------------------------------------
|
|
# General
|
|
|
|
components_general = [
|
|
ui.toprow.cb_backup,
|
|
ui.load_dataset.tb_img_directory,
|
|
ui.load_dataset.tb_caption_file_ext,
|
|
ui.load_dataset.cb_load_recursive,
|
|
ui.load_dataset.cb_load_caption_from_filename,
|
|
ui.load_dataset.cb_replace_new_line_with_comma,
|
|
ui.load_dataset.rb_use_interrogator,
|
|
ui.load_dataset.dd_intterogator_names,
|
|
ui.load_dataset.cb_use_custom_threshold_booru,
|
|
ui.load_dataset.sl_custom_threshold_booru,
|
|
ui.load_dataset.cb_use_custom_threshold_waifu,
|
|
ui.load_dataset.sl_custom_threshold_waifu,
|
|
ui.load_dataset.sl_custom_threshold_z3d,
|
|
ui.toprow.cb_save_kohya_metadata,
|
|
ui.toprow.tb_metadata_output,
|
|
ui.toprow.tb_metadata_input,
|
|
ui.toprow.cb_metadata_overwrite,
|
|
ui.toprow.cb_metadata_as_caption,
|
|
ui.toprow.cb_metadata_use_fullpath,
|
|
]
|
|
components_filter = [
|
|
ui.filter_by_tags.tag_filter_ui.cb_prefix,
|
|
ui.filter_by_tags.tag_filter_ui.cb_suffix,
|
|
ui.filter_by_tags.tag_filter_ui.cb_regex,
|
|
ui.filter_by_tags.tag_filter_ui.rb_sort_by,
|
|
ui.filter_by_tags.tag_filter_ui.rb_sort_order,
|
|
ui.filter_by_tags.tag_filter_ui.rb_logic,
|
|
] + [
|
|
ui.filter_by_tags.tag_filter_ui_neg.cb_prefix,
|
|
ui.filter_by_tags.tag_filter_ui_neg.cb_suffix,
|
|
ui.filter_by_tags.tag_filter_ui_neg.cb_regex,
|
|
ui.filter_by_tags.tag_filter_ui_neg.rb_sort_by,
|
|
ui.filter_by_tags.tag_filter_ui_neg.rb_sort_order,
|
|
ui.filter_by_tags.tag_filter_ui_neg.rb_logic,
|
|
]
|
|
components_batch_edit = [
|
|
ui.batch_edit_captions.cb_show_only_tags_selected,
|
|
ui.batch_edit_captions.cb_prepend_tags,
|
|
ui.batch_edit_captions.cb_use_regex,
|
|
ui.batch_edit_captions.rb_sr_replace_target,
|
|
ui.batch_edit_captions.tag_select_ui_remove.cb_prefix,
|
|
ui.batch_edit_captions.tag_select_ui_remove.cb_suffix,
|
|
ui.batch_edit_captions.tag_select_ui_remove.cb_regex,
|
|
ui.batch_edit_captions.tag_select_ui_remove.rb_sort_by,
|
|
ui.batch_edit_captions.tag_select_ui_remove.rb_sort_order,
|
|
ui.batch_edit_captions.rb_sort_by,
|
|
ui.batch_edit_captions.rb_sort_order,
|
|
ui.batch_edit_captions.nb_token_count,
|
|
]
|
|
components_edit_selected = [
|
|
ui.edit_caption_of_selected_image.cb_copy_caption_automatically,
|
|
ui.edit_caption_of_selected_image.cb_sort_caption_on_save,
|
|
ui.edit_caption_of_selected_image.cb_ask_save_when_caption_changed,
|
|
ui.edit_caption_of_selected_image.dd_intterogator_names_si,
|
|
ui.edit_caption_of_selected_image.rb_sort_by,
|
|
ui.edit_caption_of_selected_image.rb_sort_order,
|
|
]
|
|
components_move_delete = [
|
|
ui.move_or_delete_files.rb_move_or_delete_target_data,
|
|
ui.move_or_delete_files.cbg_move_or_delete_target_file,
|
|
ui.move_or_delete_files.tb_move_or_delete_caption_ext,
|
|
ui.move_or_delete_files.tb_move_or_delete_destination_dir,
|
|
]
|
|
|
|
configurable_components = (
|
|
components_general
|
|
+ components_filter
|
|
+ components_batch_edit
|
|
+ components_edit_selected
|
|
+ components_move_delete
|
|
)
|
|
|
|
def reload_config_file():
|
|
config.load()
|
|
p, n = read_filter_config()
|
|
logger.write("Reload config.json")
|
|
return (
|
|
read_general_config()
|
|
+ p
|
|
+ n
|
|
+ read_batch_edit_config()
|
|
+ read_edit_selected_config()
|
|
+ read_move_delete_config()
|
|
)
|
|
|
|
btn_reload_config_file.click(
|
|
fn=reload_config_file, outputs=configurable_components
|
|
)
|
|
|
|
def save_settings_callback(*a):
|
|
p = 0
|
|
|
|
def inc(v):
|
|
nonlocal p
|
|
p += v
|
|
return p
|
|
|
|
write_general_config(*a[p : inc(len(components_general))])
|
|
write_filter_config(*a[p : inc(len(components_filter))])
|
|
write_batch_edit_config(*a[p : inc(len(components_batch_edit))])
|
|
write_edit_selected_config(*a[p : inc(len(components_edit_selected))])
|
|
write_move_delete_config(*a[p:])
|
|
config.save()
|
|
logger.write("Current settings have been saved into config.json")
|
|
|
|
btn_save_setting_as_default.click(
|
|
fn=save_settings_callback, inputs=configurable_components
|
|
)
|
|
|
|
def restore_default_settings():
|
|
write_general_config(*CFG_GENERAL_DEFAULT)
|
|
write_filter_config(*CFG_FILTER_P_DEFAULT, *CFG_FILTER_N_DEFAULT)
|
|
write_batch_edit_config(*CFG_BATCH_EDIT_DEFAULT)
|
|
write_edit_selected_config(*CFG_EDIT_SELECTED_DEFAULT)
|
|
write_move_delete_config(*CFG_MOVE_DELETE_DEFAULT)
|
|
logger.write("Restore default settings")
|
|
return (
|
|
CFG_GENERAL_DEFAULT
|
|
+ CFG_FILTER_P_DEFAULT
|
|
+ CFG_FILTER_N_DEFAULT
|
|
+ CFG_BATCH_EDIT_DEFAULT
|
|
+ CFG_EDIT_SELECTED_DEFAULT
|
|
+ CFG_MOVE_DELETE_DEFAULT
|
|
)
|
|
|
|
btn_restore_default.click(
|
|
fn=restore_default_settings, outputs=configurable_components
|
|
)
|
|
|
|
o_update_gallery = [
|
|
ui.dataset_gallery.cbg_hidden_dataset_filter,
|
|
ui.dataset_gallery.nb_hidden_dataset_filter_apply,
|
|
ui.dataset_gallery.nb_hidden_image_index,
|
|
ui.dataset_gallery.nb_hidden_image_index_prev,
|
|
ui.edit_caption_of_selected_image.nb_hidden_image_index_save_or_not,
|
|
ui.gallery_state.txt_gallery,
|
|
]
|
|
|
|
o_update_filter_and_gallery = (
|
|
[
|
|
ui.filter_by_tags.tag_filter_ui.cbg_tags,
|
|
ui.filter_by_tags.tag_filter_ui_neg.cbg_tags,
|
|
]
|
|
+ o_update_gallery
|
|
+ [
|
|
ui.batch_edit_captions.tb_common_tags,
|
|
ui.batch_edit_captions.tb_edit_tags,
|
|
]
|
|
+ [ui.batch_edit_captions.tb_sr_selected_tags]
|
|
+ [ui.batch_edit_captions.tag_select_ui_remove.cbg_tags]
|
|
+ [
|
|
ui.edit_caption_of_selected_image.tb_caption,
|
|
ui.edit_caption_of_selected_image.tb_edit_caption,
|
|
]
|
|
)
|
|
|
|
ui.toprow.set_callbacks(ui.load_dataset)
|
|
ui.load_dataset.set_callbacks(
|
|
o_update_filter_and_gallery,
|
|
ui.toprow,
|
|
ui.dataset_gallery,
|
|
ui.filter_by_tags,
|
|
ui.filter_by_selection,
|
|
ui.batch_edit_captions,
|
|
update_filter_and_gallery,
|
|
)
|
|
ui.dataset_gallery.set_callbacks(ui.load_dataset, ui.gallery_state, get_filters)
|
|
ui.gallery_state.set_callbacks(ui.dataset_gallery)
|
|
ui.filter_by_tags.set_callbacks(
|
|
o_update_gallery,
|
|
o_update_filter_and_gallery,
|
|
ui.batch_edit_captions,
|
|
ui.move_or_delete_files,
|
|
update_gallery,
|
|
update_filter_and_gallery,
|
|
get_filters,
|
|
)
|
|
ui.filter_by_selection.set_callbacks(
|
|
o_update_filter_and_gallery,
|
|
ui.dataset_gallery,
|
|
ui.filter_by_tags,
|
|
get_filters,
|
|
update_filter_and_gallery,
|
|
)
|
|
ui.batch_edit_captions.set_callbacks(
|
|
o_update_filter_and_gallery,
|
|
ui.load_dataset,
|
|
ui.filter_by_tags,
|
|
get_filters,
|
|
update_filter_and_gallery,
|
|
)
|
|
ui.edit_caption_of_selected_image.set_callbacks(
|
|
o_update_filter_and_gallery,
|
|
ui.dataset_gallery,
|
|
ui.load_dataset,
|
|
get_filters,
|
|
update_filter_and_gallery,
|
|
)
|
|
ui.move_or_delete_files.set_callbacks(
|
|
o_update_filter_and_gallery,
|
|
ui.dataset_gallery,
|
|
ui.filter_by_tags,
|
|
ui.batch_edit_captions,
|
|
ui.filter_by_selection,
|
|
ui.edit_caption_of_selected_image,
|
|
get_filters,
|
|
update_filter_and_gallery,
|
|
)
|
|
|
|
return [
|
|
(
|
|
dataset_tag_editor_interface,
|
|
"Dataset Tag Editor",
|
|
"dataset_tag_editor_interface",
|
|
)
|
|
]
|
|
|
|
|
|
def on_ui_settings():
|
|
section = ("dataset-tag-editor", "Dataset Tag Editor")
|
|
shared.opts.add_option(
|
|
"dataset_editor_image_columns",
|
|
shared.OptionInfo(6, "Number of columns on image gallery", section=section),
|
|
)
|
|
shared.opts.add_option(
|
|
"dataset_editor_max_res",
|
|
shared.OptionInfo(0, "Max resolution of temporary files", section=section),
|
|
)
|
|
shared.opts.add_option(
|
|
"dataset_editor_use_temp_files",
|
|
shared.OptionInfo(
|
|
False, "Force image gallery to use temporary files", section=section
|
|
),
|
|
)
|
|
shared.opts.add_option(
|
|
"dataset_editor_use_raw_clip_token",
|
|
shared.OptionInfo(
|
|
True,
|
|
"Use raw CLIP token to calculate token count (without emphasis or embeddings)",
|
|
section=section,
|
|
),
|
|
)
|
|
shared.opts.add_option(
|
|
"dataset_editor_use_rating",
|
|
shared.OptionInfo(
|
|
False,
|
|
"Use rating tags",
|
|
section=section,
|
|
),
|
|
)
|
|
|
|
shared.opts.add_option(
|
|
"dataset_editor_num_cpu_workers",
|
|
shared.OptionInfo(
|
|
-1,
|
|
"Number of CPU workers when preprocessing images (set -1 to auto)",
|
|
section=section,
|
|
),
|
|
)
|
|
|
|
shared.opts.add_option(
|
|
"dataset_editor_batch_size_vit",
|
|
shared.OptionInfo(
|
|
4,
|
|
"Inference batch size for ViT taggers",
|
|
section=section,
|
|
),
|
|
)
|
|
|
|
shared.opts.add_option(
|
|
"dataset_editor_batch_size_convnext",
|
|
shared.OptionInfo(
|
|
4,
|
|
"Inference batch size for ConvNeXt taggers",
|
|
section=section,
|
|
),
|
|
)
|
|
|
|
shared.opts.add_option(
|
|
"dataset_editor_batch_size_swinv2",
|
|
shared.OptionInfo(
|
|
4,
|
|
"Inference batch size for SwinTransformerV2 taggers",
|
|
section=section,
|
|
),
|
|
)
|
|
|
|
|
|
script_callbacks.on_ui_settings(on_ui_settings)
|
|
script_callbacks.on_ui_tabs(on_ui_tabs)
|