From 91643f651d2794349876b12abbf2449cdc4f30b6 Mon Sep 17 00:00:00 2001 From: William Moorehouse Date: Sun, 25 Sep 2022 19:22:12 -0400 Subject: Add support for checkpoint merging --- modules/extras.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index 382ffa7d..2c5b1fd6 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -3,6 +3,8 @@ import os import numpy as np from PIL import Image +import torch + from modules import processing, shared, images, devices from modules.shared import opts import modules.gfpgan_model @@ -135,3 +137,25 @@ def run_pnginfo(image): info = f"

{message}

" return '', geninfo, info + + +def run_modelmerger(modelname_0, modelname_1, alpha): + model_0 = torch.load('models/' + modelname_0 + '.ckpt') + model_1 = torch.load('models/' + modelname_1 + '.ckpt') + + theta_0 = model_0['state_dict'] + theta_1 = model_1['state_dict'] + + for key in theta_0.keys(): + if 'model' in key and key in theta_1: + theta_0[key] = (1 - alpha) * theta_0[key] + alpha * theta_1[key] + + 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'; + + torch.save(model_0, output_modelname) + + return "

Model saved to " + output_modelname + "

" -- cgit v1.2.3 From dc11748dea36e7618a7cdad55526fa9d6faaa6cf Mon Sep 17 00:00:00 2001 From: William Moorehouse Date: Mon, 26 Sep 2022 10:50:21 -0400 Subject: Added smoothstep interpolation to checkpoint merging --- modules/extras.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index 2c5b1fd6..a9788e7d 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -139,16 +139,31 @@ def run_pnginfo(image): return '', geninfo, info -def run_modelmerger(modelname_0, modelname_1, alpha): +def run_modelmerger(modelname_0, modelname_1, interp_method, interp_amount): + # Linear interpolation (https://en.wikipedia.org/wiki/Linear_interpolation) + def weighted_sum(theta0, theta1, alpha): + return ((1 - alpha) * theta0) + (alpha * theta1) + + # Smoothstep (https://en.wikipedia.org/wiki/Smoothstep) + def sigmoid(theta0, theta1, alpha): + 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') 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(): if 'model' in key and key in theta_1: - theta_0[key] = (1 - alpha) * theta_0[key] + alpha * theta_1[key] + 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: -- cgit v1.2.3 From ada901ed661a717c44281d640b8fc0a275d4cb48 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Tue, 27 Sep 2022 10:44:00 +0300 Subject: added console outputs, more clear indication of progress, and ability to specify full filename to checkpoint merger restore "Loading..." text --- modules/extras.py | 48 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 15 deletions(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index a9788e7d..15873204 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 -- cgit v1.2.3 From 0bd8aa937f9172b44de30d158258c51b2174f136 Mon Sep 17 00:00:00 2001 From: Bernard Maltais Date: Tue, 27 Sep 2022 10:29:52 -0400 Subject: Add interpolation method and weight to output filename --- modules/extras.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'modules/extras.py') diff --git a/modules/extras.py b/modules/extras.py index 15873204..c4ee2b62 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -186,7 +186,7 @@ def run_modelmerger(modelname_0, modelname_1, interp_method, interp_amount): 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 + '-' + interp_method.replace(" ", "_") + '-' + str(interp_amount) + '-merged.ckpt' print(f"Saving to {output_modelname}...") torch.save(model_0, output_modelname) -- cgit v1.2.3