1580158
/
stable-diffusion-webui-images-browser
mirror of https://github.com/AlUlkesh/stable-diffusion-webui-images-browser.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Maintenance tab, #57

pull/71/head
AlUlkesh 2023-03-01 00:52:48 +01:00
parent 7eaa53ff32
commit 4cbcd1de17
3 changed files with 301 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
README.md
View File

@ -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

135
scripts/image_browser.py
View File

@ -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("&nbsp")
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"))

218
scripts/wib/wib_db.py
View File

@ -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