diff options
| author | William Moorehouse <moorehousew@gmail.com> | 2022-09-26 10:50:21 -0400 |
|---|---|---|
| committer | William Moorehouse <moorehousew@gmail.com> | 2022-09-26 10:50:21 -0400 |
| commit | dc11748dea36e7618a7cdad55526fa9d6faaa6cf (patch) | |
| tree | 27126c07306f06062522ca386d60693a35b966d1 /modules/extras.py | |
| parent | 91643f651d2794349876b12abbf2449cdc4f30b6 (diff) | |
Added smoothstep interpolation to checkpoint merging
Diffstat (limited to 'modules/extras.py')
| -rw-r--r-- | modules/extras.py | 19 |
1 files changed, 17 insertions, 2 deletions
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:
|