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