mindspore-lab / mindone

one for all, Optimal generator with No Exception
https://mindspore-lab.github.io/mindone/
Apache License 2.0
353 stars 67 forks source link

Opensora-PKU: fix fa attention mask #541

Closed wtomin closed 3 months ago

wtomin commented 3 months ago

FlashAttention in MindSpore only accepts mask with (0, 1), where 0 means to keep, and 1 means to discard.

Editing the current Opensora-PKU attention mask:

  1. comment the code to fill in mask with -ms.numpy.inf in LatteT2V.construct;
  2. Flip the mask using 1 - mask since FA treats 1 as discard, 0 as retain.
  3. Fill in with -ms.numpy.inf in vanilla Attention, and leave the mask untouched for FA

Minor changes:

  1. Print Use FA in sample_t2v.py;
  2. Force tensor to fp32 before saving to npz in sample_text_embed.py.