diff options
| author | William Moorehouse <moorehousew@gmail.com> | 2022-09-25 23:22:12 +0000 |
|---|---|---|
| committer | William Moorehouse <moorehousew@gmail.com> | 2022-09-25 23:22:12 +0000 |
| commit | 91643f651d2794349876b12abbf2449cdc4f30b6 (patch) | |
| tree | bb91ef06b902b68486fca856ccd5deb2eb020ef6 /modules/extras.py | |
| parent | ca3e5519e8b6dc020c5e7ae508738afb5dc6f3ec (diff) | |
| download | stable-diffusion-webui-gfx803-91643f651d2794349876b12abbf2449cdc4f30b6.tar.gz stable-diffusion-webui-gfx803-91643f651d2794349876b12abbf2449cdc4f30b6.tar.bz2 stable-diffusion-webui-gfx803-91643f651d2794349876b12abbf2449cdc4f30b6.zip | |
Add support for checkpoint merging
Diffstat (limited to 'modules/extras.py')
| -rw-r--r-- | modules/extras.py | 24 |
1 files changed, 24 insertions, 0 deletions
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"<div><p>{message}<p></div>"
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 "<p>Model saved to " + output_modelname + "</p>"
|