diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 766e01b..49688f3 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -238,7 +238,8 @@ class Unet(nn.Module): channels = 3, with_time_emb = True, use_convnext = False, - resnet_block_groups = 8 + resnet_block_groups = 8, + convnext_mult = 2 ): super().__init__() @@ -255,7 +256,7 @@ class Unet(nn.Module): # resnet or convnext if use_convnext: - block_klass = ConvNextBlock + block_klass = partial(ConvNextBlock, mult = convnext_mult) else: block_klass = partial(ResnetBlock, groups = resnet_block_groups) diff --git a/setup.py b/setup.py index 3002c57..6b0f4fb 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ from setuptools import setup, find_packages setup( name = 'denoising-diffusion-pytorch', packages = find_packages(), - version = '0.11.0', + version = '0.11.1', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang',