Repository: huggingface/diffusers

Issue:
Sample with Consistency Model
### Describe the bug

When I use diffusers to sample images with offered code by Consistency Model, I can't get an image that corresponds to the label at all: 
label 207 (Golden Retriever)
![0](https://github.com/huggingface/diffusers/assets/102647326/5b01b720-1c1f-4c98-afb9-89081a196e5e)


### Reproduction

I just used the officially offered code by Consistency Model:
https://github.com/openai/consistency_models

### Logs

```shell
python /home/dqwang/scratch/tutu/diffuser/check.py
unet/diffusion_pytorch_model.safetensors not found
Loading pipeline components...:   0%|                                                     | 0/2 [00:00<?, ?it/s]Some weights of the model checkpoint were not used when initializing UNet2DModel: 
 ['mid_block.attentions.0.group_norm.weight, mid_block.attentions.0.group_norm.bias']
Loading pipeline components...: 100%|█████████████████████████████████████████████| 2/2 [00:02<00:00,  1.11s/it]
/home/dqwang/mambaforge/envs/diffuser/lib/python3.7/site-packages/diffusers/configuration_utils.py:134: FutureWarning: Accessing config attribute `sigma_min` directly via 'CMStochasticIterativeScheduler' object attribute is deprecated. Please access 'sigma_min' over 'CMStochasticIterativeScheduler's config object instead, e.g. 'scheduler.config.sigma_min'.
  deprecate("direct config name access", "1.0.0", deprecation_message, standard_warn=False)
100%|█████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 14.55it/s]
```


### System Info

- `diffusers` version: 0.19.3
- Platform: Linux-5.15.0-25-generic-x86_64-with-debian-bookworm-sid
- Python version: 3.7.12
- PyTorch version (GPU?): 1.11.0+cu113 (True)
- Huggingface_hub version: 0.16.4
- Transformers version: not installed
- Accelerate version: 0.20.3
- xFormers version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

_No response_


Discussion:
I have solved this problem, diffusers' source code misses a group norm layer in consistency model's mid_block init, which leads to a bad result.
Would you be open to maybe opening a PR to fix the problem? 

Cc: @dg845 @ayushtues 
My understanding of the issue is as follows: for some consistency model checkpoints like [`openai/diffusers-cd_imagenet64_l2`](https://huggingface.co/openai/diffusers-cd_imagenet64_l2), the [`UNet2DModel`](https://huggingface.co/openai/diffusers-cd_imagenet64_l2/blob/main/unet/config.json) should both use `resnet_time_scale_shift="scale_shift"` and have group norms in the U-Net's mid block, but in the implementation of [`UNetMidBlock2D`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py#L426) a GroupNorm is added to the [`Attention`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py#L36) block only if `resnet_time_scale_shift == "default"`:

https://github.com/huggingface/diffusers/blob/816ca0048f7a38531a6f036299d2dd274abb5639/src/diffusers/models/unet_2d_blocks.py#L472-L478

The [conversion script](https://github.com/huggingface/diffusers/blob/main/scripts/convert_consistency_to_diffusers.py) correctly loads the mid block attention group norm weights from the original checkpoints, but since `unet.mid_block.attentions.0.group_norm` isn't created when the `unet` is initialized, we see the warning

```bash
Some weights of the model checkpoint were not used when initializing UNet2DModel: 
 ['mid_block.attentions.0.group_norm.weight, mid_block.attentions.0.group_norm.bias']
```

as above.
I have opened an PR: https://github.com/huggingface/diffusers/pull/4437#issue-1833173822.
While there seems to be a little problem.
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/diffusers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Resolving patch:
diff --git a/scripts/convert_consistency_to_diffusers.py b/scripts/convert_consistency_to_diffusers.py
--- a/scripts/convert_consistency_to_diffusers.py
+++ b/scripts/convert_consistency_to_diffusers.py
@@ -27,6 +27,7 @@
         "ResnetUpsampleBlock2D",
     ],
     "resnet_time_scale_shift": "scale_shift",
+    "attn_norm_num_groups": 32,
     "upsample_type": "resnet",
     "downsample_type": "resnet",
 }
@@ -52,6 +53,7 @@
         "ResnetUpsampleBlock2D",
     ],
     "resnet_time_scale_shift": "scale_shift",
+    "attn_norm_num_groups": 32,
     "upsample_type": "resnet",
     "downsample_type": "resnet",
 }
diff --git a/src/diffusers/models/unet_2d.py b/src/diffusers/models/unet_2d.py
--- a/src/diffusers/models/unet_2d.py
+++ b/src/diffusers/models/unet_2d.py
@@ -74,6 +74,10 @@ class UNet2DModel(ModelMixin, ConfigMixin):
         act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
         attention_head_dim (`int`, *optional*, defaults to `8`): The attention head dimension.
         norm_num_groups (`int`, *optional*, defaults to `32`): The number of groups for normalization.
+        attn_norm_num_groups (`int`, *optional*, defaults to `None`):
+            If set to an integer, a group norm layer will be created in the mid block's [`Attention`] layer with the
+            given number of groups. If left as `None`, the group norm layer will only be created if
+            `resnet_time_scale_shift` is set to `default`, and if created will have `norm_num_groups` groups.
         norm_eps (`float`, *optional*, defaults to `1e-5`): The epsilon for normalization.
         resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
             for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
