Mirrors
/
waifu-diffusion
mirror of https://github.com/harubaru/waifu-diffusion.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix pruning script yet again

This commit is contained in:
harubaru 2022-09-27 19:50:55 -07:00 committed by GitHub
parent aa91eaf228
commit a4628125fe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
scripts/prune.py
View File

@ -1,8 +1,14 @@
import os import os
from pathlib import Path
import torch import torch
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--input', '-I', type=str, help='Input file to prune', required = True)
args = parser.parse_args()
file = args.input
def prune_it(p, keep_only_ema=False): def prune_it(p, keep_only_ema=True):
print(f"prunin' in path: {p}") print(f"prunin' in path: {p}")
size_initial = os.path.getsize(p) size_initial = os.path.getsize(p)
nsd = dict() nsd = dict()
@ -18,14 +24,15 @@ def prune_it(p, keep_only_ema=False):
if keep_only_ema: if keep_only_ema:
sd = nsd["state_dict"].copy() sd = nsd["state_dict"].copy()
# infer ema keys # infer ema keys
ema_keys = {k: "model_ema." + k[6:].replace(".", ".") for k in sd.keys() if k.startswith("model.")} ema_keys = {k: "model_ema." + k[6:].replace(".", "") for k in sd.keys() if k.startswith('model.')}
new_sd = dict() new_sd = dict()
for k in sd: for k in sd:
if k in ema_keys: if k in ema_keys:
new_sd[k] = sd[ema_keys[k]].half() print(k, ema_keys[k])
new_sd[k] = sd[ema_keys[k]]
elif not k.startswith("model_ema.") or k in ["model_ema.num_updates", "model_ema.decay"]: elif not k.startswith("model_ema.") or k in ["model_ema.num_updates", "model_ema.decay"]:
new_sd[k] = sd[k].half() new_sd[k] = sd[k]
assert len(new_sd) == len(sd) - len(ema_keys) assert len(new_sd) == len(sd) - len(ema_keys)
nsd["state_dict"] = new_sd nsd["state_dict"] = new_sd
@ -33,7 +40,7 @@ def prune_it(p, keep_only_ema=False):
sd = nsd['state_dict'].copy() sd = nsd['state_dict'].copy()
new_sd = dict() new_sd = dict()
for k in sd: for k in sd:
new_sd[k] = sd[k].half() new_sd[k] = sd[k]
nsd['state_dict'] = new_sd nsd['state_dict'] = new_sd
fn = f"{os.path.splitext(p)[0]}-pruned.ckpt" if not keep_only_ema else f"{os.path.splitext(p)[0]}-ema-pruned.ckpt" fn = f"{os.path.splitext(p)[0]}-pruned.ckpt" if not keep_only_ema else f"{os.path.splitext(p)[0]}-ema-pruned.ckpt"
@ -48,4 +55,4 @@ def prune_it(p, keep_only_ema=False):
if __name__ == "__main__": if __name__ == "__main__":
prune_it('models/ldm/stable-diffusion-v1/model.ckpt') prune_it(file)