sd_smartprocess/scripts/main.py

110 lines
5.2 KiB
Python

import gradio as gr
from modules import script_callbacks, shared
from modules.shared import cmd_opts
from modules.ui import setup_progressbar, gr_show
from webui import wrap_gradio_gpu_call
import smartprocess
def on_ui_tabs():
with gr.Blocks() as sp_interface:
with gr.Row(equal_height=True):
with gr.Column(variant="panel"):
with gr.Tab("Directories"):
sp_src = gr.Textbox(label='Source directory')
sp_dst = gr.Textbox(label='Destination directory')
with gr.Tab("Cropping"):
sp_size = gr.Slider(minimum=64, maximum=2048, step=64, label="Output Size", value=512)
sp_pad = gr.Checkbox(label="Pad Images")
sp_crop = gr.Checkbox(label='Crop Images')
sp_flip = gr.Checkbox(label='Create flipped copies')
sp_split = gr.Checkbox(label='Split over-sized images')
sp_split_threshold = gr.Slider(label='Split image threshold', value=0.5, minimum=0.0,
maximum=1.0,
step=0.05)
sp_overlap_ratio = gr.Slider(label='Split image overlap ratio', value=0.2, minimum=0.0,
maximum=0.9, step=0.05)
with gr.Tab("Captions"):
sp_caption = gr.Checkbox(label='Generate Captions')
sp_caption_length = gr.Number(label='Max Caption length (0=unlimited)', value=0, precision=0)
sp_txt_action = gr.Dropdown(label='Existing Caption Action', value="ignore",
choices=["ignore", "copy", "prepend", "append"])
sp_caption_append_file = gr.Checkbox(label="Append Caption to File Name", value=True)
sp_caption_save_txt = gr.Checkbox(label="Save Caption to .txt File", value=False)
sp_caption_deepbooru = gr.Checkbox(label='Append DeepDanbooru to Caption',
visible=True if cmd_opts.deepdanbooru else False)
sp_replace_class = gr.Checkbox(label='Replace Class with Subject in Caption', value=True)
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')
with gr.Tab("Post-Processing"):
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)
sp_upscale_ratio = gr.Slider(label="Upscale Ratio", value=2, step=1, minimum=2, maximum=4)
sp_scaler = gr.Radio(label='Upscaler', elem_id="sp_scaler",
choices=[x.name for x in shared.sd_upscalers],
value=shared.sd_upscalers[0].name, type="index")
# Preview/progress
with gr.Column(variant="panel"):
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_gallery = gr.Gallery(label='Output', show_label=False, elem_id='sp_gallery').style(grid=4)
sp_preview = gr.Image(elem_id='sp_preview', visible=False)
setup_progressbar(sp_progressbar, sp_preview, 'sp', textinfo=sp_progress)
with gr.Row():
sp_cancel = gr.Button(value="Cancel")
sp_run = gr.Button(value="Preprocess", variant='primary')
sp_cancel.click(
fn=lambda: shared.state.interrupt()
)
sp_run.click(
fn=wrap_gradio_gpu_call(smartprocess.preprocess, extra_outputs=[gr.update()]),
_js="start_smart_process",
inputs=[
sp_src,
sp_dst,
sp_pad,
sp_crop,
sp_size,
sp_caption_append_file,
sp_caption_save_txt,
sp_txt_action,
sp_flip,
sp_split,
sp_caption,
sp_caption_length,
sp_caption_deepbooru,
sp_split_threshold,
sp_overlap_ratio,
sp_class,
sp_subject,
sp_replace_class,
sp_restore_faces,
sp_face_model,
sp_upscale,
sp_upscale_ratio,
sp_scaler
],
outputs=[
sp_progress,
sp_outcome
],
)
return (sp_interface, "Smart Preprocess", "smartsp_interface"),
script_callbacks.on_ui_tabs(on_ui_tabs)