Skip to content

Commit 571e406

Browse files
patrickvonplatenpatrickvonplaten
committed
merge from master
2 parents 14bd356 + c2bc59d commit 571e406

6 files changed

Lines changed: 174 additions & 40 deletions

File tree

‎scripts/conversion_bddm.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
2+
import argparse
3+
import torch
4+
5+
from diffusers.pipelines.bddm import DiffWave, BDDMPipeline
6+
from diffusers import DDPMScheduler
7+
8+
9+
def convert_bddm_orginal(checkpoint_path, noise_scheduler_checkpoint_path, output_path):
10+
sd = torch.load(checkpoint_path, map_location="cpu")["model_state_dict"]
11+
noise_scheduler_sd = torch.load(noise_scheduler_checkpoint_path, map_location="cpu")
12+
13+
model = DiffWave()
14+
model.load_state_dict(sd, strict=False)
15+
16+
ts, _, betas, _ = noise_scheduler_sd
17+
ts, betas = list(ts.numpy().tolist()), list(betas.numpy().tolist())
18+
19+
noise_scheduler = DDPMScheduler(
20+
timesteps=12,
21+
trained_betas=betas,
22+
timestep_values=ts,
23+
clip_sample=False,
24+
tensor_format="np",
25+
)
26+
27+
pipeline = BDDMPipeline(model, noise_scheduler)
28+
pipeline.save_pretrained(output_path)
29+
30+
31+
if __name__ == "__main__":
32+
parser = argparse.ArgumentParser()
33+
parser.add_argument("--checkpoint_path", type=str, required=True)
34+
parser.add_argument("--noise_scheduler_checkpoint_path", type=str, required=True)
35+
parser.add_argument("--output_path", type=str, required=True)
36+
args = parser.parse_args()
37+
38+
convert_bddm_orginal(args.checkpoint_path, args.noise_scheduler_checkpoint_path, args.output_path)
39+
40+

‎scripts/conversion_ldm_uncond.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import argparse
2+
3+
import OmegaConf
4+
import torch
5+
6+
from diffusers import UNetLDMModel, VQModel, LatentDiffusionUncondPipeline, DDIMScheduler
7+
8+
def convert_ldm_original(checkpoint_path, config_path, output_path):
9+
config = OmegaConf.load(config_path)
10+
state_dict = torch.load(checkpoint_path, map_location="cpu")["model"]
11+
keys = list(state_dict.keys())
12+
13+
# extract state_dict for VQVAE
14+
first_stage_dict = {}
15+
first_stage_key = "first_stage_model."
16+
for key in keys:
17+
if key.startswith(first_stage_key):
18+
first_stage_dict[key.replace(first_stage_key, "")] = state_dict[key]
19+
20+
# extract state_dict for UNetLDM
21+
unet_state_dict = {}
22+
unet_key = "model.diffusion_model."
23+
for key in keys:
24+
if key.startswith(unet_key):
25+
unet_state_dict[key.replace(unet_key, "")] = state_dict[key]
26+
27+
vqvae_init_args = config.model.params.first_stage_config.params
28+
unet_init_args = config.model.params.unet_config.params
29+
30+
vqvae = VQModel(**vqvae_init_args).eval()
31+
vqvae.load_state_dict(first_stage_dict)
32+
33+
unet = UNetLDMModel(**unet_init_args).eval()
34+
unet.load_state_dict(unet_state_dict)
35+
36+
noise_scheduler = DDIMScheduler(
37+
timesteps=config.model.params.timesteps,
38+
beta_schedule="scaled_linear",
39+
beta_start=config.model.params.linear_start,
40+
beta_end=config.model.params.linear_end,
41+
clip_sample=False,
42+
)
43+
44+
pipeline = LatentDiffusionUncondPipeline(vqvae, unet, noise_scheduler)
45+
pipeline.save_pretrained(output_path)
46+
47+
48+
if __name__ == "__main__":
49+
parser = argparse.ArgumentParser()
50+
parser.add_argument("--checkpoint_path", type=str, required=True)
51+
parser.add_argument("--config_path", type=str, required=True)
52+
parser.add_argument("--output_path", type=str, required=True)
53+
args = parser.parse_args()
54+
55+
convert_ldm_original(args.checkpoint_path, args.config_path, args.output_path)
56+

‎src/diffusers/models/resnet.py‎

