From 4cbcd1de1710dfeaf1ce809042b40412dc1d2e38 Mon Sep 17 00:00:00 2001 From: AlUlkesh <99896447+AlUlkesh@users.noreply.github.com> Date: Wed, 1 Mar 2023 00:52:48 +0100 Subject: [PATCH] Maintenance tab, #57 --- README.md | 1 + scripts/image_browser.py | 135 ++++++++++++++++++------ scripts/wib/wib_db.py | 218 +++++++++++++++++++++++++++++++++++---- 3 files changed, 301 insertions(+), 53 deletions(-) diff --git a/README.md b/README.md index 5e23cf9..6be91d7 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ and restart your stable-diffusion-webui, then you can see the new tab "Image Bro Please be aware that when scanning a directory for the first time, the png-cache will be built. This can take several minutes, depending on the amount of images. ## Recent updates +- Maintenance tab - Custom tabs - Copy/Move to directory - Keybindings diff --git a/scripts/image_browser.py b/scripts/image_browser.py index f37e0fd..15fb9db 100644 --- a/scripts/image_browser.py +++ b/scripts/image_browser.py @@ -2,8 +2,8 @@ import gradio as gr import csv import logging import os -import random import platform +import random import re import shutil import stat @@ -38,10 +38,15 @@ yappi_do = False favorite_tab_name = "Favorites" default_tab_options = ["txt2img", "img2img", "txt2img-grids", "img2img-grids", "Extras", favorite_tab_name, "Others"] tabs_list = [tab.strip() for tab in chain.from_iterable(csv.reader(StringIO(opts.image_browser_active_tabs))) if tab] if hasattr(opts, "image_browser_active_tabs") else default_tab_options +try: + if opts.image_browser_enable_maint: + tabs_list.append("Maintenance") # mandatory tab +except AttributeError: + tabs_list.append("Maintenance") # mandatory tab + num_of_imgs_per_page = 0 loads_files_num = 0 image_ext_list = [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"] -cur_ranking_value="0" finfo_aes = {} exif_cache = {} finfo_exif = {} @@ -50,7 +55,7 @@ none_select = "Nothing selected" refresh_symbol = '\U0001f504' # 🔄 up_symbol = '\U000025b2' # â–² down_symbol = '\U000025bc' # â–¼ -rebuild_symbol = '\U0001f50e\U0001f4c2' # 🔎📂 +caution_symbol = '\U000026a0' # âš  current_depth = 0 init = True copy_move = ["Move", "Copy"] @@ -71,13 +76,13 @@ class ImageBrowserTab(): def __init__(self, name: str): self.name: str = os.path.basename(name) if os.path.isdir(name) else name - self.path: str = path_maps.get(name, name) + self.path: str = os.path.abspath(path_maps.get(name, name)) self.base_tag: str = f"image_browser_tab_{self.get_unique_base_tag(self.remove_invalid_html_tag_chars(self.name).lower())}" def remove_invalid_html_tag_chars(self, tag: str) -> str: - # Matches any character that is not a letter, a digit, a hyphen, or an underscore - pattern = re.compile(r'[^a-zA-Z0-9\-_]') - return re.sub(pattern, '', tag) + # Removes any character that is not a letter, a digit, a hyphen, or an underscore + removed = re.sub(r'[^a-zA-Z0-9\-_]', '', tag) + return removed def get_unique_base_tag(self, base_tag: str) -> str: counter = 1 @@ -147,6 +152,7 @@ def img_path_subdirs_get(img_path): return gr.update(choices=subdirs) def img_path_add_remove(img_dir, path_recorder, add_remove, img_path_depth): + img_dir = os.path.abspath(img_dir) if add_remove == "add" or (add_remove == "remove" and img_dir in path_recorder): if add_remove == "add": path_recorder[img_dir] = { @@ -173,7 +179,7 @@ def sort_order_flip(turn_page_switch, sort_order): sort_order = up_symbol return 1, -turn_page_switch, sort_order -def read_path_recorder(path_recorder): +def read_path_recorder(): path_recorder = wib_db.load_path_recorder() path_recorder_formatted = [value.get("path_display") for key, value in path_recorder.items()] path_recorder_formatted = sorted(path_recorder_formatted, key=lambda x: natural_keys(x.lower())) @@ -191,6 +197,7 @@ def pure_path(path): depth = int(match.group(1)) else: depth = 0 + path = os.path.abspath(path) return path, depth def browser2path(img_path_browser): @@ -203,8 +210,8 @@ def totxt(file): return file_txt -def tab_select(path_recorder): - path_recorder, path_recorder_formatted, path_recorder_unformatted = read_path_recorder(path_recorder) +def tab_select(): + path_recorder, path_recorder_formatted, path_recorder_unformatted = read_path_recorder() return path_recorder, gr.update(choices=path_recorder_unformatted) def reduplicative_file_move(src, dst): @@ -417,10 +424,10 @@ def cache_exif(fileinfos): cache_exif_end = time.time() logger.debug(f"cache_exif: {new_exif}/{len(fileinfos)} cache_aes: {new_aes}/{len(fileinfos)} {round(cache_exif_end - cache_exif_start, 1)} seconds") -def exif_rebuild(exif_rebuild_button): +def exif_rebuild(maint_wait): global finfo_exif, exif_cache, finfo_aes, aes_cache if opts.image_browser_scan_exif: - print("Rebuild start") + logger.debug("Rebuild start") exif_dirs = wib_db.get_exif_dirs() finfo_aes = {} exif_cache = {} @@ -431,9 +438,37 @@ def exif_rebuild(exif_rebuild_button): print(f"Rebuilding {key}") fileinfos = traverse_all_files(key, [], "", 0) cache_exif(fileinfos) - print("Rebuild end") + logger.debug("Rebuild end") + maint_last_msg = "Rebuild finished" + else: + maint_last_msg = "Exif cache not enabled in settings" - return exif_rebuild_button + return maint_wait, maint_last_msg + +def exif_update_dirs(maint_update_dirs_from, maint_update_dirs_to, maint_wait): + global exif_cache, aes_cache + if maint_update_dirs_from == "": + maint_last_msg = "From is empty" + elif maint_update_dirs_to == "": + maint_last_msg = "To is empty" + else: + maint_update_dirs_from = os.path.abspath(maint_update_dirs_from) + maint_update_dirs_to = os.path.abspath(maint_update_dirs_to) + rows = 0 + conn, cursor = wib_db.transaction_begin() + wib_db.update_path_recorder_mult(cursor, maint_update_dirs_from, maint_update_dirs_to) + rows = rows + cursor.rowcount + wib_db.update_exif_data_mult(cursor, maint_update_dirs_from, maint_update_dirs_to) + rows = rows + cursor.rowcount + wib_db.update_ranking_mult(cursor, maint_update_dirs_from, maint_update_dirs_to) + rows = rows + cursor.rowcount + wib_db.transaction_end(conn, cursor) + if rows == 0: + maint_last_msg = "No rows updated" + else: + maint_last_msg = f"{rows} rows updated. Please reload UI!" + + return maint_wait, maint_last_msg def atof(text): try: @@ -633,24 +668,31 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): global init, exif_cache, aes_cache dir_name = None others_dir = False + maint = False + standard_ui = True path_recorder = {} path_recorder_formatted = [] path_recorder_unformatted = [] if init: - wib_db.check() + db_version = wib_db.check() + logger.debug(f"db_version: {db_version}") exif_cache = wib_db.load_exif_data(exif_cache) aes_cache = wib_db.load_aes_data(aes_cache) init = False - path_recorder, path_recorder_formatted, path_recorder_unformatted = read_path_recorder(path_recorder) + path_recorder, path_recorder_formatted, path_recorder_unformatted = read_path_recorder() if tab.name == "Others": others_dir = True + standard_ui = False + elif tab.name == "Maintenance": + maint = True + standard_ui = False else: dir_name = tab.path - if not others_dir: + if standard_ui: dir_name = str(Path(dir_name)) if not os.path.exists(dir_name): os.makedirs(dir_name) @@ -681,14 +723,8 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): img_path_subdirs = gr.Dropdown(choices=[none_select], value=none_select, label="Sub directories", interactive=True, elem_id=f"{tab.base_tag}_img_path_subdirs") with gr.Column(scale=1): img_path_subdirs_button = gr.Button(value="Get sub directories") - - with gr.Row(visible=others_dir): - with gr.Column(scale=10): - gr.Dropdown(visible=False) - with gr.Column(scale=1): - exif_rebuild_button = gr.Button(value=f"{rebuild_symbol} Rebuild exif cache") - - with gr.Row(visible=not others_dir, elem_id=f"{tab.base_tag}_image_browser") as main_panel: + + with gr.Row(visible=standard_ui, elem_id=f"{tab.base_tag}_image_browser") as main_panel: with gr.Column(): with gr.Row(): with gr.Column(scale=2): @@ -774,6 +810,34 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): mod_keys = f"{mod_keys}S" image_browser_mod_keys = gr.Textbox(value=mod_keys, elem_id=f"{tab.base_tag}_image_browser_mod_keys") + # Maintenance tab + with gr.Row(visible=maint): + with gr.Column(scale=4): + gr.HTML(f"{caution_symbol} Caution: You should only use these options if you know what you are doing. {caution_symbol}") + with gr.Column(scale=3): + maint_wait = gr.HTML("Status:") + with gr.Column(scale=7): + gr.HTML(" ") + with gr.Row(visible=maint): + maint_last_msg = gr.Textbox(label="Last message", interactive=False) + with gr.Row(visible=maint): + with gr.Column(scale=1): + maint_rebuild = gr.Button(value="Rebuild exif cache") + with gr.Column(scale=10): + gr.HTML(visible=False) + with gr.Row(visible=maint): + with gr.Column(scale=1): + maint_update_dirs = gr.Button(value="Update directory names in database") + with gr.Column(scale=10): + maint_update_dirs_from = gr.Textbox(label="From (full path)") + with gr.Column(scale=10): + maint_update_dirs_to = gr.Textbox(label="to (full path)") + with gr.Row(visible=False): + with gr.Column(scale=1): + maint_rebuild_ranking = gr.Button(value="Rebuild ranking from exif info") + with gr.Column(scale=10): + gr.HTML(visible=False) + change_dir_outputs = [warning_box, main_panel, img_path_browser, path_recorder, load_switch, img_path, img_path_depth] img_path.submit(change_dir, inputs=[img_path, path_recorder, load_switch, img_path_browser, img_path_depth, img_path], outputs=change_dir_outputs) img_path_browser.change(change_dir, inputs=[img_path_browser, path_recorder, load_switch, img_path_browser, img_path_depth, img_path], outputs=change_dir_outputs) @@ -841,10 +905,17 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): inputs=[img_path, path_recorder, img_path_remove, img_path_depth], outputs=[path_recorder, img_path_browser] ) - exif_rebuild_button.click( - fn=exif_rebuild, - inputs=[exif_rebuild_button], - outputs=[exif_rebuild_button] + maint_rebuild.click( + fn=exif_rebuild, + show_progress=True, + inputs=[maint_wait], + outputs=[maint_wait, maint_last_msg] + ) + maint_update_dirs.click( + fn=exif_update_dirs, + show_progress=True, + inputs=[maint_update_dirs_from, maint_update_dirs_to, maint_wait], + outputs=[maint_wait, maint_last_msg] ) # other functions @@ -865,10 +936,10 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): except: pass - if not others_dir: + if standard_ui: current_gr_tab.select( fn=tab_select, - inputs=[path_recorder], + inputs=[], outputs=[path_recorder, to_dir_saved] ) @@ -949,6 +1020,8 @@ def on_ui_settings(): image_browser_options.append(("image_browser_scan_exif", True, "Scan Exif-/.txt-data (slower, but required for exif-keyword-search)", "images_scan_exif")) image_browser_options.append(("image_browser_mod_shift", False, "Change CTRL keybindings to SHIFT", None)) image_browser_options.append(("image_browser_mod_ctrl_shift", False, "or to CTRL+SHIFT", None)) + #image_browser_options.append(("image_browser_ranking_pnginfo", False, "Save ranking in image's pnginfo", None)) + image_browser_options.append(("image_browser_enable_maint", True, "Enable Maintenance tab", None)) image_browser_options.append(("image_browser_page_columns", 6, "Number of columns on the page", "images_history_page_columns")) image_browser_options.append(("image_browser_page_rows", 6, "Number of rows on the page", "images_history_page_rows")) diff --git a/scripts/wib/wib_db.py b/scripts/wib/wib_db.py index e20298d..600e45e 100644 --- a/scripts/wib/wib_db.py +++ b/scripts/wib/wib_db.py @@ -3,6 +3,8 @@ import os import sqlite3 from modules import scripts +version = 2 + path_recorder_file = os.path.join(scripts.basedir(), "path_recorder.txt") aes_cache_file = os.path.join(scripts.basedir(), "aes_scores.json") exif_cache_file = os.path.join(scripts.basedir(), "exif_data.json") @@ -65,12 +67,17 @@ def create_db(cursor): cursor.execute(''' CREATE TABLE IF NOT EXISTS ranking ( file TEXT PRIMARY KEY, + name TEXT, ranking TEXT, created TIMESTAMP DEFAULT CURRENT_TIMESTAMP, updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) ''') + cursor.execute(''' + CREATE INDEX IF NOT EXISTS ranking_name ON ranking (name) + ''') + cursor.execute(''' CREATE TRIGGER ranking_tr AFTER UPDATE ON ranking @@ -88,8 +95,9 @@ def migrate_path_recorder(cursor): # json-version path_recorder = json.load(f) for path, values in path_recorder.items(): + path = os.path.abspath(path) depth = values["depth"] - path_display = values["path_display"] + path_display = f"{path} [{depth}]" cursor.execute(''' INSERT INTO path_recorder (path, depth, path_display) VALUES (?, ?, ?) @@ -99,6 +107,7 @@ def migrate_path_recorder(cursor): # old txt-version path = f.readline().rstrip("\n") while len(path) > 0: + path = os.path.abspath(path) cursor.execute(''' INSERT INTO path_recorder (path, depth, path_display) VALUES (?, ?, ?) @@ -180,10 +189,10 @@ def update_exif_data(cursor, file, info): def migrate_exif_data(cursor): if os.path.exists(exif_cache_file): with open(exif_cache_file, 'r') as file: - #with open("c:\\tools\\ai\\stable-diffusion-webui\\extensions\\stable-diffusion-webui-images-browser\\test.json ", 'r') as file: exif_cache = json.load(file) for file, info in exif_cache.items(): + file = os.path.abspath(file) update_exif_data(cursor, file, info) return @@ -194,10 +203,12 @@ def migrate_ranking(cursor): ranking = json.load(file) for file, info in ranking.items(): if info != "None": + file = os.path.abspath(file) + name = os.path.basename(file) cursor.execute(''' - INSERT INTO ranking (file, ranking) - VALUES (?, ?) - ''', (file, info)) + INSERT INTO ranking (file, name, ranking) + VALUES (?, ?, ?) + ''', (file, name, info)) return @@ -210,17 +221,122 @@ def update_db_data(cursor, key, value): return +def get_version(): + with sqlite3.connect(db_file, timeout=timeout) as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT value + FROM db_data + WHERE key = 'version' + ''',) + db_version = cursor.fetchone() + + return db_version + +def migrate_path_recorder_dirs(cursor): + cursor.execute(''' + SELECT path, path_display + FROM path_recorder + ''') + for (path, path_display) in cursor.fetchall(): + abs_path = os.path.abspath(path) + if path != abs_path: + update_from = path + update_to = abs_path + cursor.execute(''' + UPDATE path_recorder + SET path = ?, + path_display = ? || SUBSTR(path_display, LENGTH(?) + 1) + WHERE path = ? + ''', (update_to, update_to, update_from, update_from)) + + return + +def migrate_exif_data_dirs(cursor): + cursor.execute(''' + SELECT file + FROM exif_data + ''') + for (filepath,) in cursor.fetchall(): + (path, file) = os.path.split(filepath) + abs_path = os.path.abspath(path) + if path != abs_path: + update_from = filepath + update_to = os.path.join(abs_path, file) + try: + cursor.execute(''' + UPDATE exif_data + SET file = ? + WHERE file = ? + ''', (update_to, update_from)) + except sqlite3.IntegrityError as e: + # these are double keys, because the same file can be in the db with different path notations + (e_msg,) = e.args + if e_msg.startswith("UNIQUE constraint"): + cursor.execute(''' + DELETE FROM exif_data + WHERE file = ? + ''', (update_from,)) + else: + raise + + return + +def migrate_ranking_dirs(cursor): + cursor.execute(''' + ALTER TABLE ranking + ADD COLUMN name TEXT + ''') + + cursor.execute(''' + CREATE INDEX IF NOT EXISTS ranking_name ON ranking (name) + ''') + + cursor.execute(''' + SELECT file, ranking + FROM ranking + ''') + for (filepath, ranking) in cursor.fetchall(): + if filepath == "" or ranking == "None": + cursor.execute(''' + DELETE FROM ranking + WHERE file = ? + ''', (filepath,)) + else: + (path, file) = os.path.split(filepath) + abs_path = os.path.abspath(path) + name = file + update_from = filepath + update_to = os.path.join(abs_path, file) + cursor.execute(''' + UPDATE ranking + SET file = ?, + name = ? + WHERE file = ? + ''', (update_to, name, update_from)) + + return + def check(): if not os.path.exists(db_file): conn, cursor = transaction_begin() create_db(cursor) - update_db_data(cursor, "version", "1") + update_db_data(cursor, "version", version) migrate_path_recorder(cursor) migrate_exif_data(cursor) migrate_ranking(cursor) transaction_end(conn, cursor) + db_version = get_version() + if db_version[0] == "1": + # version 1 database had mixed path notations, this will change them all to abspath + conn, cursor = transaction_begin() + update_db_data(cursor, "version", version) + migrate_path_recorder_dirs(cursor) + migrate_exif_data_dirs(cursor) + migrate_ranking_dirs(cursor) + transaction_end(conn, cursor) - return + return version def load_path_recorder(): with sqlite3.connect(db_file, timeout=timeout) as conn: @@ -234,28 +350,54 @@ def load_path_recorder(): return path_recorder def select_ranking(filename): - with sqlite3.connect(db_file, timeout=timeout) as conn: - cursor = conn.cursor() + conn, cursor = transaction_begin() + cursor.execute(''' + SELECT ranking + FROM ranking + WHERE file = ? + ''', (filename,)) + ranking_value = cursor.fetchone() + + # if ranking not found search again, without path (moved?) + if ranking_value is None: + name = os.path.basename(filename) cursor.execute(''' SELECT ranking FROM ranking - WHERE file = ? - ''', (filename,)) + WHERE name = ? + ''', (name,)) ranking_value = cursor.fetchone() - + # and insert with current filepath + if ranking_value is not None: + (insert_ranking,) = ranking_value + cursor.execute(''' + INSERT INTO ranking (file, name, ranking) + VALUES (?, ?, ?) + ''', (filename, name, insert_ranking)) + if ranking_value is None: - ranking_value = ["0"] - return ranking_value[0] + return_ranking = "None" + else: + (return_ranking,) = ranking_value + transaction_end(conn, cursor) + + return return_ranking def update_ranking(file, ranking): - if ranking != "0": - with sqlite3.connect(db_file, timeout=timeout) as conn: - cursor = conn.cursor() - cursor.execute(''' - INSERT OR REPLACE - INTO ranking (file, ranking) - VALUES (?, ?) - ''', (file, ranking)) + name = os.path.basename(file) + with sqlite3.connect(db_file, timeout=timeout) as conn: + cursor = conn.cursor() + if ranking == "None": + cursor.execute(''' + DELETE FROM ranking + WHERE name = ? + ''', (name,)) + else: + cursor.execute(''' + INSERT OR REPLACE + INTO ranking (file, name, ranking) + VALUES (?, ?, ?) + ''', (file, name, ranking)) return @@ -280,6 +422,38 @@ def delete_path_recorder(path): return +def update_path_recorder_mult(cursor, update_from, update_to): + cursor.execute(''' + UPDATE path_recorder + SET path = ?, + path_display = ? || SUBSTR(path_display, LENGTH(?) + 1) + WHERE path = ? + ''', (update_to, update_to, update_from, update_from)) + + return + +def update_exif_data_mult(cursor, update_from, update_to): + update_from = update_from + os.path.sep + update_to = update_to + os.path.sep + cursor.execute(''' + UPDATE exif_data + SET file = ? || SUBSTR(file, LENGTH(?) + 1) + WHERE file like ? || '%' + ''', (update_to, update_from, update_from)) + + return + +def update_ranking_mult(cursor, update_from, update_to): + update_from = update_from + os.path.sep + update_to = update_to + os.path.sep + cursor.execute(''' + UPDATE ranking + SET file = ? || SUBSTR(file, LENGTH(?) + 1) + WHERE file like ? || '%' + ''', (update_to, update_from, update_from)) + + return + def transaction_begin(): conn = sqlite3.connect(db_file, timeout=timeout) conn.isolation_level = None