Failed to run on M1Mac

[KohakuBlueleaf]

I'm using M1Mac and sd-webui.
I try to add apply_patch in load_model method and it works normally.
But when I want to generate image, it jump out some error and I'm not sure if it is caused by M1 Mac environment/ my fault or some bugs in your repo
So I decide to post it here.

You can just close this issue if you think this is caused by myself.

error message:

/AppleInternal/Library/BuildRoots/c651a45f-806e-11ed-a221-7ef33c48bc85/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSNDArray/Kernels/MPSNDArrayGatherND.mm:234: failed assertion `Rank of updates array (1) must be greater than or equal to inner-most dimension of indices array (2048)'
[1]    22333 abort      ./webui.sh
/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

[dbolya]

Hmm seems the mac implementation of gather might have different requirements. Can you find the exact line in tome's merge.py that causes this? As well as the tensor sizes for that operation.

[KohakuBlueleaf]

@dbolya i can do some inspect on it this night
Wait me few hrs

[KohakuBlueleaf]

@dbolya I found the reason

tomesd/tomesd/merge.py

Line 68 in 113ddc3

dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx)

This looks like mps doesn't accept index array size < input array

[KohakuBlueleaf]

I tried the shape I get when error occured and get this kind result:

[KohakuBlueleaf]

@dbolya I run it on m1 mac successfully now
And it should not affect other device type
I will do more experiment and make pr for it

[dvttran]

@dbolya I run it on m1 mac successfully now
And it should not affect other device type
I will do more experiment and make pr for it

@KohakuBlueleaf
I am new to the community. Can you please show me how to plug it into StableDiffusion WebUI and how to fix the error with M1 Chip?

Thank you.

[KohakuBlueleaf]

@tvdtran For sd webui
You can just add some line in sd_model.py's load_model method
or install the tome extension

And if you want to fix m1 problem, you can wait me to upload the fix into my fork(

[KohakuBlueleaf]

And since my method is to convert the x to cpu when we need to to merge/unmerge
So actually it will comsume a little bit more memory (than invoke AI optimization only)
but the speed is increased after apply patch

maybe we need some more efficient way(

[dbolya]

@dbolya I found the reason

tomesd/tomesd/merge.py

Line 68 in 113ddc3

dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx)

This looks like mps doesn't accept index array size < input array

Hmm, this is a big issue because that's required behavior for gather there. I wonder if computing the merge on the CPU is okay? Does performance suffer a lot? If not, I could just wrap those operations in conversions to CPU if using MPS.

[KohakuBlueleaf]

@dbolya merge on cpu is ok but the ram usage is increased