Skip to content

Commit fa7443c

Browse files
patrickvonplatenpatrickvonplaten
committed
finish resnet
1 parent 8d7771d commit fa7443c

2 files changed

Lines changed: 92 additions & 25 deletions

File tree

‎src/diffusers/models/resnet.py‎

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -380,7 +380,7 @@ def __init__(
380380
eps=1e-6,
381381
non_linearity="swish",
382382
time_embedding_norm="default",
383-
fir_kernel=(1, 3, 3, 1),
383+
kernel=None,
384384
output_scale_factor=1.0,
385385
use_nin_shortcut=None,
386386
up=False,
@@ -433,8 +433,18 @@ def __init__(
433433
# elif down:
434434
# self.h_upd = Downsample(in_channels, use_conv=False, dims=2, padding=1, name="op")
435435
# self.x_upd = Downsample(in_channels, use_conv=False, dims=2, padding=1, name="op")
436-
self.upsample = Upsample(in_channels, use_conv=False, dims=2) if self.up else None
437-
self.downsample = Downsample(in_channels, use_conv=False, dims=2, padding=1, name="op") if self.down else None
436+
437+
self.upsample = self.downsample = None
438+
if self.up and kernel == "fir":
439+
fir_kernel = (1, 3, 3, 1)
440+
self.upsample = lambda x: upsample_2d(x, k=fir_kernel)
441+
elif self.up and kernel is None:
442+
self.upsample = Upsample(in_channels, use_conv=False, dims=2)
443+
elif self.down and kernel == "fir":
444+
fir_kernel = (1, 3, 3, 1)
445+
self.downsample = lambda x: downsample_2d(x, k=fir_kernel)
446+
elif self.down and kernel is None:
447+
self.downsample = Downsample(in_channels, use_conv=False, dims=2, padding=1, name="op")
438448

439449
self.use_nin_shortcut = self.in_channels != self.out_channels if use_nin_shortcut is None else use_nin_shortcut
440450

@@ -505,8 +515,6 @@ def __init__(
505515
self.GroupNorm_0 = nn.GroupNorm(num_groups=num_groups, num_channels=in_ch, eps=eps)
506516
self.up = up
507517
self.down = down
508-
self.fir_kernel = fir_kernel
509-
510518
self.Conv_0 = conv2d(in_ch, out_ch, kernel_size=3, padding=1)
511519
if temb_dim is not None:
512520
self.Dense_0 = nn.Linear(temb_dim, out_ch)
@@ -525,11 +533,6 @@ def __init__(
525533
self.out_ch = out_ch
526534

527535
# TODO(Patrick) - move to main init
528-
if self.up:
529-
self.upsample = functools.partial(upsample_2d, k=self.fir_kernel)
530-
if self.down:
531-
self.downsample = functools.partial(downsample_2d, k=self.fir_kernel)
532-
533536
self.is_overwritten = False
534537

535538
def set_weights_grad_tts(self):

‎src/diffusers/models/unet_sde_score_estimation.py‎

Lines changed: 79 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -348,24 +348,40 @@ def __init__(
348348
for i_block in range(num_res_blocks):
349349
out_ch = nf * ch_mult[i_level]
350350
# modules.append(ResnetBlock(in_ch=in_ch, out_ch=out_ch))
351-
modules.append(ResnetNew(
352-
in_channels=in_ch,
353-
out_channels=out_ch,
354-
temb_channels=4 * nf,
355-
output_scale_factor=np.sqrt(2.0),
356-
non_linearity="silu",
357-
groups=min(in_ch // 4, 32),
358-
groups_out=min(out_ch // 4, 32),
359-
overwrite_for_score_vde=True,
360-
))
351+
modules.append(
352+
ResnetNew(
353+
in_channels=in_ch,
354+
out_channels=out_ch,
355+
temb_channels=4 * nf,
356+
output_scale_factor=np.sqrt(2.0),
357+
non_linearity="silu",
358+
groups=min(in_ch // 4, 32),
359+
groups_out=min(out_ch // 4, 32),
360+
overwrite_for_score_vde=True,
361+
)
362+
)
361363
in_ch = out_ch
362364

363365
if all_resolutions[i_level] in attn_resolutions:
364366
modules.append(AttnBlock(channels=in_ch))
365367
hs_c.append(in_ch)
366368

367369
if i_level != self.num_resolutions - 1:
368-
modules.append(ResnetBlock(down=True, in_ch=in_ch))
370+
# modules.append(ResnetBlock(down=True, in_ch=in_ch))
371+
modules.append(
372+
ResnetNew(
373+
in_channels=in_ch,
374+
temb_channels=4 * nf,
375+
output_scale_factor=np.sqrt(2.0),
376+
non_linearity="silu",
377+
groups=min(in_ch // 4, 32),
378+
groups_out=min(out_ch // 4, 32),
379+
overwrite_for_score_vde=True,
380+
down=True,
381+
kernel="fir", # TODO(Patrick) - it seems like both fir and non-fir kernels are fine
382+
use_nin_shortcut=True,
383+
)
384+
)
369385

370386
if progressive_input == "input_skip":
371387
modules.append(combiner(dim1=input_pyramid_ch, dim2=in_ch))
@@ -379,16 +395,50 @@ def __init__(
379395
hs_c.append(in_ch)
380396

381397
in_ch = hs_c[-1]
382-
modules.append(ResnetBlock(in_ch=in_ch))
398+
# modules.append(ResnetBlock(in_ch=in_ch))
399+
modules.append(
400+
ResnetNew(
401+
in_channels=in_ch,
402+
temb_channels=4 * nf,
403+
output_scale_factor=np.sqrt(2.0),
404+
non_linearity="silu",
405+
groups=min(in_ch // 4, 32),
406+
groups_out=min(out_ch // 4, 32),
407+
overwrite_for_score_vde=True,
408+
)
409+
)
383410
modules.append(AttnBlock(channels=in_ch))
384-
modules.append(ResnetBlock(in_ch=in_ch))
411+
# modules.append(ResnetBlock(in_ch=in_ch))
412+
modules.append(
413+
ResnetNew(
414+
in_channels=in_ch,
415+
temb_channels=4 * nf,
416+
output_scale_factor=np.sqrt(2.0),
417+
non_linearity="silu",
418+
groups=min(in_ch // 4, 32),
419+
groups_out=min(out_ch // 4, 32),
420+
overwrite_for_score_vde=True,
421+
)
422+
)
385423

386424
pyramid_ch = 0
387425
# Upsampling block
388426
for i_level in reversed(range(self.num_resolutions)):
389427
for i_block in range(num_res_blocks + 1):
390428
out_ch = nf * ch_mult[i_level]
391-
modules.append(ResnetBlock(in_ch=in_ch + hs_c.pop(), out_ch=out_ch))
429+
# modules.append(ResnetBlock(in_ch=in_ch + hs_c.pop(), out_ch=out_ch))
430+
modules.append(
431+
ResnetNew(
432+
in_channels=in_ch + hs_c.pop(),
433+
out_channels=out_ch,
434+
temb_channels=4 * nf,
435+
output_scale_factor=np.sqrt(2.0),
436+
non_linearity="silu",
437+
groups=min(in_ch // 4, 32),
438+
groups_out=min(out_ch // 4, 32),
439+
overwrite_for_score_vde=True,
440+
)
441+
)
392442
in_ch = out_ch
393443

394444
if all_resolutions[i_level] in attn_resolutions:
@@ -420,7 +470,21 @@ def __init__(
420470
raise ValueError(f"{progressive} is not a valid name")
421471

422472
if i_level != 0:
423-
modules.append(ResnetBlock(in_ch=in_ch, up=True))
473+
# modules.append(ResnetBlock(in_ch=in_ch, up=True))
474+
modules.append(
475+
ResnetNew(
476+
in_channels=in_ch,
477+
temb_channels=4 * nf,
478+
output_scale_factor=np.sqrt(2.0),
479+
non_linearity="silu",
480+
groups=min(in_ch // 4, 32),
481+
groups_out=min(out_ch // 4, 32),
482+
overwrite_for_score_vde=True,
483+
up=True,
484+
kernel="fir",
485+
use_nin_shortcut=True,
486+
)
487+
)
424488

425489
assert not hs_c
426490

0 commit comments

Comments
 (0)