From bd644bdee08fef56ce85fb3b1f3c131196b9c544 Mon Sep 17 00:00:00 2001 From: AlUlkesh <99896447+AlUlkesh@users.noreply.github.com> Date: Fri, 8 Aug 2025 20:29:23 +0200 Subject: [PATCH] Use ui_common's save_files, #196, small cleanup --- scripts/image_browser.py | 225 ++++++++++++++++++++++++++++----------- 1 file changed, 164 insertions(+), 61 deletions(-) diff --git a/scripts/image_browser.py b/scripts/image_browser.py index db7f796..bae3f7c 100644 --- a/scripts/image_browser.py +++ b/scripts/image_browser.py @@ -1,10 +1,25 @@ -import gradio as gr +from PIL import Image, ImageOps, UnidentifiedImageError, ImageDraw +from datetime import datetime +from io import StringIO +from itertools import chain +from modules import paths, shared, script_callbacks, scripts, images +from modules.shared import opts, cmd_opts +from modules.ui_common import plaintext_to_html, save_files +from modules.ui_components import ToolButton, DropdownMulti +from packaging import version +from pathlib import Path +from typing import List, Tuple import codecs import csv +import gradio as gr +import hashlib import importlib import json import logging import math +import modules.extras +import modules.images +import modules.ui import os import platform import random @@ -17,21 +32,6 @@ import tempfile import time import torch import traceback -import hashlib -import modules.extras -import modules.images -import modules.ui -from datetime import datetime -from modules import paths, shared, script_callbacks, scripts, images -from modules.shared import opts, cmd_opts -from modules.ui_common import plaintext_to_html -from modules.ui_components import ToolButton, DropdownMulti -from PIL import Image, ImageOps, UnidentifiedImageError, ImageDraw -from packaging import version -from pathlib import Path -from typing import List, Tuple -from itertools import chain -from io import StringIO try: from scripts.wib import wib_db @@ -80,12 +80,12 @@ if not opencv_installed: exif_cache = {} aes_cache = {} none_select = "Nothing selected" -refresh_symbol = '\U0001f504' # 🔄 -up_symbol = '\U000025b2' # ▲ -down_symbol = '\U000025bc' # ▼ -caution_symbol = '\U000026a0' # ⚠ -folder_symbol = '\U0001f4c2' # 📂 -play_symbol = '\U000023f5' # ⏵ +refresh_symbol = "\U0001f504" # 🔄 +up_symbol = "\U000025b2" # ▲ +down_symbol = "\U000025bc" # ▼ +caution_symbol = "\U000026a0" # ⚠ +folder_symbol = "\U0001f4c2" # 📂 +play_symbol = "\U000023f5" # ⏵ current_depth = 0 init = True copy_move = ["Move", "Copy"] @@ -143,16 +143,16 @@ class ImageBrowserTab(): def remove_invalid_html_tag_chars(self, tag: str) -> str: # Removes any character that is not a letter, a digit, a hyphen, or an underscore - removed = re.sub(r'[^a-zA-Z0-9\-_]', '', tag) + removed = re.sub(r"[^a-zA-Z0-9\-_]", "", tag) return removed def get_unique_base_tag(self, base_tag: str) -> str: counter = 1 while base_tag in self.seen_base_tags: - match = re.search(r'_(\d+)$', base_tag) + match = re.search(r"_(\d+)$", base_tag) if match: counter = int(match.group(1)) + 1 - base_tag = re.sub(r'_(\d+)$', f"_{counter}", base_tag) + base_tag = re.sub(r"_(\d+)$", f"_{counter}", base_tag) else: base_tag = f"{base_tag}_{counter}" counter += 1 @@ -208,7 +208,7 @@ def setup_file_handler(): if not handler_active: file_handler = logging.FileHandler(log_file) file_handler.setLevel(logger_mode) - formatter = logging.Formatter(f'%(asctime)s image_browser.py: %(message)s', datefmt='%Y-%m-%d-%H:%M:%S') + formatter = logging.Formatter(f"%(asctime)s image_browser.py: %(message)s", datefmt="%Y-%m-%d-%H:%M:%S") file_handler.setFormatter(formatter) logger.addHandler(file_handler) @@ -237,7 +237,7 @@ def setup_debug(): if not common_logger: console_handler = logging.StreamHandler() console_handler.setLevel(logger_mode) - formatter = logging.Formatter(f'%(asctime)s image_browser.py: %(message)s', datefmt='%Y-%m-%d-%H:%M:%S') + formatter = logging.Formatter(f"%(asctime)s image_browser.py: %(message)s", datefmt="%Y-%m-%d-%H:%M:%S") console_handler.setFormatter(formatter) logger.addHandler(console_handler) if level_value >= capture_level_value: @@ -252,7 +252,7 @@ def setup_debug(): logger.debug(f"{sys.executable} {sys.version}") logger.debug(f"{platform.system()} {platform.version()}") try: - git = os.environ.get('GIT', "git") + git = os.environ.get("GIT", "git") webui_commit_hash = os.popen(f"{git} rev-parse HEAD").read().strip() sm_hashes = os.popen(f"{git} submodule").read() sm_hashes_lines = sm_hashes.splitlines() @@ -350,7 +350,7 @@ def browser2path(img_path_browser): def totxt(file): base, _ = os.path.splitext(file) - file_txt = base + '.txt' + file_txt = base + ".txt" return file_txt @@ -415,12 +415,114 @@ def reduplicative_file_move(src, dst): def save_image(file_name, filenames, page_index, turn_page_switch, dest_path): if file_name is not None and os.path.exists(file_name): - reduplicative_file_move(file_name, dest_path) - message = f"
{copied_moved[opts.image_browser_copy_image]} to {dest_path}
" - if not opts.image_browser_copy_image: - # Force page refresh with checking filenames - filenames = [] - turn_page_switch += 1 + try: + # Check if txt file exists and should be handled + src_txt_exists = False + src_txt = None + if opts.image_browser_txt_files: + src_txt = totxt(file_name) + if os.path.exists(src_txt): + src_txt_exists = True + + # Use the existing save_files function from ui_common.py + # First, we need to prepare the data in the format expected by save_files + + # Read image metadata + from PIL import Image + try: + image = Image.open(file_name) + geninfo, items = images.read_info_from_image(image) + + # Parse generation parameters if available + if geninfo: + # Use the already imported sendto module for parameter parsing + parameters = sendto.parse_generation_parameters(geninfo, []) + else: + parameters = {} + + # Create data structure expected by save_files + js_data = { + "infotexts": [geninfo or ""], + "index_of_first_image": 0, + "width": image.width, + "height": image.height, + "sampler_name": parameters.get("Sampler", "Unknown"), + "cfg_scale": parameters.get("CFG scale", 7.0), + "steps": parameters.get("Steps", 20), + "sd_model_name": parameters.get("Model", "Unknown"), + "sd_model_hash": parameters.get("Model hash", "Unknown") + } + + # Prepare image data as expected by save_files + images_data = [(image, None)] # (image, info) tuple + + # Temporarily change the output directory to our destination + original_outdir = shared.opts.outdir_save + shared.opts.outdir_save = dest_path + + try: + # Call the existing save_files function + file_result, message_result = save_files( + json.dumps(js_data), + images_data, + False, # do_make_zip + 0 # index + ) + message = message_result + + # Handle .txt file copying/moving manually since save_files doesn't know about it + if src_txt_exists and file_result and "value" in file_result and file_result["value"]: + # Get the saved image filename from the result + saved_files = file_result["value"] if isinstance(file_result["value"], list) else [file_result["value"]] + if saved_files: + # Find the image file (not zip) + saved_image_path = None + for saved_file in saved_files: + if isinstance(saved_file, str) and not saved_file.endswith(".zip"): + saved_image_path = saved_file + break + + if saved_image_path: + # Calculate destination txt file path + dest_txt = totxt(saved_image_path) + try: + if opts.image_browser_copy_image: + shutil.copy2(src_txt, dest_txt) + else: + shutil.move(src_txt, dest_txt) + except Exception as e: + logger.warning(f"Failed to handle txt file: {e}") + + finally: + # Restore original output directory + shared.opts.outdir_save = original_outdir + + # If we're moving (not copying), remove the original file + if not opts.image_browser_copy_image: + try: + os.remove(file_name) + # txt file should already be moved above, but clean up if it wasn't + if src_txt_exists and os.path.exists(src_txt): + os.remove(src_txt) + except Exception as e: + logger.warning(f"Failed to remove original file: {e}") + + # Force page refresh with checking filenames + filenames = [] + turn_page_switch += 1 + + except Exception as e: + logger.warning(f"Failed to use save_files function, falling back to reduplicative_file_move: {e}") + # Fallback to original method + reduplicative_file_move(file_name, dest_path) + message = f"
{copied_moved[opts.image_browser_copy_image]} to {dest_path}
" + if not opts.image_browser_copy_image: + filenames = [] + turn_page_switch += 1 + + except Exception as e: + logger.error(f"Error in save_image: {e}") + message = "
Error occurred during save operation
" else: message = "
Image not found (may have been already moved)
" @@ -591,10 +693,10 @@ def cache_exif(fileinfos): if yappi_do: yappi.stop() - pd.set_option('display.float_format', lambda x: '%.6f' % x) + pd.set_option("display.float_format", lambda x: "%.6f" % x) yappi_stats = yappi.get_func_stats().strip_dirs() data = [(s.name, s.ncall, s.tsub, s.ttot, s.ttot/s.ncall) for s in yappi_stats] - df = pd.DataFrame(data, columns=['name', 'ncall', 'tsub', 'ttot', 'tavg']) + df = pd.DataFrame(data, columns=["name", "ncall", "tsub", "ttot", "tavg"]) print(df.to_string(index=False)) yappi.get_thread_stats().print_all() @@ -728,7 +830,7 @@ def natural_keys(text): (See Toothy's implementation in the comments) float regex comes from https://stackoverflow.com/a/12643073/190597 ''' - return [ atof(c) for c in re.split(r'[+-]?([0-9]+(?:[.][0-9]*)?|[.][0-9]+)', text) ] + return [ atof(c) for c in re.split(r"[+-]?([0-9]+(?:[.][0-9]*)?|[.][0-9]+)", text) ] def open_folder(path): if os.path.exists(path): @@ -766,13 +868,13 @@ def exif_search(needle, haystack, use_regex): # Function to parse and evaluate the logical expression def parse_expression(expression): # Split the expression by 'or' to handle OR logic first - or_parts = expression.split(' or ') + or_parts = expression.split(" or ") for part in or_parts: # Split each part by 'and' to handle AND logic - and_parts = part.split(' and ') + and_parts = part.split(" and ") # Check if all conditions in the AND part are satisfied if all( - (subpart[4:] not in haystack) if subpart.startswith('not ') + (subpart[4:] not in haystack) if subpart.startswith("not ") else (subpart in haystack) for subpart in and_parts ): @@ -849,7 +951,8 @@ def get_all_images(dir_name, sort_by, sort_order, keyword, tab_base_tag_box, img exif_info = dict(exif_cache) if exif_info: for k, v in exif_info.items(): - match = re.search(r'(?<='+ sort_by + ":" ').*?(?=(,|$))', v, flags=re.DOTALL|re.IGNORECASE) + pattern = "(?<=" + re.escape(sort_by + ":") + ").*?(?=(,|$))" + match = re.search(pattern, v, flags=re.DOTALL | re.IGNORECASE) if match: sort_values[k] = match.group().strip() else: @@ -995,14 +1098,14 @@ def get_thumbnail(image_video, image_list): thumbnail.thumbnail((opts.image_browser_thumbnail_size, opts.image_browser_thumbnail_size)) if image_video == "video": - play_button_img = Image.new('RGBA', (100, 100), (0, 0, 0, 0)) + play_button_img = Image.new("RGBA", (100, 100), (0, 0, 0, 0)) play_button_draw = ImageDraw.Draw(play_button_img) - play_button_draw.polygon([(20, 20), (80, 50), (20, 80)], fill='white') + play_button_draw.polygon([(20, 20), (80, 50), (20, 80)], fill="white") play_button_img = play_button_img.resize((50, 50)) - button_for_img = Image.new('RGBA', thumbnail.size, (0, 0, 0, 0)) + button_for_img = Image.new("RGBA", thumbnail.size, (0, 0, 0, 0)) button_for_img.paste(play_button_img, (thumbnail.width - play_button_img.width, thumbnail.height - play_button_img.height), mask=play_button_img) - thumbnail_play = Image.alpha_composite(thumbnail.convert('RGBA'), button_for_img) + thumbnail_play = Image.alpha_composite(thumbnail.convert("RGBA"), button_for_img) thumbnail.close() thumbnail = thumbnail_play if thumbnail.mode != "RGB": @@ -1068,9 +1171,9 @@ def get_current_file(tab_base_tag_box, num, page_index, filenames): return file def pnginfo2html(pnginfo, items): - items = {**{'parameters': pnginfo}, **items} + items = {**{"parameters": pnginfo}, **items} - info = '' + info = "" for key, text in items.items(): info += f"""
@@ -1172,7 +1275,7 @@ def change_dir(change_dir_type, img_dir, path_recorder, load_switch, img_path_br try: f = os.listdir(img_dir) except: - warning = f"'{img_dir} is not a directory" + warning = f"{img_dir} is not a directory" else: warning = "The directory does not exist" except: @@ -1206,11 +1309,11 @@ def update_exif(img_file_name, key, value): if geninfo is not None: if f"{key}: " in geninfo: if value == "None": - geninfo = re.sub(f', {key}: \d+(\.\d+)*', '', geninfo) + geninfo = re.sub(f", {key}: \d+(\.\d+)*", "", geninfo) else: - geninfo = re.sub(f'{key}: \d+(\.\d+)*', f'{key}: {value}', geninfo) + geninfo = re.sub(f"{key}: \d+(\.\d+)*", f"{key}: {value}", geninfo) else: - geninfo = f'{geninfo}, {key}: {value}' + geninfo = f"{geninfo}, {key}: {value}" original_time = os.path.getmtime(img_file_name) images.save_image(image, os.path.dirname(img_file_name), "", extension=os.path.splitext(img_file_name)[1][1:], info=geninfo, forced_filename=os.path.splitext(os.path.basename(img_file_name))[0], save_to_dirs=False) @@ -1254,7 +1357,7 @@ def img_file_info_do_format(img_file_info): key_value_pairs.append(("Prompt", prompt)) key_value_pairs.append(("Negative prompt", negative_prompt)) # Sort the key_value_pairs by the order in img_file_info_order - img_file_info_order = opts.image_browser_info_order.split(',') + img_file_info_order = opts.image_browser_info_order.split(",") key_value_pairs = sorted(key_value_pairs, key=lambda pair: (img_file_info_order.index(pair[0]) if pair[0] in img_file_info_order else len(img_file_info_order), pair[0])) img_file_info_formatted = key_value_pairs return img_file_info_formatted @@ -1374,7 +1477,7 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): delete_num = gr.Number(value=1, interactive=True, label="delete next", elem_id=f"{tab.base_tag}_image_browser_del_num") delete_confirm = gr.Checkbox(value=False, label="also delete off-screen images") with gr.Column(scale=3): - delete = gr.Button('Delete', elem_id=f"{tab.base_tag}_image_browser_del_img_btn") + delete = gr.Button("Delete", elem_id=f"{tab.base_tag}_image_browser_del_img_btn") with gr.Row() as info_add_panel: with gr.Box(visible=opts.image_browser_info_add): gr.HTML("

Additional Generation Info

") @@ -1432,7 +1535,7 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): favorites_btn_show = False else: favorites_btn_show = True - favorites_btn = gr.Button(f'{copy_move[opts.image_browser_copy_image]} to favorites', elem_id=f"{tab.base_tag}_image_browser_favorites_btn", visible=favorites_btn_show) + favorites_btn = gr.Button(f"{copy_move[opts.image_browser_copy_image]} to favorites", elem_id=f"{tab.base_tag}_image_browser_favorites_btn", visible=favorites_btn_show) try: send_to_buttons = sendto.create_buttons(["txt2img", "img2img", "inpaint", "extras"]) except: @@ -1460,7 +1563,7 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): with gr.Row(): to_dir_saved = gr.Dropdown(choices=path_recorder_unformatted, label="Saved directories") with gr.Row(): - to_dir_btn = gr.Button(f'{copy_move[opts.image_browser_copy_image]} to directory', elem_id=f"{tab.base_tag}_image_browser_to_dir_btn") + to_dir_btn = gr.Button(f"{copy_move[opts.image_browser_copy_image]} to directory", elem_id=f"{tab.base_tag}_image_browser_to_dir_btn") with gr.Row(): collected_warning = gr.HTML() @@ -1470,7 +1573,7 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): visible_img_num = gr.Number() tab_base_tag_box = gr.Textbox(tab.base_tag) image_index = gr.Textbox(value=-1, elem_id=f"{tab.base_tag}_image_browser_image_index") - set_index = gr.Button('set_index', elem_id=f"{tab.base_tag}_image_browser_set_index") + set_index = gr.Button("set_index", elem_id=f"{tab.base_tag}_image_browser_set_index") filenames = gr.State([]) hidden = gr.Image(type="pil", elem_id=f"{tab.base_tag}_image_browser_hidden_image") image_page_list = gr.Textbox(elem_id=f"{tab.base_tag}_image_browser_image_page_list") @@ -1809,7 +1912,7 @@ def create_tab(tab: ImageBrowserTab, current_gr_tab: gr.Tab): def run_pnginfo(image, image_path, image_file_name): if image is None or os.path.splitext(image_file_name)[1] not in image_ext_list: - return '', '', '', '', '' + return "", "", "", "", "" try: geninfo, items = images.read_info_from_image(image) info = pnginfo2html(geninfo, items) @@ -1837,7 +1940,7 @@ def run_pnginfo(image, image_path, image_file_name): prompt = "" neg_prompt = "" - return '', geninfo, info, prompt, neg_prompt + return "", geninfo, info, prompt, neg_prompt def on_ui_tabs(): @@ -1932,7 +2035,7 @@ def on_ui_settings(): ("image_browser_video_y", None, 640, "Video player height (px)"), ] - section = ('image-browser', "Image Browser") + section = ("image-browser", "Image Browser") # Move historic setting names to current names added = 0 for cur_setting_name, old_setting_name, *option_info in image_browser_options: @@ -1945,4 +2048,4 @@ def on_ui_settings(): shared.opts.add_option(cur_setting_name, shared.OptionInfo(*option_info, section=section)) script_callbacks.on_ui_settings(on_ui_settings) -script_callbacks.on_ui_tabs(on_ui_tabs) +script_callbacks.on_ui_tabs(on_ui_tabs) \ No newline at end of file