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}
|
||||
}
|
||||
```
|
||||
|
||||
```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)
|
||||
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):
|
||||
def __init__(self, dim, heads = 4, dim_head = 32):
|
||||
super().__init__()
|
||||
@@ -237,9 +204,7 @@ class Unet(nn.Module):
|
||||
dim_mults=(1, 2, 4, 8),
|
||||
channels = 3,
|
||||
with_time_emb = True,
|
||||
use_convnext = False,
|
||||
resnet_block_groups = 8,
|
||||
convnext_mult = 2
|
||||
resnet_block_groups = 8
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
@@ -253,12 +218,7 @@ class Unet(nn.Module):
|
||||
dims = [init_dim, *map(lambda m: dim * m, dim_mults)]
|
||||
in_out = list(zip(dims[:-1], dims[1:]))
|
||||
|
||||
# resnet or convnext
|
||||
|
||||
if use_convnext:
|
||||
block_klass = partial(ConvNextBlock, mult = convnext_mult)
|
||||
else:
|
||||
block_klass = partial(ResnetBlock, groups = resnet_block_groups)
|
||||
block_klass = partial(ResnetBlock, groups = resnet_block_groups)
|
||||
|
||||
# time embeddings
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
|
||||
setup(
|
||||
name = 'denoising-diffusion-pytorch',
|
||||
packages = find_packages(),
|
||||
version = '0.11.2',
|
||||
version = '0.12.0',
|
||||
license='MIT',
|
||||
description = 'Denoising Diffusion Probabilistic Models - Pytorch',
|
||||
author = 'Phil Wang',
|
||||
|
||||
Reference in New Issue
Block a user