Skip to content

Commit 321f979

Browse files
patrickvonplatenpatrickvonplaten
authored
Downsample / Upsample - clean to 1D and 2D (huggingface#68)
* make unet rl work * uploaad files / code * upload files * make style correct * finish
1 parent c524244 commit 321f979

10 files changed

Lines changed: 254 additions & 232 deletions

File tree

‎src/diffusers/models/resnet.py‎

Lines changed: 219 additions & 64 deletions
Large diffs are not rendered by default.

‎src/diffusers/models/unet.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
from ..modeling_utils import ModelMixin
2323
from .attention import AttentionBlock
2424
from .embeddings import get_timestep_embedding
25-
from .resnet import Downsample, ResnetBlock2D, Upsample
25+
from .resnet import Downsample2D, ResnetBlock2D, Upsample2D
2626

2727

2828
def nonlinearity(x):
@@ -100,7 +100,7 @@ def __init__(
100100
down.block = block
101101
down.attn = attn
102102
if i_level != self.num_resolutions - 1:
103-
down.downsample = Downsample(block_in, use_conv=resamp_with_conv, padding=0)
103+
down.downsample = Downsample2D(block_in, use_conv=resamp_with_conv, padding=0)
104104
curr_res = curr_res // 2
105105
self.down.append(down)
106106

@@ -139,7 +139,7 @@ def __init__(
139139
up.block = block
140140
up.attn = attn
141141
if i_level != 0:
142-
up.upsample = Upsample(block_in, use_conv=resamp_with_conv)
142+
up.upsample = Upsample2D(block_in, use_conv=resamp_with_conv)
143143
curr_res = curr_res * 2
144144
self.up.insert(0, up) # prepend to get consistent order
145145

‎src/diffusers/models/unet_glide.py‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from ..modeling_utils import ModelMixin
77
from .attention import AttentionBlock
88
from .embeddings import get_timestep_embedding
9-
from .resnet import Downsample, ResnetBlock2D, Upsample
9+
from .resnet import Downsample2D, ResnetBlock2D, Upsample2D
1010

1111

1212
def convert_module_to_f16(l):
@@ -218,9 +218,7 @@ def __init__(
218218
down=True,
219219
)
220220
if resblock_updown
221-
else Downsample(
222-
ch, use_conv=conv_resample, dims=dims, out_channels=out_ch, padding=1, name="op"
223-
)
221+
else Downsample2D(ch, use_conv=conv_resample, out_channels=out_ch, padding=1, name="op")
224222
)
225223
)
226224
ch = out_ch
@@ -299,7 +297,7 @@ def __init__(
299297
up=True,
300298
)
301299
if resblock_updown
302-
else Upsample(ch, use_conv=conv_resample, dims=dims, out_channels=out_ch)
300+
else Upsample2D(ch, use_conv=conv_resample, out_channels=out_ch)
303301
)
304302
ds //= 2
305303
self.output_blocks.append(TimestepEmbedSequential(*layers))

‎src/diffusers/models/unet_grad_tts.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from ..modeling_utils import ModelMixin
55
from .attention import LinearAttention
66
from .embeddings import get_timestep_embedding
7-
from .resnet import Downsample, ResnetBlock2D, Upsample
7+
from .resnet import Downsample2D, ResnetBlock2D, Upsample2D
88

99

1010
class Mish(torch.nn.Module):
@@ -105,7 +105,7 @@ def __init__(self, dim, dim_mults=(1, 2, 4), groups=8, n_spks=None, spk_emb_dim=
105105
overwrite_for_grad_tts=True,
106106
),
107107
Residual(Rezero(LinearAttention(dim_out))),
108-
Downsample(dim_out, use_conv=True, padding=1) if not is_last else torch.nn.Identity(),
108+
Downsample2D(dim_out, use_conv=True, padding=1) if not is_last else torch.nn.Identity(),
109109
]
110110
)
111111
)
@@ -158,7 +158,7 @@ def __init__(self, dim, dim_mults=(1, 2, 4), groups=8, n_spks=None, spk_emb_dim=
158158
overwrite_for_grad_tts=True,
159159
),
160160
Residual(Rezero(LinearAttention(dim_in))),
161-
Upsample(dim_in, use_conv_transpose=True),
161+
Upsample2D(dim_in, use_conv_transpose=True),
162162
]
163163
)
164164
)

