Skip to content

Commit a2b72fa

Browse files
patrickvonplatenpatrickvonplaten
committed
Merge branch 'main' of https://github.com/huggingface/diffusers into main
2 parents c9504bb + 26ea58d commit a2b72fa

6 files changed

Lines changed: 91 additions & 98 deletions

File tree

‎src/diffusers/models/resnet.py‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -101,26 +101,37 @@ class Downsample(nn.Module):
101101
downsampling occurs in the inner-two dimensions.
102102
"""
103103

104-
def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
104+
def __init__(self, channels, use_conv=False, dims=2, out_channels=None, padding=1, name="conv"):
105105
super().__init__()
106106
self.channels = channels
107107
self.out_channels = out_channels or channels
108108
self.use_conv = use_conv
109109
self.dims = dims
110110
self.padding = padding
111111
stride = 2 if dims != 3 else (1, 2, 2)
112+
self.name = name
113+
112114
if use_conv:
113-
self.down = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=padding)
115+
conv = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=padding)
114116
else:
115117
assert self.channels == self.out_channels
116-
self.down = avg_pool_nd(dims, kernel_size=stride, stride=stride)
118+
conv = avg_pool_nd(dims, kernel_size=stride, stride=stride)
119+
120+
if name == "conv":
121+
self.conv = conv
122+
else:
123+
self.op = conv
117124

118125
def forward(self, x):
119126
assert x.shape[1] == self.channels
120127
if self.use_conv and self.padding == 0 and self.dims == 2:
121128
pad = (0, 1, 0, 1)
122129
x = F.pad(x, pad, mode="constant", value=0)
123-
return self.down(x)
130+
131+
if self.name == "conv":
132+
return self.conv(x)
133+
else:
134+
return self.op(x)
124135

125136

126137
class UNetUpsample(nn.Module):

‎src/diffusers/models/unet.py‎

Lines changed: 2 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@
3131
from ..configuration_utils import ConfigMixin
3232
from ..modeling_utils import ModelMixin
3333
from .embeddings import get_timestep_embedding
34-
from .resnet import Upsample
34+
from .resnet import Downsample, Upsample
3535

3636

3737
def nonlinearity(x):
@@ -43,24 +43,6 @@ def Normalize(in_channels):
4343
return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
4444

4545

46-
class Downsample(nn.Module):
47-
def __init__(self, in_channels, with_conv):
48-
super().__init__()
49-
self.with_conv = with_conv
50-
if self.with_conv:
51-
# no asymmetric padding in torch conv, must do it ourselves
52-
self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
53-
54-
def forward(self, x):
55-
if self.with_conv:
56-
pad = (0, 1, 0, 1)
57-
x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
58-
x = self.conv(x)
59-
else:
60-
x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
61-
return x
62-
63-
6446
class ResnetBlock(nn.Module):
6547
def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout, temb_channels=512):
6648
super().__init__()
@@ -207,7 +189,7 @@ def __init__(
207189
down.block = block
208190
down.attn = attn
209191
if i_level != self.num_resolutions - 1:
210-
down.downsample = Downsample(block_in, resamp_with_conv)
192+
down.downsample = Downsample(block_in, use_conv=resamp_with_conv, padding=0)
211193
curr_res = curr_res // 2
212194
self.down.append(down)
213195

‎src/diffusers/models/unet_glide.py‎

Lines changed: 6 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from ..configuration_utils import ConfigMixin
99
from ..modeling_utils import ModelMixin
1010
from .embeddings import get_timestep_embedding
11-
from .resnet import Upsample
11+
from .resnet import Downsample, Upsample
1212

1313

1414
def convert_module_to_f16(l):
@@ -124,33 +124,6 @@ def forward(self, x, emb, encoder_out=None):
124124
return x
125125

126126

127-
class Downsample(nn.Module):
128-
"""
129-
A downsampling layer with an optional convolution.
130-
131-
:param channels: channels in the inputs and outputs. :param use_conv: a bool determining if a convolution is
132-
applied. :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
133-
downsampling occurs in the inner-two dimensions.
134-
"""
135-
136-
def __init__(self, channels, use_conv, dims=2, out_channels=None):
137-
super().__init__()
138-
self.channels = channels
139-
self.out_channels = out_channels or channels
140-
self.use_conv = use_conv
141-
self.dims = dims
142-
stride = 2 if dims != 3 else (1, 2, 2)
143-
if use_conv:
144-
self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=1)
145-
else:
146-
assert self.channels == self.out_channels
147-
self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
148-
149-
def forward(self, x):
150-
assert x.shape[1] == self.channels
151-
return self.op(x)
152-
153-
154127
class ResBlock(TimestepBlock):
155128
"""
156129
A residual block that can optionally change the number of channels.
@@ -198,8 +171,8 @@ def __init__(
198171
self.h_upd = Upsample(channels, use_conv=False, dims=dims)
199172
self.x_upd = Upsample(channels, use_conv=False, dims=dims)
200173
elif down:
201-
self.h_upd = Downsample(channels, False, dims)
202-
self.x_upd = Downsample(channels, False, dims)
174+
self.h_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op")
175+
self.x_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op")
203176
else:
204177
self.h_upd = self.x_upd = nn.Identity()
205178

@@ -450,7 +423,9 @@ def __init__(
450423
down=True,
451424
)
452425
if resblock_updown
453-
else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
426+
else Downsample(
427+
ch, use_conv=conv_resample, dims=dims, out_channels=out_ch, padding=1, name="op"
428+
)
454429
)
455430
)
456431
ch = out_ch

