|
import gradio as gr
|
|
import os
|
|
import argparse
|
|
from kohya_gui.class_gui_config import KohyaSSGUIConfig
|
|
from kohya_gui.dreambooth_gui import dreambooth_tab
|
|
from kohya_gui.finetune_gui import finetune_tab
|
|
from kohya_gui.textual_inversion_gui import ti_tab
|
|
from kohya_gui.utilities import utilities_tab
|
|
from kohya_gui.lora_gui import lora_tab
|
|
from kohya_gui.class_lora_tab import LoRATools
|
|
|
|
from kohya_gui.custom_logging import setup_logging
|
|
from kohya_gui.localization_ext import add_javascript
|
|
|
|
|
|
def UI(**kwargs):
|
|
add_javascript(kwargs.get("language"))
|
|
css = ""
|
|
|
|
headless = kwargs.get("headless", False)
|
|
log.info(f"headless: {headless}")
|
|
|
|
if os.path.exists("./assets/style.css"):
|
|
with open(os.path.join("./assets/style.css"), "r", encoding="utf8") as file:
|
|
log.debug("Load CSS...")
|
|
css += file.read() + "\n"
|
|
|
|
if os.path.exists("./.release"):
|
|
with open(os.path.join("./.release"), "r", encoding="utf8") as file:
|
|
release = file.read()
|
|
|
|
if os.path.exists("./README.md"):
|
|
with open(os.path.join("./README.md"), "r", encoding="utf8") as file:
|
|
README = file.read()
|
|
|
|
interface = gr.Blocks(
|
|
css=css, title=f"Kohya_ss GUI {release}", theme=gr.themes.Default()
|
|
)
|
|
|
|
config = KohyaSSGUIConfig(config_file_path=kwargs.get("config"))
|
|
|
|
if config.is_config_loaded():
|
|
log.info(f"Loaded default GUI values from '{kwargs.get('config')}'...")
|
|
|
|
use_shell_flag = True
|
|
|
|
|
|
|
|
use_shell_flag = config.get("settings.use_shell", use_shell_flag)
|
|
|
|
if kwargs.get("do_not_use_shell", False):
|
|
use_shell_flag = False
|
|
|
|
if use_shell_flag:
|
|
log.info("Using shell=True when running external commands...")
|
|
|
|
with interface:
|
|
with gr.Tab("Dreambooth"):
|
|
(
|
|
train_data_dir_input,
|
|
reg_data_dir_input,
|
|
output_dir_input,
|
|
logging_dir_input,
|
|
) = dreambooth_tab(
|
|
headless=headless, config=config, use_shell_flag=use_shell_flag
|
|
)
|
|
with gr.Tab("LoRA"):
|
|
lora_tab(headless=headless, config=config, use_shell_flag=use_shell_flag)
|
|
with gr.Tab("Textual Inversion"):
|
|
ti_tab(headless=headless, config=config, use_shell_flag=use_shell_flag)
|
|
with gr.Tab("Finetuning"):
|
|
finetune_tab(
|
|
headless=headless, config=config, use_shell_flag=use_shell_flag
|
|
)
|
|
with gr.Tab("Utilities"):
|
|
utilities_tab(
|
|
train_data_dir_input=train_data_dir_input,
|
|
reg_data_dir_input=reg_data_dir_input,
|
|
output_dir_input=output_dir_input,
|
|
logging_dir_input=logging_dir_input,
|
|
headless=headless,
|
|
config=config,
|
|
)
|
|
with gr.Tab("LoRA"):
|
|
_ = LoRATools(headless=headless)
|
|
with gr.Tab("About"):
|
|
gr.Markdown(f"kohya_ss GUI release {release}")
|
|
with gr.Tab("README"):
|
|
gr.Markdown(README)
|
|
|
|
htmlStr = f"""
|
|
<html>
|
|
<body>
|
|
<div class="ver-class">{release}</div>
|
|
</body>
|
|
</html>
|
|
"""
|
|
gr.HTML(htmlStr)
|
|
|
|
launch_kwargs = {}
|
|
username = kwargs.get("username")
|
|
password = kwargs.get("password")
|
|
server_port = kwargs.get("server_port", 0)
|
|
inbrowser = kwargs.get("inbrowser", False)
|
|
share = kwargs.get("share", False)
|
|
do_not_share = kwargs.get("do_not_share", False)
|
|
server_name = kwargs.get("listen")
|
|
root_path = kwargs.get("root_path", None)
|
|
|
|
launch_kwargs["server_name"] = server_name
|
|
if username and password:
|
|
launch_kwargs["auth"] = (username, password)
|
|
if server_port > 0:
|
|
launch_kwargs["server_port"] = server_port
|
|
if inbrowser:
|
|
launch_kwargs["inbrowser"] = inbrowser
|
|
if do_not_share:
|
|
launch_kwargs["share"] = False
|
|
else:
|
|
if share:
|
|
launch_kwargs["share"] = share
|
|
if root_path:
|
|
launch_kwargs["root_path"] = root_path
|
|
launch_kwargs["debug"] = True
|
|
interface.launch(**launch_kwargs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--config",
|
|
type=str,
|
|
default="./config.toml",
|
|
help="Path to the toml config file for interface defaults",
|
|
)
|
|
parser.add_argument("--debug", action="store_true", help="Debug on")
|
|
parser.add_argument(
|
|
"--listen",
|
|
type=str,
|
|
default="127.0.0.1",
|
|
help="IP to listen on for connections to Gradio",
|
|
)
|
|
parser.add_argument(
|
|
"--username", type=str, default="", help="Username for authentication"
|
|
)
|
|
parser.add_argument(
|
|
"--password", type=str, default="", help="Password for authentication"
|
|
)
|
|
parser.add_argument(
|
|
"--server_port",
|
|
type=int,
|
|
default=0,
|
|
help="Port to run the server listener on",
|
|
)
|
|
parser.add_argument("--inbrowser", action="store_true", help="Open in browser")
|
|
parser.add_argument("--share", action="store_true", help="Share the gradio UI")
|
|
parser.add_argument(
|
|
"--headless", action="store_true", help="Is the server headless"
|
|
)
|
|
parser.add_argument(
|
|
"--language", type=str, default=None, help="Set custom language"
|
|
)
|
|
|
|
parser.add_argument("--use-ipex", action="store_true", help="Use IPEX environment")
|
|
parser.add_argument("--use-rocm", action="store_true", help="Use ROCm environment")
|
|
|
|
parser.add_argument(
|
|
"--do_not_use_shell", action="store_true", help="Enforce not to use shell=True when running external commands"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--do_not_share", action="store_true", help="Do not share the gradio UI"
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--root_path", type=str, default=None, help="`root_path` for Gradio to enable reverse proxy support. e.g. /kohya_ss"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
log = setup_logging(debug=args.debug)
|
|
|
|
UI(**vars(args))
|
|
|
