From 3f61b9f2997d2b043e3ac806d4c841bd8d3d4fb5 Mon Sep 17 00:00:00 2001
From: AlUlkesh <99896447+AlUlkesh@users.noreply.github.com>
Date: Mon, 20 Feb 2023 00:28:30 +0100
Subject: [PATCH] rebuild button, fix multiline prompt exif, #42

---
 scripts/image_browser.py | 30 ++++++++++++++++++++++++++++++
 scripts/wib/wib_db.py    | 23 ++++++++++++++++++++++-
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/scripts/image_browser.py b/scripts/image_browser.py
index e6aaff7..42d2fe6 100644
--- a/scripts/image_browser.py
+++ b/scripts/image_browser.py
@@ -45,6 +45,7 @@ none_select = "Nothing selected"
 refresh_symbol = '\U0001f504'  # 🔄
 up_symbol = '\U000025b2'  # ▲
 down_symbol = '\U000025bc'  # ▼
+rebuild_symbol = '\U0001f50e\U0001f4c2'  # 🔎📂
 current_depth = 0
 init = True
 copy_move = ["Move", "Copy"]
@@ -367,6 +368,24 @@ 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):
+    global finfo_exif, exif_cache, finfo_aes, aes_cache
+    if opts.image_browser_scan_exif:
+        print("Rebuild start")
+        exif_dirs = wib_db.get_exif_dirs()
+        finfo_aes = {}
+        exif_cache = {}
+        finfo_exif = {}
+        aes_cache = {}
+        for key, value in exif_dirs.items():
+            if os.path.exists(key):
+                print(f"Rebuilding {key}")
+                fileinfos = traverse_all_files(key, [], "", 0)
+                cache_exif(fileinfos)
+        print("Rebuild end")
+
+    return exif_rebuild_button
+
 def atof(text):
     try:
         retval = float(text)
@@ -623,6 +642,12 @@ def create_tab(tabname, current_gr_tab):
             img_path_subdirs = gr.Dropdown(choices=[none_select], value=none_select, label="Sub directories", interactive=True, elem_id=f"{tabname}_img_path_subdirs")
         with gr.Column(scale=1):
             img_path_subdirs_button = gr.Button(value="Get sub directories")
+
+    with gr.Row(visible=custom_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 custom_dir, elem_id=f"{tabname}_image_browser") as main_panel:
         with gr.Column():  
@@ -777,6 +802,11 @@ def create_tab(tabname, current_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]
+    )
 
     # other functions
     set_index.click(show_image_info, _js="image_browser_get_current_img", inputs=[tabname_box, image_index, page_index, filenames, turn_page_switch], outputs=[img_file_name, img_file_time, image_index, hidden, turn_page_switch])
diff --git a/scripts/wib/wib_db.py b/scripts/wib/wib_db.py
index 125bfed..e20298d 100644
--- a/scripts/wib/wib_db.py
+++ b/scripts/wib/wib_db.py
@@ -122,7 +122,11 @@ def update_exif_data(cursor, file, info):
             elif info_item.startswith(np):
                 negative_prompt = info_item.replace(np, "")
             else:
-                prompt = info_item
+                if prompt == "":
+                    prompt = info_item
+                else:
+                    # multiline prompts
+                    prompt = f"{prompt}\n{info_item}"
     if key_values != "":
         key_value_pairs = []
         key_value = ""
@@ -339,3 +343,20 @@ def load_aes_data(aes_cache):
         aes_cache[row[0]] = row[1]
 
     return aes_cache
+
+def get_exif_dirs():
+    with sqlite3.connect(db_file, timeout=timeout) as conn:
+        cursor = conn.cursor()
+        cursor.execute('''
+        SELECT file
+        FROM exif_data
+        ''')
+
+    rows = cursor.fetchall()
+
+    dirs = {}
+    for row in rows:
+        dir = os.path.dirname(row[0])
+        dirs[dir] = dir
+
+    return dirs