* Fix SDP attention mask by creating a diagonal block mask in the right shape.

city96 / ComfyUI_ExtraModels

Support for miscellaneous image models. Currently supports: DiT, PixArt, HunYuanDiT, MiaoBi, and a few VAEs.

Apache License 2.0

393 stars 34 forks source link

* Fix SDP attention mask by creating a diagonal block mask in the right shape. #36

Closed GavChap closed 6 months ago

GavChap commented 6 months ago

This corrects the diagonal mask so it's the equivalent mask that xformers uses.

GavChap commented 6 months ago

Batch of 3, SDP attention: Batch of 4: Batch of 8:

city96 commented 6 months ago

Amazing job! Thank you for digging into this.

GavChap commented 6 months ago

Amazing job! Thank you for digging into this.

You're welcome, I'm glad the fix turned out to be a few lines of code, and I learned things along the way! Additionally It seems that SDP is actually a bit faster than xformers now at least from my testing on my 4060 anyway.