[Performance 2/6] Replace einops.rearrange with torch native ops

[huchenlei]

trafficstars

Description

According to https://github.com/lllyasviel/stable-diffusion-webui-forge/discussions/716#discussioncomment-9336461, einops.rearrange calls in crossattn is causing extra overhead. Replacing it with torch native ops can save ~55ms/it.

Screenshots/videos:

TODO

There are other places where einops.rearrange can be replaced by torch native ops, but this one in CrossAttn is the most critical one. Instrument the usage of einops.rearrange elsewhere might also yield some improvements.

Checklist:

• [x] I have read contributing wiki page
• [x] I have performed a self-review of my own code
• [x] My code follows the style guidelines
• [x] My code passes tests