diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index de628db..f67e404 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -23,6 +23,8 @@ from ema_pytorch import EMA from accelerate import Accelerator +from denoising_diffusion_pytorch.version import __version__ + # constants ModelPrediction = namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start']) @@ -828,7 +830,8 @@ class Trainer(object): 'model': self.accelerator.get_state_dict(self.model), 'opt': self.opt.state_dict(), 'ema': self.ema.state_dict(), - 'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None + 'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None, + 'version': __version__ } torch.save(data, str(self.results_folder / f'model-{milestone}.pt')) @@ -846,6 +849,9 @@ class Trainer(object): self.opt.load_state_dict(data['opt']) self.ema.load_state_dict(data['ema']) + if 'version' in data: + print(f"loading from version {data['version']}") + if exists(self.accelerator.scaler) and exists(data['scaler']): self.accelerator.scaler.load_state_dict(data['scaler']) diff --git a/denoising_diffusion_pytorch/version.py b/denoising_diffusion_pytorch/version.py new file mode 100644 index 0000000..b794fd4 --- /dev/null +++ b/denoising_diffusion_pytorch/version.py @@ -0,0 +1 @@ +__version__ = '0.1.0' diff --git a/setup.py b/setup.py index 1937242..cdb095e 100644 --- a/setup.py +++ b/setup.py @@ -1,9 +1,11 @@ from setuptools import setup, find_packages +exec(open('denoising_diffusion_pytorch/version.py').read()) + setup( name = 'denoising-diffusion-pytorch', packages = find_packages(), - version = '0.32.0', + version = __version__, license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',