Raphael Walker 61b50720d0
Add support for attention masking in Flux (#5942)
* fix attention OOM in xformers

* allow passing attention mask in flux attention

* allow an attn_mask in flux

* attn masks can be done using replace patches instead of a separate dict

* fix return types

* fix return order

* enumerate

* patch the right keys

* arg names

* fix a silly bug

* fix xformers masks

* replace match with if, elif, else

* mask with image_ref_size

* remove unused import

* remove unused import 2

* fix pytorch/xformers attention

This corrects a weird inconsistency with skip_reshape.
It also allows masks of various shapes to be passed, which will be
automtically expanded (in a memory-efficient way) to a size that is
compatible with xformers or pytorch sdpa respectively.

* fix mask shapes
2024-12-16 18:21:17 -05:00
..
2024-12-12 17:59:16 -05:00
2023-01-16 22:37:14 -05:00
2024-06-27 18:43:11 -04:00
2024-11-21 08:38:23 -05:00
2024-06-10 16:00:03 -04:00
2024-12-12 17:59:16 -05:00
2024-11-28 06:46:15 -05:00
2024-12-07 14:00:31 -05:00
2024-12-12 17:59:16 -05:00
2024-12-12 17:59:16 -05:00
2024-12-12 17:59:16 -05:00