‎src/diffusers/models/unet_ldm.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from ..modeling_utils import ModelMixin
1111
from .attention import AttentionBlock
1212
from .embeddings import get_timestep_embedding
13-
from .resnet import Downsample, ResnetBlock2D, Upsample
13+
from .resnet import Downsample2D, ResnetBlock2D, Upsample2D
1414

1515

1616
# from .resnet import ResBlock
@@ -350,7 +350,7 @@ def __init__(
350350
out_ch = ch
351351
self.input_blocks.append(
352352
TimestepEmbedSequential(
353-
Downsample(ch, use_conv=conv_resample, dims=dims, out_channels=out_ch, padding=1, name="op")
353+
Downsample2D(ch, use_conv=conv_resample, out_channels=out_ch, padding=1, name="op")
354354
)
355355
)
356356
ch = out_ch
@@ -437,7 +437,7 @@ def __init__(
437437
)
438438
if level and i == num_res_blocks:
439439
out_ch = ch
440-
layers.append(Upsample(ch, use_conv=conv_resample, dims=dims, out_channels=out_ch))
440+
layers.append(Upsample2D(ch, use_conv=conv_resample, out_channels=out_ch))
441441
ds //= 2
442442
self.output_blocks.append(TimestepEmbedSequential(*layers))
443443
self._feature_size += ch

‎src/diffusers/models/unet_rl.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from ..configuration_utils import ConfigMixin
77
from ..modeling_utils import ModelMixin
88
from .embeddings import get_timestep_embedding
9-
from .resnet import Downsample, ResidualTemporalBlock, Upsample
9+
from .resnet import Downsample1D, ResidualTemporalBlock, Upsample1D
1010

1111

1212
class SinusoidalPosEmb(nn.Module):
@@ -96,7 +96,7 @@ def __init__(
9696
[
9797
ResidualTemporalBlock(dim_in, dim_out, embed_dim=time_dim, horizon=training_horizon),
9898
ResidualTemporalBlock(dim_out, dim_out, embed_dim=time_dim, horizon=training_horizon),
99-
Downsample(dim_out, use_conv=True, dims=1) if not is_last else nn.Identity(),
99+
Downsample1D(dim_out, use_conv=True) if not is_last else nn.Identity(),
100100
]
101101
)
102102
)
@@ -116,7 +116,7 @@ def __init__(
116116
[
117117
ResidualTemporalBlock(dim_out * 2, dim_in, embed_dim=time_dim, horizon=training_horizon),
118118
ResidualTemporalBlock(dim_in, dim_in, embed_dim=time_dim, horizon=training_horizon),
119-
Upsample(dim_in, use_conv_transpose=True, dims=1) if not is_last else nn.Identity(),
119+
Upsample1D(dim_in, use_conv_transpose=True) if not is_last else nn.Identity(),
120120
]
121121
)
122122
)

‎src/diffusers/models/unet_sde_score_estimation.py‎

Lines changed: 5 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,12 @@
2121
import numpy as np
2222
import torch
2323
import torch.nn as nn
24-
import torch.nn.functional as F
2524

2625
from ..configuration_utils import ConfigMixin
2726
from ..modeling_utils import ModelMixin
2827
from .attention import AttentionBlock
2928
from .embeddings import GaussianFourierProjection, get_timestep_embedding
30-
from .resnet import Downsample, ResnetBlock2D, Upsample, downsample_2d, upfirdn2d, upsample_2d
29+
from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, ResnetBlock2D, Upsample2D
3130

3231

3332
def _setup_kernel(k):
@@ -40,96 +39,6 @@ def _setup_kernel(k):
4039
return k
4140

4241

