1580158
/
sd-webui-controlnet
mirror of https://github.com/Mikubill/sd-webui-controlnet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Extraction add half option

pull/742/head
RcINS 2023-04-13 17:50:33 +08:00
parent b1c0bc62e0
commit 2cca805b2a
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
extract_controlnet.py
View File

@ -6,6 +6,7 @@ if __name__ == "__main__":
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--src", default=None, type=str, required=True, help="Path to the model to convert.") parser.add_argument("--src", default=None, type=str, required=True, help="Path to the model to convert.")
parser.add_argument("--dst", default=None, type=str, required=True, help="Path to the output model.") parser.add_argument("--dst", default=None, type=str, required=True, help="Path to the output model.")
parser.add_argument("--half", action="store_true", help="Cast to FP16.")
args = parser.parse_args() args = parser.parse_args()
assert args.src is not None, "Must provide a model path!" assert args.src is not None, "Must provide a model path!"
@ -15,9 +16,10 @@ if __name__ == "__main__":
state_dict = load_file(args.src) state_dict = load_file(args.src)
else: else:
state_dict = torch.load(args.src) state_dict = torch.load(args.src)
if any([k.startswith("control_model.") for k, v in state_dict.items()]): if any([k.startswith("control_model.") for k, v in state_dict.items()]):
state_dict = {k.replace("control_model.", ""): v for k, v in state_dict.items() if k.startswith("control_model.")} dtype = torch.float16 if args.half else torch.float32
state_dict = {k.replace("control_model.", ""): v.to(dtype) for k, v in state_dict.items() if k.startswith("control_model.")}
if args.dst.endswith(".safetensors"): if args.dst.endswith(".safetensors"):
save_file(state_dict, args.dst) save_file(state_dict, args.dst)