attention mask for transformer Flux

[christopher5106]

Describe the bug

Is it possible to get back the attention_mask argument in the flux attention processor

hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False,attn_mask=attention_mask)

https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py#L1910

in order to tweak things a bit ? otherwise the argument attention_mask is unused.

Thanks a lot

Reproduction

pip install diffusers

Logs

No response

System Info

Ubuntu

Who can help?

@yiyixuxu @sayakpaul @DN6 @asomoza

[sayakpaul]

Cc: @yiyixuxu

[yiyixuxu]

hey @christopher5106
yes! do you want to open a PR?

[rootonchair]

Hi @yiyixuxu, I am working on this issue and it seems like attention_mask is not being used by all the pipelines. Could you help me finding a case that an attention mask is being used and passed to attention processsor?

[sayakpaul]

Thanks, @rootonchair! The reason why that is not the case is because the original Flux implementation doesn't really use any mask so the Flux related pipelines don't use them. So, if we were actually use the attention mask in the Flux attention processor, users will have to make sure to pass them accordingly in their implementations.

[bghira]

it was never done, so it can't be "brought back". you will have to figure out how to handle the image inputs as well. but this probably should not be done, since the original model wasn't trained with it and not even large-scale finetuning on 32x H100 has been able to reliably introduce attention masking.

finetuning from eg. simpletuner or kohya scripts offer the option to reintroduce the attn masking but no inference engine (comfyUI, Forge, A1111, Diffusers) support the attn masking, and no one has really complained about that yet.

the few finetunes that lightly introduced attn masking did not go long enough to have much of an impact from it.

[christopher5106]

i implemented a door then you can do whatever you want with the extra args. it makes no sense not to accept that attention_mask argument to be linked to the attention mecanism

[bonlime]

I can give 2 motivations for introducing attention masking:

1. You have added Flux Redux which in default implementation only works as image variation method, if you try to combine it with text prompts it completely ignores text, but this can be easily fixed by passing attn masks to put less weight on image condition tokens, here is an example reference image on the right, scales are written on the bottom, prompt is "anime style drawing, dog jumping in the field, castle background", you can see that using scales ~0.06-0.08 achieves decent prompt following, but this must pass attn_mask to work
   
   

2. There is a recent paper OminiControl that implements a reference based image generation, which requires passing scale through attn mask, even if you wouldn't add it to diffusers, there are people out there (me for example), who would like to use native diffusers components instead of copy-pasted ones

[bghira]

what do you mean scale the attention mask? it is binary, false/true and the masked positions have a -inf softmax score. pytorch requires the values be 0 or 1. there is no float mask support for performance reasons.

i was hoping OP would share their code and implementation details because it isnt quite lining up for me. i would still rather encoder attn masking be implemented fully and be more than just "a door".

[bonlime]

it is binary, false/true and the masked positions have a -inf softmax score. pytorch requires the values be 0 or 1. there is no float mask support for performance reasons.

torch does support float mask which is treated as attention bias and added before softmax. In my implementation above bias is equal to torch.log(scale.clamp_min_(1e-5)) for example and only added to some parts of the attention

the problem is that you're trying to enable passing argument & adding some implementations for where it's used, why not to just fix the issue of "there is an argument that is not being forwarded correctly", not arguing on what's the "correct" way to have attn masking

[bghira]

if you wanted to just fix it you could copy-paste it but the discussion here is about what's best for the Diffusers project and all of its users, not just a select few

would also be nice to see some performance numbers for the use of sparse attention masks, which are notoriously slow.