mirror of
https://github.com/wassname/denoising-diffusion-pytorch.git
synced 2026-09-10 12:01:08 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
989f0fcb8e |
@@ -99,14 +99,3 @@ Samples and model checkpoints will be logged to `./results` periodically
|
|||||||
note = {under review}
|
note = {under review}
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
```bibtex
|
|
||||||
@misc{liu2022convnet,
|
|
||||||
title = {A ConvNet for the 2020s},
|
|
||||||
author = {Zhuang Liu and Hanzi Mao and Chao-Yuan Wu and Christoph Feichtenhofer and Trevor Darrell and Saining Xie},
|
|
||||||
year = {2022},
|
|
||||||
eprint = {2201.03545},
|
|
||||||
archivePrefix = {arXiv},
|
|
||||||
primaryClass = {cs.CV}
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|||||||
@@ -142,39 +142,6 @@ class ResnetBlock(nn.Module):
|
|||||||
h = self.block2(h)
|
h = self.block2(h)
|
||||||
return h + self.res_conv(x)
|
return h + self.res_conv(x)
|
||||||
|
|
||||||
class ConvNextBlock(nn.Module):
|
|
||||||
""" https://arxiv.org/abs/2201.03545 """
|
|
||||||
|
|
||||||
def __init__(self, dim, dim_out, *, time_emb_dim = None, mult = 2, norm = True):
|
|
||||||
super().__init__()
|
|
||||||
self.mlp = nn.Sequential(
|
|
||||||
nn.GELU(),
|
|
||||||
nn.Linear(time_emb_dim, dim)
|
|
||||||
) if exists(time_emb_dim) else None
|
|
||||||
|
|
||||||
self.ds_conv = nn.Conv2d(dim, dim, 7, padding = 3, groups = dim)
|
|
||||||
|
|
||||||
self.net = nn.Sequential(
|
|
||||||
LayerNorm(dim) if norm else nn.Identity(),
|
|
||||||
nn.Conv2d(dim, dim_out * mult, 3, padding = 1),
|
|
||||||
nn.GELU(),
|
|
||||||
LayerNorm(dim_out * mult),
|
|
||||||
nn.Conv2d(dim_out * mult, dim_out, 3, padding = 1)
|
|
||||||
)
|
|
||||||
|
|
||||||
self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity()
|
|
||||||
|
|
||||||
def forward(self, x, time_emb = None):
|
|
||||||
h = self.ds_conv(x)
|
|
||||||
|
|
||||||
if exists(self.mlp) and exists(time_emb):
|
|
||||||
assert exists(time_emb), 'time emb must be passed in'
|
|
||||||
condition = self.mlp(time_emb)
|
|
||||||
h = h + rearrange(condition, 'b c -> b c 1 1')
|
|
||||||
|
|
||||||
h = self.net(h)
|
|
||||||
return h + self.res_conv(x)
|
|
||||||
|
|
||||||
class LinearAttention(nn.Module):
|
class LinearAttention(nn.Module):
|
||||||
def __init__(self, dim, heads = 4, dim_head = 32):
|
def __init__(self, dim, heads = 4, dim_head = 32):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
@@ -237,9 +204,7 @@ class Unet(nn.Module):
|
|||||||
dim_mults=(1, 2, 4, 8),
|
dim_mults=(1, 2, 4, 8),
|
||||||
channels = 3,
|
channels = 3,
|
||||||
with_time_emb = True,
|
with_time_emb = True,
|
||||||
use_convnext = False,
|
resnet_block_groups = 8
|
||||||
resnet_block_groups = 8,
|
|
||||||
convnext_mult = 2
|
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
|
||||||
@@ -253,12 +218,7 @@ class Unet(nn.Module):
|
|||||||
dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
|
dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
|
||||||
in_out = list(zip(dims[:-1], dims[1:]))
|
in_out = list(zip(dims[:-1], dims[1:]))
|
||||||
|
|
||||||
# resnet or convnext
|
block_klass = partial(ResnetBlock, groups = resnet_block_groups)
|
||||||
|
|
||||||
if use_convnext:
|
|
||||||
block_klass = partial(ConvNextBlock, mult = convnext_mult)
|
|
||||||
else:
|
|
||||||
block_klass = partial(ResnetBlock, groups = resnet_block_groups)
|
|
||||||
|
|
||||||
# time embeddings
|
# time embeddings
|
||||||
|
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
|
|||||||
setup(
|
setup(
|
||||||
name = 'denoising-diffusion-pytorch',
|
name = 'denoising-diffusion-pytorch',
|
||||||
packages = find_packages(),
|
packages = find_packages(),
|
||||||
version = '0.11.2',
|
version = '0.12.0',
|
||||||
license='MIT',
|
license='MIT',
|
||||||
description = 'Denoising Diffusion Probabilistic Models - Pytorch',
|
description = 'Denoising Diffusion Probabilistic Models - Pytorch',
|
||||||
author = 'Phil Wang',
|
author = 'Phil Wang',
|
||||||
|
|||||||
Reference in New Issue
Block a user