r/stable-diffusion-finetune
1
0
Fork 0
Code Issues Pull requests Projects Releases Wiki Activity

prune more

Browse source
This commit is contained in:
Robin Rombach 2022-08-01 00:13:50 +02:00
parent d878c53f13
commit 4358de0a9a
2 changed files with 22 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

0
configs/stable-diffusion/txt2img-1p4B-multinode-clip-encoder-high-res-512-inference.yaml → configs/stable-diffusion/v1-inference.yaml
View file

26
scripts/prune-ckpt.py
View file

@ -3,7 +3,7 @@ import torch
import fire import fire
def prune_it(p): def prune_it(p, keep_only_ema=False):
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()
@ -16,12 +16,30 @@ def prune_it(p):
print(f"removing optimizer states for path {p}") print(f"removing optimizer states for path {p}")
if "global_step" in sd: if "global_step" in sd:
print(f"This is global step {sd['global_step']}.") print(f"This is global step {sd['global_step']}.")
fn = f"{os.path.splitext(p)[0]}-pruned.ckpt" if keep_only_ema:
sd = nsd["state_dict"].copy()
# infer ema keys
ema_keys = {k: "model_ema." + k[6:].replace(".", "") for k in sd.keys() if k.startswith("model.")}
new_sd = dict()
for k in sd:
if k in ema_keys:
new_sd[k] = sd[ema_keys[k]]
elif not k.startswith("model_ema.") or k in ["model_ema.num_updates", "model_ema.decay"]:
new_sd[k] = sd[k]
assert len(new_sd) == len(sd) - len(ema_keys)
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"
print(f"saving pruned checkpoint at: {fn}") print(f"saving pruned checkpoint at: {fn}")
torch.save(nsd, fn) torch.save(nsd, fn)
newsize = os.path.getsize(fn) newsize = os.path.getsize(fn)
print(f"New ckpt size: {newsize*1e-9:.2f} GB. " MSG = f"New ckpt size: {newsize*1e-9:.2f} GB. " + \
f"Saved {(size_initial - newsize)*1e-9:.2f} GB by removing optimizer states") f"Saved {(size_initial - newsize)*1e-9:.2f} GB by removing optimizer states"
if keep_only_ema:
MSG += " and non-EMA weights"
print(MSG)
if __name__ == "__main__": if __name__ == "__main__":