@@ -107,6 +111,7 @@ def __init__(
         act_fn: str = "silu",
         attention_head_dim: Optional[int] = 8,
         norm_num_groups: int = 32,
+        attn_norm_num_groups: Optional[int] = None,
         norm_eps: float = 1e-5,
         resnet_time_scale_shift: str = "default",
         add_attention: bool = True,
@@ -192,6 +197,7 @@ def __init__(
             resnet_time_scale_shift=resnet_time_scale_shift,
             attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1],
             resnet_groups=norm_num_groups,
+            attn_groups=attn_norm_num_groups,
             add_attention=add_attention,
         )
 
diff --git a/src/diffusers/models/unet_2d_blocks.py b/src/diffusers/models/unet_2d_blocks.py
--- a/src/diffusers/models/unet_2d_blocks.py
+++ b/src/diffusers/models/unet_2d_blocks.py
@@ -485,6 +485,7 @@ def __init__(
         resnet_time_scale_shift: str = "default",  # default, spatial
         resnet_act_fn: str = "swish",
         resnet_groups: int = 32,
+        attn_groups: Optional[int] = None,
         resnet_pre_norm: bool = True,
         add_attention: bool = True,
         attention_head_dim=1,
@@ -494,6 +495,9 @@ def __init__(
         resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
         self.add_attention = add_attention
 
+        if attn_groups is None:
+            attn_groups = resnet_groups if resnet_time_scale_shift == "default" else None
+
         # there is always at least one resnet
         resnets = [
             ResnetBlock2D(
@@ -526,7 +530,7 @@ def __init__(
                         dim_head=attention_head_dim,
                         rescale_output_factor=output_scale_factor,
                         eps=resnet_eps,
-                        norm_num_groups=resnet_groups if resnet_time_scale_shift == "default" else None,
+                        norm_num_groups=attn_groups,
                         spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
                         residual_connection=True,
                         bias=True,


Test patch:
diff --git a/tests/models/test_models_unet_2d.py b/tests/models/test_models_unet_2d.py
--- a/tests/models/test_models_unet_2d.py
+++ b/tests/models/test_models_unet_2d.py
@@ -74,6 +74,36 @@ def prepare_init_args_and_inputs_for_common(self):
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
 
+    def test_mid_block_attn_groups(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+
+        init_dict["norm_num_groups"] = 16
+        init_dict["add_attention"] = True
+        init_dict["attn_norm_num_groups"] = 8
+
+        model = self.model_class(**init_dict)
+        model.to(torch_device)
+        model.eval()
+
+        self.assertIsNotNone(
+            model.mid_block.attentions[0].group_norm, "Mid block Attention group norm should exist but does not."
+        )
+        self.assertEqual(
+            model.mid_block.attentions[0].group_norm.num_groups,
+            init_dict["attn_norm_num_groups"],
+            "Mid block Attention group norm does not have the expected number of groups.",
+        )
+
+        with torch.no_grad():
+            output = model(**inputs_dict)
+
+            if isinstance(output, dict):
+                output = output.to_tuple()[0]
+
+        self.assertIsNotNone(output)
+        expected_shape = inputs_dict["sample"].shape
+        self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match")
+
 
 class UNetLDMModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase):
     model_class = UNet2DModel
diff --git a/tests/pipelines/consistency_models/test_consistency_models.py b/tests/pipelines/consistency_models/test_consistency_models.py
--- a/tests/pipelines/consistency_models/test_consistency_models.py
+++ b/tests/pipelines/consistency_models/test_consistency_models.py
@@ -216,9 +216,9 @@ def test_consistency_model_cd_multistep(self):
 
         image_slice = image[0, -3:, -3:, -1]
 
-        expected_slice = np.array([0.0888, 0.0881, 0.0666, 0.0479, 0.0292, 0.0195, 0.0201, 0.0163, 0.0254])
+        expected_slice = np.array([0.0146, 0.0158, 0.0092, 0.0086, 0.0000, 0.0000, 0.0000, 0.0000, 0.0058])
 
-        assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2
+        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
 
     def test_consistency_model_cd_onestep(self):
         unet = UNet2DModel.from_pretrained("diffusers/consistency_models", subfolder="diffusers_cd_imagenet64_l2")
@@ -239,9 +239,9 @@ def test_consistency_model_cd_onestep(self):
 
         image_slice = image[0, -3:, -3:, -1]
 
-        expected_slice = np.array([0.0340, 0.0152, 0.0063, 0.0267, 0.0221, 0.0107, 0.0416, 0.0186, 0.0217])
+        expected_slice = np.array([0.0059, 0.0003, 0.0000, 0.0023, 0.0052, 0.0007, 0.0165, 0.0081, 0.0095])
 
-        assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2
+        assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
 
     @require_torch_2
     def test_consistency_model_cd_multistep_flash_attn(self):
@@ -263,7 +263,7 @@ def test_consistency_model_cd_multistep_flash_attn(self):
 
         image_slice = image[0, -3:, -3:, -1]
 
-        expected_slice = np.array([0.1875, 0.1428, 0.1289, 0.2151, 0.2092, 0.1477, 0.1877, 0.1641, 0.1353])
+        expected_slice = np.array([0.1845, 0.1371, 0.1211, 0.2035, 0.1954, 0.1323, 0.1773, 0.1593, 0.1314])
 
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3
 
@@ -289,6 +289,6 @@ def test_consistency_model_cd_onestep_flash_attn(self):
 
         image_slice = image[0, -3:, -3:, -1]
 
-        expected_slice = np.array([0.1663, 0.1948, 0.2275, 0.1680, 0.1204, 0.1245, 0.1858, 0.1338, 0.2095])
+        expected_slice = np.array([0.1623, 0.2009, 0.2387, 0.1731, 0.1168, 0.1202, 0.2031, 0.1327, 0.2447])
 
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3