‎src/diffusers/models/unet_grad_tts.py‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,17 @@
11
import torch
2+
from numpy import pad
23

34
from ..configuration_utils import ConfigMixin
45
from ..modeling_utils import ModelMixin
56
from .embeddings import get_timestep_embedding
6-
from .resnet import Upsample
7+
from .resnet import Downsample, Upsample
78

89

910
class Mish(torch.nn.Module):
1011
def forward(self, x):
1112
return x * torch.tanh(torch.nn.functional.softplus(x))
1213

1314

14-
class Downsample(torch.nn.Module):
15-
def __init__(self, dim):
16-
super(Downsample, self).__init__()
17-
self.conv = torch.nn.Conv2d(dim, dim, 3, 2, 1)
18-
19-
def forward(self, x):
20-
return self.conv(x)
21-
22-
2315
class Rezero(torch.nn.Module):
2416
def __init__(self, fn):
2517
super(Rezero, self).__init__()
@@ -141,7 +133,7 @@ def __init__(self, dim, dim_mults=(1, 2, 4), groups=8, n_spks=None, spk_emb_dim=
141133
ResnetBlock(dim_in, dim_out, time_emb_dim=dim),
142134
ResnetBlock(dim_out, dim_out, time_emb_dim=dim),
143135
Residual(Rezero(LinearAttention(dim_out))),
144-
Downsample(dim_out) if not is_last else torch.nn.Identity(),
136+
Downsample(dim_out, use_conv=True, padding=1) if not is_last else torch.nn.Identity(),
145137
]
146138
)
147139
)

‎src/diffusers/models/unet_ldm.py‎

Lines changed: 9 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from ..configuration_utils import ConfigMixin
1111
from ..modeling_utils import ModelMixin
1212
from .embeddings import get_timestep_embedding
13-
from .resnet import Upsample
13+
from .resnet import Downsample, Upsample
1414

1515

1616
def exists(val):
@@ -392,32 +392,6 @@ def forward(self, x, emb, context=None):
392392
return x
393393

394394

