Fix Consistency Models UNet2DMidBlock2D Attention GroupNorm Bug (#4863) · sepal/diffusers@4c8a05f · GitHub
Skip to content

Commit 4c8a05f

Browse files
dg845sayakpaul
dg845
and
sayakpaul
authored
Fix Consistency Models UNet2DMidBlock2D Attention GroupNorm Bug (huggingface#4863)
* Add attn_groups argument to UNet2DMidBlock2D to control theinternal Attention block's GroupNorm. * Add docstring for attn_norm_num_groups in UNet2DModel. * Since the test UNet config uses resnet_time_scale_shift == 'scale_shift', also set attn_norm_num_groups to 32. * Add test for attn_norm_num_groups to UNet2DModelTests. * Fix expected slices for slow tests. * Also fix tolerances for slow tests. --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
1 parent 5fd42e5 commit 4c8a05f

5 files changed

Lines changed: 49 additions & 7 deletions

File tree

‎scripts/convert_consistency_to_diffusers.py‎

Lines changed: 2 additions & 0 deletions

‎src/diffusers/models/unet_2d.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,10 @@ class UNet2DModel(ModelMixin, ConfigMixin):
7474
act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
7575
attention_head_dim (`int`, *optional*, defaults to `8`): The attention head dimension.
7676
norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups for normalization.
77+
attn_norm_num_groups (`int`, *optional*, defaults to `None`):
78+
If set to an integer, a group norm layer will be created in the mid block's [`Attention`] layer with the
79+
given number of groups. If left as `None`, the group norm layer will only be created if
80+
`resnet_time_scale_shift` is set to `default`, and if created will have `norm_num_groups` groups.
7781
norm_eps (`float`, *optional*, defaults to `1e-5`): The epsilon for normalization.
7882
resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
7983
for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
@@ -107,6 +111,7 @@ def __init__(
107111
act_fn: str = "silu",
108112
attention_head_dim: Optional[int] = 8,
109113
norm_num_groups: int = 32,
114+
attn_norm_num_groups: Optional[int] = None,
110115
norm_eps: float = 1e-5,
111116
resnet_time_scale_shift: str = "default",
112117
add_attention: bool = True,
@@ -192,6 +197,7 @@ def __init__(
192197
resnet_time_scale_shift=resnet_time_scale_shift,
193198
attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1],
194199
resnet_groups=norm_num_groups,
200+
attn_groups=attn_norm_num_groups,
195201
add_attention=add_attention,
196202
)
197203

‎src/diffusers/models/unet_2d_blocks.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -485,6 +485,7 @@ def __init__(
485485
resnet_time_scale_shift: str = "default", # default, spatial
486486
resnet_act_fn: str = "swish",
487487
resnet_groups: int = 32,
488+
attn_groups: Optional[int] = None,
488489
resnet_pre_norm: bool = True,
489490
add_attention: bool = True,
490491
attention_head_dim=1,
@@ -494,6 +495,9 @@ def __init__(
494495
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
495496
self.add_attention = add_attention
496497

498+
if attn_groups is None:
499+
attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None
500+
497501
# there is always at least one resnet
498502
resnets = [
499503
ResnetBlock2D(
@@ -526,7 +530,7 @@ def __init__(
526530
dim_head=attention_head_dim,
527531
rescale_output_factor=output_scale_factor,
528532
eps=resnet_eps,
529-
norm_num_groups=resnet_groups if resnet_time_scale_shift == "default" else None,
533+
norm_num_groups=attn_groups,
530534
spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
531535
residual_connection=True,
532536
bias=True,

‎tests/models/test_models_unet_2d.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,36 @@ def prepare_init_args_and_inputs_for_common(self):
7474
inputs_dict = self.dummy_input
7575
return init_dict, inputs_dict
7676

77+
def test_mid_block_attn_groups(self):
78+
init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
79+
80+
init_dict["norm_num_groups"] = 16
81+
init_dict["add_attention"] = True
82+
init_dict["attn_norm_num_groups"] = 8
83+
84+
model = self.model_class(**init_dict)
85+
model.to(torch_device)
86+
model.eval()
87+
88+
self.assertIsNotNone(
89+
model.mid_block.attentions[0].group_norm, "Mid block Attention group norm should exist but does not."
90+
)
91+
self.assertEqual(
92+
model.mid_block.attentions[0].group_norm.num_groups,
93+
init_dict["attn_norm_num_groups"],
94+
"Mid block Attention group norm does not have the expected number of groups.",
95+
)
96+
97+
with torch.no_grad():
98+
output = model(**inputs_dict)
99+
100+
if isinstance(output, dict):
101+
output = output.to_tuple()[0]
102+
103+
self.assertIsNotNone(output)
104+
expected_shape = inputs_dict["sample"].shape
105+
self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
106+
77107

78108
class UNetLDMModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
79109
model_class = UNet2DModel

‎tests/pipelines/consistency_models/test_consistency_models.py‎

Lines changed: 6 additions & 6 deletions

0 commit comments

Comments
 (0)