Apple Silicon Compatibility with torch > 2.4.1

[ShinChven]

Custom Node Testing

• [x] I have tried disabling custom nodes and the issue persists (see how to disable custom nodes if you need help)

Expected Behavior

Generate image larger than 2048x2048 on Apple Silicon macOS with latest requirements.txt

Actual Behavior

ComfyUI crashes with error:

/AppleInternal/Library/BuildRoots/.../MPSNDArray.mm:829: failed assertion ... Error: NDArray dimension length > INT_MAX'

Steps to Reproduce

1. Install latest version of ComfyUI with torch > 2.4.1 on Apple Silicon macOS
2. Generate an image with dimensions larger than 2048 x 2048

Debug Logs

/AppleInternal/Library/BuildRoots/.../MPSNDArray.mm:829: failed assertion ... Error: NDArray dimension length > INT_MAX'

Other

I have tried to downgrade Torch to 2.4.1, which fixed the issue temporarily.

pip install torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1

However, the issue remains with any later version of torch.

[otecmichal]

I have similar issues with the desktop Comfy-UI. It currently ships with torch 2.5.1 and it crashes at vae decode nodes when the latent image is somewhat larger (i did not try to pin the exact size where this happens). The same workflow does not crash on comfy-ui installed from github with torch==2.4.1 as indicated by ShinChven