123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348 |
- from typing import NamedTuple, Type, Dict, Any
- from modules import shared, script_callbacks, scripts
- from modules.shared import opts
- import gradio as gr
- import json
- from pathlib import Path
- from collections import namedtuple
- import scripts.ui as ui
- from scripts.dte_instance import dte_instance
- # ================================================================
- # General Callbacks
- # ================================================================
- CONFIG_PATH = Path(scripts.basedir(), 'config.json')
- SortBy = dte_instance.SortBy
- SortOrder = dte_instance.SortOrder
- GeneralConfig = namedtuple('GeneralConfig', [
- 'backup',
- 'dataset_dir',
- 'caption_ext',
- 'load_recursive',
- 'load_caption_from_filename',
- 'replace_new_line',
- 'use_interrogator',
- 'use_interrogator_names',
- 'use_custom_threshold_booru',
- 'custom_threshold_booru',
- 'use_custom_threshold_waifu',
- 'custom_threshold_waifu',
- 'save_kohya_metadata',
- 'meta_output_path',
- 'meta_input_path',
- 'meta_overwrite',
- 'meta_save_as_caption',
- 'meta_use_full_path'
- ])
- FilterConfig = namedtuple('FilterConfig', ['sw_prefix', 'sw_suffix', 'sw_regex','sort_by', 'sort_order', 'logic'])
- BatchEditConfig = namedtuple('BatchEditConfig', ['show_only_selected', 'prepend', 'use_regex', 'target', 'sw_prefix', 'sw_suffix', 'sw_regex', 'sory_by', 'sort_order', 'batch_sort_by', 'batch_sort_order', 'token_count'])
- EditSelectedConfig = namedtuple('EditSelectedConfig', ['auto_copy', 'sort_on_save', 'warn_change_not_saved', 'use_interrogator_name', 'sort_by', 'sort_order'])
- MoveDeleteConfig = namedtuple('MoveDeleteConfig', ['range', 'target', 'caption_ext', 'destination'])
- CFG_GENERAL_DEFAULT = GeneralConfig(True, '', '.txt', False, True, False, 'No', [], False, 0.7, False, 0.35, False, '', '', True, False, False)
- CFG_FILTER_P_DEFAULT = FilterConfig(False, False, False, SortBy.ALPHA.value, SortOrder.ASC.value, 'AND')
- CFG_FILTER_N_DEFAULT = FilterConfig(False, False, False, SortBy.ALPHA.value, SortOrder.ASC.value, 'OR')
- CFG_BATCH_EDIT_DEFAULT = BatchEditConfig(True, False, False, 'Only Selected Tags', False, False, False, SortBy.ALPHA.value, SortOrder.ASC.value, SortBy.ALPHA.value, SortOrder.ASC.value, 75)
- CFG_EDIT_SELECTED_DEFAULT = EditSelectedConfig(False, False, False, '', SortBy.ALPHA.value, SortOrder.ASC.value)
- CFG_MOVE_DELETE_DEFAULT = MoveDeleteConfig('Selected One', [], '.txt', '')
- class Config:
- def __init__(self):
- self.config = dict()
- def load(self):
- if not CONFIG_PATH.is_file():
- self.config = dict()
- return
- try:
- self.config = json.loads(CONFIG_PATH.read_text('utf8'))
- except:
- print('[tag-editor] Error on loading config.json. Default settings will be loaded.')
- self.config = dict()
- else:
- print('[tag-editor] Settings has been read from config.json')
- def save(self):
- CONFIG_PATH.write_text(json.dumps(self.config, indent=4), 'utf8')
-
- def read(self, name: str):
- return self.config.get(name)
- def write(self, cfg: dict, name: str):
- self.config[name] = cfg
- 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()
- 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().style(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.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()
- print('[tag-editor] 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()
- print('[tag-editor] 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)
- print('[tag-editor] 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))
- script_callbacks.on_ui_settings(on_ui_settings)
- script_callbacks.on_ui_tabs(on_ui_tabs)
|