395-
class Downsample(nn.Module):
396-
"""
397-
A downsampling layer with an optional convolution. :param channels: channels in the inputs and outputs. :param
398-
use_conv: a bool determining if a convolution is applied. :param dims: determines if the signal is 1D, 2D, or 3D.
399-
If 3D, then
400-
downsampling occurs in the inner-two dimensions.
401-
"""
402-
403-
def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
404-
super().__init__()
405-
self.channels = channels
406-
self.out_channels = out_channels or channels
407-
self.use_conv = use_conv
408-
self.dims = dims
409-
stride = 2 if dims != 3 else (1, 2, 2)
410-
if use_conv:
411-
self.op = conv_nd(dims, self.channels, self.out_channels, 3, stride=stride, padding=padding)
412-
else:
413-
assert self.channels == self.out_channels
414-
self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
415-
416-
def forward(self, x):
417-
assert x.shape[1] == self.channels
418-
return self.op(x)
419-
420-
421395
class ResBlock(TimestepBlock):
422396
"""
423397
A residual block that can optionally change the number of channels. :param channels: the number of input channels.
@@ -464,8 +438,8 @@ def __init__(
464438
self.h_upd = Upsample(channels, use_conv=False, dims=dims)
465439
self.x_upd = Upsample(channels, use_conv=False, dims=dims)
466440
elif down:
467-
self.h_upd = Downsample(channels, False, dims)
468-
self.x_upd = Downsample(channels, False, dims)
441+
self.h_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op")
442+
self.x_upd = Downsample(channels, use_conv=False, dims=dims, padding=1, name="op")
469443
else:
470444
self.h_upd = self.x_upd = nn.Identity()
471445

@@ -820,7 +794,9 @@ def __init__(
820794
down=True,
821795
)
822796
if resblock_updown
823-
else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
797+
else Downsample(
798+
ch, use_conv=conv_resample, dims=dims, out_channels=out_ch, padding=1, name="op"
799+
)
824800
)
825801
)
826802
ch = out_ch
@@ -1089,7 +1065,9 @@ def __init__(
10891065
down=True,
10901066
)
10911067
if resblock_updown
1092-
else Downsample(ch, conv_resample, dims=dims, out_channels=out_ch)
1068+
else Downsample(
1069+
ch, use_conv=conv_resample, dims=dims, out_channels=out_ch, padding=1, name="op"
1070+
)
10931071
)
10941072
)
10951073
ch = out_ch

‎tests/test_layers_utils.py‎

Lines changed: 56 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
import torch
2323

2424
from diffusers.models.embeddings import get_timestep_embedding
25-
from diffusers.models.resnet import Upsample
25+
from diffusers.models.resnet import Downsample, Upsample
2626
from diffusers.testing_utils import floats_tensor, slow, torch_device
2727

2828

@@ -164,3 +164,58 @@ def test_upsample_with_transpose(self):
164164
output_slice = upsampled[0, -1, -3:, -3:]
165165
expected_slice = torch.tensor([-0.3028, -0.1582, 0.0071, 0.0350, -0.4799, -0.1139, 0.1056, -0.1153, -0.1046])
166166
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3)
167+
168+
169+
class DownsampleBlockTests(unittest.TestCase):
170+
def test_downsample_default(self):
171+
torch.manual_seed(0)
172+
sample = torch.randn(1, 32, 64, 64)
173+
downsample = Downsample(channels=32, use_conv=False)
174+
with torch.no_grad():
175+
downsampled = downsample(sample)
176+
177+
assert downsampled.shape == (1, 32, 32, 32)
178+
output_slice = downsampled[0, -1, -3:, -3:]
179+
expected_slice = torch.tensor([-0.0513, -0.3889, 0.0640, 0.0836, -0.5460, -0.0341, -0.0169, -0.6967, 0.1179])
180+
max_diff = (output_slice.flatten() - expected_slice).abs().sum().item()
181+
assert max_diff <= 1e-3
182+
# assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-1)
183+
184+
def test_downsample_with_conv(self):
185+
torch.manual_seed(0)
186+
sample = torch.randn(1, 32, 64, 64)
187+
downsample = Downsample(channels=32, use_conv=True)
188+
with torch.no_grad():
189+
downsampled = downsample(sample)
190+
191+
assert downsampled.shape == (1, 32, 32, 32)
192+
output_slice = downsampled[0, -1, -3:, -3:]
193+
194+
expected_slice = torch.tensor(
195+
[0.9267, 0.5878, 0.3337, 1.2321, -0.1191, -0.3984, -0.7532, -0.0715, -0.3913],
196+
)
197+
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3)
198+
199+
def test_downsample_with_conv_pad1(self):
200+
torch.manual_seed(0)
201+
sample = torch.randn(1, 32, 64, 64)
202+
downsample = Downsample(channels=32, use_conv=True, padding=1)
203+
with torch.no_grad():
204+
downsampled = downsample(sample)
205+
206+
assert downsampled.shape == (1, 32, 32, 32)
207+
output_slice = downsampled[0, -1, -3:, -3:]
208+
expected_slice = torch.tensor([0.9267, 0.5878, 0.3337, 1.2321, -0.1191, -0.3984, -0.7532, -0.0715, -0.3913])
209+
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3)
210+
211+
def test_downsample_with_conv_out_dim(self):
212+
torch.manual_seed(0)
213+
sample = torch.randn(1, 32, 64, 64)
214+
downsample = Downsample(channels=32, use_conv=True, out_channels=16)
215+
with torch.no_grad():
216+
downsampled = downsample(sample)
217+
218+
assert downsampled.shape == (1, 16, 32, 32)
219+
output_slice = downsampled[0, -1, -3:, -3:]
220+
expected_slice = torch.tensor([-0.6586, 0.5985, 0.0721, 0.1256, -0.1492, 0.4436, -0.2544, 0.5021, 1.1522])
221+
assert torch.allclose(output_slice.flatten(), expected_slice, atol=1e-3)

0 commit comments

Comments
 (0)