Maintenance tab, #57

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

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
+        # Removes any character that is not a letter, a digit, a hyphen, or an underscore
-        pattern = re.compile(r'[^a-zA-Z0-9\-_]')
+        removed = re.sub(r'[^a-zA-Z0-9\-_]', '', tag)
-        return re.sub(pattern, '', 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):
+def tab_select():
-    path_recorder, path_recorder_formatted, path_recorder_unformatted = read_path_recorder(path_recorder)
+    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.Row(visible=standard_ui, elem_id=f"{tab.base_tag}_image_browser") as main_panel:
-        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.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(
+    maint_rebuild.click(
-        fn=exif_rebuild, 
+        fn=exif_rebuild,
-        inputs=[exif_rebuild_button],
+        show_progress=True,
-        outputs=[exif_rebuild_button]
+        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"))

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)
+                INSERT INTO ranking (file, name, ranking)
-                VALUES (?, ?)
+                VALUES (?, ?, ?)
-                ''', (file, info))
+                ''', (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:
+    conn, cursor = transaction_begin()
-        cursor = conn.cursor()
+    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 = ?
+        WHERE name = ?
-        ''', (filename,))
+        ''', (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 = "None"
-    return ranking_value[0]
+    else:
+        (return_ranking,) = ranking_value
+    transaction_end(conn, cursor)
+    
+    return return_ranking
 
 def update_ranking(file, ranking):
-    if ranking != "0":
+    name = os.path.basename(file)
-        with sqlite3.connect(db_file, timeout=timeout) as conn:
+    with sqlite3.connect(db_file, timeout=timeout) as conn:
-            cursor = conn.cursor()
+        cursor = conn.cursor()
-            cursor.execute('''
+        if ranking == "None":
-            INSERT OR REPLACE
+                cursor.execute('''
-            INTO ranking (file, ranking)
+                DELETE FROM ranking
-            VALUES (?, ?)
+                WHERE name = ?
-            ''', (file, ranking))
+                ''', (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