43-
def _upsample_conv_2d(x, w, k=None, factor=2, gain=1):
44-
"""Fused `upsample_2d()` followed by `Conv2d()`.
45-
46-
Args:
47-
Padding is performed only once at the beginning, not between the operations. The fused op is considerably more
48-
efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of arbitrary
49-
order.
50-
x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W,
51-
C]`.
52-
w: Weight tensor of the shape `[filterH, filterW, inChannels,
53-
outChannels]`. Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`.
54-
k: FIR filter of the shape `[firH, firW]` or `[firN]`
55-
(separable). The default is `[1] * factor`, which corresponds to nearest-neighbor upsampling.
56-
factor: Integer upsampling factor (default: 2). gain: Scaling factor for signal magnitude (default: 1.0).
57-
58-
Returns:
59-
Tensor of the shape `[N, C, H * factor, W * factor]` or `[N, H * factor, W * factor, C]`, and same datatype as
60-
`x`.
61-
"""
62-
63-
assert isinstance(factor, int) and factor >= 1
64-
65-
# Check weight shape.
66-
assert len(w.shape) == 4
67-
convH = w.shape[2]
68-
convW = w.shape[3]
69-
inC = w.shape[1]
70-
71-
assert convW == convH
72-
73-
# Setup filter kernel.
74-
if k is None:
75-
k = [1] * factor
76-
k = _setup_kernel(k) * (gain * (factor**2))
77-
p = (k.shape[0] - factor) - (convW - 1)
78-
79-
stride = (factor, factor)
80-
81-
# Determine data dimensions.
82-
stride = [1, 1, factor, factor]
83-
output_shape = ((x.shape[2] - 1) * factor + convH, (x.shape[3] - 1) * factor + convW)
84-
output_padding = (
85-
output_shape[0] - (x.shape[2] - 1) * stride[0] - convH,
86-
output_shape[1] - (x.shape[3] - 1) * stride[1] - convW,
87-
)
88-
assert output_padding[0] >= 0 and output_padding[1] >= 0
89-
num_groups = x.shape[1] // inC
90-
91-
# Transpose weights.
92-
w = torch.reshape(w, (num_groups, -1, inC, convH, convW))
93-
w = w[..., ::-1, ::-1].permute(0, 2, 1, 3, 4)
94-
w = torch.reshape(w, (num_groups * inC, -1, convH, convW))
95-
96-
x = F.conv_transpose2d(x, w, stride=stride, output_padding=output_padding, padding=0)
97-
98-
return upfirdn2d(x, torch.tensor(k, device=x.device), pad=((p + 1) // 2 + factor - 1, p // 2 + 1))
99-
100-
101-
def _conv_downsample_2d(x, w, k=None, factor=2, gain=1):
102-
"""Fused `Conv2d()` followed by `downsample_2d()`.
103-
104-
Args:
105-
Padding is performed only once at the beginning, not between the operations. The fused op is considerably more
106-
efficient than performing the same calculation using standard TensorFlow ops. It supports gradients of arbitrary
107-
order.
108-
x: Input tensor of the shape `[N, C, H, W]` or `[N, H, W,
109-
C]`.
110-
w: Weight tensor of the shape `[filterH, filterW, inChannels,
111-
outChannels]`. Grouped convolution can be performed by `inChannels = x.shape[0] // numGroups`.
112-
k: FIR filter of the shape `[firH, firW]` or `[firN]`
113-
(separable). The default is `[1] * factor`, which corresponds to average pooling.
114-
factor: Integer downsampling factor (default: 2). gain: Scaling factor for signal magnitude (default: 1.0).
115-
116-
Returns:
117-
Tensor of the shape `[N, C, H // factor, W // factor]` or `[N, H // factor, W // factor, C]`, and same datatype
118-
as `x`.
119-
"""
120-
121-
assert isinstance(factor, int) and factor >= 1
122-
_outC, _inC, convH, convW = w.shape
123-
assert convW == convH
124-
if k is None:
125-
k = [1] * factor
126-
k = _setup_kernel(k) * gain
127-
p = (k.shape[0] - factor) + (convW - 1)
128-
s = [factor, factor]
129-
x = upfirdn2d(x, torch.tensor(k, device=x.device), pad=((p + 1) // 2, p // 2))
130-
return F.conv2d(x, w, stride=s, padding=0)
131-
132-
13342
def _variance_scaling(scale=1.0, in_axis=1, out_axis=0, dtype=torch.float32, device="cpu"):
13443
"""Ported from JAX."""
13544
scale = 1e-10 if scale == 0 else scale
@@ -183,46 +92,6 @@ def forward(self, x, y):
18392
raise ValueError(f"Method {self.method} not recognized.")
18493

18594

186-
class FirUpsample(nn.Module):
187-
def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)):
188-
super().__init__()
189-
out_channels = out_channels if out_channels else channels
190-
if use_conv:
191-
self.Conv2d_0 = Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)
192-
self.use_conv = use_conv
193-
self.fir_kernel = fir_kernel
194-
self.out_channels = out_channels
195-
196-
def forward(self, x):
197-
if self.use_conv:
198-
h = _upsample_conv_2d(x, self.Conv2d_0.weight, k=self.fir_kernel)
199-
h = h + self.Conv2d_0.bias.reshape(1, -1, 1, 1)
200-
else:
201-
h = upsample_2d(x, self.fir_kernel, factor=2)
202-
203-
return h
204-
205-
206-
class FirDownsample(nn.Module):
207-
def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)):
208-
super().__init__()
209-
out_channels = out_channels if out_channels else channels
210-
if use_conv:
211-
self.Conv2d_0 = self.Conv2d_0 = Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)
212-
self.fir_kernel = fir_kernel
213-
self.use_conv = use_conv
214-
self.out_channels = out_channels
215-
216-
def forward(self, x):
217-
if self.use_conv:
218-
x = _conv_downsample_2d(x, self.Conv2d_0.weight, k=self.fir_kernel)
219-
x = x + self.Conv2d_0.bias.reshape(1, -1, 1, 1)
220-
else:
221-
x = downsample_2d(x, self.fir_kernel, factor=2)
222-
223-
return x
224-
225-
22695
class NCSNpp(ModelMixin, ConfigMixin):
22796
"""NCSN++ model"""
22897

