mirror of https://github.com/bmaltais/kohya_ss
422 lines
15 KiB
Python
422 lines
15 KiB
Python
# Standard library imports
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import json
|
|
|
|
# Third-party imports
|
|
import gradio as gr
|
|
from easygui import msgbox
|
|
|
|
# Local module imports
|
|
from .common_gui import get_saveasfilename_path, get_file_path, scriptdir, list_files, create_refresh_button
|
|
from .custom_logging import setup_logging
|
|
|
|
# Set up logging
|
|
log = setup_logging()
|
|
|
|
folder_symbol = '\U0001f4c2' # 📂
|
|
refresh_symbol = '\U0001f504' # 🔄
|
|
save_style_symbol = '\U0001f4be' # 💾
|
|
document_symbol = '\U0001F4C4' # 📄
|
|
|
|
PYTHON = sys.executable
|
|
|
|
|
|
def check_model(model):
|
|
if not model:
|
|
return True
|
|
if not os.path.isfile(model):
|
|
msgbox(f'The provided {model} is not a file')
|
|
return False
|
|
return True
|
|
|
|
|
|
def verify_conditions(sd_model, lora_models):
|
|
lora_models_count = sum(1 for model in lora_models if model)
|
|
if sd_model and lora_models_count >= 1:
|
|
return True
|
|
elif not sd_model and lora_models_count >= 2:
|
|
return True
|
|
return False
|
|
|
|
|
|
class GradioMergeLoRaTab:
|
|
def __init__(self, headless=False):
|
|
self.headless = headless
|
|
self.build_tab()
|
|
|
|
def save_inputs_to_json(self, file_path, inputs):
|
|
with open(file_path, 'w') as file:
|
|
json.dump(inputs, file)
|
|
log.info(f'Saved inputs to {file_path}')
|
|
|
|
def load_inputs_from_json(self, file_path):
|
|
with open(file_path, 'r') as file:
|
|
inputs = json.load(file)
|
|
log.info(f'Loaded inputs from {file_path}')
|
|
return inputs
|
|
|
|
def build_tab(self):
|
|
current_sd_model_dir = os.path.join(scriptdir, "outputs")
|
|
current_save_dir = os.path.join(scriptdir, "outputs")
|
|
current_a_model_dir = current_sd_model_dir
|
|
current_b_model_dir = current_sd_model_dir
|
|
current_c_model_dir = current_sd_model_dir
|
|
current_d_model_dir = current_sd_model_dir
|
|
|
|
def list_sd_models(path):
|
|
nonlocal current_sd_model_dir
|
|
current_sd_model_dir = path
|
|
return list(list_files(path, exts=[".ckpt", ".safetensors"], all=True))
|
|
|
|
def list_a_models(path):
|
|
nonlocal current_a_model_dir
|
|
current_a_model_dir = path
|
|
return list(list_files(path, exts=[".pt", ".safetensors"], all=True))
|
|
|
|
def list_b_models(path):
|
|
nonlocal current_b_model_dir
|
|
current_b_model_dir = path
|
|
return list(list_files(path, exts=[".pt", ".safetensors"], all=True))
|
|
|
|
def list_c_models(path):
|
|
nonlocal current_c_model_dir
|
|
current_c_model_dir = path
|
|
return list(list_files(path, exts=[".pt", ".safetensors"], all=True))
|
|
|
|
def list_d_models(path):
|
|
nonlocal current_d_model_dir
|
|
current_d_model_dir = path
|
|
return list(list_files(path, exts=[".pt", ".safetensors"], all=True))
|
|
|
|
def list_save_to(path):
|
|
nonlocal current_save_dir
|
|
current_save_dir = path
|
|
return list(list_files(path, exts=[".ckpt", ".safetensors"], all=True))
|
|
|
|
with gr.Tab('Merge LoRA'):
|
|
gr.Markdown(
|
|
'This utility can merge up to 4 LoRA together or alternatively merge up to 4 LoRA into a SD checkpoint.'
|
|
)
|
|
|
|
lora_ext = gr.Textbox(value='*.safetensors *.pt', visible=False)
|
|
lora_ext_name = gr.Textbox(value='LoRA model types', visible=False)
|
|
ckpt_ext = gr.Textbox(value='*.safetensors *.ckpt', visible=False)
|
|
ckpt_ext_name = gr.Textbox(value='SD model types', visible=False)
|
|
|
|
with gr.Group(), gr.Row():
|
|
sd_model = gr.Dropdown(
|
|
label='SD Model (Optional. Stable Diffusion model path, if you want to merge it with LoRA files)',
|
|
interactive=True,
|
|
choices=[""] + list_sd_models(current_sd_model_dir),
|
|
value="",
|
|
allow_custom_value=True,
|
|
)
|
|
create_refresh_button(sd_model, lambda: None, lambda: {"choices": list_sd_models(current_sd_model_dir)}, "open_folder_small")
|
|
sd_model_file = gr.Button(
|
|
folder_symbol,
|
|
elem_id='open_folder_small',
|
|
elem_classes=['tool'],
|
|
visible=(not self.headless),
|
|
)
|
|
sd_model_file.click(
|
|
get_file_path,
|
|
inputs=[sd_model, ckpt_ext, ckpt_ext_name],
|
|
outputs=sd_model,
|
|
show_progress=False,
|
|
)
|
|
sdxl_model = gr.Checkbox(label='SDXL model', value=False)
|
|
|
|
sd_model.change(
|
|
fn=lambda path: gr.Dropdown(choices=[""] + list_sd_models(path)),
|
|
inputs=sd_model,
|
|
outputs=sd_model,
|
|
show_progress=False,
|
|
)
|
|
|
|
with gr.Group(), gr.Row():
|
|
lora_a_model = gr.Dropdown(
|
|
label='LoRA model "A" (path to the LoRA A model)',
|
|
interactive=True,
|
|
choices=[""] + list_a_models(current_a_model_dir),
|
|
value="",
|
|
allow_custom_value=True,
|
|
)
|
|
create_refresh_button(lora_a_model, lambda: None, lambda: {"choices": list_a_models(current_a_model_dir)}, "open_folder_small")
|
|
button_lora_a_model_file = gr.Button(
|
|
folder_symbol,
|
|
elem_id='open_folder_small',
|
|
elem_classes=['tool'],
|
|
visible=(not self.headless),
|
|
)
|
|
button_lora_a_model_file.click(
|
|
get_file_path,
|
|
inputs=[lora_a_model, lora_ext, lora_ext_name],
|
|
outputs=lora_a_model,
|
|
show_progress=False,
|
|
)
|
|
|
|
lora_b_model = gr.Dropdown(
|
|
label='LoRA model "B" (path to the LoRA B model)',
|
|
interactive=True,
|
|
choices=[""] + list_b_models(current_b_model_dir),
|
|
value="",
|
|
allow_custom_value=True,
|
|
)
|
|
create_refresh_button(lora_b_model, lambda: None, lambda: {"choices": list_b_models(current_b_model_dir)}, "open_folder_small")
|
|
button_lora_b_model_file = gr.Button(
|
|
folder_symbol,
|
|
elem_id='open_folder_small',
|
|
elem_classes=['tool'],
|
|
visible=(not self.headless),
|
|
)
|
|
button_lora_b_model_file.click(
|
|
get_file_path,
|
|
inputs=[lora_b_model, lora_ext, lora_ext_name],
|
|
outputs=lora_b_model,
|
|
show_progress=False,
|
|
)
|
|
|
|
lora_a_model.change(
|
|
fn=lambda path: gr.Dropdown(choices=[""] + list_a_models(path)),
|
|
inputs=lora_a_model,
|
|
outputs=lora_a_model,
|
|
show_progress=False,
|
|
)
|
|
lora_b_model.change(
|
|
fn=lambda path: gr.Dropdown(choices=[""] + list_b_models(path)),
|
|
inputs=lora_b_model,
|
|
outputs=lora_b_model,
|
|
show_progress=False,
|
|
)
|
|
|
|
with gr.Row():
|
|
ratio_a = gr.Slider(
|
|
label='Model A merge ratio (eg: 0.5 mean 50%)',
|
|
minimum=0,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0.0,
|
|
interactive=True,
|
|
)
|
|
|
|
ratio_b = gr.Slider(
|
|
label='Model B merge ratio (eg: 0.5 mean 50%)',
|
|
minimum=0,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0.0,
|
|
interactive=True,
|
|
)
|
|
|
|
with gr.Group(), gr.Row():
|
|
lora_c_model = gr.Dropdown(
|
|
label='LoRA model "C" (path to the LoRA C model)',
|
|
interactive=True,
|
|
choices=[""] + list_c_models(current_c_model_dir),
|
|
value="",
|
|
allow_custom_value=True,
|
|
)
|
|
create_refresh_button(lora_c_model, lambda: None, lambda: {"choices": list_c_models(current_c_model_dir)}, "open_folder_small")
|
|
button_lora_c_model_file = gr.Button(
|
|
folder_symbol,
|
|
elem_id='open_folder_small',
|
|
elem_classes=['tool'],
|
|
visible=(not self.headless),
|
|
)
|
|
button_lora_c_model_file.click(
|
|
get_file_path,
|
|
inputs=[lora_c_model, lora_ext, lora_ext_name],
|
|
outputs=lora_c_model,
|
|
show_progress=False,
|
|
)
|
|
|
|
lora_d_model = gr.Dropdown(
|
|
label='LoRA model "D" (path to the LoRA D model)',
|
|
interactive=True,
|
|
choices=[""] + list_d_models(current_d_model_dir),
|
|
value="",
|
|
allow_custom_value=True,
|
|
)
|
|
create_refresh_button(lora_d_model, lambda: None, lambda: {"choices": list_d_models(current_d_model_dir)}, "open_folder_small")
|
|
button_lora_d_model_file = gr.Button(
|
|
folder_symbol,
|
|
elem_id='open_folder_small',
|
|
elem_classes=['tool'],
|
|
visible=(not self.headless),
|
|
)
|
|
button_lora_d_model_file.click(
|
|
get_file_path,
|
|
inputs=[lora_d_model, lora_ext, lora_ext_name],
|
|
outputs=lora_d_model,
|
|
show_progress=False,
|
|
)
|
|
lora_c_model.change(
|
|
fn=lambda path: gr.Dropdown(choices=[""] + list_c_models(path)),
|
|
inputs=lora_c_model,
|
|
outputs=lora_c_model,
|
|
show_progress=False,
|
|
)
|
|
lora_d_model.change(
|
|
fn=lambda path: gr.Dropdown(choices=[""] + list_d_models(path)),
|
|
inputs=lora_d_model,
|
|
outputs=lora_d_model,
|
|
show_progress=False,
|
|
)
|
|
|
|
with gr.Row():
|
|
ratio_c = gr.Slider(
|
|
label='Model C merge ratio (eg: 0.5 mean 50%)',
|
|
minimum=0,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0.0,
|
|
interactive=True,
|
|
)
|
|
|
|
ratio_d = gr.Slider(
|
|
label='Model D merge ratio (eg: 0.5 mean 50%)',
|
|
minimum=0,
|
|
maximum=1,
|
|
step=0.01,
|
|
value=0.0,
|
|
interactive=True,
|
|
)
|
|
|
|
with gr.Group(), gr.Row():
|
|
save_to = gr.Dropdown(
|
|
label='Save to (path for the file to save...)',
|
|
interactive=True,
|
|
choices=[""] + list_save_to(current_d_model_dir),
|
|
value="",
|
|
allow_custom_value=True,
|
|
)
|
|
create_refresh_button(save_to, lambda: None, lambda: {"choices": list_save_to(current_save_dir)}, "open_folder_small")
|
|
button_save_to = gr.Button(
|
|
folder_symbol,
|
|
elem_id='open_folder_small',
|
|
elem_classes=['tool'],
|
|
visible=(not self.headless),
|
|
)
|
|
button_save_to.click(
|
|
get_saveasfilename_path,
|
|
inputs=[save_to, lora_ext, lora_ext_name],
|
|
outputs=save_to,
|
|
show_progress=False,
|
|
)
|
|
precision = gr.Radio(
|
|
label='Merge precision',
|
|
choices=['fp16', 'bf16', 'float'],
|
|
value='float',
|
|
interactive=True,
|
|
)
|
|
save_precision = gr.Radio(
|
|
label='Save precision',
|
|
choices=['fp16', 'bf16', 'float'],
|
|
value='fp16',
|
|
interactive=True,
|
|
)
|
|
|
|
save_to.change(
|
|
fn=lambda path: gr.Dropdown(choices=[""] + list_save_to(path)),
|
|
inputs=save_to,
|
|
outputs=save_to,
|
|
show_progress=False,
|
|
)
|
|
|
|
merge_button = gr.Button('Merge model')
|
|
|
|
merge_button.click(
|
|
self.merge_lora,
|
|
inputs=[
|
|
sd_model,
|
|
sdxl_model,
|
|
lora_a_model,
|
|
lora_b_model,
|
|
lora_c_model,
|
|
lora_d_model,
|
|
ratio_a,
|
|
ratio_b,
|
|
ratio_c,
|
|
ratio_d,
|
|
save_to,
|
|
precision,
|
|
save_precision,
|
|
],
|
|
show_progress=False,
|
|
)
|
|
|
|
def merge_lora(
|
|
self,
|
|
sd_model,
|
|
sdxl_model,
|
|
lora_a_model,
|
|
lora_b_model,
|
|
lora_c_model,
|
|
lora_d_model,
|
|
ratio_a,
|
|
ratio_b,
|
|
ratio_c,
|
|
ratio_d,
|
|
save_to,
|
|
precision,
|
|
save_precision,
|
|
):
|
|
|
|
log.info('Merge model...')
|
|
models = [
|
|
sd_model,
|
|
lora_a_model,
|
|
lora_b_model,
|
|
lora_c_model,
|
|
lora_d_model,
|
|
]
|
|
lora_models = models[1:]
|
|
ratios = [ratio_a, ratio_b, ratio_c, ratio_d]
|
|
|
|
if not verify_conditions(sd_model, lora_models):
|
|
log.info(
|
|
'Warning: Either provide at least one LoRa model along with the sd_model or at least two LoRa models if no sd_model is provided.'
|
|
)
|
|
return
|
|
|
|
for model in models:
|
|
if not check_model(model):
|
|
return
|
|
|
|
if not sdxl_model:
|
|
run_cmd = fr'"{PYTHON}" "{scriptdir}/sd-scripts/networks/merge_lora.py"'
|
|
else:
|
|
run_cmd = (
|
|
fr'"{PYTHON}" "{scriptdir}/sd-scripts/networks/sdxl_merge_lora.py"'
|
|
)
|
|
if sd_model:
|
|
run_cmd += fr' --sd_model "{sd_model}"'
|
|
run_cmd += f' --save_precision {save_precision}'
|
|
run_cmd += f' --precision {precision}'
|
|
run_cmd += fr' --save_to "{save_to}"'
|
|
|
|
# Create a space-separated string of non-empty models (from the second element onwards), enclosed in double quotes
|
|
models_cmd = ' '.join([fr'"{model}"' for model in lora_models if model])
|
|
|
|
# Create a space-separated string of non-zero ratios corresponding to non-empty LoRa models
|
|
valid_ratios = [
|
|
ratios[i] for i, model in enumerate(lora_models) if model
|
|
]
|
|
ratios_cmd = ' '.join([str(ratio) for ratio in valid_ratios])
|
|
|
|
if models_cmd:
|
|
run_cmd += f' --models {models_cmd}'
|
|
run_cmd += f' --ratios {ratios_cmd}'
|
|
|
|
log.info(run_cmd)
|
|
|
|
env = os.environ.copy()
|
|
env['PYTHONPATH'] = fr"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
|
|
|
|
# Run the command
|
|
subprocess.run(run_cmd, shell=True, env=env)
|
|
|
|
log.info('Done merging...')
|