UNet2DConditionModel: add support for QK Normalization by propagating qk_norm value from config through to child attention modules

[damian0815]

What does this PR do?

Fixes #12050

QK Normalization was already implemented in Attention __init__, but adding eg "qk_norm": "rms_norm" to the config.json for a UNet2DConfitionModel had no effect.

This PR makes config qk_norm have an effect by propagating its value through the various UNet2DConfitionModel block initialization logic.

Without this PR:

from diffusers import UNet2DConditionModel
config_minimal = {"qk_norm": "rms_norm"}
model = UNet2DConditionModel.from_config(config_minimal)
print([n for n, _ in model.named_modules()
       if 'attn1.norm_' in n])
# output: []

With this PR:

from diffusers import UNet2DConditionModel
config_minimal = {"qk_norm": "rms_norm"}
model = UNet2DConditionModel.from_config(config_minimal)
print([n for n, _ in model.named_modules()
       if 'attn1.norm_' in n])
# output: ['down_blocks.0.attentions.0.transformer_blocks.0.attn1.norm_q', 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.norm_k', ...]

(I have successfully finetuned a model after making this change)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Pings based on git blame:
@yiyixuxu @gnobitab @sayakpaul

[damian0815]

For the tests I find myself writing code like this:

for down_block_type in ["DownBlock2D", "ResnetDownsampleBlock2D",
                     "AttnDownBlock2D", "CrossAttnDownBlock2D",
                     "SimpleCrossAttnDownBlock2D", "SkipDownBlock2D",
                     "AttnSkipDownBlock2D", "DownEncoderBlock2D",
                     "AttnDownEncoderBlock2D", "KDownBlock2D",
                     "KCrossAttnDownBlock2D"]:
    block = get_down_block(
        down_block_type=down_block_type,
        ...

Is there a canonical way of obtaining the list of block types so it doesn't have to be hardcoded in the test?

[damian0815]

I wasn't able to run the full test suite - significant components seem to be broken on macOS/mps