[Feature Request]: Implementation of Token Merging for fast SD

[GreenLandisaLie]

Is there an existing issue for this?

• I have searched the existing issues and checked the recent builds/commits

What would your feature do ?

Assuming its possible - an implementation of https://github.com/dbolya/tomesd would be great!
They claim significant memory savings and inference speed improvements with very little quality loss.
However, it requires pytorch >= 1.12.1 (for scatter_reduce) so I'm not sure this can be ported to DirectML though.

Proposed workflow

.

Additional information

No response

[lshqqytiger]

DirectML seems to be weak on handling partially modified dimensions. (dim < 0)
I added some fallback (to cpu) for it. Then, it was slower and generated some broken images.

[GreenLandisaLie]

Thank you for taking a look into it.
Since its not something solvable on your end, should I close this issue or keep it open in hopes the DirectML team fixes this someday?

[lshqqytiger]

I think they'll fix it someday, but I don't know when it will be. I don't think we need to maintain this issue because an extension for Token Merging is already out.