Lines changed: 36 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from abc import abstractmethod
2+
from functools import partial
23

34
import numpy as np
45
import torch
@@ -78,18 +79,25 @@ class Upsample(nn.Module):
7879
upsampling occurs in the inner-two dimensions.
7980
"""
8081

81-
def __init__(self, channels, use_conv=False, use_conv_transpose=False, dims=2, out_channels=None):
82+
def __init__(self, channels, use_conv=False, use_conv_transpose=False, dims=2, out_channels=None, name="conv"):
8283
super().__init__()
8384
self.channels = channels
8485
self.out_channels = out_channels or channels
8586
self.use_conv = use_conv
8687
self.dims = dims
8788
self.use_conv_transpose = use_conv_transpose
89+
self.name = name
8890

91+
conv = None
8992
if use_conv_transpose:
90-
self.conv = conv_transpose_nd(dims, channels, self.out_channels, 4, 2, 1)
93+
conv = conv_transpose_nd(dims, channels, self.out_channels, 4, 2, 1)
9194
elif use_conv:
92-
self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1)
95+
conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=1)
96+
97+
if name == "conv":
98+
self.conv = conv
99+
else:
100+
self.Conv2d_0 = conv
93101

94102
def forward(self, x):
95103
assert x.shape[1] == self.channels
@@ -102,7 +110,10 @@ def forward(self, x):
102110
x = F.interpolate(x, scale_factor=2.0, mode="nearest")
103111

104112
if self.use_conv:
105-
x = self.conv(x)
113+
if self.name == "conv":
114+
x = self.conv(x)
115+
else:
116+
x = self.Conv2d_0(x)
106117

107118
return x
108119

@@ -134,6 +145,8 @@ def __init__(self, channels, use_conv=False, dims=2, out_channels=None, padding=
134145

135146
if name == "conv":
136147
self.conv = conv
148+
elif name == "Conv2d_0":
149+
self.Conv2d_0 = conv
137150
else:
138151
self.op = conv
139152

@@ -145,6 +158,8 @@ def forward(self, x):
145158

146159
if self.name == "conv":
147160
return self.conv(x)
161+
elif self.name == "Conv2d_0":
162+
return self.Conv2d_0(x)
148163
else:
149164
return self.op(x)
150165

@@ -469,6 +484,7 @@ def __init__(
469484
up=False,
470485
down=False,
471486
dropout=0.1,
487+
fir=False,
472488
fir_kernel=(1, 3, 3, 1),
473489
skip_rescale=True,
474490
init_scale=0.0,
@@ -479,8 +495,20 @@ def __init__(
479495
self.GroupNorm_0 = nn.GroupNorm(num_groups=min(in_ch // 4, 32), num_channels=in_ch, eps=1e-6)
480496
self.up = up
481497
self.down = down
498+
self.fir = fir
482499
self.fir_kernel = fir_kernel
483500

501+
if self.up:
502+
if self.fir:
503+
self.upsample = partial(upsample_2d, k=self.fir_kernel, factor=2)
504+
else:
505+
self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
506+
elif self.down:
507+
if self.fir:
508+
self.downsample = partial(downsample_2d, k=self.fir_kernel, factor=2)
509+
else:
510+
self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2)
511+
484512
self.Conv_0 = conv2d(in_ch, out_ch, kernel_size=3, padding=1)
485513
if temb_dim is not None:
486514
self.Dense_0 = nn.Linear(temb_dim, out_ch)
@@ -503,11 +531,11 @@ def forward(self, x, temb=None):
503531
h = self.act(self.GroupNorm_0(x))
504532

505533
if self.up:
506-
h = upsample_2d(h, self.fir_kernel, factor=2)
507-
x = upsample_2d(x, self.fir_kernel, factor=2)
534+
h = self.upsample(h)
535+
x = self.upsample(x)
508536
elif self.down:
509-
h = downsample_2d(h, self.fir_kernel, factor=2)
510-
x = downsample_2d(x, self.fir_kernel, factor=2)
537+
h = self.downsample(h)
538+
x = self.downsample(x)
511539

512540
h = self.Conv_0(h)
513541
# Add bias to each feature map conditioned on the time embedding

‎src/diffusers/models/unet_sde_score_estimation.py‎

Lines changed: 38 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
from ..modeling_utils import ModelMixin
2828
from .attention import AttentionBlock
2929
from .embeddings import GaussianFourierProjection, get_timestep_embedding
30-
from .resnet import downsample_2d, upfirdn2d, upsample_2d
30+
from .resnet import downsample_2d, upfirdn2d, upsample_2d, Downsample, Upsample
3131
from .resnet import ResnetBlock
3232

3333

@@ -185,37 +185,39 @@ def forward(self, x, y):
185185

186186

187187
class FirUpsample(nn.Module):
188-
def __init__(self, in_ch=None, out_ch=None, with_conv=False, fir_kernel=(1, 3, 3, 1)):
188+
def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)):
189189
super().__init__()
190-
out_ch = out_ch if out_ch else in_ch
191-
if with_conv:
192-
self.Conv2d_0 = Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1)
193-
self.with_conv = with_conv
190+
out_channels = out_channels if out_channels else channels
191+
if use_conv:
192+
self.Conv2d_0 = Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)
193+
self.use_conv = use_conv
194194
self.fir_kernel = fir_kernel
195-
self.out_ch = out_ch
195+
self.out_channels = out_channels
196196

197197
def forward(self, x):
198-
if self.with_conv:
198+
if self.use_conv:
199199
h = _upsample_conv_2d(x, self.Conv2d_0.weight, k=self.fir_kernel)
200+
h = h + self.Conv2d_0.bias.reshape(1, -1, 1, 1)
200201
else:
201202
h = upsample_2d(x, self.fir_kernel, factor=2)
202203

203204
return h
204205

205206

206207
class FirDownsample(nn.Module):
207-
def __init__(self, in_ch=None, out_ch=None, with_conv=False, fir_kernel=(1, 3, 3, 1)):
208+
def __init__(self, channels=None, out_channels=None, use_conv=False, fir_kernel=(1, 3, 3, 1)):
208209
super().__init__()
209-
out_ch = out_ch if out_ch else in_ch
210-
if with_conv:
211-
self.Conv2d_0 = self.Conv2d_0 = Conv2d(in_ch, out_ch, kernel_size=3, stride=1, padding=1)
210+
out_channels = out_channels if out_channels else channels
211+
if use_conv:
212+
self.Conv2d_0 = self.Conv2d_0 = Conv2d(channels, out_channels, kernel_size=3, stride=1, padding=1)
212213
self.fir_kernel = fir_kernel
213-
self.with_conv = with_conv
214-
self.out_ch = out_ch
214+
self.use_conv = use_conv
215+
self.out_channels = out_channels
215216

216217
def forward(self, x):
217-
if self.with_conv:
218+
if self.use_conv:
218219
x = _conv_downsample_2d(x, self.Conv2d_0.weight, k=self.fir_kernel)
220+
x = x + self.Conv2d_0.bias.reshape(1, -1, 1, 1)
219221
else:
220222
x = downsample_2d(x, self.fir_kernel, factor=2)
221223

@@ -229,13 +231,14 @@ def __init__(
229231
self,
230232
image_size=1024,
231233
num_channels=3,
234+
centered=False,
232235
attn_resolutions=(16,),
233236
ch_mult=(1, 2, 4, 8, 16, 32, 32, 32),
234237
conditional=True,
235238
conv_size=3,
236239
dropout=0.0,
237240
embedding_type="fourier",
238-
fir=True, # TODO (patil-suraj) remove this option from here and pre-trained model configs
241+
fir=True,
239242
fir_kernel=(1, 3, 3, 1),
240243
fourier_scale=16,
241244
init_scale=0.0,
@@ -253,12 +256,14 @@ def __init__(
253256
self.register_to_config(
254257
image_size=image_size,
255258
num_channels=num_channels,
259+
centered=centered,
256260
attn_resolutions=attn_resolutions,
257261
ch_mult=ch_mult,
258262
conditional=conditional,
259263
conv_size=conv_size,
260264
dropout=dropout,
261265
embedding_type=embedding_type,
266+
fir=fir,
262267
fir_kernel=fir_kernel,
263268
fourier_scale=fourier_scale,
264269
init_scale=init_scale,
@@ -308,21 +313,26 @@ def __init__(
308313
modules.append(Linear(nf * 4, nf * 4))
309314

310315
AttnBlock = functools.partial(AttentionBlock, overwrite_linear=True, rescale_output_factor=math.sqrt(2.0))
311-
Up_sample = functools.partial(FirUpsample, with_conv=resamp_with_conv, fir_kernel=fir_kernel)
316+
317+
if self.fir:
318+
Up_sample = functools.partial(FirUpsample, fir_kernel=fir_kernel, use_conv=resamp_with_conv)
319+
else:
320+
Up_sample = functools.partial(Upsample, name="Conv2d_0")
312321

313322
if progressive == "output_skip":
314-
self.pyramid_upsample = Up_sample(fir_kernel=fir_kernel, with_conv=False)
323+
self.pyramid_upsample = Up_sample(channels=None, use_conv=False)
315324
elif progressive == "residual":
316-
pyramid_upsample = functools.partial(Up_sample, fir_kernel=fir_kernel, with_conv=True)
325+
pyramid_upsample = functools.partial(Up_sample, use_conv=True)
317326

318-
Down_sample = functools.partial(FirDownsample, with_conv=resamp_with_conv, fir_kernel=fir_kernel)
327+
if self.fir:
328+
Down_sample = functools.partial(FirDownsample, fir_kernel=fir_kernel, use_conv=resamp_with_conv)
329+
else:
330+
Down_sample = functools.partial(Downsample, padding=0, name="Conv2d_0")
319331

320332
if progressive_input == "input_skip":
321-
self.pyramid_downsample = Down_sample(fir_kernel=fir_kernel, with_conv=False)
333+
self.pyramid_downsample = Down_sample(channels=None, use_conv=False)
322334
elif progressive_input == "residual":
323-
pyramid_downsample = functools.partial(Down_sample, fir_kernel=fir_kernel, with_conv=True)
324-
325-
# Downsampling block
335+
pyramid_downsample = functools.partial(Down_sample, use_conv=True)
326336

327337
channels = num_channels
328338
if progressive_input != "none":
@@ -376,7 +386,7 @@ def __init__(
376386
in_ch *= 2
377387

378388
elif progressive_input == "residual":
379-
modules.append(pyramid_downsample(in_ch=input_pyramid_ch, out_ch=in_ch))
389+
modules.append(pyramid_downsample(channels=input_pyramid_ch, out_channels=in_ch))
380390
input_pyramid_ch = in_ch
381391

382392
hs_c.append(in_ch)
@@ -448,7 +458,7 @@ def __init__(
448458
)
449459
pyramid_ch = channels
450460
elif progressive == "residual":
451-
modules.append(pyramid_upsample(in_ch=pyramid_ch, out_ch=in_ch))
461+
modules.append(pyramid_upsample(channels=pyramid_ch, out_channels=in_ch))
452462
pyramid_ch = in_ch
453463
else:
454464
raise ValueError(f"{progressive} is not a valid name")
@@ -505,7 +515,8 @@ def forward(self, x, timesteps, sigmas=None):
505515
temb = None
506516

507517
# If input data is in [0, 1]
508-
x = 2 * x - 1.0
518+
if not self.config.centered:
519+
x = 2 * x - 1.0
509520

510521
# Downsampling block
511522
input_pyramid = None

‎src/diffusers/pipelines/latent_diffusion_uncond/pipeline_latent_diffusion_uncond.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,6 @@ def __call__(
6363
# 4. set current image to prev_image: x_t -> x_t-1
6464
image = pred_prev_image + variance
6565

66-
# scale and decode image with vae
67-
image = 1 / 0.18215 * image
66+
# decode image with vae
6867
image = self.vqvae.decode(image)
69-
image = torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
70-
7168
return image

‎tests/test_modeling_utils.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1159,7 +1159,9 @@ def test_ldm_uncond(self):
11591159
image_slice = image[0, -1, -3:, -3:].cpu()
11601160

11611161
assert image.shape == (1, 3, 256, 256)
1162-
expected_slice = torch.tensor([0.5025, 0.4121, 0.3851, 0.4806, 0.3996, 0.3745, 0.4839, 0.4559, 0.4293])
1162+
expected_slice = torch.tensor(
1163+
[-0.1202, -0.1005, -0.0635, -0.0520, -0.1282, -0.0838, -0.0981, -0.1318, -0.1106]
1164+
)
11631165
assert (image_slice.flatten() - expected_slice).abs().max() < 1e-2
11641166

11651167
def test_module_from_pipeline(self):

0 commit comments

Comments
 (0)