diff --git a/modules/extras.py b/modules/extras.py index a9788e7dec668b534f0151438ee425323a4ffde0..158732048ae01eeb01c98cdeacf9b81272044fc7 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -4,6 +4,7 @@ import numpy as np from PIL import Image import torch +import tqdm from modules import processing, shared, images, devices from modules.shared import opts @@ -149,28 +150,45 @@ def run_modelmerger(modelname_0, modelname_1, interp_method, interp_amount): alpha = alpha * alpha * (3 - (2 * alpha)) return theta0 + ((theta1 - theta0) * alpha) - model_0 = torch.load('models/' + modelname_0 + '.ckpt') - model_1 = torch.load('models/' + modelname_1 + '.ckpt') + if os.path.exists(modelname_0): + model0_filename = modelname_0 + modelname_0 = os.path.splitext(os.path.basename(modelname_0))[0] + else: + model0_filename = 'models/' + modelname_0 + '.ckpt' + + if os.path.exists(modelname_1): + model1_filename = modelname_1 + modelname_1 = os.path.splitext(os.path.basename(modelname_1))[0] + else: + model1_filename = 'models/' + modelname_1 + '.ckpt' + + print(f"Loading {model0_filename}...") + model_0 = torch.load(model0_filename, map_location='cpu') + + print(f"Loading {model1_filename}...") + model_1 = torch.load(model1_filename, map_location='cpu') theta_0 = model_0['state_dict'] theta_1 = model_1['state_dict'] - theta_func = weighted_sum - - if interp_method == "Weighted Sum": - theta_func = weighted_sum - if interp_method == "Sigmoid": - theta_func = sigmoid - - for key in theta_0.keys(): + + theta_funcs = { + "Weighted Sum": weighted_sum, + "Sigmoid": sigmoid, + } + theta_func = theta_funcs[interp_method] + + print(f"Merging...") + for key in tqdm.tqdm(theta_0.keys()): if 'model' in key and key in theta_1: theta_0[key] = theta_func(theta_0[key], theta_1[key], interp_amount) for key in theta_1.keys(): if 'model' in key and key not in theta_0: theta_0[key] = theta_1[key] - - output_modelname = 'models/' + modelname_0 + '-' + modelname_1 + '-merged.ckpt'; - + + output_modelname = 'models/' + modelname_0 + '-' + modelname_1 + '-merged.ckpt' + print(f"Saving to {output_modelname}...") torch.save(model_0, output_modelname) - - return "

Model saved to " + output_modelname + "

" + + print(f"Checkpoint saved.") + return "Checkpoint saved to " + output_modelname diff --git a/modules/ui.py b/modules/ui.py index 5476c32fc297e3e0f004e99945016d17e4712153..e96109c9c59a1bcbd2ba68ddc2342bf9741d3a6b 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -49,6 +49,7 @@ sample_img2img = sample_img2img if os.path.exists(sample_img2img) else None css_hide_progressbar = """ .wrap .m-12 svg { display:none!important; } +.wrap .m-12::before { content:"Loading..." } .progress-bar { display:none!important; } .meta-text { display:none!important; } """ @@ -865,7 +866,7 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger): submit = gr.Button(elem_id="modelmerger_merge", label="Merge", variant='primary') with gr.Column(variant='panel'): - submit_result = gr.HTML(elem_id="modelmerger_result") + submit_result = gr.Textbox(elem_id="modelmerger_result", show_label=False) submit.click( fn=run_modelmerger,