diff options
| author | Liam <liamthekerr@gmail.com> | 2022-09-27 20:37:24 +0000 |
|---|---|---|
| committer | Liam <liamthekerr@gmail.com> | 2022-09-27 20:37:24 +0000 |
| commit | 981fe9c4a3994bb42ea5ff5212e4fe53b748bdd9 (patch) | |
| tree | 80fd53962ccafeb773b2d43b178e1ee39ac03ca3 /modules/extras.py | |
| parent | 5034f7d7597685aaa4779296983be0f49f4f991f (diff) | |
| parent | f2a4a2c3a672e22f088a7455d6039557370dd3f2 (diff) | |
| download | stable-diffusion-webui-gfx803-981fe9c4a3994bb42ea5ff5212e4fe53b748bdd9.tar.gz stable-diffusion-webui-gfx803-981fe9c4a3994bb42ea5ff5212e4fe53b748bdd9.tar.bz2 stable-diffusion-webui-gfx803-981fe9c4a3994bb42ea5ff5212e4fe53b748bdd9.zip | |
Merge remote-tracking branch 'upstream/master' into token_count
Diffstat (limited to 'modules/extras.py')
| -rw-r--r-- | modules/extras.py | 57 |
1 files changed, 57 insertions, 0 deletions
diff --git a/modules/extras.py b/modules/extras.py index 382ffa7d..c4ee2b62 100644 --- a/modules/extras.py +++ b/modules/extras.py @@ -3,6 +3,9 @@ import os import numpy as np
from PIL import Image
+import torch
+import tqdm
+
from modules import processing, shared, images, devices
from modules.shared import opts
import modules.gfpgan_model
@@ -135,3 +138,57 @@ def run_pnginfo(image): info = f"<div><p>{message}<p></div>"
return '', geninfo, info
+
+
+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)
+
+ 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_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 + '-' + interp_method.replace(" ", "_") + '-' + str(interp_amount) + '-merged.ckpt'
+ print(f"Saving to {output_modelname}...")
+ torch.save(model_0, output_modelname)
+
+ print(f"Checkpoint saved.")
+ return "Checkpoint saved to " + output_modelname
|