mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-02-24 00:33:57 +00:00
92 lines
3.7 KiB
Python
92 lines
3.7 KiB
Python
import argparse
|
|
import torch
|
|
from safetensors.torch import load_file, save_file
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--sd15", default=None, type=str, required=True, help="Path to the original sd15.")
|
|
parser.add_argument("--control", default=None, type=str, required=True, help="Path to the sd15 with control.")
|
|
parser.add_argument("--dst", default=None, type=str, required=True, help="Path to the output difference model.")
|
|
parser.add_argument("--fp16", action="store_true", help="Save as fp16.")
|
|
parser.add_argument("--bf16", action="store_true", help="Save as bf16.")
|
|
args = parser.parse_args()
|
|
|
|
assert args.sd15 is not None, "Must provide a original sd15 model path!"
|
|
assert args.control is not None, "Must provide a sd15 with control model path!"
|
|
assert args.dst is not None, "Must provide a output path!"
|
|
|
|
# make differences: copy from https://github.com/lllyasviel/ControlNet/blob/main/tool_transfer_control.py
|
|
|
|
def get_node_name(name, parent_name):
|
|
if len(name) <= len(parent_name):
|
|
return False, ''
|
|
p = name[:len(parent_name)]
|
|
if p != parent_name:
|
|
return False, ''
|
|
return True, name[len(parent_name):]
|
|
|
|
# remove first/cond stage from sd to reduce memory usage
|
|
def remove_first_and_cond(sd):
|
|
keys = list(sd.keys())
|
|
for key in keys:
|
|
is_first_stage, _ = get_node_name(key, 'first_stage_model')
|
|
is_cond_stage, _ = get_node_name(key, 'cond_stage_model')
|
|
if is_first_stage or is_cond_stage:
|
|
sd.pop(key, None)
|
|
return sd
|
|
|
|
print(f"loading: {args.sd15}")
|
|
if args.sd15.endswith(".safetensors"):
|
|
sd15_state_dict = load_file(args.sd15)
|
|
else:
|
|
sd15_state_dict = torch.load(args.sd15)
|
|
sd15_state_dict = sd15_state_dict.pop("state_dict", sd15_state_dict)
|
|
sd15_state_dict = remove_first_and_cond(sd15_state_dict)
|
|
|
|
print(f"loading: {args.control}")
|
|
if args.control.endswith(".safetensors"):
|
|
control_state_dict = load_file(args.control)
|
|
else:
|
|
control_state_dict = torch.load(args.control)
|
|
control_state_dict = remove_first_and_cond(control_state_dict)
|
|
|
|
# make diff of original and control
|
|
print(f"create difference")
|
|
keys = list(control_state_dict.keys())
|
|
final_state_dict = {"difference": torch.tensor(1.0)} # indicates difference
|
|
for key in keys:
|
|
p = control_state_dict.pop(key)
|
|
|
|
is_control, node_name = get_node_name(key, 'control_')
|
|
if not is_control:
|
|
continue
|
|
|
|
sd15_key_name = 'model.diffusion_' + node_name
|
|
if sd15_key_name in sd15_state_dict: # part of U-Net
|
|
# print("in sd15", key, sd15_key_name)
|
|
p_new = p - sd15_state_dict.pop(sd15_key_name)
|
|
if torch.max(torch.abs(p_new)) < 1e-6: # no difference?
|
|
print("no diff", key, sd15_key_name)
|
|
continue
|
|
else:
|
|
# print("not in sd15", key, sd15_key_name)
|
|
p_new = p # hint or zero_conv
|
|
|
|
final_state_dict[key] = p_new
|
|
|
|
save_dtype = None
|
|
if args.fp16:
|
|
save_dtype = torch.float16
|
|
elif args.bf16:
|
|
save_dtype = torch.bfloat16
|
|
if save_dtype is not None:
|
|
for key in final_state_dict.keys():
|
|
final_state_dict[key] = final_state_dict[key].to(save_dtype)
|
|
|
|
print("saving difference.")
|
|
if args.dst.endswith(".safetensors"):
|
|
save_file(final_state_dict, args.dst)
|
|
else:
|
|
torch.save({"state_dict": final_state_dict}, args.dst)
|
|
print("done!")
|