AttentionProcessor.group_norm num_channels should be query_dim

[williamberman]

group norm channels fix

The group_norm on the attention processor should really norm the number of channels in the query not the inner dim. This wasn't caught before because the group_norm is only used by the added kv attention processors and the added kv attention processors are only used by the karlo models which are configured such that the inner dim is the same as the query dim.

I separately ran the integration tests on the unclip/karlo models to confirm they all still pass.

See here that hidden states is normed before the projection to inner dim

diffusers/src/diffusers/models/attention_processor.py

Line 406 in 67c3518

hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

add_{k,v}_proj projection to inner dim fix

Similarly, add_{k,v}_proj should be projecting to inner_dim. This is similarly not caught in karlo/unclip because the cross attention dimension is the same as the inner dim.

See here that in order to concat the projections along dimension 1, they must have the same hidden dimension

diffusers/src/diffusers/models/attention_processor.py

Lines 411 to 422 in 67c3518

	key = attn.to_k(hidden_states)
	value = attn.to_v(hidden_states)
	key = attn.head_to_batch_dim(key)
	value = attn.head_to_batch_dim(value)
	encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
	encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
	encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)
	encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)
	key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
	value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[yiyixuxu]

Thanks for catching this!

[patrickvonplaten]

Ah nice I see! Cool nice catch!