import gradio as gr import json import logging import os import random import re import shutil import stat import time import modules.extras import modules.ui from modules import script_callbacks from modules import shared, scripts, images from modules.shared import opts, cmd_opts from modules.ui_common import plaintext_to_html from modules.ui_components import ToolButton from PIL import Image from PIL.ExifTags import TAGS from PIL.JpegImagePlugin import JpegImageFile from PIL.PngImagePlugin import PngImageFile from pathlib import Path from send2trash import send2trash from typing import List, Tuple favorite_tab_name = "Favorites" tabs_list = ["txt2img", "img2img", "instruct-pix2pix", "txt2img-grids", "img2img-grids", "Extras", favorite_tab_name, "Others"] #txt2img-grids and img2img-grids added by HaylockGrant num_of_imgs_per_page = 0 loads_files_num = 0 path_recorder_filename = os.path.join(scripts.basedir(), "path_recorder.txt") path_recorder_filename_tmp = f"{path_recorder_filename}.tmp" aes_cache_file = os.path.join(scripts.basedir(), "aes_scores.json") exif_cache_file = os.path.join(scripts.basedir(), "exif_data.json") ranking_file = os.path.join(scripts.basedir(), "ranking.json") image_ext_list = [".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp"] cur_ranking_value="0" finfo_aes = {} finfo_exif = {} none_select = "Nothing selected" refresh_symbol = '\U0001f504' # 🔄 up_symbol = '\U000025b2' # ▲ down_symbol = '\U000025bc' # ▼ warning_permission = "You have no permission to visit {}. If you want to visit all directories, add command line argument option '--administrator', More details here" current_depth = 0 logger = logging.getLogger(__name__) logger_mode = logging.INFO if hasattr(opts, "images_logger_warning"): if opts.images_logger_warning: logger_mode = logging.WARNING logger.setLevel(logger_mode) def delete_recycle(filename): if opts.images_delete_recycle: send2trash(filename) else: os.remove(filename) return def img_path_subdirs_get(img_path): subdirs = [] subdirs.append(none_select) for item in os.listdir(img_path): item_path = os.path.join(img_path, item) if os.path.isdir(item_path): subdirs.append(item_path) return gr.update(choices=subdirs) def img_path_add_remove(img_dir, path_recorder, add_remove, img_path_depth): if add_remove == "add" or (add_remove == "remove" and img_dir in path_recorder): if os.path.exists(path_recorder_filename_tmp): os.remove(path_recorder_filename_tmp) if add_remove == "add": path_recorder[img_dir] = { "depth": int(img_path_depth), "path_display": f"{img_dir} [{int(img_path_depth)}]" } else: del path_recorder[img_dir] path_recorder = {key: value for key, value in sorted(path_recorder.items(), key=lambda x: x[0].lower())} with open(path_recorder_filename_tmp, "w") as f: json.dump(path_recorder, f, indent=4) os.replace(path_recorder_filename_tmp, path_recorder_filename) path_recorder_formatted = [value.get("path_display") for key, value in path_recorder.items()] if add_remove == "remove": selected = None else: selected = path_recorder[img_dir]["path_display"] return path_recorder, gr.update(choices=path_recorder_formatted, value=selected) def sort_order_flip(turn_page_switch, sort_order): if sort_order == up_symbol: sort_order = down_symbol else: sort_order = up_symbol return 1, -turn_page_switch, sort_order def read_path_recorder(path_recorder, path_recorder_formatted): if os.path.exists(path_recorder_filename): try: with open(path_recorder_filename) as f: path_recorder = json.load(f) except json.JSONDecodeError: with open(path_recorder_filename) as f: path = f.readline().rstrip("\n") while len(path) > 0: path_recorder[path] = { "depth": 0, "path_display": f"{path} [0]" } path = f.readline().rstrip("\n") path_recorder_formatted = [value.get("path_display") for key, value in path_recorder.items()] path_recorder_formatted = sorted(path_recorder_formatted, key=lambda x: x.lower()) return path_recorder, path_recorder_formatted def pure_path(path): match = re.search(r" \[(\d+)\]$", path) if match: path = path[:match.start()] depth = int(match.group(1)) else: depth = 0 return path, depth def history2path(img_path_history): img_path, _ = pure_path(img_path_history) return img_path def reduplicative_file_move(src, dst): def same_name_file(basename, path): name, ext = os.path.splitext(basename) f_list = os.listdir(path) max_num = 0 for f in f_list: if len(f) <= len(basename): continue f_ext = f[-len(ext):] if len(ext) > 0 else "" if f[:len(name)] == name and f_ext == ext: if f[len(name)] == "(" and f[-len(ext)-1] == ")": number = f[len(name)+1:-len(ext)-1] if number.isdigit(): if int(number) > max_num: max_num = int(number) return f"{name}({max_num + 1}){ext}" name = os.path.basename(src) save_name = os.path.join(dst, name) if not os.path.exists(save_name): if opts.images_copy_image: shutil.copy2(src, dst) else: shutil.move(src, dst) else: name = same_name_file(name, dst) if opts.images_copy_image: shutil.copy2(src, os.path.join(dst, name)) else: shutil.move(src, os.path.join(dst, name)) def save_image(file_name): if file_name is not None and os.path.exists(file_name): reduplicative_file_move(file_name, opts.outdir_save) return "
 ") with gr.Row(visible= custom_dir): with gr.Column(scale=10): img_path = gr.Textbox(dir_name, label="Images directory", placeholder="Input images directory", interactive=custom_dir) with gr.Column(scale=1): img_path_depth = gr.Number(value="0", label="Sub directory depth") with gr.Column(scale=1): img_path_save_button = gr.Button(value="Add to / replace in saved directories") with gr.Row(visible= custom_dir): with gr.Column(scale=10): img_path_history = gr.Dropdown(choices=path_recorder_formatted, label="Saved directories") with gr.Column(scale=1): img_path_remove_button = gr.Button(value="Remove from saved directories") path_recorder = gr.State(path_recorder) with gr.Row(visible= custom_dir): with gr.Column(scale=10): img_path_subdirs = gr.Dropdown(choices=[none_select], value=none_select, label="Sub directories", interactive=True, elem_id="img_path_subdirs") with gr.Column(scale=1): img_path_subdirs_button = gr.Button(value="Get sub directories") with gr.Row(visible= not custom_dir, elem_id=tabname + "_images_history") as main_panel: with gr.Column(): with gr.Row(): with gr.Column(scale=2): with gr.Row(): first_page = gr.Button('First Page') prev_page = gr.Button('Prev Page') page_index = gr.Number(value=1, label="Page Index") refresh_index_button = ToolButton(value=refresh_symbol) next_page = gr.Button('Next Page') end_page = gr.Button('End Page') with gr.Column(scale=10): ranking = gr.Radio(value="None", choices=["1", "2", "3", "4", "5", "None"], label="ranking", interactive=True) auto_next = gr.Checkbox(label="Next Image After Ranking (To be implemented)", interactive=False, visible=False) history_gallery = gr.Gallery(show_label=False, elem_id=tabname + "_images_history_gallery").style(grid=opts.images_history_page_columns) with gr.Row() as delete_panel: with gr.Column(scale=1): delete_num = gr.Number(value=1, interactive=True, label="delete next") with gr.Column(scale=3): delete = gr.Button('Delete', elem_id=tabname + "_images_history_del_button") with gr.Column(scale=1): with gr.Row(scale=0.5): sort_by = gr.Dropdown(value="date", choices=["path name", "date", "aesthetic_score", "random", "cfg scale", "steps", "seed", "sampler", "size", "model", "model hash", "ranking"], label="sort by") sort_order = ToolButton(value=down_symbol) with gr.Row(): keyword = gr.Textbox(value="", label="filename keyword") exif_keyword = gr.Textbox(value="", label="exif keyword") with gr.Column(): ranking_filter = gr.Radio(value="All", choices=["All", "1", "2", "3", "4", "5", "None"], label="ranking filter", interactive=True) with gr.Row(): aes_filter = gr.Textbox(value="", label="minimum aesthetic_score") with gr.Row(): with gr.Column(): img_file_info = gr.Textbox(label="Generate Info", interactive=False, lines=6) img_file_name = gr.Textbox(value="", label="File Name", interactive=False) img_file_time= gr.HTML() with gr.Row(elem_id=tabname + "_images_history_button_panel") as button_panel: if tabname != favorite_tab_name: save_btn = gr.Button(f'{"Move" if not opts.images_copy_image else "Copy"} to favorites') try: send_to_buttons = modules.generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "inpaint", "extras"]) except: pass with gr.Row(): collected_warning = gr.HTML() # hidden items with gr.Row(visible=False): renew_page = gr.Button("Renew Page", elem_id=tabname + "_images_history_renew_page") visible_img_num = gr.Number() tabname_box = gr.Textbox(tabname) image_index = gr.Textbox(value=-1) set_index = gr.Button('set_index', elem_id=tabname + "_images_history_set_index") filenames = gr.State([]) all_images_list = gr.State() hidden = gr.Image(type="pil") info1 = gr.Textbox() info2 = gr.Textbox() load_switch = gr.Textbox(value="load_switch", label="load_switch") turn_page_switch = gr.Number(value=1, label="turn_page_switch") img_path_add = gr.Textbox(value="add") img_path_remove = gr.Textbox(value="remove") change_dir_outputs = [warning_box, main_panel, img_path_history, path_recorder, load_switch, img_path, img_path_depth] img_path.submit(change_dir, inputs=[img_path, path_recorder, load_switch, img_path_history, img_path_depth, img_path], outputs=change_dir_outputs) img_path_history.change(change_dir, inputs=[img_path_history, path_recorder, load_switch, img_path_history, img_path_depth, img_path], outputs=change_dir_outputs) img_path_history.change(history2path, inputs=[img_path_history], outputs=[img_path]) #delete delete.click(delete_image, inputs=[delete_num, img_file_name, filenames, image_index, visible_img_num], outputs=[filenames, delete_num, visible_img_num]) delete.click(fn=None, _js="images_history_delete", inputs=[delete_num, tabname_box, image_index], outputs=None) if tabname != favorite_tab_name: save_btn.click(save_image, inputs=[img_file_name], outputs=[collected_warning]) img_file_name.change(fn=update_move_text, inputs=[img_file_name], outputs=[save_btn]) #turn page first_page.click(lambda s:(1, -s) , inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) next_page.click(lambda p, s: (p + 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch]) prev_page.click(lambda p, s: (p - 1, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch]) end_page.click(lambda s: (-1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) load_switch.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) keyword.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) exif_keyword.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) aes_filter.submit(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) sort_by.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) ranking_filter.change(lambda s:(1, -s), inputs=[turn_page_switch], outputs=[page_index, turn_page_switch]) page_index.submit(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch]) renew_page.click(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch]) refresh_index_button.click(lambda p, s:(p, -s), inputs=[page_index, turn_page_switch], outputs=[page_index, turn_page_switch]) img_path_depth.change(lambda s: -s, inputs=[turn_page_switch], outputs=[turn_page_switch]) turn_page_switch.change( fn=get_image_page, inputs=[img_path, page_index, filenames, keyword, sort_by, sort_order, tabname_box, img_path_depth, ranking_filter, aes_filter, exif_keyword], outputs=[filenames, page_index, history_gallery, img_file_name, img_file_time, img_file_info, visible_img_num, warning_box] ) turn_page_switch.change(fn=None, inputs=[tabname_box], outputs=None, _js="images_history_turnpage") turn_page_switch.change(fn=lambda:(gr.update(visible=False), gr.update(visible=False)), inputs=None, outputs=[delete_panel, button_panel]) sort_order.click( fn=sort_order_flip, inputs=[turn_page_switch, sort_order], outputs=[page_index, turn_page_switch, sort_order] ) # Others img_path_subdirs_button.click( fn=img_path_subdirs_get, inputs=[img_path], outputs=[img_path_subdirs] ) img_path_subdirs.change( fn=change_dir, inputs=[img_path_subdirs, path_recorder, load_switch, img_path_history, img_path_depth, img_path], outputs=change_dir_outputs ) img_path_save_button.click( fn=img_path_add_remove, inputs=[img_path, path_recorder, img_path_add, img_path_depth], outputs=[path_recorder, img_path_history] ) img_path_remove_button.click( fn=img_path_add_remove, inputs=[img_path, path_recorder, img_path_remove], outputs=[path_recorder, img_path_history] ) # other functions set_index.click(show_image_info, _js="images_history_get_current_img", inputs=[tabname_box, image_index, page_index, filenames], outputs=[img_file_name, img_file_time, image_index, hidden]) set_index.click(fn=lambda:(gr.update(visible=True), gr.update(visible=True)), inputs=None, outputs=[delete_panel, button_panel]) img_file_name.change(fn=lambda : "", inputs=None, outputs=[collected_warning]) img_file_name.change(get_ranking, inputs=img_file_name, outputs=ranking) hidden.change(fn=run_pnginfo, inputs=[hidden, img_path, img_file_name], outputs=[info1, img_file_info, info2]) #ranking ranking.change(create_ranked_file, inputs=[img_file_name, ranking]) #ranking.change(show_next_image_info, _js="images_history_get_current_img", inputs=[tabname_box, image_index, page_index, auto_next], outputs=[img_file_name, img_file_time, image_index, hidden]) try: modules.generation_parameters_copypaste.bind_buttons(send_to_buttons, hidden, img_file_info) except: pass def run_pnginfo(image, image_path, image_file_name): if image is None: return '', '', '' geninfo, items = images.read_info_from_image(image) items = {**{'parameters': geninfo}, **items} info = '' for key, text in items.items(): info += f"""
{plaintext_to_html(str(key))}
{plaintext_to_html(str(text))}