From a4628125feacca98af1ecd8d9bdd50afcab4e122 Mon Sep 17 00:00:00 2001 From: harubaru Date: Tue, 27 Sep 2022 19:50:55 -0700 Subject: [PATCH] fix pruning script yet again --- scripts/prune.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/scripts/prune.py b/scripts/prune.py index ed4dc43..05d7861 100644 --- a/scripts/prune.py +++ b/scripts/prune.py @@ -1,8 +1,14 @@ import os +from pathlib import Path 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}") size_initial = os.path.getsize(p) nsd = dict() @@ -18,14 +24,15 @@ def prune_it(p, keep_only_ema=False): 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.")} + 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]].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"]: - new_sd[k] = sd[k].half() + new_sd[k] = sd[k] assert len(new_sd) == len(sd) - len(ema_keys) nsd["state_dict"] = new_sd @@ -33,7 +40,7 @@ def prune_it(p, keep_only_ema=False): sd = nsd['state_dict'].copy() new_sd = dict() for k in sd: - new_sd[k] = sd[k].half() + new_sd[k] = sd[k] 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" @@ -48,4 +55,4 @@ def prune_it(p, keep_only_ema=False): if __name__ == "__main__": - prune_it('models/ldm/stable-diffusion-v1/model.ckpt') + prune_it(file)