Error Running Pull #25 Version Files On Mac

[jrittvo]

I just downloaded and tried the updated merge.py and patch.py files from the merged pull #25 in a simple python diffusers CLI pipeline. Tomesd worked fine on my Mac with the prior versions. The new versions give the traceback below. Would it help you troubleshoot if I swapped out one file at a time, or are they dependent on changes in each other to work at all?

Traceback (most recent call last):
File "/Users/jrittvo/miniconda3/envs/python_playground/diffusers/gen-sd21.py", line 25, in
image = pipe(args.prompt, generator=generator, num_inference_steps=24, guidance_scale=7.5, height=768, width=768).images[0]
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 676, in call
noise_pred = self.unet(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
result = forward_call(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/models/unet_2d_condition.py", line 685, in forward
sample, res_samples = downsample_block(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/models/unet_2d_blocks.py", line 867, in forward
hidden_states = attn(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/models/transformer_2d.py", line 265, in forward
hidden_states = block(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/patch.py", line 105, in forward
m_a, m_c, m_m, u_a, u_c, u_m = compute_merge(hidden_states, self._tome_info)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/patch.py", line 38, in compute_merge
args["generator"] = init_generator(x.device)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/patch.py", line 19, in init_generator
return torch.Generator(device=device).set_state(torch.mps.get_rng_state())
AttributeError: module 'torch' has no attribute 'mps'

[dbolya]

Huh, that mps is there specifically for macs, but maybe it requires a newer torch version. What torch version are you on? I can probably just default it to use cpu if mps isn't available.

[dbolya]

Was able to grab a hold of a friend with an M1 mac and I think I found a solution to the issue.

Does the most recent commit work for you?

[jrittvo]

I set an environment variable in my python diffusers pipeline, that is necessary when I run the full Automatic1111 wrbui:

export PYTORCH_ENABLE_MPS_FALLBACK=1

And this is the full terminal output on a successful run:

(python_playground) jrittvo@M1PRO diffusers % python gen-sd21.py --prompt "photograph of a mouse" --seed "123" --save "images/mouse123.jpg"
0%| | 0/24 [00:00<?, ?it/s]/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/merge.py:66: UserWarning: MPS: no support for int64 min/max ops, casting it to int32 (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/Sort.mm:39.)
rand_idx = idx_buffer.reshape(1, -1, 1).argsort(dim=1)
/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/merge.py:104: UserWarning: The operator 'aten::scatter_reduce.two_out' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode)
100%|███████████████████████████████████████████| 24/24 [00:43<00:00, 1.82s/it]

Changing ONLY the merge.py and patch.py files, I get the traceback in my first post above. You may be right that it is not actually using mps for some or all of the pipeline, but the prior tomesd file versions don't trigger any traceback.

If I try without the environment variable set, the new versions give the same traceback. The old versions give this traceback:

NotImplementedError: The operator 'aten::scatter_reduce.two_out' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on pytorch/pytorch#77764. As a temporary fix, you can set the environment variable PYTORCH_ENABLE_MPS_FALLBACK=1 to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS.

I'm using torch 2.0.0.dist. I have torch 2.1.0.dev20230407.dist on my Automatic1111. I ca try swapping that newer package over to my python diffusers pipeline if that would be useful to you. I also have a special torch build that brkirch bundles with his experimental Mac build of Automatic1111. I am going to try that one in a bit, because his pieces have fixes for Macs that have not been upstreamed yet.

[jrittvo]

Sit tight for a few minutes and I'll try them.

[jrittvo]

Nope on the new new build. I get the traceback below with both PYTORCH_ENABLE_MPS_FALLBACK=1 and PYTORCH_ENABLE_MPS_FALLBACK=0. I assume that =0 is taking it in and out of the mix??

(python_playground) jrittvo@M1PRO diffusers % python gen-sd21.py --prompt "photograph of a mouse" --seed "123" --save "images/mouse123.jpg"
0%| | 0/24 [00:00<?, ?it/s]
Traceback (most recent call last):
File "/Users/jrittvo/miniconda3/envs/python_playground/diffusers/gen-sd21.py", line 25, in
image = pipe(args.prompt, generator=generator, num_inference_steps=24, guidance_scale=7.5, height=768, width=768).images[0]
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
return func(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 676, in call
noise_pred = self.unet(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
result = forward_call(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/models/unet_2d_condition.py", line 685, in forward
sample, res_samples = downsample_block(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/models/unet_2d_blocks.py", line 867, in forward
hidden_states = attn(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/diffusers/models/transformer_2d.py", line 265, in forward
hidden_states = block(
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(args, **kwargs)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/patch.py", line 97, in forward
m_a, m_c, m_m, u_a, u_c, u_m = compute_merge(hidden_states, self._tome_info)
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/patch.py", line 33, in compute_merge
m, u = merge.bipartite_soft_matching_random2d(x, w, h, args["sx"], args["sy"], r,
File "/Users/jrittvo/miniconda3/envs/python_playground/lib/python3.10/site-packages/tomesd/merge.py", line 52, in bipartite_soft_matching_random2d
rand_idx = torch.randint(sysx, size=(hsy, wsx, 1), device=metric.device, generator=generator)
RuntimeError: Expected a 'mps' device type for generator but found 'cpu'

[dbolya]

Thanks for all the information. I'll have direct access to an M1 mac in a week or two, and then I can create a natively compatible version that doesn't require PYTORCH_ENABLE_MPS_FALLBACK=1.

Otherwise, looks like the fix did make some progress. Can you try the latest commit? I think that should be all that's needed.

[jrittvo]

Yes, I'll grab it in a few minutes and let you know. I tried the brkirch build of torch and it didn't help with our without MPS_FALLBACK enabled. I do think that is something he has been working on though.

[jrittvo]

Your latest set works again (with MPS_FALLBACK=1).

Do you you know who brkirch is? He seems to be an upper level developer working with Automatic1111 on Mac issues. He has an experimantal build that has had tomesd in it for a few weeks now. This is his repo: https://github.com/brkirch/stable-diffusion-webui/actions. If you look at the "Actions" there, you can see what he has been working on. I think he has spent a fair amount of time recently on tomesd and Macs.

[dbolya]

Oh huh, no I was not aware of it. I'll see if I can reach out once I start adding native support.

Thanks for all the testing! Closing the issue, as it seems to work now.

[jrittvo]

Yes, issue resolved. Glad to be able to help in a small way.