933 lines
44 KiB
Python
933 lines
44 KiB
Python
|
|
import os
|
|
import shutil
|
|
import sys
|
|
from importlib import import_module
|
|
|
|
import gradio as gr
|
|
|
|
from extensions.sd_smartprocess import smartprocess
|
|
from extensions.sd_smartprocess.file_manager import FileManager, ImageData
|
|
from extensions.sd_smartprocess.interrogators.interrogator import InterrogatorRegistry
|
|
from extensions.sd_smartprocess.process_params import ProcessParams
|
|
from extensions.sd_smartprocess.smartprocess import is_image, get_backup_path
|
|
from modules import script_callbacks, shared
|
|
from modules.call_queue import wrap_gradio_gpu_call, wrap_gradio_call
|
|
from modules.ui import setup_progressbar
|
|
from modules.upscaler import Upscaler
|
|
|
|
refresh_symbol = "\U0001f4c2" # 📂
|
|
delete_symbol = "\U0001F5D1" # 🗑️
|
|
|
|
file_manager = FileManager()
|
|
# This holds all the inputs for the Smart Process tab
|
|
inputs_dict = {}
|
|
|
|
# This holds the accordions for the captioners
|
|
accordion_keys = []
|
|
captioner_accordions = []
|
|
|
|
# Sort all the tags by frequency
|
|
tags_dict = {}
|
|
|
|
current_image: ImageData = None
|
|
current_caption = None
|
|
|
|
# The list of processors to run
|
|
registry = InterrogatorRegistry()
|
|
|
|
int_dict = registry.list_interrogators()
|
|
print(f"Found {len(int_dict.keys())} interrogators: {int_dict}")
|
|
natural_captioner_names = ["CLIP", "BLIP", "MPLUG2"]
|
|
default_captioners = ["Swin"]
|
|
|
|
natural_captioners = {}
|
|
tag_captioners = {}
|
|
|
|
for interrogator, params in int_dict.items():
|
|
enable = False
|
|
display_name = interrogator.replace("Interrogator", "")
|
|
if display_name in default_captioners:
|
|
enable = True
|
|
if display_name in natural_captioner_names:
|
|
natural_captioners[display_name] = enable
|
|
else:
|
|
tag_captioners[display_name] = enable
|
|
|
|
wolf_captioners = ["Moat", "Swin", "Conv", "Conv2", "Vit"]
|
|
|
|
|
|
def list_scalers():
|
|
scaler_dir = os.path.join(shared.script_path, "extensions", "sd_smartprocess", "upscalers")
|
|
scalers = []
|
|
|
|
for root, dirs, files in os.walk(scaler_dir):
|
|
for file in files:
|
|
if file.endswith("_model.py"):
|
|
relative_path = os.path.relpath(os.path.join(root, file), scaler_dir)
|
|
module_name = "extensions.sd_smartprocess.upscalers." + relative_path.replace(os.sep, '.').replace('.py', '')
|
|
imported_module = import_module(module_name)
|
|
|
|
module_dir = os.path.dirname(os.path.join(root, file))
|
|
if module_dir not in sys.path:
|
|
sys.path.append(module_dir)
|
|
|
|
for name, obj in vars(imported_module).items():
|
|
# Check if the object is a class and a subclass of Upscaler
|
|
if isinstance(obj, type) and (issubclass(obj, Upscaler) or any(issubclass(base, Upscaler) for base in obj.__bases__)):
|
|
# Create an instance of the class and get the "scalers" attribute
|
|
instance = obj()
|
|
if hasattr(instance, 'scalers'):
|
|
for scaler_data in instance.scalers:
|
|
scalers.append(scaler_data)
|
|
|
|
system_scalers = shared.sd_upscalers
|
|
|
|
# Add unique scalers to the system_scalers list
|
|
for scaler in scalers:
|
|
if scaler.name not in [x.name for x in system_scalers]:
|
|
system_scalers.append(scaler)
|
|
|
|
def generate_caption_section():
|
|
global captioner_accordions
|
|
global inputs_dict
|
|
wolf_params = {}
|
|
for gator, gator_params in int_dict.items():
|
|
accordion_name = gator.replace("Interrogator", "")
|
|
|
|
if accordion_name in wolf_captioners:
|
|
wolf_params = gator_params
|
|
continue
|
|
show_accordion = False
|
|
if gator in natural_captioners.keys():
|
|
show_accordion = natural_captioners.get(gator, False)
|
|
elif gator in tag_captioners.keys():
|
|
show_accordion = tag_captioners.get(gator, False)
|
|
if len(gator_params.keys()) == 0:
|
|
continue
|
|
print(f"Adding {accordion_name} to captioners")
|
|
with gr.Accordion(label=accordion_name, open=False, visible=show_accordion) as temp_accordion:
|
|
with gr.Column():
|
|
for param, param_value in gator_params.items():
|
|
temp_element = None
|
|
label = param.replace("_", " ").title()
|
|
if isinstance(param_value, bool):
|
|
temp_element = gr.Checkbox(label=label, value=param_value, interactive=True)
|
|
elif isinstance(param_value, int):
|
|
max_value = 100
|
|
if "tokens" in param:
|
|
max_value = 300
|
|
temp_element = gr.Slider(label=label, value=param_value, step=1, minimum=0, maximum=max_value,
|
|
interactive=True)
|
|
elif isinstance(param_value, float):
|
|
temp_element = gr.Slider(label=label, value=param_value, step=0.01, minimum=0, maximum=1,
|
|
interactive=True)
|
|
elif isinstance(param_value, str):
|
|
if param == "group":
|
|
continue
|
|
temp_element = gr.Textbox(label=label, value=param_value, interactive=True)
|
|
if temp_element:
|
|
inputs_dict[param] = temp_element
|
|
captioner_accordions.append(temp_accordion)
|
|
accordion_keys.append(accordion_name)
|
|
show_accordion = False
|
|
for wc in wolf_captioners:
|
|
if wc in default_captioners:
|
|
show_accordion = True
|
|
accordion_name = "Wolf"
|
|
print(f"Adding {accordion_name} to captioners")
|
|
|
|
with gr.Accordion(label=accordion_name, open=False, visible=show_accordion) as temp_accordion:
|
|
with gr.Column():
|
|
for param, param_value in wolf_params.items():
|
|
temp_element = None
|
|
label = param.replace("_", " ").title()
|
|
if isinstance(param_value, bool):
|
|
temp_element = gr.Checkbox(label=label, value=param_value, interactive=True)
|
|
elif isinstance(param_value, int):
|
|
temp_element = gr.Slider(label=label, value=param_value, step=1, minimum=0, maximum=100,
|
|
interactive=True)
|
|
elif isinstance(param_value, float):
|
|
temp_element = gr.Slider(label=label, value=param_value, step=0.01, minimum=0, maximum=1,
|
|
interactive=True)
|
|
elif isinstance(param_value, str):
|
|
if param == "group":
|
|
continue
|
|
temp_element = gr.Textbox(label=label, value=param_value, interactive=True)
|
|
if temp_element:
|
|
setattr(temp_element, "do_not_save_to_config", True)
|
|
inputs_dict[param] = temp_element
|
|
# Sort all_inputs and input_keys
|
|
captioner_accordions.append(temp_accordion)
|
|
accordion_keys.append(accordion_name)
|
|
|
|
|
|
def sort_tags_dict():
|
|
global tags_dict
|
|
global file_manager
|
|
# Sort the tags into a list by frequency
|
|
sorted_tags = sorted(tags_dict.items(), key=lambda x: x[1], reverse=True)
|
|
tag_options = []
|
|
for tag in sorted_tags:
|
|
if tag[0] in file_manager.included_tags or tag[0] in file_manager.excluded_tags:
|
|
continue
|
|
tag_name = tag[0]
|
|
tag_label = f"{tag_name} ({tag[1]})"
|
|
tag_options.append(tag_label)
|
|
return tag_options
|
|
|
|
|
|
def create_process_ui():
|
|
"""Create the UI for the Smart Process tab"""
|
|
global tag_captioners
|
|
global natural_captioners
|
|
global inputs_dict
|
|
global captioner_accordions
|
|
global file_manager
|
|
list_scalers()
|
|
with gr.Blocks() as sp_interface:
|
|
with gr.Row(equal_height=True):
|
|
with gr.Column():
|
|
with gr.Row():
|
|
sp_src = gr.Textbox(label='Source directory', elem_id="sp_src")
|
|
sp_load_src = gr.Button(value=refresh_symbol, size="sm", variant='primary', elem_id="sp_load_src")
|
|
sp_clear_src = gr.Button(value=delete_symbol, size="sm", variant='secondary',
|
|
elem_id="sp_clear_src")
|
|
with gr.Column():
|
|
with gr.Row(elem_id="sp_top_btn_row"):
|
|
sp_process_selected = gr.Button(value="Process Selected")
|
|
sp_cancel = gr.Button(value="Cancel", visible=False)
|
|
sp_process_all = gr.Button(value="Process All", variant='primary')
|
|
|
|
with gr.Column():
|
|
with gr.Row():
|
|
with gr.Accordion(label="Filtering", open=False):
|
|
with gr.Row():
|
|
with gr.Column():
|
|
sp_filter = gr.Textbox(label="Filter", value="")
|
|
with gr.Column():
|
|
with gr.Row():
|
|
sp_filter_include = gr.Button(value="Include", elem_classes=["inc_exc_btn"])
|
|
sp_filter_exclude = gr.Button(value="Exclude", elem_classes=["inc_exc_btn"])
|
|
with gr.Row():
|
|
with gr.Column():
|
|
sp_tag_dropdown = gr.Dropdown(label="Tags", choices=sort_tags_dict(), value="")
|
|
with gr.Column():
|
|
with gr.Row():
|
|
sp_tag_include = gr.Button(value="Include", elem_classes=["inc_exc_btn"])
|
|
sp_tag_exclude = gr.Button(value="Exclude", elem_classes=["inc_exc_btn"])
|
|
with gr.Row(variant="panel"):
|
|
with gr.Row():
|
|
with gr.Column():
|
|
sp_filter_string_include_group = gr.CheckboxGroup(label="Included Strings",
|
|
choices=file_manager.included_strings,
|
|
value=file_manager.included_strings,
|
|
interactive=True)
|
|
with gr.Column():
|
|
sp_filter_string_exclude_group = gr.CheckboxGroup(label="Excluded Strings",
|
|
choices=file_manager.excluded_strings,
|
|
value=file_manager.excluded_strings,
|
|
interactive=True)
|
|
|
|
with gr.Row():
|
|
with gr.Column():
|
|
sp_include_tag_group = gr.CheckboxGroup(label="Included Tags",
|
|
choices=file_manager.included_tags,
|
|
value=file_manager.included_tags,
|
|
interactive=True)
|
|
with gr.Column():
|
|
sp_exclude_tag_group = gr.CheckboxGroup(label="Excluded Tags",
|
|
choices=file_manager.excluded_tags,
|
|
value=file_manager.excluded_tags,
|
|
interactive=True)
|
|
|
|
with gr.Row(equal_height=False):
|
|
with gr.Column(variant="panel", scale=1):
|
|
with gr.Tab("Pre-Process") as sp_crop_tab:
|
|
sp_size = gr.Slider(minimum=64, maximum=4096, step=64, label="Output Size", value=1024)
|
|
with gr.Row():
|
|
sp_pad = gr.Checkbox(label="Pad Images", elem_classes=["short_check"])
|
|
sp_crop = gr.Checkbox(label="Crop Images", elem_classes=["short_check"])
|
|
sp_do_rename = gr.Checkbox(label="Rename Images", elem_classes=["short_check"])
|
|
sp_crop_mode = gr.Dropdown(label='Crop Mode', value="smart",
|
|
choices=["smart", "square", "empty", "contain"],
|
|
visible=False)
|
|
|
|
def toggle_mode_dropdown(evt: gr.SelectData):
|
|
return gr.update(visible=evt.selected)
|
|
|
|
sp_crop.select(
|
|
fn=toggle_mode_dropdown,
|
|
outputs=[sp_crop_mode]
|
|
)
|
|
|
|
with gr.Tab("Captions") as sp_cap_tab:
|
|
sp_caption = gr.Checkbox(label='Generate Captions')
|
|
with gr.Column(visible=False) as cap_settings:
|
|
sp_nl_captioners = gr.CheckboxGroup(label='Natural Language Captioners',
|
|
choices=natural_captioners.keys(),
|
|
value=[x for x in natural_captioners.keys() if
|
|
natural_captioners[x]],
|
|
interactive=True)
|
|
sp_captioners = gr.CheckboxGroup(label='Captioners',
|
|
choices=tag_captioners.keys(),
|
|
value=[x for x in tag_captioners.keys() if tag_captioners[x]],
|
|
interactive=True)
|
|
sp_txt_action = gr.Dropdown(label='Existing Caption Action', value="ignore",
|
|
choices=["ignore", "include"])
|
|
generate_caption_section()
|
|
sp_tags_to_ignore = gr.Textbox(label="Tags To Ignore", value="")
|
|
sp_replace_class = gr.Checkbox(label='Replace Class with Subject in Caption', value=False)
|
|
sp_class = gr.Textbox(label='Subject Class', placeholder='Subject class to crop (leave '
|
|
'blank to auto-detect)')
|
|
sp_subject = gr.Textbox(label='Subject Name', placeholder='Subject Name to replace class '
|
|
'with in captions')
|
|
|
|
def update_caption_groups(evt: gr.SelectData):
|
|
global tag_captioners
|
|
global natural_captioners
|
|
if evt.value in natural_captioners.keys():
|
|
natural_captioners[evt.value] = evt.selected
|
|
elif evt.value in tag_captioners.keys():
|
|
tag_captioners[evt.value] = evt.selected
|
|
outputs = []
|
|
show_wolf = False
|
|
for cap_name, accordion in zip(accordion_keys, captioner_accordions):
|
|
if cap_name in tag_captioners.keys():
|
|
outputs.append(gr.update(visible=tag_captioners.get(cap_name, False)))
|
|
elif cap_name in natural_captioners.keys():
|
|
outputs.append(gr.update(visible=natural_captioners.get(cap_name, False)))
|
|
for wc in wolf_captioners:
|
|
if tag_captioners.get(wc, False):
|
|
show_wolf = True
|
|
outputs.append(gr.update(visible=show_wolf))
|
|
return outputs
|
|
|
|
sp_captioners.select(
|
|
fn=update_caption_groups,
|
|
outputs=captioner_accordions
|
|
)
|
|
|
|
sp_nl_captioners.select(
|
|
fn=update_caption_groups,
|
|
outputs=captioner_accordions
|
|
)
|
|
|
|
def toggle_cap_settings(evt: gr.SelectData):
|
|
return gr.update(visible=evt.selected)
|
|
|
|
sp_caption.select(
|
|
fn=toggle_cap_settings,
|
|
outputs=[cap_settings]
|
|
)
|
|
|
|
with gr.Tab("Post-Process") as sp_post_tab:
|
|
sp_restore_faces = gr.Checkbox(label='Restore Faces', value=False)
|
|
sp_face_model = gr.Dropdown(label="Face Restore Model", choices=["GFPGAN", "Codeformer"],
|
|
value="GFPGAN")
|
|
sp_upscale = gr.Checkbox(label='Upscale and Resize', value=False)
|
|
with gr.Column(visible=False) as sp_upscale_settings:
|
|
sp_upscale_mode = gr.Radio(label="Upscale Mode", choices=["Size", "Ratio"], value="Size")
|
|
sp_upscale_ratio = gr.Slider(label="Upscale Ratio", value=2, step=1, minimum=2, maximum=8,
|
|
visible=False)
|
|
sp_upscale_size = gr.Slider(label="Upscale Size", value=1024, step=64, minimum=512,
|
|
maximum=8192)
|
|
sp_upscaler_1 = gr.Dropdown(label='Upscaler', elem_id="sp_scaler_1",
|
|
choices=[x.name for x in shared.sd_upscalers],
|
|
value=shared.sd_upscalers[0].name)
|
|
sp_upscaler_2 = gr.Dropdown(label='Upscaler', elem_id="sp_scaler_2",
|
|
choices=[x.name for x in shared.sd_upscalers],
|
|
value=shared.sd_upscalers[0].name)
|
|
|
|
def toggle_upscale_mode(evt: gr.SelectData):
|
|
return gr.update(visible=evt.value == "Ratio"), gr.update(visible=evt.value == "Size")
|
|
|
|
def toggle_upscale_settings(evt: gr.SelectData):
|
|
return gr.update(visible=evt.selected)
|
|
|
|
sp_upscale_mode.select(
|
|
fn=toggle_upscale_mode,
|
|
outputs=[sp_upscale_ratio, sp_upscale_size]
|
|
)
|
|
|
|
sp_upscale.select(
|
|
fn=toggle_upscale_settings,
|
|
outputs=[sp_upscale_settings]
|
|
)
|
|
|
|
# Preview/progress
|
|
with gr.Column(variant="panel", scale=1):
|
|
sp_gallery_all = gr.Gallery(label='All', elem_id='sp_gallery', rows=2, columns=4,
|
|
allow_preview=False)
|
|
with gr.Row():
|
|
sp_add_all = gr.Button(value="Add All", size="sm")
|
|
sp_add_selected = gr.Button(value="Add Selected", size="sm")
|
|
sp_remove_selected = gr.Button(value="Remove Selected", size="sm")
|
|
sp_clear_selected = gr.Button(value="Remove All", size="sm")
|
|
sp_gallery_selected = gr.Gallery(label='Selected', elem_id='sp_selected', rows=2, columns=4,
|
|
allow_preview=False)
|
|
|
|
with gr.Column(variant="panel", scale=1):
|
|
with gr.Row(elem_id="sp_file_options"):
|
|
sp_do_backup = gr.Checkbox(label="Backup Files", elem_classes=["short_check"])
|
|
sp_auto_save = gr.Checkbox(label="Auto-Save", elem_classes=["short_check"])
|
|
with gr.Row() as sp_crop_row:
|
|
sp_pre_current = gr.Button(value="PreProcess Current", size="sm")
|
|
sp_save_current_pre = gr.Button(value="Save Pre", size="sm")
|
|
with gr.Row(visible=False) as sp_cap_row:
|
|
with gr.Column():
|
|
with gr.Row():
|
|
sp_current_caption = gr.Textbox(label="Caption", value="", lines=2)
|
|
with gr.Row():
|
|
sp_caption_current = gr.Button(value="Caption Current", size="sm")
|
|
sp_save_current_caption = gr.Button(value="Save Caption", size="sm")
|
|
with gr.Row(visible=False) as sp_post_row:
|
|
sp_post_current = gr.Button(value="PostProcess Current", size="sm")
|
|
sp_save_current_post = gr.Button(value="Save Post", size="sm")
|
|
sp_progress = gr.HTML(elem_id="sp_progress", value="")
|
|
sp_outcome = gr.HTML(elem_id="sp_error", value="")
|
|
sp_progressbar = gr.HTML(elem_id="sp_progressbar")
|
|
sp_preview = gr.Image(label='Preview', elem_id='sp_preview', interactive=False)
|
|
setup_progressbar(sp_progressbar, sp_preview, 'sp', textinfo=sp_progress)
|
|
|
|
def toggle_group_buttons(gd: gr.SelectData):
|
|
value = gd.value
|
|
selected = gd.selected
|
|
show_cap = False
|
|
show_pre = False
|
|
show_post = False
|
|
if value == "Pre-Process":
|
|
show_pre = selected
|
|
elif value == "Captions":
|
|
show_cap = selected
|
|
elif value == "Post-Process":
|
|
show_post = selected
|
|
return gr.update(visible=show_pre), gr.update(visible=show_cap), gr.update(visible=show_post)
|
|
|
|
for elem in [sp_crop_tab, sp_cap_tab, sp_post_tab]:
|
|
elem.select(
|
|
fn=toggle_group_buttons,
|
|
outputs=[sp_crop_row, sp_cap_row, sp_post_row]
|
|
)
|
|
|
|
def process_outputs(params: ProcessParams, current=False, selected=False):
|
|
global file_manager
|
|
global current_image
|
|
image_data, output = smartprocess.do_process(params)
|
|
file_manager.update_files(image_data)
|
|
if current:
|
|
if len(image_data) > 0:
|
|
current_image = image_data[0]
|
|
return gr.update(value=current_image.get_image()), gr.update(
|
|
value=current_image.caption), gr.update(value=output)
|
|
else:
|
|
return gr.update(value=None), gr.update(value=None), gr.update(value=output)
|
|
elif selected:
|
|
images = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=images), gr.update(value=output)
|
|
else:
|
|
images = file_manager.filtered_files(True)
|
|
selected_files = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=images), gr.update(value=output), gr.update(
|
|
value=selected_files)
|
|
|
|
def pre_process_current(*args):
|
|
global current_image
|
|
params = params_to_dict(*args)
|
|
params.src_files = [current_image]
|
|
params.pre_only()
|
|
params.save_image = False
|
|
return process_outputs(params, current=True)
|
|
|
|
def caption_current(*args):
|
|
global current_image
|
|
cap_params = params_to_dict(*args)
|
|
cap_params.src_files = [current_image]
|
|
cap_params.save_caption = False
|
|
cap_params.cap_only()
|
|
return process_outputs(cap_params, current=True)
|
|
|
|
def post_process_current(*args):
|
|
global current_image
|
|
post_params = params_to_dict(*args)
|
|
post_params.src_files = [current_image]
|
|
post_params.post_only()
|
|
post_params.save_image = False
|
|
return process_outputs(post_params, current=True)
|
|
|
|
def process_selected(*args):
|
|
global file_manager
|
|
params = params_to_dict(*args)
|
|
params.src_files = file_manager.filtered_and_selected_files()
|
|
return process_outputs(params, selected=True)
|
|
|
|
def process_all(*args):
|
|
global file_manager
|
|
params = params_to_dict(*args)
|
|
params.src_files = file_manager.filtered_files()
|
|
return process_outputs(params)
|
|
|
|
def start_caption(params: ProcessParams, current=False, selected=False):
|
|
params.captioners = tag_captioners
|
|
params.crop = False
|
|
params.pad = False
|
|
params.scale = False
|
|
params.restore_faces = False
|
|
params.caption = True
|
|
if len(params.src_files) > 0:
|
|
return process_outputs(params, current=current, selected=selected)
|
|
else:
|
|
return gr.update(value=None), gr.update(value=None), gr.update(value="No images selected")
|
|
|
|
def process_tags(captions):
|
|
global tags_dict
|
|
tags_dict = {}
|
|
for caption in captions:
|
|
tags = caption.split(",")
|
|
for tag in tags:
|
|
tag = tag.strip()
|
|
if tag not in tags_dict.keys():
|
|
tags_dict[tag] = 1
|
|
else:
|
|
tags_dict[tag] += 1
|
|
|
|
def clear_globals():
|
|
global file_manager
|
|
global current_image
|
|
file_manager.clear()
|
|
current_image = None
|
|
|
|
def load_src(src_path):
|
|
# Enumerate all files recursively in src_path, and if they're an image, add them to the gallery
|
|
global file_manager
|
|
file_manager.file_path = src_path
|
|
file_manager.load_files()
|
|
process_tags(file_manager.all_captions())
|
|
tag_group = gr.update(choices=sort_tags_dict(), value="")
|
|
all_files = file_manager.filtered_files(True)
|
|
print(f"Found {len(all_files)} images in {src_path}")
|
|
outputs = [gr.update(value=all_files), tag_group]
|
|
# Enumerate all elements in elements_to_clear except for the first two
|
|
for _ in elements_to_clear[2:]:
|
|
outputs.append(gr.update(value=None))
|
|
return outputs
|
|
|
|
def set_all_current(evt: gr.SelectData):
|
|
global current_image
|
|
global file_manager
|
|
if evt.selected:
|
|
files = file_manager.filtered_files(False)
|
|
current_image = files[evt.index]
|
|
return gr.update(value=current_image.image_path), gr.update(value=current_image.caption)
|
|
else:
|
|
return gr.update(value=None), gr.update(value=None)
|
|
|
|
def set_selected_current(evt: gr.SelectData):
|
|
global current_image
|
|
global file_manager
|
|
if evt.selected:
|
|
files = file_manager.filtered_files(False)
|
|
current_image = files[evt.index]
|
|
return gr.update(value=current_image.image_path), gr.update(value=current_image.caption)
|
|
else:
|
|
return gr.update(value=None), gr.update(value=None)
|
|
|
|
def add_selected():
|
|
global current_image
|
|
global file_manager
|
|
if current_image and not current_image.selected:
|
|
current_image.selected = True
|
|
file_manager.update_file(current_image)
|
|
selected_files = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=selected_files)
|
|
|
|
def remove_selected():
|
|
global file_manager
|
|
global current_image
|
|
if current_image and current_image.selected:
|
|
current_image.selected = False
|
|
file_manager.update_file(current_image)
|
|
selected_files = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=selected_files)
|
|
|
|
def clear_selected():
|
|
global file_manager
|
|
global current_image
|
|
current_image = None
|
|
all_files = file_manager.all_files()
|
|
for file in all_files:
|
|
file.selected = False
|
|
file_manager.update_files(all_files)
|
|
selected_files = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=selected_files)
|
|
|
|
def select_all():
|
|
global file_manager
|
|
all_files = file_manager.filtered_files()
|
|
for file in all_files:
|
|
file.selected = True
|
|
file_manager.update_files(all_files)
|
|
all_files = file_manager.filtered_files(True)
|
|
return gr.update(value=all_files)
|
|
|
|
def params_to_dict(*args):
|
|
global inputs_dict
|
|
input_keys = list(inputs_dict.keys())
|
|
pp = list(args)
|
|
params_dict = dict(zip(input_keys, pp))
|
|
pp = ProcessParams().from_dict(params_dict)
|
|
return pp
|
|
|
|
def include_tag(tag_dropdown_value):
|
|
global file_manager
|
|
included_tags = file_manager.included_tags
|
|
tag_dropdown_value = tag_dropdown_value.split(" (")[0]
|
|
if tag_dropdown_value not in included_tags:
|
|
included_tags.append(tag_dropdown_value)
|
|
# Sort filter_tags_include by frequency, then alphabetically, ensuring we only add the values in filter_tags_include
|
|
included_tags = [x[0] for x in sorted(tags_dict.items(), key=lambda x: (-x[1], x[0])) if
|
|
x[0] in included_tags]
|
|
file_manager.included_tags = included_tags
|
|
file_manager.update_filters()
|
|
tag_group = gr.update(choices=sort_tags_dict(), value="")
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_files(True)
|
|
return tag_group, gr.update(value=included_tags, choices=included_tags), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items)
|
|
|
|
def exclude_tag(tag_dropdown_value):
|
|
global file_manager
|
|
excluded_tags = file_manager.excluded_tags
|
|
tag_dropdown_value = tag_dropdown_value.split(" (")[0]
|
|
if tag_dropdown_value not in excluded_tags:
|
|
excluded_tags.append(tag_dropdown_value)
|
|
# Sort filter_tags_exclude by frequency, then alphabetically
|
|
excluded_tags = [x[0] for x in sorted(tags_dict.items(), key=lambda x: (-x[1], x[0])) if
|
|
x[0] in excluded_tags]
|
|
file_manager.excluded_tags = excluded_tags
|
|
file_manager.update_filters()
|
|
tag_group = gr.update(choices=sort_tags_dict(), value="")
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_and_selected_files(True)
|
|
return tag_group, gr.update(value=excluded_tags, choices=excluded_tags), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items)
|
|
|
|
def remove_include_tag(select_data: gr.SelectData):
|
|
global file_manager
|
|
included_tags = file_manager.included_tags
|
|
|
|
if not select_data.selected:
|
|
included_tags.remove(select_data.value)
|
|
included_tags = [x[0] for x in sorted(tags_dict.items(), key=lambda x: (-x[1], x[0])) if
|
|
x[0] in included_tags]
|
|
file_manager.included_tags = included_tags
|
|
file_manager.update_filters()
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(choices=sort_tags_dict(), value=""), gr.update(value=included_tags,
|
|
choices=included_tags), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items)
|
|
|
|
def remove_exclude_tag(select_data: gr.SelectData):
|
|
global file_manager
|
|
filter_tags_exclude = file_manager.excluded_tags
|
|
if select_data.selected:
|
|
filter_tags_exclude.append(select_data.value)
|
|
else:
|
|
filter_tags_exclude.remove(select_data.value)
|
|
# Sort filter tags alphabetically
|
|
filter_tags_exclude = [x[0] for x in sorted(tags_dict.items(), key=lambda x: (-x[1], x[0])) if
|
|
x[0] in filter_tags_exclude]
|
|
file_manager.excluded_tags = filter_tags_exclude
|
|
file_manager.update_filters()
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(choices=sort_tags_dict(), value=""), gr.update(value=filter_tags_exclude,
|
|
choices=filter_tags_exclude), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items)
|
|
|
|
def include_string(string_group_value):
|
|
global file_manager
|
|
filter_string_include = file_manager.included_strings
|
|
if string_group_value not in filter_string_include:
|
|
filter_string_include.append(string_group_value)
|
|
# Sort filter_string_include alphabetically
|
|
filter_string_include = sorted(filter_string_include)
|
|
file_manager.included_strings = filter_string_include
|
|
file_manager.update_filters()
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=filter_string_include, choices=filter_string_include), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items), gr.update(value="")
|
|
|
|
def exclude_string(string_group_value):
|
|
global file_manager
|
|
filter_string_exclude = file_manager.excluded_strings
|
|
if string_group_value not in filter_string_exclude:
|
|
filter_string_exclude.append(string_group_value)
|
|
# Sort filter_string_exclude alphabetically
|
|
filter_string_exclude = sorted(filter_string_exclude)
|
|
file_manager.excluded_strings = filter_string_exclude
|
|
file_manager.update_filters()
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=filter_string_exclude, choices=filter_string_exclude), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items), gr.update(value="")
|
|
|
|
def remove_include_string(select_data: gr.SelectData):
|
|
global file_manager
|
|
filter_string_include = file_manager.included_strings
|
|
if select_data.selected:
|
|
filter_string_include.append(select_data.value)
|
|
else:
|
|
filter_string_include.remove(select_data.value)
|
|
filter_string_include = sorted(filter_string_include)
|
|
file_manager.included_strings = filter_string_include
|
|
file_manager.update_filters()
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=filter_string_include, choices=filter_string_include), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items), gr.update(value="")
|
|
|
|
def remove_exclude_string(select_data: gr.SelectData):
|
|
global file_manager
|
|
filter_string_exclude = file_manager.excluded_strings
|
|
if select_data.selected:
|
|
filter_string_exclude.append(select_data.value)
|
|
else:
|
|
filter_string_exclude.remove(select_data.value)
|
|
filter_string_exclude = sorted(filter_string_exclude)
|
|
file_manager.excluded_strings = filter_string_exclude
|
|
file_manager.update_filters()
|
|
all_gallery_items = file_manager.filtered_files(True)
|
|
selected_gallery_items = file_manager.filtered_and_selected_files(True)
|
|
return gr.update(value=filter_string_exclude, choices=filter_string_exclude), gr.update(
|
|
value=all_gallery_items), gr.update(value=selected_gallery_items), gr.update(value="")
|
|
|
|
def save_current_caption(caption, do_backup=True):
|
|
temp_params = ProcessParams()
|
|
temp_params.do_backup = do_backup
|
|
global current_image
|
|
if current_image:
|
|
image_path = current_image.image_path
|
|
caption_txt_filename = os.path.splitext(image_path)[0] + '.txt'
|
|
caption_txt_filename, caption_backup_path = get_backup_path(caption_txt_filename, temp_params)
|
|
if caption_txt_filename != caption_backup_path:
|
|
shutil.copy(caption_txt_filename, caption_backup_path)
|
|
with open(caption_txt_filename, 'w', encoding="utf8") as f:
|
|
f.write(caption)
|
|
return gr.update(value="Caption saved to " + caption_txt_filename)
|
|
else:
|
|
return gr.update(value="No image selected")
|
|
|
|
def save_current_image(do_backup, do_rename):
|
|
temp_params = ProcessParams()
|
|
temp_params.do_backup = do_backup
|
|
temp_params.do_rename = do_rename
|
|
temp_params.save_image = True
|
|
global current_image
|
|
if current_image:
|
|
image_path = current_image.image_path
|
|
save_path, backup_path = get_backup_path(image_path, temp_params)
|
|
if save_path != backup_path:
|
|
shutil.copy(image_path, backup_path)
|
|
current_image_element = current_image.get_image()
|
|
current_image_element.save(save_path)
|
|
return gr.update(value=save_path), gr.update(value="Image saved to " + save_path)
|
|
else:
|
|
return gr.update(value=None), gr.update(value="No image selected")
|
|
|
|
inputs_dict = {
|
|
"auto_save": sp_auto_save,
|
|
"caption": sp_caption,
|
|
"captioners": sp_captioners,
|
|
"class": sp_class,
|
|
"clear_src": sp_clear_src,
|
|
"crop": sp_crop,
|
|
"crop_mode": sp_crop_mode,
|
|
"do_backup": sp_do_backup,
|
|
"do_rename": sp_do_rename,
|
|
"face_model": sp_face_model,
|
|
"nl_captioners": sp_nl_captioners,
|
|
"pad": sp_pad,
|
|
"replace_class": sp_replace_class,
|
|
"restore_faces": sp_restore_faces,
|
|
"size": sp_size,
|
|
"src": sp_src,
|
|
"subject": sp_subject,
|
|
"tags_to_ignore": sp_tags_to_ignore,
|
|
"txt_action": sp_txt_action,
|
|
"upscale": sp_upscale,
|
|
"upscale_ratio": sp_upscale_ratio,
|
|
"upscaler_1": sp_upscaler_1,
|
|
"upscaler_2": sp_upscaler_2
|
|
}
|
|
|
|
def all_inputs():
|
|
return [x for x in inputs_dict.values()]
|
|
|
|
sp_tag_include.click(
|
|
fn=include_tag,
|
|
inputs=[sp_tag_dropdown],
|
|
outputs=[sp_tag_dropdown, sp_include_tag_group, sp_gallery_all, sp_gallery_selected]
|
|
)
|
|
|
|
sp_tag_exclude.click(
|
|
fn=exclude_tag,
|
|
inputs=[sp_tag_dropdown],
|
|
outputs=[sp_tag_dropdown, sp_exclude_tag_group, sp_gallery_all, sp_gallery_selected]
|
|
)
|
|
|
|
sp_exclude_tag_group.select(
|
|
fn=remove_exclude_tag,
|
|
outputs=[sp_tag_dropdown, sp_exclude_tag_group, sp_gallery_all, sp_gallery_selected]
|
|
)
|
|
|
|
sp_include_tag_group.select(
|
|
fn=remove_include_tag,
|
|
outputs=[sp_tag_dropdown, sp_include_tag_group, sp_gallery_all, sp_gallery_selected]
|
|
)
|
|
|
|
sp_filter_string_include_group.select(
|
|
fn=remove_include_string,
|
|
outputs=[sp_filter_string_include_group, sp_gallery_all, sp_gallery_selected, sp_filter]
|
|
)
|
|
|
|
sp_filter_string_exclude_group.select(
|
|
fn=remove_exclude_string,
|
|
outputs=[sp_filter_string_exclude_group, sp_gallery_all, sp_gallery_selected, sp_filter]
|
|
)
|
|
|
|
sp_filter_include.click(
|
|
fn=include_string,
|
|
inputs=[sp_filter],
|
|
outputs=[sp_filter_string_include_group, sp_gallery_all, sp_gallery_selected, sp_filter]
|
|
)
|
|
|
|
sp_filter_exclude.click(
|
|
fn=exclude_string,
|
|
inputs=[sp_filter],
|
|
outputs=[sp_filter_string_exclude_group, sp_gallery_all, sp_gallery_selected, sp_filter]
|
|
)
|
|
|
|
sp_pre_current.click(
|
|
fn=wrap_gradio_call(pre_process_current),
|
|
_js="start_smart_process",
|
|
inputs=all_inputs(),
|
|
outputs=[sp_preview, sp_current_caption, sp_outcome]
|
|
)
|
|
|
|
sp_save_current_pre.click(
|
|
fn=save_current_image,
|
|
inputs=[sp_do_backup, sp_do_rename],
|
|
outputs=[sp_preview, sp_outcome]
|
|
)
|
|
|
|
sp_caption_current.click(
|
|
fn=wrap_gradio_call(caption_current),
|
|
_js="start_smart_process",
|
|
inputs=all_inputs(),
|
|
outputs=[sp_preview, sp_current_caption, sp_outcome]
|
|
)
|
|
|
|
sp_save_current_caption.click(
|
|
fn=save_current_caption,
|
|
inputs=[sp_current_caption, sp_do_backup],
|
|
outputs=[sp_outcome]
|
|
)
|
|
|
|
sp_post_current.click(
|
|
fn=wrap_gradio_call(post_process_current),
|
|
_js="start_smart_process",
|
|
inputs=all_inputs(),
|
|
outputs=[sp_preview, sp_current_caption, sp_outcome]
|
|
)
|
|
|
|
sp_save_current_post.click(
|
|
fn=save_current_image,
|
|
inputs=[sp_do_backup, sp_do_rename],
|
|
outputs=[sp_preview, sp_outcome]
|
|
)
|
|
|
|
sp_process_selected.click(
|
|
fn=wrap_gradio_gpu_call(process_selected, extra_outputs=[gr.update()]),
|
|
_js="start_smart_process",
|
|
inputs=all_inputs(),
|
|
outputs=[sp_current_caption]
|
|
)
|
|
|
|
sp_process_all.click(
|
|
fn=wrap_gradio_gpu_call(process_all, extra_outputs=[gr.update()]),
|
|
_js="start_smart_process",
|
|
inputs=all_inputs(),
|
|
outputs=[
|
|
sp_progress,
|
|
sp_outcome
|
|
],
|
|
)
|
|
|
|
sp_cancel.click(
|
|
fn=lambda: shared.state.interrupt()
|
|
)
|
|
|
|
sp_gallery_all.select(fn=set_all_current, outputs=[sp_preview, sp_current_caption])
|
|
sp_gallery_selected.select(fn=set_selected_current, outputs=[sp_preview, sp_current_caption])
|
|
|
|
elements_to_clear = [sp_gallery_all,
|
|
sp_tag_dropdown,
|
|
sp_gallery_selected,
|
|
sp_filter,
|
|
sp_current_caption,
|
|
sp_preview,
|
|
sp_filter_string_include_group,
|
|
sp_filter_string_exclude_group,
|
|
sp_include_tag_group,
|
|
sp_exclude_tag_group,
|
|
]
|
|
|
|
elements_to_load = [sp_gallery_all, sp_tag_dropdown, *elements_to_clear[2:]]
|
|
sp_load_src.click(fn=load_src, inputs=[sp_src],
|
|
outputs=elements_to_clear)
|
|
|
|
def clear_elements():
|
|
clear_globals()
|
|
outputs = []
|
|
for _ in elements_to_clear:
|
|
# If the element has a choices attribute, set it to None
|
|
if isinstance(_, gr.Dropdown) or isinstance(_, gr.CheckboxGroup):
|
|
print(f"Clearing {_}")
|
|
outputs.append(gr.update(choices=[], value=None))
|
|
else:
|
|
print(f"Element type is {type(_)}")
|
|
outputs.append(gr.update(value=None))
|
|
return outputs
|
|
|
|
sp_clear_src.click(fn=clear_elements, outputs=elements_to_clear)
|
|
|
|
sp_add_all.click(
|
|
fn=select_all,
|
|
outputs=[sp_gallery_selected]
|
|
)
|
|
|
|
sp_add_selected.click(
|
|
fn=add_selected,
|
|
inputs=[],
|
|
outputs=[sp_gallery_selected]
|
|
)
|
|
|
|
sp_remove_selected.click(
|
|
fn=remove_selected,
|
|
inputs=[],
|
|
outputs=[sp_gallery_selected]
|
|
)
|
|
|
|
sp_clear_selected.click(
|
|
fn=clear_selected,
|
|
outputs=[sp_gallery_selected]
|
|
)
|
|
|
|
return (sp_interface, "SmartProcess", "smartsp_interface"),
|
|
|
|
|
|
script_callbacks.on_ui_tabs(create_process_ui)
|