KohyaSS/kohya_gui.py

111 lines
3.4 KiB
Python
Raw Normal View History

2022-12-22 16:51:34 +00:00
import gradio as gr
import os
import argparse
from dreambooth_gui import dreambooth_tab
from finetune_gui import finetune_tab
2023-01-26 21:22:58 +00:00
from textual_inversion_gui import ti_tab
2022-12-22 16:51:34 +00:00
from library.utilities import utilities_tab
from library.extract_lora_gui import gradio_extract_lora_tab
from library.extract_lycoris_locon_gui import gradio_extract_lycoris_locon_tab
from library.merge_lora_gui import gradio_merge_lora_tab
from library.resize_lora_gui import gradio_resize_lora_tab
2022-12-29 04:02:35 +00:00
from lora_gui import lora_tab
2022-12-22 16:51:34 +00:00
def UI(**kwargs):
2022-12-22 16:51:34 +00:00
css = ''
if os.path.exists('./style.css'):
with open(os.path.join('./style.css'), 'r', encoding='utf8') as file:
print('Load CSS...')
css += file.read() + '\n'
2023-02-06 01:07:00 +00:00
interface = gr.Blocks(css=css, title='Kohya_ss GUI')
2022-12-22 16:51:34 +00:00
with interface:
with gr.Tab('Dreambooth'):
(
train_data_dir_input,
reg_data_dir_input,
output_dir_input,
logging_dir_input,
) = dreambooth_tab()
2023-01-01 05:33:29 +00:00
with gr.Tab('Dreambooth LoRA'):
lora_tab()
2023-01-26 21:22:58 +00:00
with gr.Tab('Dreambooth TI'):
ti_tab()
2022-12-22 16:51:34 +00:00
with gr.Tab('Finetune'):
finetune_tab()
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,
enable_copy_info_button=True,
)
gradio_extract_lora_tab()
gradio_extract_lycoris_locon_tab()
gradio_merge_lora_tab()
gradio_resize_lora_tab()
2022-12-22 16:51:34 +00:00
# Show the interface
launch_kwargs = {}
username = kwargs.get('username')
password = kwargs.get('password')
server_port = kwargs.get('server_port', 0)
inbrowser = kwargs.get('inbrowser', False)
2023-02-25 01:37:51 +00:00
share = kwargs.get('share', False)
2023-03-06 02:57:06 +00:00
server_name = kwargs.get('listen')
2023-03-06 02:57:06 +00:00
launch_kwargs['server_name'] = server_name
if username and password:
2023-03-02 00:24:11 +00:00
launch_kwargs['auth'] = (username, password)
if server_port > 0:
2023-03-02 00:24:11 +00:00
launch_kwargs['server_port'] = server_port
if inbrowser:
2023-03-02 00:24:11 +00:00
launch_kwargs['inbrowser'] = inbrowser
2023-02-25 01:37:51 +00:00
if share:
2023-03-02 00:24:11 +00:00
launch_kwargs['share'] = share
interface.launch(**launch_kwargs)
2022-12-22 16:51:34 +00:00
2023-03-02 00:24:11 +00:00
2022-12-22 16:51:34 +00:00
if __name__ == '__main__':
# torch.cuda.set_per_process_memory_fraction(0.48)
parser = argparse.ArgumentParser()
2023-03-06 02:57:06 +00:00
parser.add_argument(
'--listen',
type=str,
default='127.0.0.1',
help='IP to listen on for connections to Gradio',
2023-03-06 02:57:06 +00:00
)
2022-12-22 16:51:34 +00:00
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(
2023-03-02 00:24:11 +00:00
'--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'
)
2022-12-22 16:51:34 +00:00
args = parser.parse_args()
2023-03-02 00:24:11 +00:00
UI(
username=args.username,
password=args.password,
inbrowser=args.inbrowser,
server_port=args.server_port,
share=args.share,
2023-03-06 02:57:06 +00:00
listen=args.listen,
2023-03-02 00:24:11 +00:00
)