imaginAIry/imaginairy/training_tools/prune_model.py
Bryce 316114e660 docs: add docstrings
Wrote an openai script and custom prompt to generate them.
2023-12-15 14:32:01 -08:00

42 lines
1.1 KiB
Python

"""Functions for pruning diffusion models"""
import logging
import os
import torch
logger = logging.getLogger(__name__)
def prune_diffusion_ckpt(ckpt_path, dst_path=None):
if dst_path is None:
dst_path = f"{os.path.splitext(ckpt_path)[0]}-pruned.ckpt"
data = torch.load(ckpt_path, map_location="cpu")
new_data = prune_model_data(data)
torch.save(new_data, dst_path)
size_initial = os.path.getsize(ckpt_path)
newsize = os.path.getsize(dst_path)
msg = (
f"New ckpt size: {newsize * 1e-9:.2f} GB. "
f"Saved {(size_initial - newsize) * 1e-9:.2f} GB by removing optimizer states"
)
logger.info(msg)
def prune_model_data(data, only_keep_ema=True):
data.pop("optimizer_states", None)
if only_keep_ema:
state_dict = data["state_dict"]
model_keys = [k for k in state_dict if k.startswith("model.")]
for model_key in model_keys:
ema_key = "model_ema." + model_key[6:].replace(".", "")
state_dict[model_key] = state_dict[ema_key]
del state_dict[ema_key]
return data