michael/stable-diffusion-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Added smoothstep interpolation to checkpoint merging

Browse source
This commit is contained in:
William Moorehouse 2022-09-26 10:50:21 -04:00
parent 91643f651d
commit dc11748dea
2 changed files with 21 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
modules/extras.py
View file

@ -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:

6
modules/ui.py
View file

@ -860,7 +860,8 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
modelname_0 = gr.Textbox(elem_id="modelmerger_modelname_0", label="Model Name (to)")
modelname_1 = gr.Textbox(elem_id="modelmerger_modelname_1", label="Model Name (from)")
alpha = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Alpha', value=0.3)
interp_method = gr.Radio(choices=["Weighted Sum", "Sigmoid"], value="Weighted Sum", label="Interpolation Method")
interp_amount = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Interpolation Amount', value=0.3)
submit = gr.Button(elem_id="modelmerger_merge", label="Merge", variant='primary')
with gr.Column(variant='panel'):
@ -871,7 +872,8 @@ def create_ui(txt2img, img2img, run_extras, run_pnginfo, run_modelmerger):
inputs=[
modelname_0,
modelname_1,
alpha
interp_method,
interp_amount
],
outputs=[
submit_result,