@@ -313,19 +182,19 @@ def __init__(
313182
AttnBlock = functools.partial(AttentionBlock, overwrite_linear=True, rescale_output_factor=math.sqrt(2.0))
314183

315184
if self.fir:
316-
Up_sample = functools.partial(FirUpsample, fir_kernel=fir_kernel, use_conv=resamp_with_conv)
185+
Up_sample = functools.partial(FirUpsample2D, fir_kernel=fir_kernel, use_conv=resamp_with_conv)
317186
else:
318-
Up_sample = functools.partial(Upsample, name="Conv2d_0")
187+
Up_sample = functools.partial(Upsample2D, name="Conv2d_0")
319188

320189
if progressive == "output_skip":
321190
self.pyramid_upsample = Up_sample(channels=None, use_conv=False)
322191
elif progressive == "residual":
323192
pyramid_upsample = functools.partial(Up_sample, use_conv=True)
324193

325194
if self.fir:
326-
Down_sample = functools.partial(FirDownsample, fir_kernel=fir_kernel, use_conv=resamp_with_conv)
195+
Down_sample = functools.partial(FirDownsample2D, fir_kernel=fir_kernel, use_conv=resamp_with_conv)
327196
else:
328-
Down_sample = functools.partial(Downsample, padding=0, name="Conv2d_0")
197+
Down_sample = functools.partial(Downsample2D, padding=0, name="Conv2d_0")
329198

330199
if progressive_input == "input_skip":
331200
self.pyramid_downsample = Down_sample(channels=None, use_conv=False)

‎src/diffusers/models/vae.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from ..configuration_utils import ConfigMixin
66
from ..modeling_utils import ModelMixin
77
from .attention import AttentionBlock
8-
from .resnet import Downsample, ResnetBlock2D, Upsample
8+
from .resnet import Downsample2D, ResnetBlock2D, Upsample2D
99

1010

1111
def nonlinearity(x):
@@ -65,7 +65,7 @@ def __init__(
6565
down.block = block
6666
down.attn = attn
6767
if i_level != self.num_resolutions - 1:
68-
down.downsample = Downsample(block_in, use_conv=resamp_with_conv, padding=0)
68+
down.downsample = Downsample2D(block_in, use_conv=resamp_with_conv, padding=0)
6969
curr_res = curr_res // 2
7070
self.down.append(down)
7171

@@ -179,7 +179,7 @@ def __init__(
179179
up.block = block
180180
up.attn = attn
181181
if i_level != 0:
182-
up.upsample = Upsample(block_in, use_conv=resamp_with_conv)
182+
up.upsample = Upsample2D(block_in, use_conv=resamp_with_conv)
183183
curr_res = curr_res * 2
184184
self.up.insert(0, up) # prepend to get consistent order
185185

‎src/diffusers/pipelines/bddm/pipeline_bddm.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ def forward(self, input_data):
137137
# Dilated conv layer
138138
h = self.dilated_conv_layer(h)
139139

140-
# Upsample spectrogram to size of audio
140+
# Upsample2D spectrogram to size of audio
141141
mel_spec = torch.unsqueeze(mel_spec, dim=1)
142142
mel_spec = F.leaky_relu(self.upsample_conv2d[0](mel_spec), 0.4, inplace=False)
143143
mel_spec = F.leaky_relu(self.upsample_conv2d[1](mel_spec), 0.4, inplace=False)

0 commit comments

Comments
 (0)