Merge branch 'main' of https://github.com/JoePenna/Dreambooth-Stable-Diffusion

Author: djbielejeski

Diff for: ldm/modules/pruningckptio.py

@@ -0,0 +1,15 @@
+from pytorch_lightning.plugins.io.torch_plugin import TorchCheckpointIO
+
+from typing import Any, Callable, Dict, Optional
+from pytorch_lightning.utilities.types import _PATH
+from ldm.pruner import prune_checkpoint
+ 
+class PruningCheckpointIO(TorchCheckpointIO):
+    def save_checkpoint(
+            self, 
+            checkpoint: Dict[str, Any], 
+            path: _PATH, 
+            storage_options: Optional[Any] = None
+        ) -> None:
+        pruned_checkpoint = prune_checkpoint(checkpoint)
+        TorchCheckpointIO.save_checkpoint(self, pruned_checkpoint, path, storage_options)

Diff for: ldm/pruner.py

@@ -0,0 +1,17 @@
+def prune_checkpoint(old_state):
+    print(f"Prunin' Checkpoint")
+    pruned_checkpoint = dict()
+    print(f"Checkpoint Keys: {old_state.keys()}")
+    for key in old_state.keys():
+        if key != "optimizer_states":
+            pruned_checkpoint[key] = old_state[key]
+    else:
+        print("Removing optimizer states from checkpoint")
+    if "global_step" in old_state:
+        print(f"This is global step {old_state['global_step']}.")
+    old_state = pruned_checkpoint['state_dict'].copy()
+    new_state = dict()
+    for key in old_state:
+        new_state[key] = old_state[key].half()
+    pruned_checkpoint['state_dict'] = new_state
+    return pruned_checkpoint

Diff for: main.py

@@ -1,4 +1,5 @@
 import argparse, os, sys, datetime, glob, importlib, csv
+from ldm.modules.pruningckptio import PruningCheckpointIO
 import numpy as np
 import time
 import torch
@@ -785,7 +786,8 @@ def on_train_epoch_start(self, trainer, pl_module):
 
         trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg]
         trainer_kwargs["max_steps"] = trainer_opt.max_steps
-
+        trainer_kwargs["plugins"] = PruningCheckpointIO()
+    
         trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs)
         trainer.logdir = logdir  ###
 

Diff for: prune_ckpt.py

@@ -1,58 +1,26 @@
 import os
+from ldm.pruner import prune_checkpoint
 import torch
 import argparse
-import glob
 
 
 parser = argparse.ArgumentParser(description='Pruning')
 parser.add_argument('--ckpt', type=str, default=None, help='path to model ckpt')
 args = parser.parse_args()
 ckpt = args.ckpt
 
-def prune_it(p, keep_only_ema=False):
-    print(f"prunin' in path: {p}")
-    size_initial = os.path.getsize(p)
-    nsd = dict()
-    sd = torch.load(p, map_location="cpu")
-    print(sd.keys())
-    for k in sd.keys():
-        if k != "optimizer_states":
-            nsd[k] = sd[k]
-    else:
-        print(f"removing optimizer states for path {p}")
-    if "global_step" in sd:
-        print(f"This is global step {sd['global_step']}.")
-    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]].half()
-            elif not k.startswith("model_ema.") or k in ["model_ema.num_updates", "model_ema.decay"]:
-                new_sd[k] = sd[k].half()
-
-        assert len(new_sd) == len(sd) - len(ema_keys)
-        nsd["state_dict"] = new_sd
-    else:
-        sd = nsd['state_dict'].copy()
-        new_sd = dict()
-        for k in sd:
-            new_sd[k] = sd[k].half()
-        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}")
-    torch.save(nsd, fn)
+def prune_it(checkpoint_path):
+    print(f"Prunin' checkpoint from path: {checkpoint_path}")
+    size_initial = os.path.getsize(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    pruned = prune_checkpoint(checkpoint)
+    fn = f"{os.path.splitext(checkpoint_path)[0]}-pruned.ckpt"
+    print(f"Saving pruned checkpoint at: {fn}")
+    torch.save(pruned, fn)
     newsize = os.path.getsize(fn)
     MSG = f"New ckpt size: {newsize*1e-9:.2f} GB. " + \
           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__":
     prune_it(ckpt)